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( {