new zarr version handling

Author: ilan-gold

.github/workflows/ci.yml

@@ -16,11 +16,14 @@ env:
 
 jobs:
   build_and_test:
-    name: build and test
+    name: build and test (zarr ${{ matrix.zarr_version }})
     strategy:
       fail-fast: false
       matrix:
         rust_toolchain: ["stable"]  # "nightly"
+        # "main" exercises features not yet in a zarr-python release (e.g. the
+        # sharding `subchunk_write_order`); those tests skip on "released".
+        zarr_version: ["released", "main"]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v5
@@ -53,6 +56,12 @@ jobs:
         run: |
           uv pip install --system -e . --group dev --verbose
 
+      - name: Install zarr-python from main
+        if: matrix.zarr_version == 'main'
+        run: |
+          uv pip install --system --reinstall-package zarr \
+            "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main"
+
       - name: Python Tests
         run: pytest -n auto
 

README.md

@@ -52,7 +52,7 @@ The `ZarrsCodecPipeline` specific options are:
 - `codec_pipeline.strict`: raise exceptions for unsupported operations instead of falling back to the default codec pipeline of `zarr-python`.
   - Defaults to `False`.
 
-The subchunk write order within a shard is derived from the `ShardingCodec`'s `subchunk_write_order` (`zarr-python` >= 3.2), not from a config option. `zarrs` writes `lexicographic` as row-major and falls back to unordered for `morton`/`colexicographic`/`unordered`, since it cannot reproduce those layouts (the data is identical, only the physical byte layout differs).
+The subchunk write order within a shard is derived from each `ShardingCodec`'s `subchunk_write_order` (`zarr-python` >= 3.2.2). Nested shards are handled per level: the order is read from the sharding codec at each nesting depth. `zarrs` falls back to `unordered` for `morton`/`colexicographic` but can handle `lexicographic` and `unordered` explicitly.
 
 For example:
 ```python

python/zarrs/_internal.pyi

@@ -30,9 +30,9 @@ class CodecPipelineImpl:
         chunk_concurrent_maximum: builtins.int | None = None,
         num_threads: builtins.int | None = None,
         direct_io: builtins.bool = False,
-        subchunk_write_order: typing.Literal[
-            "morton", "unordered", "lexicographic", "colexicographic"
-        ] = "unordered",
+        subchunk_write_order: builtins.list[
+            typing.Literal["morton", "unordered", "lexicographic", "colexicographic"]
+        ] = [],
     ) -> CodecPipelineImpl: ...
     def retrieve_chunks_and_apply_index(
         self,

python/zarrs/pipeline.py

@@ -3,11 +3,12 @@
 import asyncio
 import json
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, Literal, TypedDict
 from warnings import warn
 
 import numpy as np
 from zarr.abc.codec import Codec, CodecPipeline
+from zarr.codecs import ShardingCodec
 from zarr.codecs._v2 import V2Codec
 from zarr.core import BatchedCodecPipeline
 from zarr.core.config import config
@@ -32,6 +33,8 @@
     make_chunk_info_for_rust_with_indices,
 )
 
+SubchunkWriteOrder = Literal["morton", "unordered", "lexicographic", "colexicographic"]
+
 
 class UnsupportedDataTypeError(Exception):
     pass
@@ -62,7 +65,7 @@ def get_codec_pipeline_impl(
             ),
             num_threads=config.get("threading.max_workers", None),
             direct_io=config.get("codec_pipeline.direct_io", False),
-            subchunk_write_order=_subchunk_write_order(metadata),
+            subchunk_write_order=_subchunk_write_orders(metadata),
         )
     except TypeError as e:
         if strict:
@@ -90,16 +93,28 @@ class ZarrsCodecPipelineState(TypedDict):
     codecs: tuple[Codec, ...]
 
 
-def _subchunk_write_order(metadata: ArrayMetadata) -> str:
-    # Derive the shard subchunk write order from the sharding codec itself
-    # (zarr-python >=3.2). Older zarr always wrote morton order and has no
-    # attribute; the default matches that. zarrs maps morton/colexicographic
-    # to unordered since it can't reproduce those layouts.
-    for codec in array_metadata_to_codecs(metadata):
-        order = getattr(codec, "subchunk_write_order", None)
-        if order is not None:
-            return order
-    return "morton"
+def _subchunk_write_orders(
+    metadata: ArrayMetadata,
+) -> list[SubchunkWriteOrder]:
+    """Subchunk write order can be nested inside the sharding codec - return that order here.
+
+    Parameters
+    ----------
+    metadata
+        The array metadata containing the potentially nested shards.
+
+    Returns
+    -------
+        A list of subchunk write orders
+    """
+    orders: list[SubchunkWriteOrder] = []
+    codecs: Iterable[Codec] = array_metadata_to_codecs(metadata)
+    while (
+        sharding := next((c for c in codecs if isinstance(c, ShardingCodec)), None)
+    ) is not None:
+        orders.append(getattr(sharding, "subchunk_write_order", "morton"))
+        codecs = sharding.codecs
+    return orders
 
 
 def array_metadata_to_codecs(metadata: ArrayMetadata) -> list[Codec]:

src/lib.rs

@@ -18,15 +18,15 @@ use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use rayon_iter_concurrent_limit::iter_concurrent_limit;
 use unsafe_cell_slice::UnsafeCellSlice;
 use utils::is_whole_chunk;
-use zarrs::array::codec::{ShardingCodecOptions, SubchunkWriteOrder};
+use zarrs::array::codec::{ShardingCodec, ShardingCodecConfiguration, SubchunkWriteOrder};
 use zarrs::array::{
     ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayMetadata,
-    ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions,
-    CodecSpecificOptions, DataType, FillValue, StoragePartialDecoder, copy_fill_value_into,
-    update_array_bytes,
+    ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions, DataType,
+    FillValue, StoragePartialDecoder, copy_fill_value_into, update_array_bytes,
 };
 use zarrs::config::global_config;
 use zarrs::convert::array_metadata_v2_to_v3;
+use zarrs::metadata::v3::MetadataV3;
 use zarrs::plugin::ZarrVersion;
 use zarrs::storage::{ReadableWritableListableStorage, StorageHandle, StoreKey};
 
@@ -42,6 +42,59 @@ use crate::concurrency::ChunkConcurrentLimitAndCodecOptions;
 use crate::store::StoreConfig;
 use crate::utils::{PyCodecErrExt, PyErrExt as _, SubchunkWriteOrderWrapper};
 
+/// Build a codec chain from metadata, applying `orders[depth]` to the sharding
+/// codec at each nesting level (outermost = `orders[0]`). Sharding nests
+/// linearly (one array->bytes codec per chain), so the list is indexed by depth.
+/// Levels beyond the list, or an empty list, keep zarrs' default order.
+fn codec_chain_with_subchunk_write_orders(
+    codecs: &[MetadataV3],
+    orders: &[SubchunkWriteOrderWrapper],
+) -> PyResult<CodecChain> {
+    let base = CodecChain::from_metadata(codecs).map_py_err::<PyTypeError>()?;
+    let array_to_bytes = base.array_to_bytes_codec();
+    let array_to_bytes: Arc<dyn ArrayToBytesCodecTraits> =
+        if array_to_bytes.as_any().is::<ShardingCodec>() {
+            let ShardingCodecConfiguration::V1(config) = codecs
+                .iter()
+                .find_map(|m| m.to_configuration::<ShardingCodecConfiguration>().ok())
+                .ok_or_else(|| {
+                    PyErr::new::<PyTypeError, _>(
+                        "sharding codec present but its metadata was not found",
+                    )
+                })?
+            else {
+                return Err(PyErr::new::<PyTypeError, _>(
+                    "unsupported sharding codec configuration version",
+                ));
+            };
+            let order = orders
+                .first()
+                .map_or(SubchunkWriteOrder::Unordered, |o| o.0);
+            let inner = codec_chain_with_subchunk_write_orders(
+                &config.codecs,
+                orders.get(1..).unwrap_or_default(),
+            )?;
+            let index =
+                CodecChain::from_metadata(&config.index_codecs).map_py_err::<PyTypeError>()?;
+            Arc::new(
+                ShardingCodec::new(
+                    config.chunk_shape,
+                    Arc::new(inner),
+                    Arc::new(index),
+                    config.index_location,
+                )
+                .with_subchunk_write_order(order),
+            )
+        } else {
+            array_to_bytes.clone()
+        };
+    Ok(CodecChain::new(
+        base.array_to_array_codecs().to_vec(),
+        array_to_bytes,
+        base.bytes_to_bytes_codecs().to_vec(),
+    ))
+}
+
 // TODO: Use a OnceLock for store with get_or_try_init when stabilised?
 #[gen_stub_pyclass]
 #[pyclass]
@@ -212,6 +265,7 @@ impl CodecPipelineImpl {
 #[pymethods]
 impl CodecPipelineImpl {
     #[allow(clippy::too_many_arguments)] // python functions can have defaults
+    #[allow(clippy::needless_pass_by_value)] // pyo3 extracts args by value
     #[pyo3(signature = (
         array_metadata,
         store_config,
@@ -221,7 +275,7 @@ impl CodecPipelineImpl {
         chunk_concurrent_maximum=None,
         num_threads=None,
         direct_io=false,
-        subchunk_write_order=SubchunkWriteOrderWrapper(SubchunkWriteOrder::Unordered),
+        subchunk_write_order=Vec::new(),
     ))]
     #[new]
     fn new(
@@ -232,7 +286,10 @@ impl CodecPipelineImpl {
         chunk_concurrent_maximum: Option<usize>,
         num_threads: Option<usize>,
         direct_io: bool,
-        subchunk_write_order: SubchunkWriteOrderWrapper,
+        // One order per sharding-codec nesting level, outermost first. Sharding
+        // nests linearly (a codec chain has exactly one array->bytes codec), so a
+        // flat depth-indexed list is enough — no tree needed.
+        subchunk_write_order: Vec<SubchunkWriteOrderWrapper>,
     ) -> PyResult<Self> {
         store_config.direct_io(direct_io);
         let metadata = serde_json::from_str(array_metadata).map_py_err::<PyTypeError>()?;
@@ -242,16 +299,10 @@ impl CodecPipelineImpl {
             }
             ArrayMetadata::V3(v3) => Cow::Borrowed(v3),
         };
-        let codec_chain = Arc::new(
-            CodecChain::from_metadata(&metadata_v3.codecs)
-                .map_py_err::<PyTypeError>()?
-                .with_codec_specific_options(
-                    &CodecSpecificOptions::default().with_option(
-                        ShardingCodecOptions::default()
-                            .with_subchunk_write_order(subchunk_write_order.0),
-                    ),
-                ),
-        );
+        let codec_chain = Arc::new(codec_chain_with_subchunk_write_orders(
+            &metadata_v3.codecs,
+            &subchunk_write_order,
+        )?);
         let codec_options = CodecOptions::default().with_validate_checksums(validate_checksums);
 
         let chunk_concurrent_minimum =

tests/test_sharding.py

@@ -1,11 +1,13 @@
 import inspect
 import warnings
-from typing import Any, Literal
+from importlib.metadata import version
+from typing import Any, get_args
 
 import numpy as np
 import numpy.typing as npt
 import pytest
-from zarr import AsyncArray, create_array
+from packaging.version import Version
+from zarr import config, create_array
 from zarr.abc.store import Store
 from zarr.api.asynchronous import create_array as create_async_array
 from zarr.codecs import (
@@ -17,7 +19,9 @@
 from zarr.core.array import ShardsConfigParam
 from zarr.core.buffer import default_buffer_prototype
 from zarr.errors import ZarrUserWarning
-from zarr.storage import StorePath
+from zarr.storage import LocalStore, StorePath
+
+from zarrs.pipeline import SubchunkWriteOrder
 
 from .conftest import ArrayRequest
 from .test_codecs import _AsyncArrayProxy, order_from_dim
@@ -328,41 +332,63 @@ async def test_sharding_with_empty_inner_chunk(
     "subchunk_write_order" in inspect.signature(ShardingCodec).parameters
 )
 
-
-@pytest.mark.skipif(
-    not _SHARDING_HAS_WRITE_ORDER,
+# `subchunk_write_order` on the sharding codec is a zarr-python >=3.2.2 feature.
+requires_write_order = pytest.mark.skipif(
+    Version(version("zarr")) < Version("3.2.2dev0"),
     reason="zarr-python ShardingCodec has no subchunk_write_order",
 )
-@pytest.mark.parametrize(
-    "index_location", [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end]
-)
-# "lexicographic" -> C (sorted offsets); "morton" falls back to unordered in zarrs
-@pytest.mark.parametrize("subchunk_write_order", ["lexicographic", "morton"])
-async def test_sharding_subchunk_write_order(
-    store: Store,
-    index_location: ShardingCodecIndexLocation,
-    subchunk_write_order: Literal["lexicographic", "morton"],
+
+
+@requires_write_order
+@pytest.mark.parametrize("nested", [False, True], ids=["flat", "nested"])
+@pytest.mark.parametrize("subchunk_write_order", list(get_args(SubchunkWriteOrder)))
+def test_subchunk_write_order_matches_zarr_python(
+    tmp_path, *, subchunk_write_order: SubchunkWriteOrder, nested: bool
 ) -> None:
-    path = f"sharding_with_empty_inner_chunk_{index_location}"
-    spath = StorePath(store, path)
-    codec = ShardingCodec(
-        chunk_shape=(2, 2),
-        index_location=index_location,
-        subchunk_write_order=subchunk_write_order,
-    )
-    a = await AsyncArray.create(
-        spath,
-        shape=(16, 16),
-        chunk_shape=(16, 16),
-        dtype="uint32",
-        fill_value=0,
-        codecs=[codec],
+    data = np.arange(1, 32 * 32 + 1, dtype="uint32").reshape((32, 32))
+    ground_truth_subchunk_write_order = (
+        "unordered"
+        if subchunk_write_order in {"colexicographic", "unordered", "morton"}
+        else "lexicographic"
     )
-    await a.setitem(..., np.arange(16 * 16).reshape((16, 16)))
-    index = await codec._load_shard_index(a.store_path / "/c/0/0", (8, 8))
-    index_offsets = index.offsets_and_lengths[index.get_full_chunk_map()].ravel()[::2]
-    assert len(index_offsets) == 64  # 8 * 8
-    if subchunk_write_order == "lexicographic":
-        np.testing.assert_equal(np.sort(index_offsets), index_offsets)
+    if nested:
+        zarrs_codec = ShardingCodec(
+            chunk_shape=(8, 8),
+            subchunk_write_order=subchunk_write_order,
+            codecs=[
+                ShardingCodec(chunk_shape=(2, 2), subchunk_write_order="lexicographic")
+            ],
+        )
+        zarr_codec = ShardingCodec(
+            chunk_shape=(8, 8),
+            subchunk_write_order=ground_truth_subchunk_write_order,
+            codecs=[
+                ShardingCodec(chunk_shape=(2, 2), subchunk_write_order="lexicographic")
+            ],
+        )
     else:
-        assert not np.array_equal(np.sort(index_offsets), index_offsets)
+        zarrs_codec = ShardingCodec(
+            chunk_shape=(8, 8), subchunk_write_order="lexicographic"
+        )
+        zarr_codec = ShardingCodec(
+            chunk_shape=(8, 8), subchunk_write_order=ground_truth_subchunk_write_order
+        )
+
+    def write(pipeline: str) -> bytes:
+        sub = tmp_path / pipeline.rsplit(".", 1)[-1]
+        with config.set({"codec_pipeline.path": pipeline}):
+            a = create_array(
+                StorePath(LocalStore(sub)),
+                shape=(32, 32),
+                chunks=(32, 32),
+                dtype="uint32",
+                fill_value=0,
+                serializer=zarrs_codec if "zarrs" in pipeline else zarr_codec,
+                compressors=None,
+            )
+            a[:, :] = data
+        return (sub / "c" / "0" / "0").read_bytes()
+
+    zarrs_bytes = write("zarrs.ZarrsCodecPipeline")
+    zarr_bytes = write("zarr.core.codec_pipeline.BatchedCodecPipeline")
+    assert zarrs_bytes == zarr_bytes