From 7a156f247d1d71916e0895921b236394aa8c6a94 Mon Sep 17 00:00:00 2001 From: Wietze Date: Tue, 16 Dec 2025 17:54:55 +0000 Subject: [PATCH 1/2] feat(s2): direct v3 re-encode and multiscales pipeline --- .../s2_optimization/s2_converter.py | 20 ++- .../s2_optimization/s2_multiscale.py | 17 ++- src/eopf_geozarr/zarrio.py | 123 +++++++++++++++--- 3 files changed, 135 insertions(+), 25 deletions(-) diff --git a/src/eopf_geozarr/s2_optimization/s2_converter.py b/src/eopf_geozarr/s2_optimization/s2_converter.py index 834bb4a2..4a0bd5b2 100644 --- a/src/eopf_geozarr/s2_optimization/s2_converter.py +++ b/src/eopf_geozarr/s2_optimization/s2_converter.py @@ -163,7 +163,7 @@ def array_reencoder( # handle xarray datetime/timedelta encoding # If the array has time-related units, ensure the dtype attribute matches the actual dtype if attributes.get("units") in TIME_UNITS: - numpy_time_unit = _netcdf_to_numpy_timeunit(attributes["units"]) + numpy_time_unit = _netcdf_to_numpy_timeunit(str(attributes["units"])) # Check if this is a timedelta or datetime based on: # 1. The zarr dtype (if it's a native time type like TimeDelta64) # 2. The existing dtype attribute (for int64-encoded times) @@ -227,6 +227,12 @@ def array_reencoder( # don't rechunk 1D variables pass else: + prefix_shape = tuple(metadata.chunks[:-2]) + inner_spatial = ( + min(spatial_chunk, metadata.shape[-2]), + min(spatial_chunk, metadata.shape[-1]), + ) + # Check if array is too small for the requested chunk size # If any spatial dimension is smaller than spatial_chunk, don't use sharding array_too_small_for_sharding = enable_sharding and ( @@ -234,17 +240,23 @@ def array_reencoder( ) if not enable_sharding or array_too_small_for_sharding: - chunk_shape = (spatial_chunk, spatial_chunk) + chunk_shape = (*prefix_shape, *inner_spatial) subchunk_shape = None else: # Generate 2D chunking / sharding, because partitioning along # other dimensions will be set to 1 - chunk_shape, subchunk_shape = _auto_partition( + shard_spatial, subchunk_spatial = _auto_partition( array_shape=metadata.shape[-2:], - chunk_shape=(spatial_chunk,) * 2, + chunk_shape=inner_spatial, shard_shape="auto", item_size=item_size, ) + if shard_spatial is None: + chunk_shape = (*prefix_shape, *inner_spatial) + subchunk_shape = None + else: + chunk_shape = (*prefix_shape, *shard_spatial) + subchunk_shape = (*prefix_shape, *subchunk_spatial) chunk_grid = {"name": "regular", "configuration": {"chunk_shape": chunk_shape}} if enable_sharding and subchunk_shape is not None: diff --git a/src/eopf_geozarr/s2_optimization/s2_multiscale.py b/src/eopf_geozarr/s2_optimization/s2_multiscale.py index 1ea83c3b..957df31a 100644 --- a/src/eopf_geozarr/s2_optimization/s2_multiscale.py +++ b/src/eopf_geozarr/s2_optimization/s2_multiscale.py @@ -79,14 +79,25 @@ def create_multiscale_levels(group: zarr.Group, path: str) -> None: # Open the current resolution level as a dataset cur_group_path = f"{full_path}/{cur_group_name}" cur_ds = xr.open_dataset(group.store, group=cur_group_path, engine="zarr") - scale = next_factor // cur_factor to_downsample: dict[str, xr.DataArray] = {} for var_name, var in cur_ds.data_vars.items(): # Check if the variable already exists in the next level next_level_path = f"{path}/{next_group_name}" - if f"{next_level_path}/{var_name}" not in group: + next_var_key = f"{next_level_path}/{var_name}" + if next_var_key not in group: to_downsample[var_name] = var + + if not to_downsample: + next_level_key = f"{path}/{next_group_name}" + if next_level_key not in group: + log.info( + "Stopping at %s: next level missing and nothing to create", + next_group_name, + ) + break + log.info("Skipping %s: all variables already exist", next_group_name) + continue log.info("downsampling %s into %s", tuple(sorted(to_downsample.keys())), next_group_name) # Don't pass coords here - let the downsampled variables determine their own coordinates downsampled_ds = create_downsampled_resolution_group( @@ -251,7 +262,7 @@ def add_multiscales_metadata_to_parent( base_path: str, res_groups: Mapping[str, xr.Dataset], multiscales_flavor: set[MultiscalesFlavor] | None = None, -) -> xr.DataTree: +) -> xr.DataTree | None: """Add GeoZarr-compliant multiscales metadata to parent group.""" # Sort by resolution (finest to coarsest) if multiscales_flavor is None: diff --git a/src/eopf_geozarr/zarrio.py b/src/eopf_geozarr/zarrio.py index ab1a56f4..b80d38dc 100644 --- a/src/eopf_geozarr/zarrio.py +++ b/src/eopf_geozarr/zarrio.py @@ -4,6 +4,8 @@ from __future__ import annotations +import math +from itertools import product from typing import TYPE_CHECKING, Any, NotRequired, TypedDict import numcodecs @@ -15,7 +17,7 @@ from zarr.storage._common import make_store_path if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable, Iterator, Mapping from zarr.core.metadata.v2 import ArrayV2Metadata @@ -25,6 +27,69 @@ class ChunkEncodingSpec(TypedDict): read_chunks: NotRequired[tuple[int, ...]] +def _normalize_node_path(value: str) -> str: + value = value.strip() + if not value: + return "" + # reencode_group member names are relative (no leading slash) + value = value.lstrip("/") + while "//" in value: + value = value.replace("//", "/") + return value.rstrip("/") + + +def _normalize_omit_nodes(values: set[str] | None) -> set[str]: + if not values: + return set() + normalized = {_normalize_node_path(v) for v in values} + normalized.discard("") + return normalized + + +def _is_omitted(name: str, omit_nodes: set[str]) -> bool: + # Omit either the exact node, or any descendant below it. + return any(name == v or name.startswith(v + "/") for v in omit_nodes) + + +def _iter_chunk_regions( + shape: tuple[int, ...], + chunk_shape: tuple[int, ...], +) -> Iterator[tuple[slice, ...]] | None: + if len(shape) != len(chunk_shape): + return None + if any(dim_size <= 0 for dim_size in shape): + return None + + starts_per_dim = [ + range(0, dim_size, max(1, chunk_size)) + for dim_size, chunk_size in zip(shape, chunk_shape, strict=True) + ] + + def _gen() -> Iterator[tuple[slice, ...]]: + for starts in product(*starts_per_dim): + yield tuple( + slice(start, min(start + max(1, chunk), dim)) + for start, chunk, dim in zip(starts, chunk_shape, shape, strict=True) + ) + + return _gen() + + +def _dtype_itemsize(dtype: object) -> int: + itemsize = getattr(dtype, "itemsize", None) + if isinstance(itemsize, int) and itemsize > 0: + return itemsize + return 1 + + +def _estimate_nbytes(shape: tuple[int, ...], dtype: object) -> int: + try: + n = int(math.prod(shape)) + except Exception: + return 0 + return n * _dtype_itemsize(dtype) + + def convert_compression( compressor: numcodecs.abc.Codec | dict[str, object], *, compression_level: int | None = None ) -> tuple[dict[str, Any], ...]: @@ -120,16 +185,14 @@ def reencode_group( The path in the new store to use overwrite : bool, default = False Whether to overwrite contents of the new store - omit_nodes : set[str], default = {} - The names of groups or arrays to omit from re-encoding. - chunk_reencoder : Callable[[zarr.Array[Any], ChunkEncodingSpec]] | None, default = None - A function that takes a Zarr array object and returns a ChunkEncodingSpec, which is a dict - that defines a new chunk encoding. Use this parameter to define per-array chunk encoding - logic. + omit_nodes : set[str] | None + Relative group/array paths to omit (e.g., "measurements/reflectance"). + Exact matches omit that node; prefix matches omit the whole subtree. + array_reencoder : Callable[[str, ArrayV2Metadata], ArrayV3Metadata] + Maps a v2 array metadata document to v3 metadata for the destination array. """ - if omit_nodes is None: - omit_nodes = set() + omit_nodes = _normalize_omit_nodes(omit_nodes) log = structlog.get_logger() @@ -142,19 +205,21 @@ def reencode_group( ) log.info("Begin re-encoding Zarr group %s", group) + root_attrs = group.attrs.asdict() + new_members: dict[str, ArrayV3Metadata | GroupMetadata] = { - path: GroupMetadata(zarr_format=3, attributes=group.attrs.asdict()) + path: GroupMetadata(zarr_format=3, attributes=root_attrs) } - chunks_to_encode: list[str] = [] + arrays_to_copy: list[str] = [] for name in omit_nodes: - if not any(k.startswith(name) for k in members): + if not any(k == name or k.startswith(name + "/") for k in members): log.warning( "The name %s was provided in omit_nodes but no such array or group exists.", name ) for name, member in members.items(): - if any(name.startswith(v) for v in omit_nodes): + if _is_omitted(name, omit_nodes): log.info( - "Skipping node %s because it is contained in a subgroup declared in the omit_groups parameter", + "Skipping node %s because it is contained in a subtree declared in omit_nodes", name, ) continue @@ -163,7 +228,7 @@ def reencode_group( if isinstance(member, zarr.Array): new_meta = array_reencoder(member.path, member.metadata) new_members[new_path] = new_meta - chunks_to_encode.append(name) + arrays_to_copy.append(name) else: new_members[new_path] = GroupMetadata( zarr_format=3, @@ -172,12 +237,34 @@ def reencode_group( log.info("Creating new Zarr hierarchy structure at %s", f"{store}/{path}") tree = dict(zarr.create_hierarchy(store=store, nodes=new_members, overwrite=overwrite)) new_group: zarr.Group = tree[path] - for name in chunks_to_encode: - log.info("Re-encoding chunks for array %s", name) + for name in arrays_to_copy: + log.info("Copying array data %s", name) old_array = group[name] new_array = new_group[name] - new_array[...] = old_array[...] + if new_array.ndim == 0: + new_array[...] = old_array[...] + continue + + old_chunk_shape = tuple(getattr(old_array.metadata, "chunks", ())) + new_chunk_shape = tuple(new_array.metadata.chunk_grid.chunk_shape) + # If chunking differs, writing by destination chunk regions can re-read source chunks. + # Prefer eager copy for smaller arrays to minimize IO; fall back to chunk-by-chunk when + # the array is too large to comfortably materialize. + eager_copy_max_bytes = 256 * 1024 * 1024 + if old_chunk_shape != new_chunk_shape: + estimated = _estimate_nbytes(tuple(new_array.shape), getattr(old_array, "dtype", None)) + if estimated and estimated <= eager_copy_max_bytes: + new_array[...] = old_array[...] + continue + + # Iterate using the destination chunk grid to bound peak memory. + regions_iter = _iter_chunk_regions(tuple(new_array.shape), new_chunk_shape) + if regions_iter is None: + new_array[...] = old_array[...] + else: + for region in regions_iter: + new_array[region] = old_array[region] # return the root group return tree[path] From 63fcaf1a478ca21a7f742106b8cc38eecc5549c2 Mon Sep 17 00:00:00 2001 From: Wietze Date: Tue, 16 Dec 2025 16:44:59 +0000 Subject: [PATCH 2/2] test: update zarr v3 conversion expectations --- tests/test_s2_multiscale.py | 4 +--- tests/test_zarrio.py | 21 +++++++++++++++++++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/tests/test_s2_multiscale.py b/tests/test_s2_multiscale.py index 7dd0d25d..bb3f6a2f 100644 --- a/tests/test_s2_multiscale.py +++ b/tests/test_s2_multiscale.py @@ -1,15 +1,13 @@ """ Tests for S2 multiscale pyramid creation with xy-aligned sharding. """ - -from typing import TypedDict - import numpy as np import pytest import xarray as xr import zarr from pydantic_zarr.experimental.v3 import ArraySpec, GroupSpec from structlog.testing import capture_logs +from typing_extensions import TypedDict from eopf_geozarr.s2_optimization.s2_multiscale import ( calculate_aligned_chunk_size, diff --git a/tests/test_zarrio.py b/tests/test_zarrio.py index 9f1469e6..0954fc17 100644 --- a/tests/test_zarrio.py +++ b/tests/test_zarrio.py @@ -256,5 +256,22 @@ def custom_array_encoder(key: str, metadata: ArrayV2Metadata) -> ArrayV3Metadata assert new_node.attrs.asdict() == old_node.attrs.asdict() -if __name__ == "__main__": - pytest.main([__file__]) +def test_reencode_group_omit_nodes_normalizes_leading_slash() -> None: + group_v2 = zarr.create_group({}, zarr_format=2) + group_v2.create_array("data", shape=(5,), dtype="int32") + + group_v3 = reencode_group(group_v2, {}, "", omit_nodes=["/data"]) + + assert "data" not in group_v3 + + +def test_reencode_group_omit_nodes_is_prefix_safe() -> None: + group_v2 = zarr.create_group({}, zarr_format=2) + group_v2.create_array("foobar", shape=(5,), dtype="int32") + group_v2.create_group("foo").create_array("child", shape=(5,), dtype="int32") + + group_v3 = reencode_group(group_v2, {}, "", omit_nodes=["foo"]) + + assert "foo" not in group_v3 + assert "foo/child" not in group_v3 + assert "foobar" in group_v3