refactored recipe module to use FilePattern; nearly everything works
MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -20,7 +20,7 @@ import zarr
from rechunker.types import MultiStagePipeline, ParallelPipelines, Stage
from .patterns import FilePattern
-from .storage import AbstractTarget, UninitializedTarget
+from .storage import AbstractTarget, UninitializedTarget, UninitializedTargetError
from .utils import (
chunk_bounds_and_conflicts,
chunked_iterable,
@@ -226,8 +226,8 @@ class XarrayZarrRecipe(BaseRecipe):
:param target_chunks: Desired chunk structure for the targret dataset.
"""
+ file_pattern: FilePattern
inputs_per_chunk: Optional[int] = 1
- file_pattern: Optional[FilePattern] = None
target: AbstractTarget = field(default_factory=UninitializedTarget)
input_cache: AbstractTarget = field(default_factory=UninitializedTarget)
metadata_cache: AbstractTarget = field(default_factory=UninitializedTarget)
@@ -633,13 +633,14 @@ class XarrayZarrRecipe(BaseRecipe):
# get the sequence length of every file
# this line could become problematic for large (> 10_000) lists of files
input_meta = self.get_input_meta(*self._inputs_chunks)
- all_lens = [m["dims"][self._concat_dim] for m in input_meta.values()]
- all_lens.shape = self.file_pattern.shape
+ # use a numpy array to allow reshaping
+ all_lens = np.array([m["dims"][self._concat_dim] for m in input_meta.values()])
+ all_lens.shape = list(self.file_pattern.dims.values())
# check that all lens are the same along the concat dim
concat_dim_axis = list(self.file_pattern.dims).index(self._concat_dim)
selector = [0] * len(self.file_pattern.dims)
selector[concat_dim_axis] = slice(None)
- sequence_lens = all_lens(tuple(selector))
- if not all(all_lens == sequence_lens):
+ sequence_lens = all_lens[tuple(selector)]
+ if not (all_lens == sequence_lens).all(): # BROKEN! broadcasting isn't happening
raise ValueError(f"Inconsistent sequence lengths found: f{all_lens}")
- return sequence_lens
+ return sequence_lens.tolist()

MODIFY
tests/conftest.py
tests/conftest.py
@@ -109,7 +109,7 @@ def netcdf_local_paths_by_variable(daily_xarray_dataset, tmpdir_factory, request
full_paths = [tmp_path.join(fname) for fname in fnames]
xr.save_mfdataset(datasets, [str(path) for path in full_paths])
items_per_file = {"D": 1, "2D": 2}[request.param]
- path_format = str(tmp_path) + "/{variable}_{n:03d}.nc"
+ path_format = str(tmp_path) + "/{variable}_{time:03d}.nc"
return full_paths, items_per_file, fnames_by_variable, path_format
@@ -174,10 +174,8 @@ def uninitialized_target():
def netCDFtoZarr_sequential_recipe(daily_xarray_dataset, netcdf_local_paths, tmp_target, tmp_cache):
paths, items_per_file = netcdf_local_paths
file_pattern = pattern_from_file_sequence([str(path) for path in paths], "time", items_per_file)
- kwargs = dict(
- file_pattern=file_pattern, inputs_per_chunk=1, target=tmp_target, input_cache=tmp_cache,
- )
- return recipe.XarrayZarrRecipe, kwargs, daily_xarray_dataset, tmp_target
+ kwargs = dict(inputs_per_chunk=1, target=tmp_target, input_cache=tmp_cache,)
+ return recipe.XarrayZarrRecipe, file_pattern, kwargs, daily_xarray_dataset, tmp_target
@pytest.fixture
@@ -187,18 +185,16 @@ def netCDFtoZarr_sequential_multi_variable_recipe(
paths, items_per_file, fnames_by_variable, path_format = netcdf_local_paths_by_variable
time_index = list(range(len(paths) // 2))
- def format_function(variable, n):
- return path_format.format(variable=variable, n=n)
+ def format_function(variable, time):
+ return path_format.format(variable=variable, time=time)
file_pattern = FilePattern(
format_function,
ConcatDim("time", time_index, items_per_file),
MergeDim("variable", ["foo", "bar"]),
)
- kwargs = dict(
- file_pattern=file_pattern, inputs_per_chunk=1, target=tmp_target, input_cache=tmp_cache,
- )
- return recipe.XarrayZarrRecipe, kwargs, daily_xarray_dataset, tmp_target
+ kwargs = dict(inputs_per_chunk=1, target=tmp_target, input_cache=tmp_cache,)
+ return recipe.XarrayZarrRecipe, file_pattern, kwargs, daily_xarray_dataset, tmp_target
@pytest.fixture(scope="session")

MODIFY
tests/test_recipes.py
tests/test_recipes.py
@@ -16,8 +16,8 @@ all_recipes = [
def test_recipe(recipe_fixture, execute_recipe):
"""The basic recipe test. Use this as a template for other tests."""
- RecipeClass, kwargs, ds_expected, target = recipe_fixture
- rec = RecipeClass(**kwargs)
+ RecipeClass, file_pattern, kwargs, ds_expected, target = recipe_fixture
+ rec = RecipeClass(file_pattern, **kwargs)
execute_recipe(rec)
ds_actual = xr.open_zarr(target.get_mapper()).load()
xr.testing.assert_identical(ds_actual, ds_expected)
@@ -30,11 +30,14 @@ def test_recipe_caching_copying(
):
"""The basic recipe test. Use this as a template for other tests."""
- RecipeClass, kwargs, ds_expected, target = netCDFtoZarr_sequential_recipe
+ RecipeClass, file_pattern, kwargs, ds_expected, target = netCDFtoZarr_sequential_recipe
if not cache_inputs:
kwargs.pop("input_cache") # make sure recipe doesn't require input_cache
rec = RecipeClass(
- **kwargs, cache_inputs=cache_inputs, copy_input_to_local_file=copy_input_to_local_file
+ file_pattern,
+ **kwargs,
+ cache_inputs=cache_inputs,
+ copy_input_to_local_file=copy_input_to_local_file
)
execute_recipe(rec)
ds_actual = xr.open_zarr(target.get_mapper()).load()
@@ -57,10 +60,10 @@ def incr_date(ds, filename=""):
def test_process(recipe_fixture, execute_recipe, process_input, process_chunk):
"""Check that the process_chunk and process_input arguments work as expected."""
- RecipeClass, kwargs, ds_expected, target = recipe_fixture
+ RecipeClass, file_pattern, kwargs, ds_expected, target = recipe_fixture
kwargs["process_input"] = process_input
kwargs["process_chunk"] = process_chunk
- rec = RecipeClass(**kwargs)
+ rec = RecipeClass(file_pattern, **kwargs)
execute_recipe(rec)
ds_actual = xr.open_zarr(target.get_mapper()).load()
@@ -101,20 +104,21 @@ def test_chunks(
):
"""Check that chunking of datasets works as expected."""
- RecipeClass, kwargs, ds_expected, target = recipe_fixture
+ RecipeClass, file_pattern, kwargs, ds_expected, target = recipe_fixture
kwargs["target_chunks"] = target_chunks
kwargs["inputs_per_chunk"] = inputs_per_chunk
if specify_nitems_per_input:
- kwargs["nitems_per_input"] = kwargs["nitems_per_input"] # it's already there in kwargs
kwargs["metadata_cache"] = None
else:
- # file will be scanned and metadata cached
- kwargs["nitems_per_input"] = None
+ # modify file_pattern in place to remove nitems_per_file; a bit hacky
+ for cdim in file_pattern.combine_dims:
+ if hasattr(cdim, "nitems_per_file"):
+ cdim.nitems_per_file = None
kwargs["metadata_cache"] = kwargs["input_cache"]
with chunk_expectation as excinfo:
- rec = RecipeClass(**kwargs)
+ rec = RecipeClass(file_pattern, **kwargs)
if excinfo:
# don't continue if we got an exception
return
@@ -124,9 +128,8 @@ def test_chunks(
# chunk validation
ds_actual = xr.open_zarr(target.get_mapper(), consolidated=True)
sequence_chunks = ds_actual.chunks["time"]
- seq_chunk_len = target_chunks.get("time", None) or (
- kwargs["nitems_per_input"] * inputs_per_chunk
- )
+ nitems_per_input = list(file_pattern.nitems_per_input.values())[0]
+ seq_chunk_len = target_chunks.get("time", None) or (nitems_per_input * inputs_per_chunk)
# we expect all chunks but the last to have the expected size
assert all([item == seq_chunk_len for item in sequence_chunks[:-1]])
for other_dim, chunk_len in target_chunks.items():

refactored recipe module to use FilePattern; nearly everything works