Add list_bounding_boxes() for zarr arrays (#1238)

* add list_bounding_boxes() for zarr arrays

* rm test_perf

* changelog

* types

* yield

Author: normanrz

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- Added `list_bounding_boxes()` for Zarr-based datasets. [#1238](https://github.com/scalableminds/webknossos-libs/pull/1238)
 
 ### Changed
 

webknossos/tests/dataset/test_dataset.py

@@ -35,6 +35,7 @@
 from webknossos.utils import (
     copytree,
     get_executor_for_args,
+    is_remote_path,
     named_partial,
     rmtree,
     snake_to_camel_case,
@@ -1989,7 +1990,7 @@ def test_bounding_box_on_disk(data_format: DataFormat, output_path: Path) -> Non
     for offset in write_positions:
         mag.write(absolute_offset=offset * mag.mag.to_vec3_int(), data=write_data)
 
-    if data_format in (DataFormat.Zarr, DataFormat.Zarr3):
+    if is_remote_path(output_path):
         with pytest.warns(UserWarning, match=".*can be slow.*"):
             bounding_boxes_on_disk = list(mag.get_bounding_boxes_on_disk())
 
@@ -2051,22 +2052,13 @@ def test_compression(data_format: DataFormat, output_path: Path) -> None:
         compressed_dataset_path = (
             REMOTE_TESTOUTPUT_DIR / f"simple_{data_format}_dataset_compressed"
         )
-        if data_format in (DataFormat.Zarr, DataFormat.Zarr3):
-            with pytest.warns(UserWarning, match=".*can be slow.*"):
-                mag1.compress(
-                    target_path=compressed_dataset_path,
-                )
-        else:
+        with pytest.warns(UserWarning, match=".*can be slow.*"):
             mag1.compress(
                 target_path=compressed_dataset_path,
             )
         mag1 = Dataset.open(compressed_dataset_path).get_layer("color").get_mag(1)
     else:
-        if data_format in (DataFormat.Zarr, DataFormat.Zarr3):
-            with pytest.warns(UserWarning, match=".*can be slow.*"):
-                mag1.compress()
-        else:
-            mag1.compress()
+        mag1.compress()
 
     assert mag1._is_compressed()
     assert mag1.info.data_format == data_format

webknossos/webknossos/dataset/_array.py

@@ -1,6 +1,7 @@
 import re
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from dataclasses import dataclass
 from functools import lru_cache
 from os.path import relpath
@@ -11,6 +12,7 @@
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Type,
@@ -223,7 +225,7 @@ def _list_files(self) -> Iterator[Path]:
             for filename in self._wkw_dataset.list_files()
         )
 
-    def list_bounding_boxes(self) -> Iterator[NDBoundingBox]:
+    def list_bounding_boxes(self) -> Iterator[BoundingBox]:
         def _extract_num(s: str) -> int:
             match = re.search("[0-9]+", s)
             assert match is not None
@@ -592,9 +594,62 @@ def write(self, bbox: NDBoundingBox, data: np.ndarray) -> None:
         )
         array[requested_domain].write(data).result()
 
-    def list_bounding_boxes(self) -> Iterator[BoundingBox]:
+    def _chunk_key_encoding(self) -> tuple[Literal["default", "v2"], Literal["/", "."]]:
         raise NotImplementedError
 
+    def _list_bounding_boxes(
+        self, kvstore: Any, shard_shape: Vec3Int, shape: VecInt
+    ) -> Iterator[BoundingBox]:
+        _type, separator = self._chunk_key_encoding()
+
+        def _try_parse_ints(vec: Iterable[Any]) -> Optional[list[int]]:
+            output = []
+            for value in vec:
+                try:
+                    output.append(int(value))
+                except ValueError:  # noqa: PERF203
+                    return None
+            return output
+
+        keys = kvstore.list().result()
+        for key in keys:
+            key_parts = key.decode("utf-8").split(separator)
+            if _type == "default":
+                if key_parts[0] != "c":
+                    continue
+                key_parts = key_parts[1:]
+            if len(key_parts) != self._array.ndim:
+                continue
+            chunk_coords_list = _try_parse_ints(key_parts)
+            if chunk_coords_list is None:
+                continue
+
+            if shape.axes[0] == "c":
+                chunk_coords = Vec3Int(chunk_coords_list[1:])
+            else:
+                chunk_coords = Vec3Int(chunk_coords_list)
+
+            yield BoundingBox(chunk_coords * shard_shape, shard_shape)
+
+    def list_bounding_boxes(self) -> Iterator[BoundingBox]:
+        kvstore = self._array.kvstore
+
+        if kvstore.spec().to_json()["driver"] == "s3":
+            raise NotImplementedError(
+                "list_bounding_boxes() is not supported for s3 arrays."
+            )
+
+        _, _, shard_shape, _, shape = self._get_array_dimensions(self._array)
+
+        if shape.axes != ("c", "x", "y", "z") and shape.axes != ("x", "y", "z"):
+            raise NotImplementedError(
+                "list_bounding_boxes() is not supported for non 3-D arrays."
+            )
+
+        # This needs to be a separate function because we need the NotImplementedError
+        # to be raised immediately and not part of the iterator.
+        return self._list_bounding_boxes(kvstore, shard_shape, shape)
+
     def close(self) -> None:
         if self._cached_array is not None:
             self._cached_array = None
@@ -724,12 +779,10 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr3Array":
                                             "configuration": {"endian": "little"},
                                         },
                                         {
-                                            "name": "blosc",
+                                            "name": "zstd",
                                             "configuration": {
-                                                "cname": "zstd",
-                                                "clevel": 5,
-                                                "shuffle": "shuffle",
-                                                "typesize": array_info.voxel_type.itemsize,
+                                                "level": 5,
+                                                "checksum": True,
                                             },
                                         },
                                     ]
@@ -761,6 +814,17 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr3Array":
         ).result()
         return cls(path, _array)
 
+    def _chunk_key_encoding(self) -> tuple[Literal["default", "v2"], Literal["/", "."]]:
+        metadata = self._array.spec().to_json()["metadata"]
+        chunk_key_encoding = metadata["chunk_key_encoding"]
+        _type = chunk_key_encoding["name"]
+        separator = chunk_key_encoding.get("configuration", {}).get(
+            "separator", "/" if _type == "default" else "."
+        )
+        assert _type in ["default", "v2"]
+        assert separator in ["/", "."]
+        return _type, separator
+
 
 class Zarr2Array(TensorStoreArray):
     data_format = DataFormat.Zarr
@@ -812,10 +876,8 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr2Array":
                     "order": "F",
                     "compressor": (
                         {
-                            "id": "blosc",
-                            "cname": "zstd",
-                            "clevel": 5,
-                            "shuffle": 1,
+                            "id": "zstd",
+                            "level": 5,
                         }
                         if array_info.compression_mode
                         else None
@@ -827,3 +889,9 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr2Array":
             }
         ).result()
         return cls(path, _array)
+
+    def _chunk_key_encoding(self) -> tuple[Literal["default", "v2"], Literal["/", "."]]:
+        metadata = self._array.spec().to_json()["metadata"]
+        separator = metadata.get("dimension_separator", ".")
+        assert separator in ["/", "."]
+        return "v2", separator

webknossos/webknossos/dataset/mag_view.py

@@ -10,8 +10,6 @@
 from cluster_tools import Executor
 from upath import UPath
 
-from webknossos.dataset.data_format import DataFormat
-
 from ..geometry import Mag, NDBoundingBox, Vec3Int, Vec3IntLike, VecInt
 from ..utils import (
     LazyPath,
@@ -679,9 +677,13 @@ def compress(
             )
         )
         with get_executor_for_args(args, executor) as executor:
-            if self.layer.data_format == DataFormat.WKW:
+            try:
+                bbox_iterator = self._array.list_bounding_boxes()
+            except NotImplementedError:
+                bbox_iterator = None
+            if bbox_iterator is not None:
                 job_args = []
-                for i, bbox in enumerate(self._array.list_bounding_boxes()):
+                for i, bbox in enumerate(bbox_iterator):
                     bbox = bbox.from_mag_to_mag1(self._mag).intersected_with(
                         self.layer.bounding_box, dont_assert=True
                     )