Configurable Zarr3 codecs (#1343)

* adds MagView.rechunk method

* make zarr3 codecs configurable

* types

* tests

* fixes

* fixes

* types

* types

* changelog

* test

* changelog

* add zarr3_codecs to more methods

* types

* changed signature to Zarr3Config

* pr feedback

* types

* changelog + rechunk

* docsrting

* Update webknossos/webknossos/dataset/mag_view.py

Co-authored-by: Tobias <[email protected]>

* compress defaults to True

---------

Co-authored-by: Tobias <[email protected]>

Author: normanrz

webknossos/Changelog.md

@@ -16,9 +16,11 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 
 ### Added
 - Added `MagView.rechunk` methods to allow for rechunking of datasets. [#1342](https://github.com/scalableminds/webknossos-libs/pull/1342)
+- Added the option to configure the codecs of Zarr3 datasets. Supply a `Zarr3Config` to the `compress` argument in `Layer.add_mag` or similar methods. `codecs` and `chunk_key_encoding` can be customized. [#1343](https://github.com/scalableminds/webknossos-libs/pull/1343)
 
 ### Changed
 - Enforces that `chunk_shape` and `shard_shape` have power-of-two values. This assumptions was used in the code previously, but not explicitly enforced. [#1342](https://github.com/scalableminds/webknossos-libs/pull/1342)
+- The `sharding_indexed` codec is no longer used if `chunk_shape == shard_shape`. [#1343](https://github.com/scalableminds/webknossos-libs/pull/1343)
 
 ### Fixed
 

webknossos/tests/dataset/test_dataset.py

@@ -27,6 +27,7 @@
     RemoteDataset,
     View,
 )
+from webknossos.dataset._array import Zarr3ArrayInfo, Zarr3Config
 from webknossos.dataset.data_format import AttachmentDataFormat, DataFormat
 from webknossos.dataset.dataset import PROPERTIES_FILE_NAME
 from webknossos.dataset.defaults import DEFAULT_DATA_FORMAT
@@ -2484,6 +2485,9 @@ def test_invalid_chunk_shard_shape(output_path: UPath) -> None:
     with pytest.raises(ValueError, match=".*must be a multiple.*"):
         layer.add_mag("1", chunk_shape=(16, 16, 16), shard_shape=(8, 16, 16))
 
+    with pytest.raises(ValueError, match=".*must be a multiple.*"):
+        layer.add_mag("1", chunk_shape=(16, 16, 16), shard_shape=(8, 8, 8))
+
     with pytest.raises(ValueError, match=".*must be a multiple.*"):
         # also not a power-of-two shard shape
         layer.add_mag("1", chunk_shape=(16, 16, 16), shard_shape=(53, 16, 16))
@@ -2683,6 +2687,103 @@ def test_rechunking(data_format: DataFormat, output_path: UPath) -> None:
     assure_exported_properties(mag1.layer.dataset)
 
 
+@pytest.mark.parametrize("output_path", OUTPUT_PATHS)
+def test_zarr3_config(output_path: UPath) -> None:
+    new_dataset_path = prepare_dataset_path(DataFormat.Zarr3, output_path)
+    ds = Dataset(new_dataset_path, voxel_size=(2, 2, 1))
+    mag1 = ds.add_layer(
+        "color", COLOR_CATEGORY, num_channels=3, data_format=DataFormat.Zarr3
+    ).add_mag(
+        1,
+        compress=Zarr3Config(
+            codecs=(
+                {"name": "bytes"},
+                {"name": "gzip", "configuration": {"level": 3}},
+            ),
+            chunk_key_encoding={
+                "name": "default",
+                "configuration": {"separator": "."},
+            },
+        ),
+    )
+
+    # writing unaligned data to an uncompressed dataset
+    write_data = (np.random.rand(3, 10, 20, 30) * 255).astype(np.uint8)
+    mag1.write(write_data, absolute_offset=(60, 80, 100), allow_resize=True)
+
+    assert isinstance(mag1.info, Zarr3ArrayInfo)
+    assert mag1.info.codecs == (
+        {"name": "bytes"},
+        {"name": "gzip", "configuration": {"level": 3}},
+    )
+    assert mag1.info.chunk_key_encoding == {
+        "name": "default",
+        "configuration": {"separator": "."},
+    }
+    assert (mag1.path / "c.0.0.0.0").exists()
+    assert json.loads((mag1.path / "zarr.json").read_bytes())["codecs"][0][
+        "configuration"
+    ]["codecs"] == [
+        {"name": "bytes"},
+        {"name": "gzip", "configuration": {"level": 3}},
+    ]
+
+    np.testing.assert_array_equal(
+        write_data, mag1.read(absolute_offset=(60, 80, 100), size=(10, 20, 30))
+    )
+
+    assure_exported_properties(mag1.layer.dataset)
+
+
+@pytest.mark.parametrize("output_path", OUTPUT_PATHS)
+def test_zarr3_sharding(output_path: UPath) -> None:
+    new_dataset_path = prepare_dataset_path(DataFormat.Zarr3, output_path)
+    ds = Dataset(new_dataset_path, voxel_size=(2, 2, 1))
+    mag1 = ds.add_layer(
+        "color", COLOR_CATEGORY, num_channels=3, data_format=DataFormat.Zarr3
+    ).add_mag(1, chunk_shape=(32, 32, 32), shard_shape=(64, 64, 64))
+
+    # writing unaligned data to an uncompressed dataset
+    write_data = (np.random.rand(3, 10, 20, 30) * 255).astype(np.uint8)
+    mag1.write(write_data, absolute_offset=(60, 80, 100), allow_resize=True)
+
+    assert (
+        json.loads((mag1.path / "zarr.json").read_bytes())["codecs"][0]["name"]
+        == "sharding_indexed"
+    )
+
+    np.testing.assert_array_equal(
+        write_data, mag1.read(absolute_offset=(60, 80, 100), size=(10, 20, 30))
+    )
+
+    assure_exported_properties(mag1.layer.dataset)
+
+
+@pytest.mark.parametrize("output_path", OUTPUT_PATHS)
+def test_zarr3_no_sharding(output_path: UPath) -> None:
+    new_dataset_path = prepare_dataset_path(DataFormat.Zarr3, output_path)
+    ds = Dataset(new_dataset_path, voxel_size=(2, 2, 1))
+    mag1 = ds.add_layer(
+        "color", COLOR_CATEGORY, num_channels=3, data_format=DataFormat.Zarr3
+    ).add_mag(1, chunk_shape=(32, 32, 32), shard_shape=(32, 32, 32))
+
+    # writing unaligned data to an uncompressed dataset
+    write_data = (np.random.rand(3, 10, 20, 30) * 255).astype(np.uint8)
+    mag1.write(write_data, absolute_offset=(60, 80, 100), allow_resize=True)
+
+    # Don't set up a sharding codec, if no sharding is necessary, i.e. chunk_shape == shard_shape
+    assert (
+        json.loads((mag1.path / "zarr.json").read_bytes())["codecs"][0]["name"]
+        != "sharding_indexed"
+    )
+
+    np.testing.assert_array_equal(
+        write_data, mag1.read(absolute_offset=(60, 80, 100), size=(10, 20, 30))
+    )
+
+    assure_exported_properties(mag1.layer.dataset)
+
+
 def test_dataset_view_configuration() -> None:
     ds_path = prepare_dataset_path(DataFormat.WKW, TESTOUTPUT_DIR)
     ds1 = Dataset(ds_path, voxel_size=(2, 2, 1))