perf/direct zarr io

src/eopf_geozarr/cli.py

@@ -1159,6 +1159,9 @@ 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(
@@ -1189,12 +1192,15 @@ def convert_s2_optimized_command(args: argparse.Namespace) -> None:
             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,
         )

src/eopf_geozarr/s2_optimization/s2_converter.py

@@ -206,6 +206,7 @@ def convert_s2_optimized(
     spatial_chunk: int,
     compression_level: int,
     validate_output: bool,
+    omit_nodes: set[str] | None = None,
     max_retries: int = 3,
 ) -> xr.DataTree:
     """
@@ -224,6 +225,9 @@ def convert_s2_optimized(
         Optimized DataTree
     """
 
+    if omit_nodes is None:
+        omit_nodes = set()
+
     start_time = time.time()
     zg = get_zarr_group(dt_input)
     s2root_model = Sentinel2Root.from_zarr(zg)
@@ -258,17 +262,26 @@ def chunk_reencoder(array: zarr.Array[Any]) -> ChunkEncodingSpec:
             return {"write_chunks": (1,) * (array.ndim - 2) + (spatial_chunk, spatial_chunk)}
         return {"write_chunks": array.chunks}
 
-    out_group = reencode_group(zg, out_store, "", overwrite=True, chunk_reencoder=chunk_reencoder)
-
-    log.info("Adding CRS elements to datasets in measurements")
-    for _, subgroup in out_group["measurements"].groups():
-        for _, dataset in subgroup.groups():
-            add_crs_and_grid_mapping(dataset, crs=crs)
+    out_group = reencode_group(
+        zg,
+        out_store,
+        path="",
+        overwrite=True,
+        chunk_reencoder=chunk_reencoder,
+        omit_nodes=omit_nodes,
+    )
 
-    log.info("Adding CRS elements to quality datasets")
-    for _, subgroup in out_group["quality"].groups():
-        for _, dataset in subgroup.groups():
-            add_crs_and_grid_mapping(dataset, crs=crs)
+    if "measurements" not in omit_nodes:
+        log.info("Adding CRS elements to datasets in measurements")
+        for _, subgroup in out_group["measurements"].groups():
+            for _, dataset in subgroup.groups():
+                add_crs_and_grid_mapping(dataset, crs=crs)
+
+    if "quality" not in omit_nodes:
+        log.info("Adding CRS elements to quality datasets")
+        for _, subgroup in out_group["quality"].groups():
+            for _, dataset in subgroup.groups():
+                add_crs_and_grid_mapping(dataset, crs=crs)
 
     # Step 2: Create multiscale pyramids for each group in the original structure
     log.info("Adding multiscale levels")

src/eopf_geozarr/zarrio.py

@@ -105,34 +105,66 @@ def reencode_array(
 
 def reencode_group(
     group: zarr.Group,
-    store: Any,
+    store: zarr.storage.StoreLike,
     path: str,
     *,
     overwrite: bool = False,
     use_consolidated_for_children: bool = False,
+    omit_nodes: set[str] | None = None,
     chunk_reencoder: Callable[[zarr.Array[Any]], ChunkEncodingSpec] | None = None,
 ) -> zarr.Group:
     """
     Re-encode a Zarr group, applying a re-encoding to all sub-groups and sub-arrays.
+
+    Parameters
+    ----------
+    group : zarr.Group
+        The Zarr group to re-encode
+    store : zarr.storage.StoreLike
+        The store to write into
+    path : str
+        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.
+
     """
+    if omit_nodes is None:
+        omit_nodes = set()
+
     log = structlog.get_logger()
 
     # Convert store-like to a proper Store object
     store_path = sync(make_store_path(store))
     store = store_path.store
 
-    all_members = dict(
+    members = dict(
         group.members(max_depth=None, use_consolidated_for_children=use_consolidated_for_children)
     )
 
-    log = structlog.get_logger()
     log.info("Begin re-encoding Zarr group %s", group)
     new_members: dict[str, ArrayV3Metadata | GroupMetadata] = {
         path: GroupMetadata(zarr_format=3, attributes=group.attrs.asdict())
     }
     chunks_to_encode: list[str] = []
-    for name, member in all_members.items():
-        log.info("re-encoding member %s", name)
+    for name in omit_nodes:
+        if not any(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):
+            log.info(
+                "Skipping node %s because it is contained in a subgroup declared in the omit_groups parameter",
+                name,
+            )
+            continue
+        log.info("Re-encoding member %s", name)
         new_path = f"{path}/{name}"
         member_attrs = member.attrs.asdict()
         if isinstance(member, zarr.Array):