feat: serve sorted integer-array selections through zarrs

Cargo.toml

@@ -10,7 +10,11 @@ crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 pyo3 = { version = "0.27.1", features = ["abi3-py311"] }
-zarrs = { version = "0.23.6", features = ["async", "zlib", "pcodec", "bz2"] }
+# ON THE 0.24 RELEASE: replace with `zarrs = { version = "0.24", features = [...] }` AND
+# delete the [patch.crates-io] block below, in the SAME commit -- see the note there for why
+# doing only one of the two builds silently against the wrong `zarrs_storage`. src/ already
+# compiles against the 0.24 API.
+zarrs = { git = "https://github.com/zarrs/zarrs", rev = "c17fe374b1fa7df8373b6c6f6eb3f1d33c3a3bd7", 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
@@ -29,3 +33,16 @@ zarrs_object_store = "0.5.0" # object_store 0.12
 
 [profile.release]
 lto = true
+
+# ON THE 0.24 RELEASE: delete this, in the same commit that drops the git rev above.
+#
+# `zarrs_opendal` and `zarrs_object_store` do not depend on `zarrs` at all -- only on
+# `zarrs_storage`, from crates.io. While `zarrs` comes from git that is a SECOND copy of
+# `zarrs_storage`, so there are two `AsyncReadableStorageTraits` and the bounds on the async
+# stores cannot be satisfied. Patching the one crate to the same rev collapses the graph.
+#
+# Drop the git rev above WITHOUT deleting this and the build is silently wrong: released
+# `zarrs 0.24` compiled against an unpublished `zarrs_storage`. Nothing warns, because the
+# git copy carries the same version number as the published one with different contents.
+[patch.crates-io]
+zarrs_storage = { git = "https://github.com/zarrs/zarrs", rev = "c17fe374b1fa7df8373b6c6f6eb3f1d33c3a3bd7" }

python/zarrs/_internal.pyi

@@ -31,6 +31,7 @@ class CodecPipelineImpl:
         num_threads: builtins.int | None = None,
         direct_io: builtins.bool = False,
         file_handle_cache_size: builtins.int = 0,
+        store_is_read_only: builtins.bool = False,
     ) -> CodecPipelineImpl: ...
     def retrieve_chunks_and_apply_index(
         self,

python/zarrs/pipeline.py

@@ -29,7 +29,8 @@
     DiscontiguousArrayError,
     FillValueNoneError,
     UnsupportedVIndexingError,
-    make_chunk_info_for_rust_with_indices,
+    chunk_info_for_read,
+    chunk_info_for_write,
 )
 
 
@@ -65,6 +66,7 @@ def get_codec_pipeline_impl(
             file_handle_cache_size=config.get(
                 "codec_pipeline.file_handle_cache_size", 0
             ),
+            store_is_read_only=store.read_only,
         )
     except TypeError as e:
         if strict:
@@ -182,9 +184,7 @@ async def read(
             if self.impl is None:
                 raise UnsupportedMetadataError()
             self._raise_error_on_unsupported_batch_dtype(batch_info)
-            chunks_desc = make_chunk_info_for_rust_with_indices(
-                batch_info, drop_axes, out.shape
-            )
+            chunks_desc = chunk_info_for_read(batch_info, drop_axes, out.shape)
         except (
             UnsupportedMetadataError,
             DiscontiguousArrayError,
@@ -217,9 +217,7 @@ async def write(
             if self.impl is None:
                 raise UnsupportedMetadataError()
             self._raise_error_on_unsupported_batch_dtype(batch_info)
-            chunks_desc = make_chunk_info_for_rust_with_indices(
-                batch_info, drop_axes, value.shape
-            )
+            chunks_desc = chunk_info_for_write(batch_info, drop_axes, value.shape)
         except (
             UnsupportedMetadataError,
             DiscontiguousArrayError,

python/zarrs/utils.py

@@ -12,14 +12,18 @@
 from zarrs._internal import ChunkItem
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
     from types import EllipsisType
 
     from zarr.abc.store import ByteGetter, ByteSetter
     from zarr.core.array_spec import ArraySpec
     from zarr.core.indexing import SelectorTuple
     from zarr.dtype import ZDType
 
+    BatchInfo = Iterable[
+        tuple[ByteGetter | ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]
+    ]
+
 
 # adapted from https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
 def get_max_threads() -> int:
@@ -38,6 +42,28 @@ class FillValueNoneError(Exception):
     pass
 
 
+def _as_int64_batch_info(batch_info: BatchInfo) -> BatchInfo:
+    """Normalise the batch's array indices to int64 positions, lazily."""
+
+    def cast(sel: SelectorTuple) -> SelectorTuple:
+        if isinstance(sel, np.ndarray):
+            # A boolean mask is not an index array; its positions are what it means.
+            if sel.dtype.kind == "b":
+                return np.flatnonzero(sel).astype(np.int64, copy=False)
+            # "f" is required: a uint64 selection arrives as float64 (uint64 - int64 promotes).
+            if sel.dtype.kind not in "iuf":
+                raise DiscontiguousArrayError(sel.dtype)
+            return sel.astype(np.int64, copy=False)
+        if isinstance(sel, tuple) and any(isinstance(s, np.ndarray) for s in sel):
+            return tuple(map(cast, sel))
+        return sel
+
+    return (
+        (byte_getter, chunk_spec, cast(chunk_sel), cast(out_sel), is_complete)
+        for byte_getter, chunk_spec, chunk_sel, out_sel, is_complete in batch_info
+    )
+
+
 # This is a (mostly) copy of the function from zarr.core.indexing that fixes:
 #   DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated
 # TODO: Upstream this fix
@@ -53,10 +79,11 @@ def make_slice_selection(selection: tuple[np.ndarray | float]) -> list[slice]:
                     slice(int(dim_selection.item()), int(dim_selection.item()) + 1, 1)
                 )
             else:
-                diff = np.diff(dim_selection)
-                if (diff != 1).any() and (diff != 0).any():
-                    raise DiscontiguousArrayError(diff)
-                ls.append(slice(dim_selection[0], dim_selection[-1] + 1, 1))
+                # Callers must normalise to int64 first: an unsigned diff wraps a decrease into +1.
+                steps = dim_selection[1:] - dim_selection[:-1]
+                if (steps != 1).any() and (steps != 0).any():
+                    raise DiscontiguousArrayError(steps)
+                ls.append(slice(int(dim_selection[0]), int(dim_selection[-1]) + 1, 1))
         else:
             ls.append(dim_selection)
     return ls
@@ -70,6 +97,112 @@ def selector_tuple_to_slice_selection(selector_tuple: SelectorTuple) -> list[sli
     return make_slice_selection(selector_tuple)
 
 
+def _as_selector_tuples(
+    chunk_selection: SelectorTuple, out_selection: SelectorTuple
+) -> tuple[tuple, tuple]:
+    """Both selections as tuples."""
+    return (
+        chunk_selection if isinstance(chunk_selection, tuple) else (chunk_selection,),
+        out_selection if isinstance(out_selection, tuple) else (out_selection,),
+    )
+
+
+def _is_sorted_integer_axis(indices: Any, out_axis_sel: Any) -> bool:
+    """Is this one sorted 1-D integer axis written to a contiguous output slice?"""
+    return (
+        isinstance(indices, np.ndarray)
+        and indices.ndim == 1
+        # Non-decreasing only. When zarr DOES reorder for an unsorted selection the
+        # out-selection is an ndarray, which the `isinstance(..., slice)` clause below rejects
+        # first; what reaches this test is `CoordinateIndexer` with `sel_sort is None`, which
+        # hands over a contiguous slice whose indices descend. `out_start + i` would then put
+        # each element at the wrong output position.
+        and not (indices[1:] < indices[:-1]).any()
+        and isinstance(out_axis_sel, slice)
+        and out_axis_sel.step in (None, 1)
+    )
+
+
+def _output_run_matches(indices: np.ndarray, out_axis_sel: slice) -> bool:
+    """Does the output slice hold exactly one element per index."""
+    start = out_axis_sel.start or 0
+    return out_axis_sel.stop - start == indices.size
+
+
+def split_selection_runs(
+    chunk_selection: SelectorTuple,
+    out_selection: SelectorTuple,
+    chunk_shape: tuple[int, ...] | None = None,
+) -> Iterator[tuple[SelectorTuple, SelectorTuple]]:
+    """Split a selection with one non-consecutive integer-array axis into contiguous boxes.
+
+    Only one array axis is split: with two, outer and coordinate indexing disagree on what
+    the selection means. Anything not splittable is yielded unchanged.
+
+    The boxes this yields become `ArrayBytesFixedDisjointView`s, whose `unsafe` constructor
+    takes disjointness on trust. Here that holds because the output slices are consecutive
+    intervals of one run partition -- so a caller must not reorder or duplicate what this
+    yields.
+    """
+    chunk_sel, out_sel = _as_selector_tuples(chunk_selection, out_selection)
+    unsplit = ((chunk_selection, out_selection),)
+
+    array_axes = [
+        axis for axis, sel in enumerate(chunk_sel) if isinstance(sel, np.ndarray)
+    ]
+    # Equal arity means no axis was dropped, so chunk axis `axis` is output axis `axis`.
+    if len(array_axes) != 1 or len(chunk_sel) != len(out_sel):
+        yield from unsplit
+        return
+    (axis,) = array_axes
+    indices = chunk_sel[axis]
+    out_axis_sel = out_sel[axis]
+    if not _is_sorted_integer_axis(indices, out_axis_sel) or not all(
+        isinstance(sel, slice) for sel in out_sel
+    ):
+        yield from unsplit
+        return
+    # BOTH ends. Splitting made `_chunk_items`' "describable as slices" check vacuous on this
+    # axis -- a split box is already all slices, so it compares a shape against itself -- and
+    # that check is what used to reject an index outside the chunk. The low end can be reached
+    # today (zarr-developers/zarr-python#4285 wraps an unsigned decrease into a positive step
+    # and emits a negative chunk-relative index); the high end is the same failure unguarded.
+    # Indices are known non-decreasing here, so the endpoints are the extremes.
+    # `indices.size` first: the endpoints are only readable if there are any, and an empty
+    # array reaches here (`_is_sorted_integer_axis` accepts it vacuously).
+    if indices.size and (
+        indices[0] < 0 or (chunk_shape is not None and indices[-1] >= chunk_shape[axis])
+    ):
+        raise DiscontiguousArrayError(indices)
+    out_start = out_axis_sel.start or 0
+    if not _output_run_matches(indices, out_axis_sel):
+        yield from unsplit
+        return
+
+    # A single run still becomes a slice. Left as an ndarray, `resulting_shape_from_index`
+    # mis-orders a non-leading advanced index, and `_chunk_items`' drop-axis detection then
+    # inserts a phantom axis -- a subset with one dimension too many for the output.
+    boundaries = np.flatnonzero(indices[1:] != indices[:-1] + 1) + 1
+
+    for start, stop in zip(
+        np.concatenate(([0], boundaries)),
+        np.concatenate((boundaries, [indices.size])),
+        strict=True,
+    ):
+        rows = indices[start:stop]
+        # A box describes a RUN as a slice, so its two sides must hold the same count. They
+        # cannot disagree given the boundaries above -- but this is the one thing the vacuous
+        # check above used to catch, and a mismatch here would hand Rust a chunk box longer
+        # than the output box it writes into.
+        if int(rows[-1]) - int(rows[0]) + 1 != int(stop) - int(start):
+            raise DiscontiguousArrayError(rows)
+        box_chunk_sel = list(chunk_sel)
+        box_chunk_sel[axis] = slice(int(rows[0]), int(rows[-1]) + 1)
+        box_out_sel = list(out_sel)
+        box_out_sel[axis] = slice(out_start + int(start), out_start + int(stop))
+        yield tuple(box_chunk_sel), tuple(box_out_sel)
+
+
 def resulting_shape_from_index(
     array_shape: tuple[int, ...],
     index_tuple: tuple[int | slice | EllipsisType | np.ndarray],
@@ -153,13 +286,49 @@ class RustChunkInfo:
     write_empty_chunks: bool
 
 
-def make_chunk_info_for_rust_with_indices(
-    batch_info: Iterable[
-        tuple[ByteGetter | ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]
-    ],
+def chunk_info_for_write(
+    batch_info: BatchInfo,
+    drop_axes: tuple[int, ...],
+    shape: tuple[int, ...],
+) -> RustChunkInfo:
+    """Describe a write batch to Rust, one item per entry.
+
+    Never split: two items on one chunk key make the read-modify-writes race.
+    """
+    return _chunk_items(_as_int64_batch_info(batch_info), drop_axes, shape)
+
+
+def chunk_info_for_read(
+    batch_info: BatchInfo,
+    drop_axes: tuple[int, ...],
+    shape: tuple[int, ...],
+) -> RustChunkInfo:
+    """Describe a read batch to Rust, one box per run of consecutive indices."""
+    return _chunk_items(
+        [
+            (byte_getter, chunk_spec, box_chunk_sel, box_out_sel, is_complete)
+            for (
+                byte_getter,
+                chunk_spec,
+                chunk_selection,
+                out_selection,
+                is_complete,
+            ) in _as_int64_batch_info(batch_info)
+            for box_chunk_sel, box_out_sel in split_selection_runs(
+                chunk_selection, out_selection, chunk_spec.shape
+            )
+        ],
+        drop_axes,
+        shape,
+    )
+
+
+def _chunk_items(
+    batch_info: BatchInfo,
     drop_axes: tuple[int, ...],
     shape: tuple[int, ...],
 ) -> RustChunkInfo:
+    """One ChunkItem per batch entry."""
     is_constant = shape == ()
     chunk_info_with_indices: list[ChunkItem] = []
     write_empty_chunks: bool = True

src/concurrency.rs

@@ -25,7 +25,7 @@ impl ChunkConcurrentLimitAndCodecOptions for Vec<ChunkItem> {
 
         let codec_concurrency = codec_pipeline_impl
             .codec_chain
-            .recommended_concurrency(&item.shape, &codec_pipeline_impl.data_type)
+            .recommended_concurrency(&item.shape)
             .map_codec_err()?;
 
         let min_concurrent_chunks =