ConsolidatedDimensionCoords Transform -- cleaned up version of PR #556

docs/composition/styles.md

@@ -24,6 +24,14 @@ the recipe pipeline will contain at a minimum the following transforms applied t
 * {class}`pangeo_forge_recipes.transforms.OpenURLWithFSSpec`: retrieves each pattern file using the specified URLs.
 * {class}`pangeo_forge_recipes.transforms.OpenWithXarray`: load each pattern file into an [`xarray.Dataset`](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html).
 * {class}`pangeo_forge_recipes.transforms.StoreToZarr`: generate a Zarr store by combining the datasets.
+* {class}`pangeo_forge_recipes.transforms.ConsolidateDimensionCoordinates`: consolidate the Dimension Coordinates for dataset read performance.
+* {class}`pangeo_forge_recipes.transforms.ConsolidateMetadata`: calls Zarr's convinience function to consolidate metadata.
+
+```{tip}
+If using the {class}`pangeo_forge_recipes.transforms.ConsolidateDimensionCoordinates` transform, make sure to chain on the {class}`pangeo_forge_recipes.transforms.ConsolidateMetadata` transform to your recipe.
+
+```
+
 
 ## Open with Kerchunk, write to virtual Zarr
 

examples/feedstock/gpcp_from_gcs.py

@@ -3,7 +3,13 @@
 import zarr
 
 from pangeo_forge_recipes.patterns import ConcatDim, FilePattern
-from pangeo_forge_recipes.transforms import OpenURLWithFSSpec, OpenWithXarray, StoreToZarr
+from pangeo_forge_recipes.transforms import (
+    ConsolidateDimensionCoordinates,
+    ConsolidateMetadata,
+    OpenURLWithFSSpec,
+    OpenWithXarray,
+    StoreToZarr,
+)
 
 dates = [
     d.to_pydatetime().strftime("%Y%m%d")
@@ -43,5 +49,7 @@ def test_ds(store: zarr.storage.FSStore) -> zarr.storage.FSStore:
         store_name="gpcp.zarr",
         combine_dims=pattern.combine_dim_keys,
     )
+    | ConsolidateDimensionCoordinates()
+    | ConsolidateMetadata()
     | "Test dataset" >> beam.Map(test_ds)
 )

examples/feedstock/noaa_oisst.py

@@ -3,7 +3,13 @@
 import zarr
 
 from pangeo_forge_recipes.patterns import ConcatDim, FilePattern
-from pangeo_forge_recipes.transforms import OpenURLWithFSSpec, OpenWithXarray, StoreToZarr
+from pangeo_forge_recipes.transforms import (
+    ConsolidateDimensionCoordinates,
+    ConsolidateMetadata,
+    OpenURLWithFSSpec,
+    OpenWithXarray,
+    StoreToZarr,
+)
 
 dates = pd.date_range("1981-09-01", "2022-02-01", freq="D")
 
@@ -26,7 +32,7 @@ def test_ds(store: zarr.storage.FSStore) -> zarr.storage.FSStore:
     # TODO: see if --setup-file option for runner fixes this
     import xarray as xr
 
-    ds = xr.open_dataset(store, engine="zarr", chunks={})
+    ds = xr.open_dataset(store, engine="zarr", consolidated=True, chunks={})
     for var in ["anom", "err", "ice", "sst"]:
         assert var in ds.data_vars
     return store
@@ -40,5 +46,7 @@ def test_ds(store: zarr.storage.FSStore) -> zarr.storage.FSStore:
         store_name="noaa-oisst.zarr",
         combine_dims=pattern.combine_dim_keys,
     )
+    | ConsolidateDimensionCoordinates()
+    | ConsolidateMetadata()
     | beam.Map(test_ds)
 )

pangeo_forge_recipes/rechunking.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import xarray as xr
+import zarr
 
 from .aggregation import XarraySchema, determine_target_chunks
 from .chunk_grid import ChunkGrid
@@ -238,3 +239,43 @@ def _sort_by_speed_of_varying(item):
     ds_combined = xr.combine_nested(dsets_to_concat, concat_dim=concat_dims_sorted)
 
     return first_index, ds_combined
+
+
+def _gather_coordinate_dimensions(group: zarr.Group) -> List[str]:
+    return list(
+        set(itertools.chain(*(group[var].attrs.get("_ARRAY_DIMENSIONS", []) for var in group)))
+    )
+
+
+def consolidate_dimension_coordinates(
+    singleton_target_store: zarr.storage.FSStore,
+) -> zarr.storage.FSStore:
+    """Consolidate dimension coordinates chunking"""
+    group = zarr.open_group(singleton_target_store)
+
+    dims = (dim for dim in _gather_coordinate_dimensions(group) if dim in group)
+    for dim in dims:
+        arr = group[dim]
+        attrs = dict(arr.attrs)
+        data = arr[:]
+
+        # This will generally use bulk-delete API calls
+        # config.storage_config.target.rm(dim, recursive=True)
+
+        singleton_target_store.fs.rm(singleton_target_store.path + "/" + dim, recursive=True)
+
+        new = group.array(
+            dim,
+            data,
+            chunks=arr.shape,
+            dtype=arr.dtype,
+            compressor=arr.compressor,
+            fill_value=arr.fill_value,
+            order=arr.order,
+            filters=arr.filters,
+            overwrite=True,
+        )
+
+        new.attrs.update(attrs)
+
+    return singleton_target_store

pangeo_forge_recipes/transforms.py

@@ -21,7 +21,7 @@
 from .combiners import CombineMultiZarrToZarr, CombineXarraySchemas
 from .openers import open_url, open_with_kerchunk, open_with_xarray
 from .patterns import CombineOp, Dimension, FileType, Index, augment_index_with_start_stop
-from .rechunking import combine_fragments, split_fragment
+from .rechunking import combine_fragments, consolidate_dimension_coordinates, split_fragment
 from .storage import CacheFSSpecTarget, FSSpecTarget
 from .writers import (
     ZarrWriterMixin,
@@ -364,9 +364,14 @@ class PrepareZarrTarget(beam.PTransform):
         If chunking is present in the schema for a given dimension, the length of
         the first fragment will be used. Otherwise, the dimension will not be chunked.
     :param attrs: Extra group-level attributes to inject into the dataset.
-    :param consolidated_metadata: Bool controlling if xarray.to_zarr()
-    writes consolidated metadata. Default's to True.
     :param encoding: Dictionary describing encoding for xarray.to_zarr()
+    :param consolidated_metadata: Bool controlling if xarray.to_zarr()
+    writes consolidated metadata. Default's to False. In StoreToZarr,
+    always default to unconsolidated. This leaves it up to the
+    user whether or not they want to consolidate with ConsolidateMetadata(). Also,
+    it prevents a broken/inconsistent state that could arise from metadata being
+    consolidated here, and then falling out of sync with coordinates if
+    ConsolidateDimensionCoordinates() is applied to the output of StoreToZarr().
     """
 
     target: str | FSSpecTarget
@@ -386,8 +391,8 @@ def expand(self, pcoll: beam.PCollection) -> beam.PCollection:
             target_store=store,
             target_chunks=self.target_chunks,
             attrs=self.attrs,
-            consolidated_metadata=self.consolidated_metadata,
             encoding=self.encoding,
+            consolidated_metadata=False,
         )
         return initialized_target
 
@@ -430,6 +435,13 @@ def expand(self, pcoll: beam.PCollection) -> beam.PCollection:
         return new_fragments
 
 
+class ConsolidateDimensionCoordinates(beam.PTransform):
+    def expand(
+        self, pcoll: beam.PCollection[zarr.storage.FSStore]
+    ) -> beam.PCollection[zarr.storage.FSStore]:
+        return pcoll | beam.Map(consolidate_dimension_coordinates)
+
+
 @dataclass
 class CombineReferences(beam.PTransform):
     """Combines Kerchunk references into a single reference dataset.
@@ -573,6 +585,9 @@ class StoreToZarr(beam.PTransform, ZarrWriterMixin):
         `store_name` will be appended to this prefix to create a full path.
     :param target_chunks: Dictionary mapping dimension names to chunks sizes.
       If a dimension is a not named, the chunks will be inferred from the data.
+    :param consolidate_dimension_coordinates: Whether to rewrite coordinate variables as a
+      single chunk. We recommend consolidating coordinate variables to avoid
+      many small read requests to get the coordinates in xarray. Defaults to ``True``.
     :param dynamic_chunking_fn: Optionally provide a function that takes an ``xarray.Dataset``
       template dataset as its first argument and returns a dynamically generated chunking dict.
       If provided, ``target_chunks`` cannot also be passed. You can use this to determine chunking
@@ -581,8 +596,7 @@ class StoreToZarr(beam.PTransform, ZarrWriterMixin):
       out https://github.com/jbusecke/dynamic_chunks
     :param dynamic_chunking_fn_kwargs: Optional keyword arguments for ``dynamic_chunking_fn``.
     :param attrs: Extra group-level attributes to inject into the dataset.
-    :param consolidated_metadata: Bool controlling if xarray.to_zarr()
-    writes consolidated metadata. Default's to True.
+
     :param encoding: Dictionary encoding for xarray.to_zarr().
     """
 
@@ -597,7 +611,6 @@ class StoreToZarr(beam.PTransform, ZarrWriterMixin):
     dynamic_chunking_fn: Optional[Callable[[xr.Dataset], dict]] = None
     dynamic_chunking_fn_kwargs: Optional[dict] = field(default_factory=dict)
     attrs: Dict[str, str] = field(default_factory=dict)
-    consolidated_metadata: Optional[bool] = True
     encoding: Optional[dict] = field(default_factory=dict)
 
     def __post_init__(self):
@@ -624,7 +637,6 @@ def expand(
             target=self.get_full_target(),
             target_chunks=target_chunks,
             attrs=self.attrs,
-            consolidated_metadata=self.consolidated_metadata,
             encoding=self.encoding,
         )
         n_target_stores = rechunked_datasets | StoreDatasetFragments(target_store=target_store)
@@ -633,7 +645,5 @@ def expand(
             | beam.combiners.Sample.FixedSizeGlobally(1)
             | beam.FlatMap(lambda x: x)  # https://stackoverflow.com/a/47146582
         )
-        # TODO: optionally use `singleton_target_store` to
-        # consolidate metadata and/or coordinate dims here
 
         return singleton_target_store

pangeo_forge_recipes/writers.py

@@ -73,7 +73,7 @@ def consolidate_metadata(store: MutableMapping) -> MutableMapping:
     :param store: Input Store for Zarr
     :type store: MutableMapping
     :return: Output Store
-    :rtype: MutableMapping
+    :rtype: zarr.storage.FSStore
     """
 
     import zarr
@@ -83,12 +83,10 @@ def consolidate_metadata(store: MutableMapping) -> MutableMapping:
             """Creating consolidated metadata for Kerchunk references should not
             yield a performance benefit so consolidating metadata is not supported."""
         )
-
     if isinstance(store, zarr.storage.FSStore):
         zarr.convenience.consolidate_metadata(store)
 
-    zc = zarr.open_consolidated(store)
-    return zc
+    return store
 
 
 def store_dataset_fragment(

tests/test_end_to_end.py

@@ -14,6 +14,8 @@
 
 from pangeo_forge_recipes.patterns import FilePattern, pattern_from_file_sequence
 from pangeo_forge_recipes.transforms import (
+    ConsolidateDimensionCoordinates,
+    ConsolidateMetadata,
     OpenWithKerchunk,
     OpenWithXarray,
     StoreToZarr,
@@ -173,3 +175,29 @@ def test_reference_grib(
     # various inconsistencies (of dtype casting int to float, etc.). With the right combination of
     # options passed to the pipeline, seems like these should pass?
     # xr.testing.assert_equal(ds.load(), ds2)
+
+
+def test_xarray_zarr_consolidate_dimension_coordinates(
+    netcdf_local_file_pattern_sequential,
+    pipeline,
+    tmp_target,
+):
+    pattern = netcdf_local_file_pattern_sequential
+    with pipeline as p:
+        (
+            p
+            | beam.Create(pattern.items())
+            | OpenWithXarray(file_type=pattern.file_type)
+            | StoreToZarr(
+                target_root=tmp_target,
+                store_name="subpath",
+                combine_dims=pattern.combine_dim_keys,
+            )
+            | ConsolidateDimensionCoordinates()
+            | ConsolidateMetadata()
+        )
+
+    path = os.path.join(tmp_target.root_path, "subpath")
+    ds = xr.open_dataset(path, engine="zarr", consolidated=True, chunks={})
+
+    assert ds.time.encoding["chunks"][0] == ds.time.shape[0]

tests/test_rechunking.py

@@ -1,11 +1,20 @@
 import itertools
+import os
 import random
 from collections import namedtuple
+from tempfile import TemporaryDirectory
 
+import numpy as np
 import pytest
 import xarray as xr
+import zarr
 
-from pangeo_forge_recipes.rechunking import GroupKey, combine_fragments, split_fragment
+from pangeo_forge_recipes.rechunking import (
+    GroupKey,
+    combine_fragments,
+    consolidate_dimension_coordinates,
+    split_fragment,
+)
 from pangeo_forge_recipes.types import CombineOp, Dimension, Index, IndexedPosition, Position
 
 from .conftest import split_up_files_by_variable_and_day
@@ -258,3 +267,24 @@ def test_combine_fragments_errors():
     index1 = Index({Dimension("time", CombineOp.CONCAT): IndexedPosition(2)})
     with pytest.raises(ValueError, match="are not consistent"):
         _ = combine_fragments(group, [(index0, ds), (index1, ds)])
+
+
+def test_consolidate_dimension_coordinates():
+    td = TemporaryDirectory()
+    store_path = os.path.join(td.name + "tmp.zarr")
+    group = zarr.group(store=store_path, overwrite=True)
+    group.create(name="data", shape=100, chunks=10, dtype="i4")
+    group.create(name="time", shape=100, chunks=10, dtype="i4")
+    group.data[:] = np.random.randn(*group.data.shape)
+    group.time[:] = np.arange(100)
+
+    # If you don't provide these attrs,
+    # consolidate_dimension_coordinates does not
+    # raise an error, while Xarray does
+    group.data.attrs["_ARRAY_DIMENSIONS"] = ["time"]
+    group.time.attrs["_ARRAY_DIMENSIONS"] = ["time"]
+
+    consolidated_zarr = consolidate_dimension_coordinates(zarr.storage.FSStore(store_path))
+    store = zarr.open(consolidated_zarr)
+    assert store.time.chunks[0] == 100
+    assert store.data.chunks[0] == 10

tests/test_transforms.py

@@ -310,7 +310,6 @@ def dynamic_chunking_fn(template_ds: xr.Dataset, divisor: int = 1):
             combine_dims=pattern.combine_dim_keys,
             attrs={},
             dynamic_chunking_fn=dynamic_chunking_fn,
-            consolidated_metadata=True,
             **kws,
         )
         open_store = target_store | OpenZarrStore()

tests/test_writers.py

@@ -169,7 +169,6 @@ def test_zarr_consolidate_metadata(
                 target_root=tmp_target,
                 store_name="store",
                 combine_dims=pattern.combine_dim_keys,
-                consolidated_metadata=False,
             )
             | ConsolidateMetadata()
         )