Make v2→v3 re-encode safer: prefix-safe omit_nodes and lower peak memory

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,24 +227,36 @@ 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 (
                 metadata.shape[-2] < spatial_chunk or metadata.shape[-1] < spatial_chunk
             )
 
             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:

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:

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]

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,

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