Add read_shards option to xbeam.open_zarr() and Dataset.from_zarr().

docs/read-write.ipynb

@@ -128,7 +128,7 @@
         "The easiest way to get good performance from Zarr into Xarray-Beam is to use {py:func}`xarray_beam.open_zarr`. This function returns a pair of values:\n",
         "\n",
         "1. A lazily indexed `xarray.Dataset` corresponding to the Zarr store, but not using Dask. This is exactly what you would get from `xarray.open_zarr` with `chunks=None`.\n",
-        "2. A dictionary mapping from dimension names to integer chunk sizes. Obtaining this information without using Dask to chunk the array requires looking at Xarray's `encoding` dictionaries or directly inspecting the Zarr store."
+        "2. A dictionary mapping from dimension names to integer chunk sizes. Obtaining this information without using Dask to chunk the array requires looking at Xarray's `encoding` dictionaries or directly inspecting the Zarr store. For sharded Zarr v3 stores, these are the inner chunk sizes by default; pass `read_shards=True` to get shard sizes instead, so that each Beam chunk covers whole shards."
       ]
     },
     {

xarray_beam/_src/dataset.py

@@ -665,6 +665,7 @@ def from_zarr(
       path: str,
       *,
       chunks: UnnormalizedChunks | None = None,
+      read_shards: bool = False,
       split_vars: bool = False,
       pipeline: beam.Pipeline | None = None,
       label: str | None = None,
@@ -676,6 +677,9 @@ def from_zarr(
       chunks: optional mapping from dimension names to chunk sizes, or any value
         that can be passed to :py:func:`xarray_beam.normalize_chunks`. If not
         provided, the chunk sizes will be inferred from the Zarr file.
+      read_shards: If True, read shard sizes instead of chunk sizes for sharded
+        Zarr v3 datasets. If False (default), or if the store is not sharded,
+        use chunk sizes.
       split_vars: whether variables are split between separate elements in the
         ptransform, or all stored in the same element.
       pipeline: Beam pipeline to use for this dataset. If not provided, you will
@@ -688,7 +692,7 @@ def from_zarr(
     """
     if label is None:
       label = _get_label('from_zarr')
-    source, previous_chunks = zarr.open_zarr(path)
+    source, previous_chunks = zarr.open_zarr(path, read_shards=read_shards)
     if chunks is None:
       chunks = previous_chunks  # pyrefly: ignore[bad-assignment]
     result = cls.from_xarray(

xarray_beam/_src/dataset_test.py

@@ -633,6 +633,26 @@ def test_from_zarr_with_chunks(self):
     collected = beam_ds.collect_with_direct_runner()
     xarray.testing.assert_identical(ds, collected)
 
+  def test_from_zarr_with_shards(self):
+    temp_dir = self.create_tempdir().full_path
+    ds = xarray.Dataset({'foo': ('x', np.arange(12))})
+    beam_ds = xbeam.Dataset.from_xarray(ds, {'x': 6})
+    with beam.Pipeline() as p:
+      p |= beam_ds.to_zarr(
+          temp_dir,
+          zarr_chunks={'x': 2},
+          zarr_shards={'x': 6},
+      )
+    read_ds = xbeam.Dataset.from_zarr(temp_dir)
+    self.assertEqual(read_ds.chunks, {'x': 2})
+    collected = read_ds.collect_with_direct_runner()
+    xarray.testing.assert_identical(ds, collected)
+
+    read_ds_shards = xbeam.Dataset.from_zarr(temp_dir, read_shards=True)
+    self.assertEqual(read_ds_shards.chunks, {'x': 6})
+    collected_shards = read_ds_shards.collect_with_direct_runner()
+    xarray.testing.assert_identical(ds, collected_shards)
+
   def test_to_zarr(self):
     temp_dir = self.create_tempdir().full_path
     ds = xarray.Dataset({'foo': ('x', np.arange(12))})

xarray_beam/_src/zarr.py

@@ -45,7 +45,7 @@
 WritableStore = str | zarr_storage.StoreLike | os.PathLike[str]
 
 
-def _infer_chunks(dataset: xarray.Dataset) -> dict[str, int]:
+def _infer_chunks(dataset: xarray.Dataset, read_shards: bool) -> dict[str, int]:
   """Infer chunks for an xarray.Dataset loaded from Zarr."""
   # The original Zarr array chunks (as tuples) are stored in the "encoding"
   # dictionary on each DataArray object.
@@ -54,23 +54,29 @@ def _infer_chunks(dataset: xarray.Dataset) -> dict[str, int]:
   for name, variable in dataset.items():
     # exclude variables that are indexed, which are loaded into memory already
     if name not in dataset.indexes:
-      for dim, size in zip(variable.dims, variable.encoding['chunks']):
+      shards = variable.encoding.get('shards') if read_shards else None
+      var_chunks = shards or variable.encoding['chunks']
+      for dim, size in zip(variable.dims, var_chunks):
         chunks_sets[dim].add(size)
 
   chunks = {}
   for dim, sizes in chunks_sets.items():
     if len(sizes) > 1:
+      hint = ' (shard sizes used where present)' if read_shards else ''
       raise ValueError(
           f'inconsistent chunk sizes on Zarr dataset for dimension {dim!r}: '
-          f'{sizes}'
+          f'{sizes}{hint}'
       )
     (chunks[dim],) = sizes
   return chunks
 
 
 @core.export
 def open_zarr(
-    store: ReadableStore, **kwargs: Any
+    store: ReadableStore,
+    *,
+    read_shards: bool = False,
+    **kwargs: Any,
 ) -> tuple[xarray.Dataset, dict[str, int]]:
   """Returns a lazily indexable xarray.Dataset and chunks from a Zarr store.
 
@@ -80,6 +86,9 @@ def open_zarr(
 
   Args:
     store: Xarray compatible Zarr store to open.
+    read_shards: If True, read shard sizes instead of chunk sizes for sharded
+      Zarr v3 datasets. If False (default), or if the store is not sharded,
+      use chunk sizes.
     **kwargs: passed on to xarray.open_zarr. The "chunks" keyword argument is
       not supported.
 
@@ -92,7 +101,7 @@ def open_zarr(
         'xarray_beam.open_zarr does not support the `chunks` argument'
     )
   dataset = xarray.open_zarr(store, **kwargs, chunks=None)
-  chunks = _infer_chunks(dataset)
+  chunks = _infer_chunks(dataset, read_shards=read_shards)
   return dataset, chunks
 
 

xarray_beam/_src/zarr_test.py

@@ -43,6 +43,31 @@ def test_open_zarr(self):
     self.assertEqual(roundtripped_ds.chunks, {})
     self.assertEqual(chunks, {'x': 2})
 
+  def test_open_zarr_with_shards(self):
+    dataset = xarray.Dataset(
+        {'foo': ('x', np.arange(0, 60, 10))},
+        coords={'x': np.arange(6)},
+    )
+    template = xbeam.make_template(dataset)
+    inputs = [
+        (xbeam.Key({'x': 0}), dataset.isel(x=slice(0, 3))),
+        (xbeam.Key({'x': 3}), dataset.isel(x=slice(3, 6))),
+    ]
+    temp_dir = self.create_tempdir().full_path
+    inputs | xbeam.ChunksToZarr(
+        temp_dir,
+        template=template,
+        zarr_chunks={'x': 1},
+        zarr_shards={'x': 3},
+        zarr_format=3,
+    )
+    roundtripped_ds, chunks = xbeam.open_zarr(temp_dir)
+    xarray.testing.assert_identical(roundtripped_ds, dataset)
+    self.assertEqual(chunks, {'x': 1})
+
+    _, shard_chunks = xbeam.open_zarr(temp_dir, read_shards=True)
+    self.assertEqual(shard_chunks, {'x': 3})
+
   def test_open_zarr_inconsistent(self):
     source_ds = xarray.Dataset(
         {
@@ -58,6 +83,28 @@ def test_open_zarr_inconsistent(self):
     ):
       xbeam.open_zarr(temp_dir)
 
+  def test_open_zarr_inconsistent_shards(self):
+    source_ds = xarray.Dataset(
+        {'foo': ('x', np.arange(6)), 'bar': ('x', np.arange(6))}
+    )
+    temp_dir = self.create_tempdir().full_path
+    source_ds.to_zarr(
+        temp_dir,
+        zarr_format=3,
+        encoding={
+            'foo': {'chunks': (1,), 'shards': (3,)},
+            'bar': {'chunks': (1,)},
+        },
+    )
+    _, chunks = xbeam.open_zarr(temp_dir)
+    self.assertEqual(chunks, {'x': 1})
+    with self.assertRaisesRegex(
+        ValueError,
+        r"inconsistent chunk sizes on Zarr dataset for dimension 'x': {1, 3} "
+        r'\(shard sizes used where present\)',
+    ):
+      xbeam.open_zarr(temp_dir, read_shards=True)
+
   def test_make_template(self):
     source = xarray.Dataset(
         {