[Xarray-Beam] Add support for Zarr v3 sharding in ChunkToZarr

PiperOrigin-RevId: 810057956

Author: shoyer

examples/xbeam_rechunk.py

@@ -30,6 +30,21 @@
         'chunksize of -1 indicates not to chunk a dimension.'
     ),
 )
+TARGET_SHARDS = flags.DEFINE_string(
+    'target_shards',
+    None,
+    help=(
+        'Desired shards for each dimension in the output Zarr dataset, in the '
+        'same format as --target_chunks. If omitted, sharding is not used. '
+        'Shards should be multiples of corresponding chunk sizes. Only valid '
+        'with Zarr v3.'
+    ),
+)
+ZARR_FORMAT = flags.DEFINE_integer(
+    'zarr_format',
+    None,
+    help='Zarr format to use for the output.',
+)
 RUNNER = flags.DEFINE_string('runner', None, help='beam.runners.Runner')
 
 
@@ -48,7 +63,14 @@ def _parse_chunks_str(chunks_str: str) -> dict[str, int]:
 def main(argv):
   source_dataset, source_chunks = xbeam.open_zarr(INPUT_PATH.value)
   template = xbeam.make_template(source_dataset)
-  target_chunks = dict(source_chunks, **_parse_chunks_str(TARGET_CHUNKS.value))
+
+  target_chunks = source_chunks | _parse_chunks_str(TARGET_CHUNKS.value)
+
+  if TARGET_SHARDS.value is not None:
+    target_shards = source_chunks | _parse_chunks_str(TARGET_SHARDS.value)
+  else:
+    target_shards = None
+
   itemsize = max(variable.dtype.itemsize for variable in template.values())
 
   with beam.Pipeline(runner=RUNNER.value, argv=argv) as root:
@@ -58,10 +80,16 @@ def main(argv):
         | xbeam.Rechunk(  # pytype: disable=wrong-arg-types
             source_dataset.sizes,
             source_chunks,
-            target_chunks,
+            target_chunks if target_shards is None else target_shards,
             itemsize=itemsize,
         )
-        | xbeam.ChunksToZarr(OUTPUT_PATH.value, template, target_chunks)
+        | xbeam.ChunksToZarr(
+            OUTPUT_PATH.value,
+            template,
+            zarr_chunks=target_chunks,
+            zarr_shards=target_shards,
+            zarr_format=ZARR_FORMAT.value,
+        )
     )
 
 

examples/xbeam_rechunk_test.py

@@ -23,7 +23,7 @@
 
 class Era5RechunkTest(test_util.TestCase):
 
-  def test(self):
+  def test_chunks_only(self):
     input_path = self.create_tempdir('source').full_path
     output_path = self.create_tempdir('destination').full_path
 
@@ -44,6 +44,32 @@ def test(self):
     )
     xarray.testing.assert_identical(input_ds, output_ds)
 
+  def test_chunks_and_shards(self):
+    input_path = self.create_tempdir('source').full_path
+    output_path = self.create_tempdir('destination').full_path
+
+    input_ds = test_util.dummy_era5_surface_dataset(times=365)
+    input_ds.chunk({'time': 31}).to_zarr(input_path)
+
+    with flagsaver.flagsaver(
+        input_path=input_path,
+        output_path=output_path,
+        target_chunks='latitude=5,longitude=5,time=-1',
+        target_shards='latitude=10,longitude=10,time=-1',
+        zarr_format=3,
+    ):
+      xbeam_rechunk.main([])
+
+    output_ds = xarray.open_zarr(output_path)
+    self.assertEqual(
+        {k: v[0] for k, v in output_ds.chunks.items()},
+        {'latitude': 5, 'longitude': 5, 'time': 365}
+    )
+    actual_shards = {k: v.encoding['shards'] for k, v in output_ds.items()}
+    expected_shards = {k: (365, 10, 10) for k, v in output_ds.items()}
+    self.assertEqual(actual_shards, expected_shards)
+    xarray.testing.assert_identical(input_ds, output_ds)
+
 
 if __name__ == '__main__':
   absltest.main()

pyproject.toml

@@ -35,6 +35,7 @@ tests = [
     "hypothesis",
     "pandas",
     "pytest",
+    'zarr>=3',
 ]
 docs = [
     'myst-nb',

xarray_beam/_src/core.py

@@ -136,11 +136,14 @@ def __setstate__(self, state):
     self.__init__(*state)
 
 
+K = TypeVar("K")
+
+
 def offsets_to_slices(
-    offsets: Mapping[str, int],
-    sizes: Mapping[str, int],
-    base: Mapping[str, int] | None = None,
-) -> dict[str, slice]:
+    offsets: Mapping[K, int],
+    sizes: Mapping[K, int],
+    base: Mapping[K, int] | None = None,
+) -> dict[K, slice]:
   """Convert offsets into slices with an optional base offset.
 
   Args: