perf/direct zarr io

.vscode/launch.json

@@ -198,6 +198,7 @@
         "convert-s2-optimized",
         // "https://objects.eodc.eu/e05ab01a9d56408d82ac32d69a5aae2a:202509-s02msil2a/08/products/cpm_v256/S2A_MSIL2A_20250908T100041_N0511_R122_T32TQM_20250908T115116.zarr",
         "https://objects.eodc.eu/e05ab01a9d56408d82ac32d69a5aae2a:202511-s02msil2a-eu/15/products/cpm_v262/S2B_MSIL2A_20251115T091139_N0511_R050_T35SLU_20251115T111807.zarr",
+        // "https://objects.eodc.eu:443/e05ab01a9d56408d82ac32d69a5aae2a:202512-s02msil2a-eu/16/products/cpm_v262/S2B_MSIL2A_20251216T102339_N0511_R065_T32TNS_20251216T123617.zarr",
         // "https://objects.eodc.eu/e05ab01a9d56408d82ac32d69a5aae2a:202511-s02msil2a-eu/16/products/cpm_v262/S2A_MSIL2A_20251116T085431_N0511_R107_T35SQD_20251116T103813.zarr",
         // "s3://esa-zarr-sentinel-explorer-fra/tests-output/sentinel-2-l2a-opt/S2A_MSIL2A_20250908T100041_N0511_R122_T32TQM_20250908T115116.zarr",
         // "s3://esa-zarr-sentinel-explorer-fra/tests-output/sentinel-2-l2a-staging/S2B_MSIL2A_20251115T091139_N0511_R050_T35SLU_20251115T111807.zarr",
@@ -211,7 +212,7 @@
         // "--omit-nodes",
         // "quality/l2a_quicklook",
         "--dask-cluster",
-        "--verbose"
+        // "--verbose"
       ],
       "cwd": "${workspaceFolder}",
       "justMyCode": false,

pyproject.toml

@@ -27,9 +27,9 @@ classifiers = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "pydantic-zarr>=0.8.0",
+    "pydantic-zarr>=0.9.1",
     "pydantic>=2.12",
-    "zarr>=3.1.1",
+    "zarr>=3.1.4",
     "xarray>=2025.7.1",
     "dask[array,distributed]>=2025.5.1",
     "numpy>=2.3.1",
@@ -172,6 +172,8 @@ module = ["zarr.*", "xarray.*", "rioxarray.*", "cf_xarray.*", "dask.*"]
 ignore_missing_imports = true
 
 [tool.pytest.ini_options]
+filterwarnings = "error"
+log_level = "WARNING"
 minversion = "7.0"
 addopts = "-ra -q --strict-markers --strict-config"
 testpaths = ["tests"]

src/eopf_geozarr/cli.py

@@ -1155,8 +1155,16 @@ def add_s2_optimization_commands(subparsers: argparse._SubParsersAction) -> None
         choices=range(1, 10),
         help="Compression level 1-9 (default: 3)",
     )
+    s2_parser.add_argument(
+        "--omit-nodes", help="The names of groups or arrays to skip.", default="", type=str
+    )
     s2_parser.add_argument("--skip-validation", action="store_true", help="Skip output validation")
     s2_parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
+    s2_parser.add_argument(
+        "--allow-json-nan",
+        action="store_true",
+        help="Allow invalid float values (nan, inf) in output JSON",
+    )
     s2_parser.add_argument(
         "--keep-scale-offset",
         action="store_true",
@@ -1184,18 +1192,26 @@ def convert_s2_optimized_command(args: argparse.Namespace) -> None:
         # Load input dataset
         log.info("Loading Sentinel-2 dataset from", input_path=args.input_path)
         storage_options = get_storage_options(str(args.input_path))
+        store = args.input_path
         dt_input = xr.open_datatree(
-            str(args.input_path), engine="zarr", chunks="auto", storage_options=storage_options
+            store,
+            engine="zarr",
+            chunks="auto",
+            storage_options=storage_options,
         )
 
+        omit_nodes = set(args.omit_nodes.split())
+
         # Convert
         convert_s2_optimized(
             dt_input=dt_input,
             output_path=args.output_path,
             enable_sharding=args.enable_sharding,
             spatial_chunk=args.spatial_chunk,
+            omit_nodes=omit_nodes,
             compression_level=args.compression_level,
             validate_output=not args.skip_validation,
+            allow_json_nan=args.allow_json_nan,
             keep_scale_offset=args.keep_scale_offset,
         )
 

src/eopf_geozarr/conversion/geozarr.py

@@ -1189,6 +1189,7 @@ def cleanup_prefix(prefix: str) -> None:
                         engine="zarr",
                         decode_coords="all",
                         chunks="auto",
+                        consolidated=False,
                         storage_options=store_storage_options,
                     )
                 break

src/eopf_geozarr/data_api/geozarr/common.py

@@ -5,6 +5,7 @@
 import io
 import urllib
 import urllib.request
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
@@ -24,12 +25,18 @@
 from pydantic.experimental.missing_sentinel import MISSING
 from typing_extensions import Protocol, TypedDict, runtime_checkable
 
-from eopf_geozarr.data_api.geozarr.projjson import ProjJSON  # noqa: TC001
+from eopf_geozarr.data_api.geozarr.projjson import ProjJSON
+from eopf_geozarr.data_api.geozarr.types import (
+    CF_SCALE_OFFSET_KEYS,
+    CFScaleOffset,
+    EmptyDict,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
 
+
 @dataclass(frozen=True)
 class UNSET_TYPE:
     """
@@ -289,3 +296,38 @@ def check_grid_mapping(model: TDataSetLike) -> TDataSetLike:
 
 def is_none(data: object) -> TypeGuard[None]:
     return data is None
+
+
+def extract_scale_offset(
+    data: Mapping[str, object],
+) -> tuple[dict[str, object], CFScaleOffset | EmptyDict]:
+    """
+    Extract scale/offset information from a mapping, returning the remaining data and the scale/offset info.
+
+    Parameters
+    ----------
+    data : Mapping[[str, object]]
+        The input mapping from which to extract scale/offset information.
+
+    Returns
+    -------
+    tuple[Mapping[str, object], CFScaleOffset]
+        A tuple containing the remaining data (with scale/offset keys removed) and the extracted scale/offset info.
+    """
+    scale_offset: CFScaleOffset = {}  # type: ignore[typeddict-item]
+    remaining_data: dict[str, object] = {}
+
+    if set(data.keys()).isdisjoint(CF_SCALE_OFFSET_KEYS):
+        return dict(data), {}
+
+    if set(data.keys()).issuperset(CF_SCALE_OFFSET_KEYS):
+        for key, value in data.items():
+            if key in CF_SCALE_OFFSET_KEYS:
+                scale_offset[key] = value  # type: ignore[literal-required]
+            else:
+                remaining_data[key] = value
+        return remaining_data, scale_offset
+
+    raise ValueError(
+        "Incomplete scale/offset information: all of 'scale_factor', 'add_offset', must be present."
+    )

src/eopf_geozarr/data_api/geozarr/multiscales/geozarr.py

@@ -28,7 +28,7 @@ def valid_zcm(self) -> Self:
         Ensure that the ZCM metadata, if present, is valid
         """
         if self.layout is not MISSING:
-            zcm.Multiscales(**self.model_dump())
+            zcm.Multiscales(layout=self.layout, resampling_method=self.resampling_method)
 
         return self
 
@@ -38,7 +38,11 @@ def valid_tms(self) -> Self:
         Ensure that the TMS metadata, if present, is valid
         """
         if self.tile_matrix_set is not MISSING:
-            tms.Multiscales(**self.model_dump())
+            tms.Multiscales(
+                tile_matrix_set=self.tile_matrix_set,
+                tile_matrix_limits=self.tile_matrix_limits,
+                resampling_method=self.resampling_method,  # type: ignore[arg-type]
+            )
 
         return self
 

src/eopf_geozarr/data_api/geozarr/multiscales/zcm.py

@@ -2,7 +2,7 @@
 
 from typing import Final, Literal, NotRequired
 
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, field_validator, model_serializer
 from pydantic.experimental.missing_sentinel import MISSING
 from typing_extensions import TypedDict
 
@@ -73,6 +73,15 @@ class Transform(BaseModel):
     scale: tuple[float, ...] | MISSING = MISSING
     translation: tuple[float, ...] | MISSING = MISSING
 
+    @model_serializer
+    def serialize_model(self) -> dict[str, tuple[float, ...]]:
+        result: dict[str, tuple[float, ...]] = {}
+        if self.scale is not MISSING:
+            result["scale"] = self.scale
+        if self.translation is not MISSING:
+            result["translation"] = self.translation
+        return result
+
 
 class TransformJSON(TypedDict):
     scale: NotRequired[tuple[float, ...]]

src/eopf_geozarr/data_api/geozarr/types.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Literal, NotRequired, TypedDict
+from typing import TYPE_CHECKING, Final, Literal, NotRequired
+
+from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -16,7 +18,7 @@ class TileMatrixLimitJSON(TypedDict):
     maxTileRow: int
 
 
-CF_SCALE_OFFSET_KEYS: Final[set[str]] = {"scale_factor", "add_offset", "dtype"}
+CF_SCALE_OFFSET_KEYS: Final[set[str]] = {"scale_factor", "add_offset"}
 
 XARRAY_ENCODING_KEYS: Final[set[str]] = {
     "chunks",
@@ -25,16 +27,32 @@ class TileMatrixLimitJSON(TypedDict):
     "filters",
     "shards",
     "_FillValue",
+    "dtype",
 } | CF_SCALE_OFFSET_KEYS
 
 
+class CFScaleOffset(TypedDict):
+    """
+    Metadata defining scale/offset encoding for array values. Defined by the CF
+    conventions and found in EOPF Sentinel products in Zarr array attributes.
+    """
+
+    scale_factor: float
+    add_offset: float
+    dtype: str
+
+
+class EmptyDict(TypedDict, closed=True):  # type: ignore[call-arg]
+    """A dict with no keys."""
+
+
 class XarrayDataArrayEncoding(TypedDict):
     """
     The dict form of the encoding for xarray.DataArray
     """
 
     chunks: NotRequired[tuple[int, ...]]
-    preferred_chunks: NotRequired[tuple[int, ...]]
+    preferred_chunks: NotRequired[dict[str, int]]
     compressors: NotRequired[tuple[object, ...] | None]
     filters: NotRequired[tuple[object, ...]]
     shards: NotRequired[tuple[int, ...] | None]