Consolidate dimension coordinates
This consolidates dimension coordinates to address the performance
issues from having many small coordinate chunks.
Closes
MODIFY
pangeo_forge_recipes/recipes/xarray_zarr.py
pangeo_forge_recipes/recipes/xarray_zarr.py
@@ -621,9 +621,25 @@ def store_chunk(
zarr_array[zarr_region] = data
-def finalize_target(target: CacheFSSpecTarget, consolidate_zarr: bool) -> None:
+def finalize_target(
+ target: CacheFSSpecTarget,
+ consolidate_zarr: bool,
+ consolidate_dimension_coordinates: bool = True,
+) -> None:
if target is None:
raise ValueError("target has not been set.")
+
+ if consolidate_dimension_coordinates:
+ logger.info("Consolidating dimension coordinate arrays")
+ target_mapper = target.get_mapper()
+ ds = xr.open_zarr(target_mapper) # Probably a better way to get the dimension coords?
+ group = zarr.open(target_mapper)
+ for dim in ds.dims:
+ attrs = dict(group[dim].attrs)
+ data = group[dim][:]
+ group[dim] = data
+ group[dim].attrs.update(attrs)
+
if consolidate_zarr:
logger.info("Consolidating Zarr metadata")
target_mapper = target.get_mapper()
@@ -661,6 +677,9 @@ class XarrayZarrRecipe(BaseRecipe, FilePatternRecipeMixin):
``xr.open_dataset``. This is required for engines that can't open
file-like objects (e.g. pynio).
:param consolidate_zarr: Whether to consolidate the resulting Zarr dataset.
+ :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.
:param xarray_open_kwargs: Extra options for opening the inputs with Xarray.
:param xarray_concat_kwargs: Extra options to pass to Xarray when concatenating
the inputs to form a chunk.
@@ -685,6 +704,7 @@ class XarrayZarrRecipe(BaseRecipe, FilePatternRecipeMixin):
cache_inputs: Optional[bool] = None
copy_input_to_local_file: bool = False
consolidate_zarr: bool = True
+ consolidate_dimension_coordinates: bool = True
xarray_open_kwargs: dict = field(default_factory=dict)
xarray_concat_kwargs: dict = field(default_factory=dict)
delete_input_encoding: bool = True
@@ -851,7 +871,10 @@ class XarrayZarrRecipe(BaseRecipe, FilePatternRecipeMixin):
@property
def finalize_target(self) -> Callable[[], None]:
return functools.partial(
- finalize_target, target=self.target, consolidate_zarr=self.consolidate_zarr
+ finalize_target,
+ target=self.target,
+ consolidate_zarr=self.consolidate_zarr,
+ consolidate_dimension_coordinates=self.consolidate_dimension_coordinates,
)
def iter_inputs(self) -> Iterator[InputKey]:

MODIFY
tests/recipe_tests/test_XarrayZarrRecipe.py
tests/recipe_tests/test_XarrayZarrRecipe.py
@@ -4,6 +4,7 @@ from unittest.mock import patch
import pytest
import xarray as xr
+import zarr
# need to import this way (rather than use pytest.lazy_fixture) to make it work with dask
from pytest_lazyfixture import lazy_fixture
@@ -274,6 +275,10 @@ def do_actual_chunks_test(
assert all([item == chunk_len for item in ds_actual.chunks[other_dim][:-1]])
ds_actual.load()
+ store = zarr.open_consolidated(target.get_mapper())
+ for dim in ds_actual.dims:
+ assert store[dim].chunks == ds_actual[dim].shape
+
xr.testing.assert_identical(ds_actual, ds_expected)

Consolidate dimension coordinates

This consolidates dimension coordinates to address the performance
issues from having many small coordinate chunks.

Closes https://github.com/pangeo-forge/pangeo-forge-recipes/issues/209