From 5c7e1275f03c1725eaec8384a6966dabea5b846e Mon Sep 17 00:00:00 2001 From: Xarray-Beam authors Date: Fri, 7 Aug 2026 12:43:36 -0700 Subject: [PATCH] Add read_shards option to xbeam.open_zarr() and Dataset.from_zarr(). Sharded zarr v3 stores can pack many small inner chunks into each shard, so inferring Beam chunks from the inner chunks dispatches too many read operations and can cause race conditions on writing. Introduces `read_shards=True` so the unit of work is the shard size of the source dataset. `read_shards` defaults to `False`, so existing callers remain unchanged. PiperOrigin-RevId: 961070996 --- docs/read-write.ipynb | 2 +- xarray_beam/_src/dataset.py | 6 +++- xarray_beam/_src/dataset_test.py | 20 ++++++++++++++ xarray_beam/_src/zarr.py | 19 +++++++++---- xarray_beam/_src/zarr_test.py | 47 ++++++++++++++++++++++++++++++++ 5 files changed, 87 insertions(+), 7 deletions(-) diff --git a/docs/read-write.ipynb b/docs/read-write.ipynb index 1666aa9..066268e 100644 --- a/docs/read-write.ipynb +++ b/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." ] }, { diff --git a/xarray_beam/_src/dataset.py b/xarray_beam/_src/dataset.py index 3b0a051..638a585 100644 --- a/xarray_beam/_src/dataset.py +++ b/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( diff --git a/xarray_beam/_src/dataset_test.py b/xarray_beam/_src/dataset_test.py index 22aecae..98ee5ce 100644 --- a/xarray_beam/_src/dataset_test.py +++ b/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))}) diff --git a/xarray_beam/_src/zarr.py b/xarray_beam/_src/zarr.py index 56d0f63..25ec123 100644 --- a/xarray_beam/_src/zarr.py +++ b/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,15 +54,18 @@ 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 @@ -70,7 +73,10 @@ def _infer_chunks(dataset: xarray.Dataset) -> dict[str, int]: @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 diff --git a/xarray_beam/_src/zarr_test.py b/xarray_beam/_src/zarr_test.py index a7d9cf9..0f78842 100644 --- a/xarray_beam/_src/zarr_test.py +++ b/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( {