Add local staging to Zarr setup in xarray_beam.

Fixes #122

This change introduces a `stage_locally` parameter to `setup_zarr`, `ChunksToZarr` and `Dataset.to_zarr`. When enabled, Zarr metadata is first written to a local temporary directory and then copied to the final destination in parallel using `fsspec`. This can significantly speed up the setup process on high-latency filesystems, e.g., in one example, I found it sped up Zarr setup by a factor of 25x, from 100 seconds to 4 seconds.

This adds a hard dependency on fsspec in Xarray-Beam.

Hopefully in the future Xarray will have concurrent writing to stores built in (see pydata/xarray#10622), which will eliminate the primary need for this.

Alternatively, we might be able to eventually leverage Zarr's built-in stores to do this copying rather than fsspec. Zarr has all the necessary functionality (including atomic writes, which would be nice) but does not expose the required public APIs for copying store objects from a synchronous function.

PiperOrigin-RevId: 817684876

Author: shoyer

docs/requirements.txt

@@ -10,6 +10,7 @@ matplotlib==3.10.6
 # xarray-beam requirements
 apache-beam==2.67.0
 dask==2025.9.1
+fsspec==2025.9.0
 immutabledict==4.2.1
 numpy==2.2.6
 pandas==2.3.2

pyproject.toml

@@ -15,10 +15,10 @@ authors = [
 readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.10"
-# TODO(shoyer): thin these down
 dependencies = [
     "apache_beam>=2.31.0",
     "dask",
+    "fsspec",
     "immutabledict",
     "zarr",
     "xarray",

xarray_beam/__init__.py

@@ -55,4 +55,4 @@
     DatasetToZarr as DatasetToZarr,
 )
 
-__version__ = '0.10.5'  # automatically synchronized to pyproject.toml
+__version__ = '0.11.0'  # automatically synchronized to pyproject.toml

xarray_beam/_src/dataset.py

@@ -622,6 +622,7 @@ def to_zarr(
       zarr_chunks: UnnormalizedChunks | None = None,
       zarr_shards: UnnormalizedChunks | None = None,
       zarr_format: int | None = None,
+      stage_locally: bool | None = None,
   ) -> beam.PTransform:
     """Write this dataset to a Zarr file.
 
@@ -650,6 +651,11 @@ def to_zarr(
       zarr_format: optional integer specifying the explicit Zarr format to use.
         Defaults to Zarr v3 if using shards, or the default format for your
         installed version of Zarr.
+      stage_locally: If True, write Zarr metadata to a local temporary directory
+        before copying to `store` in parallel. This can significantly speed up
+        setup on high-latency filesystems. By default, uses local staging if
+        possible, which is true as long as `store` is provided as as string or
+        path.
 
     Returns:
       Beam PTransform that writes the dataset to a Zarr file.
@@ -710,6 +716,7 @@ def to_zarr(
         zarr_chunks=zarr_chunks,
         zarr_shards=zarr_shards,
         zarr_format=zarr_format,
+        stage_locally=stage_locally,
     )
 
   def collect_with_direct_runner(self) -> xarray.Dataset:

xarray_beam/_src/zarr.py

@@ -16,16 +16,19 @@
 
 import collections
 from collections.abc import Mapping, Set
+import concurrent.futures
 import dataclasses
 import logging
 import os
 import pprint
-from typing import Any, TypeVar
+import tempfile
+from typing import Any
 import warnings
 
 import apache_beam as beam
 import dask
 import dask.array
+import fsspec
 import numpy as np
 import pandas as pd
 import xarray
@@ -34,6 +37,7 @@
 from xarray_beam._src import threadmap
 from zarr import storage as zarr_storage
 
+
 # pylint: disable=logging-fstring-interpolation
 
 # Match the types accepted by xarray.open_zarr() and to_zarr().
@@ -366,13 +370,30 @@ def _get_chunk_and_shard_encoding(
   return encoding
 
 
+def _copy_zarr_store_with_fsspec(
+    source_dir: str, dest_store: str | os.PathLike[str], num_threads: int = 128
+) -> None:
+  """Copy a Zarr store from one location to another using fsspec."""
+  source_mapper = fsspec.get_mapper(source_dir)
+  dest_mapper = fsspec.get_mapper(dest_store)
+
+  def copy_item(key):
+    dest_mapper[key] = source_mapper[key]
+
+  with concurrent.futures.ThreadPoolExecutor(num_threads) as executor:
+    for _ in executor.map(copy_item, source_mapper):
+      pass
+
+
 def _setup_zarr(
     template: xarray.Dataset,
     store: WritableStore,
     zarr_chunks: Mapping[str, int] | None = None,
     zarr_shards: Mapping[str, int] | None = None,
     zarr_format: int | None = None,
     encoding: Mapping[str, Any] | None = None,
+    *,
+    stage_locally: bool | None = None,
 ) -> None:
   """setup_zarr() without finalizing args."""
   if encoding is None:
@@ -401,14 +422,39 @@ def _setup_zarr(
       f'writing Zarr metadata for template:\n{template}\n'
       f'encoding={encoding_str}'
   )
-  template.to_zarr(
-      store,
-      compute=False,
-      consolidated=True,
-      mode='w',
-      zarr_format=zarr_format,
-      encoding=encoding,
-  )
+  if stage_locally is None:
+    stage_locally = isinstance(store, (str, os.PathLike))
+    if not stage_locally:
+      logging.info(
+          'skipping local staging, because store is not a string or PathLike'
+      )
+
+  if stage_locally:
+    if not isinstance(store, (str, os.PathLike)):
+      raise ValueError(
+          'only path-like stores are supported when stage_locally=True'
+      )
+    with tempfile.TemporaryDirectory() as tmpdir:
+      logging.info(f'writing temporary copy to {tmpdir}')
+      template.to_zarr(
+          tmpdir,
+          compute=False,
+          consolidated=True,
+          mode='w',
+          zarr_format=zarr_format,
+          encoding=encoding,
+      )
+      logging.info(f'copying temporary copy to {store}')
+      _copy_zarr_store_with_fsspec(tmpdir, store)
+  else:
+    template.to_zarr(
+        store,
+        compute=False,
+        consolidated=True,
+        mode='w',
+        zarr_format=zarr_format,
+        encoding=encoding,
+    )
   logging.info('finished setting up Zarr')
 
 
@@ -419,6 +465,8 @@ def setup_zarr(
     zarr_shards: Mapping[str, int] | None = None,
     zarr_format: int | None = None,
     encoding: Mapping[str, Any] | None = None,
+    *,
+    stage_locally: bool | None = None,
 ) -> None:
   """Setup a Zarr store.
 
@@ -441,14 +489,25 @@ def setup_zarr(
       possible, otherwise defaulting to the default version used by the
       zarr-python library installed.
     encoding: Nested dictionary with variable names as keys and dictionaries of
-      variable specific encodings as values, e.g.,
-      ``{"my_variable": {"dtype": "int16", "scale_factor": 0.1,}, ...}``
+      variable specific encodings as values, e.g., ``{"my_variable": {"dtype":
+      "int16", "scale_factor": 0.1,}, ...}``
+    stage_locally: If True, write Zarr metadata to a local temporary directory
+      before copying to `store` in parallel. This can significantly speed up
+      setup on high-latency filesystems. By default, uses local staging if
+      possible, which is true as long as `store` is provided as as string or
+      path.
   """
   template, zarr_chunks, zarr_shards = _finalize_setup_zarr_args(
       template, zarr_chunks, zarr_shards
   )
   _setup_zarr(
-      template, store, zarr_chunks, zarr_shards, zarr_format, encoding
+      template,
+      store,
+      zarr_chunks,
+      zarr_shards,
+      zarr_format,
+      encoding,
+      stage_locally=stage_locally,
   )
 
 
@@ -584,6 +643,7 @@ def __init__(
       num_threads: int | None = None,
       needs_setup: bool = True,
       encoding: Mapping[str, Any] | None = None,
+      stage_locally: bool | None = None,
   ):
     # pyformat: disable
     """Initialize ChunksToZarr.
@@ -642,6 +702,11 @@ def __init__(
       encoding : Nested dictionary with variable names as keys and dictionaries
         of variable specific encodings as values, e.g.,
         ``{"my_variable": {"dtype": "int16", "scale_factor": 0.1,}, ...}``
+      stage_locally: If True, write Zarr metadata to a local temporary directory
+        before copying to `store` in parallel. This can significantly speed up
+        setup on high-latency filesystems. By default, uses local staging if
+        possible, which is true as long as `store` is provided as as string or
+        path.
     """
     # pyformat: enable
 
@@ -653,7 +718,13 @@ def __init__(
       )
       if needs_setup:
         _setup_zarr(
-            template, store, zarr_chunks, zarr_shards, zarr_format, encoding
+            template,
+            store,
+            zarr_chunks,
+            zarr_shards,
+            zarr_format,
+            encoding,
+            stage_locally=stage_locally,
         )
     elif isinstance(template, beam.pvalue.AsSingleton):
       if not needs_setup:
@@ -685,6 +756,7 @@ def __init__(
     self.num_threads = num_threads
     self.zarr_format = zarr_format
     self.encoding = encoding
+    self.stage_locally = stage_locally
 
   def _validate_zarr_chunk(self, key, chunk, template=None):
     # If template doesn't have a default value, Beam errors with "Side inputs
@@ -716,12 +788,15 @@ def expand(self, pcoll):
           template.pvalue
           | 'SetupZarr'
           >> beam.Map(
-              setup_zarr,
-              self.store,
-              self.zarr_chunks,
-              self.zarr_shards,
-              self.zarr_format,
-              self.encoding,
+              lambda t: setup_zarr(
+                  t,
+                  self.store,
+                  self.zarr_chunks,
+                  self.zarr_shards,
+                  self.zarr_format,
+                  self.encoding,
+                  stage_locally=self.stage_locally,
+              )
           )
       )
     return (

xarray_beam/_src/zarr_test.py

@@ -206,6 +206,11 @@ def test_chunks_to_zarr(self):
       inputs | xbeam.ChunksToZarr(temp_dir, chunked)
       result = xarray.open_zarr(temp_dir, consolidated=True)
       xarray.testing.assert_identical(dataset, result)
+    with self.subTest('with template and stage_locally=True'):
+      temp_dir = self.create_tempdir().full_path
+      inputs | xbeam.ChunksToZarr(temp_dir, chunked, stage_locally=True)
+      result = xarray.open_zarr(temp_dir, consolidated=True)
+      xarray.testing.assert_identical(dataset, result)
     with self.subTest('with template and needs_setup=False'):
       temp_dir = self.create_tempdir().full_path
       xbeam.setup_zarr(chunked, temp_dir)
@@ -427,6 +432,30 @@ def test_chunks_to_zarr_with_invalid_shards(self):
           zarr_format=3,
       )
 
+  @parameterized.product(
+      stage_locally=[True, False, None],
+      zarr_format=[2, 3, None],
+  )
+  def test_setup_zarr(self, stage_locally, zarr_format):
+    dataset = xarray.Dataset(
+        {'foo': ('x', np.arange(0, 60, 10))},
+        coords={'x': np.arange(6)},
+    )
+    template = xbeam.make_template(dataset)
+    temp_dir = self.create_tempdir().full_path
+    xbeam.setup_zarr(
+        template,
+        temp_dir,
+        zarr_chunks={'x': 3},
+        stage_locally=stage_locally,
+        zarr_format=zarr_format,
+    )
+    # Verify we can open it with Xarray and it has the right structure.
+    ds = xarray.open_zarr(temp_dir)
+    self.assertEqual(ds.sizes, template.sizes)
+    self.assertEqual(ds.chunks, {'x': (3, 3)})
+    xarray.testing.assert_equal(ds.coords['x'], template.coords['x'])
+
   def test_chunks_to_zarr_append(self):
     zarr_chunks = {'t': 1, 'x': 5}