feat: serve a sorted integer selection as runs, so it reaches zarrs at all

A sorted integer-array selection is mostly runs of consecutive indices, and a run is a
slice. Split into runs, each one is a contiguous box the existing read path already knows
how to serve.

Without this a fancy index that is not contiguous is declined outright and the whole read
falls back to zarr-python's pipeline -- so the shape a row-sampling data loader produces
never reached zarrs, however well zarrs could have served it.

Reads and writes stop sharing one description function, because they now want different
ones. A read is split into runs; a WRITE must not be. Splitting a write can put two items
on one chunk key, and the write path is read-modify-write, so two items on one key race.

Indices are normalised to int64 first, which two dtypes needed for different reasons. A
uint8 selection like [255, 0] differenced to 1, read as consecutive, and returned an empty
slice -- wrong data, with no error. And a BOOLEAN MASK is not an index array at all: cast
blindly it becomes [1, 1, ...], which is non-decreasing and exactly as long as the output
slice, so it passed every test here and read element 1 once per True. A mask's POSITIONS
are what it means, so that is what it is turned into -- which also lets a mask take this
path rather than falling back.

Author: selmanozleyen

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,
 )
 
 
@@ -183,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,
@@ -218,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,34 @@ 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: `zarr`'s `BoolArrayDimIndexer` hands one
+            # over with a slice out-selection. Cast blindly it becomes [1, 1, ...], which is
+            # non-decreasing and the right length, so it passes every test below and reads
+            # element 1 once per True -- silently wrong data. Its positions are what it means.
+            if sel.dtype.kind == "b":
+                return np.flatnonzero(sel).astype(np.int64, copy=False)
+            # Integers and floats only. A `uint64` selection reaches us as float64, because
+            # zarr subtracts an int64 offset from it, so "f" has to be accepted; anything
+            # else is not a selection this path can read, and declining is what the caller's
+            # fallback is for.
+            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 +85,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))
+                # int64 (see `_as_int64`): 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 +103,89 @@ 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, so an axis can be addressed by position."""
+    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?"""
+    # Negative indices and the output length are the caller's to judge: one raises where the
+    # other declines, so the order of those checks belongs at the call site.
+    return (
+        isinstance(indices, np.ndarray)
+        and indices.ndim == 1
+        # Non-decreasing only: any decrease would mean one box per element, a decode each.
+        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
+) -> Iterator[tuple[SelectorTuple, SelectorTuple]]:
+    """Split a selection with one non-consecutive integer-array axis into contiguous boxes.
+
+    zarrs describes a chunk read as a rectangular subset, so ``z[[3, 7, 8], :]`` has no
+    single-box description -- but it is a *stack* of boxes, one per run of consecutive
+    indices. Only one array axis is split: with two, outer and coordinate indexing disagree
+    on what the selection means. Anything not splittable is yielded unchanged.
+    """
+    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
+    # this line can be removed once https://github.com/zarr-developers/zarr-python/issues/4285 is fixed
+    if (indices < 0).any():
+        raise DiscontiguousArrayError(indices)
+    out_start = out_axis_sel.start or 0
+    if not _output_run_matches(indices, out_axis_sel):
+        yield from unsplit
+        return
+
+    # Always slices, even for one run: `resulting_shape_from_index` mis-orders a non-leading
+    # advanced index, and the caller's element-count check then rejects the selection.
+    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]
+        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 +269,56 @@ 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.
+
+    Neither the chunk-unit grouping nor the run splitting a READ gets: both can put two
+    items on one chunk key, and the write path is a read-modify-write, so two items on one
+    key race. A write therefore describes exactly what it was given.
+    """
+    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.
+
+    A sorted integer selection is mostly runs, and a run is a slice. Handing each run over
+    as a slice is what lets the read reach zarrs at all: a fancy index that is not
+    contiguous is otherwise declined, and the whole read falls back to zarr-python.
+    """
+    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
+            )
+        ],
+        drop_axes,
+        shape,
+    )
+
+
+def _chunk_items(
+    batch_info: BatchInfo,
     drop_axes: tuple[int, ...],
     shape: tuple[int, ...],
 ) -> RustChunkInfo:
+    """One ChunkItem per batch entry, the description both paths end at."""
     is_constant = shape == ()
     chunk_info_with_indices: list[ChunkItem] = []
     write_empty_chunks: bool = True

tests/test_index_dtype_overflow.py

@@ -0,0 +1,124 @@
+"""An index array's dtype must not change which selections are accepted.
+
+Subtracting in the incoming dtype inverts the comparison, because on an unsigned
+array a decrease wraps to a large positive step:
+
+    np.diff(np.array([255, 0], dtype="uint8"))   -> array([1], dtype=uint8)
+
+the most extreme possible decrease reads as consecutive, and the slice built from it,
+`slice(255, 0 + 1)`, is empty. uint64 is worse than wrong: mixed with a signed step it
+promotes to float64 and loses exactness above 2**53.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+import zarr
+
+from zarrs.utils import (
+    DiscontiguousArrayError,
+    _as_int64_batch_info,
+    make_slice_selection,
+    split_selection_runs,
+)
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+SETTINGS = {
+    "codec_pipeline.path": "zarrs.ZarrsCodecPipeline",
+    # No fallback to hide behind: a selection zarrs cannot serve must raise rather
+    # than be served correctly by zarr-python and look like a passing test.
+    "codec_pipeline.strict": True,
+}
+
+UNSIGNED = ["uint8", "uint16", "uint32", "uint64"]
+
+
+@pytest.fixture
+def sharded(tmp_path: Path) -> tuple[Path, np.ndarray]:
+    path = tmp_path / "a.zarr"
+    values = np.arange(32 * 40, dtype="float32").reshape(32, 40)
+    array = zarr.create_array(
+        path,
+        shape=values.shape,
+        dtype="float32",
+        chunks=(4, 5),
+        shards=(16, 20),
+        compressors=[zarr.codecs.BloscCodec(cname="lz4")],
+    )
+    array[:] = values
+    return path, values
+
+
+def test_wraparound_decrease_is_not_consecutive() -> None:
+    """[255, 0] as uint8 differences to 1. It is a decrease of 255, not a step of 1.
+
+    `make_slice_selection` differences directly, so go through the boundary that normalises:
+    neither half is a guarantee alone.
+    """
+    selection = (np.array([255, 0], dtype="uint8"),)
+    ((_, _, chunk_selection, _, _),) = _as_int64_batch_info(
+        [(None, None, selection, selection, True)]
+    )
+    with pytest.raises(DiscontiguousArrayError):
+        make_slice_selection(chunk_selection)
+
+
+@pytest.mark.parametrize("dtype", UNSIGNED)
+def test_consecutive_unsigned_still_collapses(dtype: str) -> None:
+    """The fix must not reject what was always valid."""
+    (result,) = make_slice_selection((np.array([7, 8, 9], dtype=dtype),))
+    assert result == slice(7, 10, 1)
+
+
+@pytest.mark.parametrize("dtype", UNSIGNED)
+def test_unsigned_rows_read_the_same_as_signed(dtype: str, sharded) -> None:
+    """A selection's dtype is not part of its meaning."""
+    path, values = sharded
+    rows = [3, 4, 5, 11, 12, 27]
+    with zarr.config.set(SETTINGS):
+        array = zarr.open_array(path, mode="r")
+        unsigned = array[np.array(rows, dtype=dtype), :]
+        signed = array[np.array(rows, dtype="int64"), :]
+    np.testing.assert_array_equal(unsigned, values[rows, :])
+    np.testing.assert_array_equal(unsigned, signed)
+
+
+@pytest.mark.parametrize("dtype", UNSIGNED)
+def test_unsorted_unsigned_is_still_refused(dtype: str, sharded) -> None:
+    """Descending rows must not be admitted just because the dtype hides the descent."""
+    path, _ = sharded
+    with zarr.config.set(SETTINGS), pytest.raises(DiscontiguousArrayError):
+        zarr.open_array(path, mode="r")[np.array([27, 3], dtype=dtype), :]
+
+
+def test_negative_chunk_relative_index_is_refused() -> None:
+    """A negative index must never become a slice bound.
+
+    Indices are chunk-relative, and zarr miscomputes them for unsigned dtypes -- uint8
+    `[27, 3]` arrives as `[-13]`. `slice(-13, -12)` is an empty subset, not a row near the end.
+    """
+    with pytest.raises(DiscontiguousArrayError):
+        list(
+            split_selection_runs(
+                (np.array([-13]), slice(0, 20, 1)),
+                (slice(0, 1), slice(0, 20)),
+            )
+        )
+
+
+def test_sorted_selections_never_produce_a_negative_bound(sharded) -> None:
+    """The guard above must not be firing on ordinary reads."""
+    path, values = sharded
+    rng = np.random.default_rng(0)
+    with zarr.config.set(SETTINGS):
+        array = zarr.open_array(path, mode="r")
+        for _ in range(50):
+            rows = np.sort(
+                rng.choice(values.shape[0], size=rng.integers(1, 8), replace=False)
+            )
+            np.testing.assert_array_equal(array[rows, :], values[rows, :])