feat: subchunk write order

Cargo.toml

@@ -10,7 +10,7 @@ crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 pyo3 = { version = "0.27.1", features = ["abi3-py311"] }
-zarrs = { version = "0.23.6", features = ["async", "zlib", "pcodec", "bz2"] }
+zarrs = { version = "0.23.7", features = ["async", "zlib", "pcodec", "bz2"] }
 rayon_iter_concurrent_limit = "0.2.0"
 rayon = "1.10.0"
 # fix for https://stackoverflow.com/questions/76593417/package-openssl-was-not-found-in-the-pkg-config-search-path

README.md

@@ -51,6 +51,8 @@ The `ZarrsCodecPipeline` specific options are:
   - Defaults to `False`.
 - `codec_pipeline.strict`: raise exceptions for unsupported operations instead of falling back to the default codec pipeline of `zarr-python`.
   - Defaults to `False`.
+- `codec_pipeline.subchunk_write_order`: Tells `zarrs` in what order to write subchunks within a shard. One of "C" or "random."
+  - Defaults to `random`.
 
 For example:
 ```python
@@ -63,7 +65,8 @@ zarr.config.set({
         "chunk_concurrent_maximum": None,
         "chunk_concurrent_minimum": 4,
         "direct_io": False,
-        "strict": False
+        "strict": False,
+        "subchunk_write_order": "C"
     }
 })
 ```

python/zarrs/_internal.pyi

@@ -30,6 +30,7 @@ class CodecPipelineImpl:
         chunk_concurrent_maximum: builtins.int | None = None,
         num_threads: builtins.int | None = None,
         direct_io: builtins.bool = False,
+        subchunk_write_order: builtins.str = "random",
     ) -> CodecPipelineImpl: ...
     def retrieve_chunks_and_apply_index(
         self,

python/zarrs/pipeline.py

@@ -62,6 +62,9 @@ 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=config.get(
+                "codec_pipeline.subchunk_write_order", "random"
+            ),
         )
     except TypeError as e:
         if strict:

src/lib.rs

@@ -18,10 +18,12 @@ 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::{
     ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayMetadata,
-    ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions, DataType,
-    FillValue, StoragePartialDecoder, copy_fill_value_into, update_array_bytes,
+    ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions,
+    CodecSpecificOptions, DataType, FillValue, StoragePartialDecoder, copy_fill_value_into,
+    update_array_bytes,
 };
 use zarrs::config::global_config;
 use zarrs::convert::array_metadata_v2_to_v3;
@@ -38,7 +40,7 @@ mod utils;
 
 use crate::concurrency::ChunkConcurrentLimitAndCodecOptions;
 use crate::store::StoreConfig;
-use crate::utils::{PyCodecErrExt, PyErrExt as _};
+use crate::utils::{PyCodecErrExt, PyErrExt as _, SubchunkWriteOrderWrapper};
 
 // TODO: Use a OnceLock for store with get_or_try_init when stabilised?
 #[gen_stub_pyclass]
@@ -218,6 +220,7 @@ impl CodecPipelineImpl {
         chunk_concurrent_maximum=None,
         num_threads=None,
         direct_io=false,
+        subchunk_write_order=SubchunkWriteOrderWrapper(SubchunkWriteOrder::Random),
     ))]
     #[new]
     fn new(
@@ -228,6 +231,7 @@ impl CodecPipelineImpl {
         chunk_concurrent_maximum: Option<usize>,
         num_threads: Option<usize>,
         direct_io: bool,
+        subchunk_write_order: SubchunkWriteOrderWrapper,
     ) -> PyResult<Self> {
         store_config.direct_io(direct_io);
         let metadata = serde_json::from_str(array_metadata).map_py_err::<PyTypeError>()?;
@@ -237,8 +241,16 @@ impl CodecPipelineImpl {
             }
             ArrayMetadata::V3(v3) => Cow::Borrowed(v3),
         };
-        let codec_chain =
-            Arc::new(CodecChain::from_metadata(&metadata_v3.codecs).map_py_err::<PyTypeError>()?);
+        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_options = CodecOptions::default().with_validate_checksums(validate_checksums);
 
         let chunk_concurrent_minimum =

src/tests.rs

@@ -10,8 +10,8 @@ use crate::CodecPipelineImpl;
 
 #[test]
 fn test_nparray_to_unsafe_cell_slice_empty() -> PyResult<()> {
-    pyo3::prepare_freethreaded_python();
-    Python::with_gil(|py| {
+    Python::initialize();
+    Python::attach(|py| {
         let arr: Bound<'_, PyUntypedArray> = PyModule::from_code(
             py,
             c_str!(

src/utils.rs

@@ -1,7 +1,10 @@
 use std::fmt::Display;
 
-use pyo3::{PyErr, PyResult, PyTypeInfo};
-use zarrs::array::CodecError;
+use pyo3::{
+    Borrowed, Bound, FromPyObject, IntoPyObject, PyAny, PyErr, PyResult, PyTypeInfo, Python,
+    exceptions::PyValueError, types::PyString,
+};
+use zarrs::array::{CodecError, codec::SubchunkWriteOrder};
 
 use crate::ChunkItem;
 
@@ -41,3 +44,42 @@ pub fn is_whole_chunk(item: &ChunkItem) -> bool {
     item.chunk_subset.start().iter().all(|&o| o == 0)
         && item.chunk_subset.shape() == bytemuck::must_cast_slice::<_, u64>(&item.shape)
 }
+
+#[derive(Debug, Clone)]
+pub struct SubchunkWriteOrderWrapper(pub SubchunkWriteOrder);
+
+impl<'py> IntoPyObject<'py> for SubchunkWriteOrderWrapper {
+    type Target = PyString;
+    type Output = Bound<'py, PyString>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
+        match self.0 {
+            SubchunkWriteOrder::C => Ok("C".into_pyobject(py)?),
+            SubchunkWriteOrder::Random => Ok("random".into_pyobject(py)?),
+            _ => Err(PyValueError::new_err(
+                "Unrecognized subchunk write order for converting to python object, only `C` and `random` allowed.",
+            )),
+        }
+    }
+}
+
+impl<'py> FromPyObject<'_, 'py> for SubchunkWriteOrderWrapper {
+    type Error = PyErr;
+
+    fn extract(option: Borrowed<'_, 'py, PyAny>) -> PyResult<SubchunkWriteOrderWrapper> {
+        match option.extract::<&str>()? {
+            "C" => Ok(SubchunkWriteOrderWrapper(SubchunkWriteOrder::C)),
+            "random" => Ok(SubchunkWriteOrderWrapper(SubchunkWriteOrder::Random)),
+            _ => Err(PyValueError::new_err(
+                "Unrecognized subchunk write order while extracting to rust, only `C` and `random` allowed.",
+            )),
+        }
+    }
+}
+
+impl pyo3_stub_gen::PyStubType for SubchunkWriteOrderWrapper {
+    fn type_output() -> pyo3_stub_gen::TypeInfo {
+        pyo3_stub_gen::TypeInfo::builtin("str")
+    }
+}

tests/test_sharding.py

@@ -1,8 +1,9 @@
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import numpy.typing as npt
 import pytest
+import zarr
 from zarr import Array, AsyncArray
 from zarr.abc.store import Store
 from zarr.codecs import (
@@ -367,3 +368,36 @@ async def test_sharding_with_empty_inner_chunk(
     print("read data")
     data_read = await a.getitem(...)
     assert np.array_equal(data_read, data)
+
+
+@pytest.mark.parametrize(
+    "index_location", [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end]
+)
+@pytest.mark.parametrize("subchunk_write_order", ["C", "random"])
+async def test_sharding_subchunk_write_order(
+    store: Store,
+    index_location: ShardingCodecIndexLocation,
+    subchunk_write_order: Literal["C", "random"],
+) -> None:
+    with zarr.config.set({"codec_pipeline.subchunk_write_order": subchunk_write_order}):
+        path = f"sharding_with_empty_inner_chunk_{index_location}"
+        spath = StorePath(store, path)
+        codec = ShardingCodec(chunk_shape=(2, 2), index_location=index_location)
+        a = await AsyncArray.create(
+            spath,
+            shape=(16, 16),
+            chunk_shape=(16, 16),
+            dtype="uint32",
+            fill_value=0,
+            codecs=[codec],
+        )
+        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 == "C":
+            np.testing.assert_equal(np.sort(index_offsets), index_offsets)
+        else:
+            assert not np.array_equal(np.sort(index_offsets), index_offsets)