From 824b5dee2b1aef18661172823a28054f3bc2a7ff Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 28 Aug 2026 08:39:05 +0200 Subject: [PATCH 1/3] chore: migrate to zarrs 0.24 A dependency bump and the API migration it forces, and nothing else. Kept on its own because a dependency migration reviewed together with a behaviour change means neither can be reverted alone. All of it is one upstream idea -- the #411 refactor -- which is that a codec chain is BOUND to a data type and fill value once, where it is built, instead of being handed both at every call. `CodecChain` is now unbound; `decode`, `encode`, `decode_into`, `partial_decoder` and `recommended_concurrency` live on `CodecChainBound`, reached by `with_context`. So the chain is bound in the constructor and five call sites stop passing a pair they no longer need. `zarrs::array::StoragePartialDecoder` is also gone, and nothing named replaced it: the (storage, key) TUPLE is the store-backed BytesPartialDecoderTraits implementation now. No behaviour change, with one exception worth naming rather than hiding: a codec chain that cannot BIND to its data type now fails in the constructor as a TypeError, where before it surfaced at the first read as a RuntimeError. Binding has to happen somewhere, and the constructor is the only place it can. The codec metadata is still PARSED where it was, so an array with both bad codecs and a bad fill value still reports the codecs. The dependency is a git rev rather than a version, because 0.24 is not released yet, and it carries a [patch.crates-io] block for zarrs_storage. That patch is load-bearing, not cosmetic: zarrs_opendal and zarrs_object_store track RELEASED zarrs, so they pull zarrs_storage from crates.io while zarrs comes from git -- two copies of one crate, two distinct AsyncReadableStorageTraits, and trait bounds that cannot be satisfied. Patching that single crate collapses the graph; every other zarrs crate already resolves through the git checkout, verified from a clean lockfile. Both the pin and the patch are commented with exactly what to do on release day, and nothing in src/ changes then. --- Cargo.toml | 19 ++++++++++++++++- src/concurrency.rs | 2 +- src/lib.rs | 52 ++++++++++++++++++---------------------------- 3 files changed, 39 insertions(+), 34 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e3454bec..6f611074 100644 --- a/Cargo.toml +++ b/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" } diff --git a/src/concurrency.rs b/src/concurrency.rs index 08b9dca1..633839e6 100644 --- a/src/concurrency.rs +++ b/src/concurrency.rs @@ -25,7 +25,7 @@ impl ChunkConcurrentLimitAndCodecOptions for Vec { 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 = diff --git a/src/lib.rs b/src/lib.rs index 3ce7eb2e..08403ea9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,10 +18,11 @@ 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::api::BytesPartialDecoderTraits; use zarrs::array::{ ArrayBytes, ArrayBytesDecodeIntoTarget, ArrayBytesFixedDisjointView, ArrayMetadata, - ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecOptions, DataType, - FillValue, StoragePartialDecoder, copy_fill_value_into, update_array_bytes, + ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, CodecChain, CodecChainBound, CodecOptions, + DataType, FillValue, copy_fill_value_into, update_array_bytes, }; use zarrs::config::global_config; use zarrs::convert::array_metadata_v2_to_v3; @@ -45,7 +46,7 @@ use crate::utils::{PyCodecErrExt, PyErrExt as _}; #[pyclass] pub struct CodecPipelineImpl { pub(crate) store: ReadableWritableListableStorage, - pub(crate) codec_chain: Arc, + pub(crate) codec_chain: Arc, pub(crate) codec_options: CodecOptions, pub(crate) chunk_concurrent_minimum: usize, pub(crate) chunk_concurrent_maximum: usize, @@ -58,20 +59,14 @@ impl CodecPipelineImpl { fn retrieve_chunk_bytes<'a>( &self, item: &ChunkItem, - codec_chain: &CodecChain, + codec_chain: &CodecChainBound, codec_options: &CodecOptions, ) -> PyResult> { let value_encoded = self.store.get(&item.key).map_py_err::()?; let value_decoded = if let Some(value_encoded) = value_encoded { let value_encoded: Vec = value_encoded.into(); // zero-copy in this case codec_chain - .decode( - value_encoded.into(), - &item.shape, - &self.data_type, - &self.fill_value, - codec_options, - ) + .decode(value_encoded.into(), &item.shape, codec_options) .map_codec_err()? } else { ArrayBytes::new_fill_value(&self.data_type, item.num_elements, &self.fill_value) @@ -83,7 +78,7 @@ impl CodecPipelineImpl { fn store_chunk_bytes( &self, item: &ChunkItem, - codec_chain: &CodecChain, + codec_chain: &CodecChainBound, value_decoded: ArrayBytes, codec_options: &CodecOptions, ) -> PyResult<()> { @@ -95,13 +90,7 @@ impl CodecPipelineImpl { self.store.erase(&item.key).map_py_err::() } else { let value_encoded = codec_chain - .encode( - value_decoded, - &item.shape, - &self.data_type, - &self.fill_value, - codec_options, - ) + .encode(value_decoded, &item.shape, codec_options) .map(Cow::into_owned) .map_codec_err()?; @@ -115,7 +104,7 @@ impl CodecPipelineImpl { fn store_chunk_subset_bytes( &self, item: &ChunkItem, - codec_chain: &CodecChain, + codec_chain: &CodecChainBound, chunk_subset_bytes: ArrayBytes, codec_options: &CodecOptions, ) -> PyResult<()> { @@ -240,8 +229,10 @@ impl CodecPipelineImpl { } ArrayMetadata::V3(v3) => Cow::Borrowed(v3), }; + // Parsed before binding, so an array with bad codecs and a bad fill value still + // reports the codecs. let codec_chain = - Arc::new(CodecChain::from_metadata(&metadata_v3.codecs).map_py_err::()?); + CodecChain::from_metadata(&metadata_v3.codecs).map_py_err::()?; let codec_options = CodecOptions::default().with_validate_checksums(validate_checksums); let chunk_concurrent_minimum = @@ -271,6 +262,10 @@ impl CodecPipelineImpl { }) .map_py_err::()?; + let codec_chain = codec_chain + .with_context(data_type.clone(), fill_value.clone()) + .map_py_err::()?; + Ok(Self { store, codec_chain, @@ -310,18 +305,13 @@ impl CodecPipelineImpl { if !partial_chunk_items.is_empty() { let key_decoder_pairs = iter_concurrent_limit!(chunk_concurrent_limit, partial_chunk_items, map, |item| { - let storage_handle = Arc::new(StorageHandle::new(self.store.clone())); - let input_handle = StoragePartialDecoder::new(storage_handle, item.key.clone()); + // The (storage, key) tuple IS the store-backed `BytesPartialDecoderTraits`. + let input_handle: Arc = + Arc::new((StorageHandle::new(self.store.clone()), item.key.clone())); let partial_decoder = self .codec_chain .clone() - .partial_decoder( - Arc::new(input_handle), - &item.shape, - &self.data_type, - &self.fill_value, - &codec_options, - ) + .partial_decoder(input_handle, &item.shape, &codec_options) .map_codec_err()?; Ok((item.key.clone(), partial_decoder)) }) @@ -362,8 +352,6 @@ impl CodecPipelineImpl { self.codec_chain.decode_into( Cow::Owned(chunk_encoded), &item.shape, - &self.data_type, - &self.fill_value, target, &codec_options, ) From 8b90ba56b10ec98b2ec1ca76a980026369e2c4e3 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 28 Aug 2026 09:59:03 +0200 Subject: [PATCH 2/3] fix: refuse writes to a store zarr-python opened read-only zarr-python enforces read-only in the store itself -- `Store._check_writable`, reached from the concrete store's `_set`. This pipeline never gets there: it is handed a `StoreConfig` and builds its own Rust store, writable whatever mode the array was opened in, so a write to a `mode="r"` array succeeds here and raises through the default pipeline. Silent divergence, and silent data loss for anyone relying on the mode. The flag has to be carried from Python because nothing on the Rust side can infer it: the mode is zarr-python's, not the store's. One guard, at the only entry point Python can write through. `store_chunk_bytes` and `store_chunk_subset_bytes` are reached only from it. The tests open STRICT. zarr's own refusal message is byte-identical to the guard's, so with a fallback available they would pass whether the guard fired or zarr-python served the write. --- python/zarrs/_internal.pyi | 1 + python/zarrs/pipeline.py | 1 + src/lib.rs | 13 ++++++++ tests/test_read_only_store.py | 56 +++++++++++++++++++++++++++++++++++ 4 files changed, 71 insertions(+) create mode 100644 tests/test_read_only_store.py diff --git a/python/zarrs/_internal.pyi b/python/zarrs/_internal.pyi index 2f676818..d324408a 100644 --- a/python/zarrs/_internal.pyi +++ b/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, diff --git a/python/zarrs/pipeline.py b/python/zarrs/pipeline.py index b3368219..9bb3ccbb 100644 --- a/python/zarrs/pipeline.py +++ b/python/zarrs/pipeline.py @@ -65,6 +65,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: diff --git a/src/lib.rs b/src/lib.rs index 3ce7eb2e..39bb1263 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,6 +52,11 @@ pub struct CodecPipelineImpl { pub(crate) num_threads: usize, pub(crate) fill_value: FillValue, pub(crate) data_type: DataType, + /// Whether zarr-python opened this store read-only. + /// + /// Not inferable here: `StoreConfig` builds a writable Rust store whatever mode the + /// array was opened in. + pub(crate) store_is_read_only: bool, } impl CodecPipelineImpl { @@ -219,6 +224,7 @@ impl CodecPipelineImpl { num_threads=None, direct_io=false, file_handle_cache_size=0, + store_is_read_only=false, ))] #[new] fn new( @@ -230,6 +236,7 @@ impl CodecPipelineImpl { num_threads: Option, direct_io: bool, file_handle_cache_size: usize, + store_is_read_only: bool, ) -> PyResult { store_config.direct_io(direct_io); store_config.file_handle_cache_size(file_handle_cache_size); @@ -280,6 +287,7 @@ impl CodecPipelineImpl { num_threads, fill_value, data_type, + store_is_read_only, }) } @@ -399,6 +407,11 @@ impl CodecPipelineImpl { value: &Bound<'_, PyUntypedArray>, write_empty_chunks: bool, ) -> PyResult<()> { + if self.store_is_read_only { + return Err(PyValueError::new_err( + "store was opened in read-only mode and does not support writing", + )); + } enum InputValue<'a> { Array(ArrayBytes<'a>), Constant(FillValue), diff --git a/tests/test_read_only_store.py b/tests/test_read_only_store.py new file mode 100644 index 00000000..24fd65b9 --- /dev/null +++ b/tests/test_read_only_store.py @@ -0,0 +1,56 @@ +"""A store opened read-only must refuse writes, as zarr-python's own pipeline does. + +zarr-python enforces this in the store itself -- `Store._check_writable`, reached from the +concrete store's `_set`. This pipeline never gets there: it is handed a `StoreConfig` and +builds its own Rust store, writable whatever mode the array was opened in. Without the guard +a write to a `mode="r"` array SUCCEEDS here and raises through the default pipeline. + +Opened STRICT throughout, and that is what makes the assertion mean anything: zarr's own +refusal message is byte-identical to the Rust guard's, so with a fallback available these +tests would pass whether the guard fired or zarr-python served the write. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import zarr + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.fixture +def array(tmp_path: Path) -> tuple[Path, np.ndarray]: + values = np.arange(64, dtype=np.float32) + path = tmp_path / "a" + zarr.create_array(path, dtype=values.dtype, shape=values.shape, chunks=(16,))[:] = ( + values + ) + return path, values + + +def open_strict(path: Path, mode: str) -> zarr.Array: + """Open with no fallback. `strict` has to be set BEFORE the open: that is when the + pipeline decides whether it has one.""" + with zarr.config.set({"codec_pipeline.strict": True}): + return zarr.open_array(path, mode=mode) + + +def test_write_to_a_read_only_array_raises(array: tuple[Path, np.ndarray]) -> None: + path, values = array + z = open_strict(path, "r") + with pytest.raises(ValueError, match="read-only"): + z[0:16] = -1.0 + np.testing.assert_array_equal(zarr.open_array(path, mode="r")[:], values) + + +def test_a_writable_array_still_writes(array: tuple[Path, np.ndarray]) -> None: + path, values = array + z = open_strict(path, "r+") + z[0:16] = -1.0 + expected = values.copy() + expected[0:16] = -1.0 + np.testing.assert_array_equal(zarr.open_array(path, mode="r")[:], expected) From ee7d47e264f893cbd2f10aed2419dc2898a84560 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Fri, 28 Aug 2026 12:24:14 +0200 Subject: [PATCH 3/3] 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. --- python/zarrs/pipeline.py | 11 +- python/zarrs/utils.py | 187 +++++++++++++++++++++++++-- tests/test_index_dtype_overflow.py | 113 +++++++++++++++++ tests/test_sorted_fancy_indexing.py | 189 ++++++++++++++++++++++++++++ 4 files changed, 484 insertions(+), 16 deletions(-) create mode 100644 tests/test_index_dtype_overflow.py create mode 100644 tests/test_sorted_fancy_indexing.py diff --git a/python/zarrs/pipeline.py b/python/zarrs/pipeline.py index 9bb3ccbb..19de23b2 100644 --- a/python/zarrs/pipeline.py +++ b/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, diff --git a/python/zarrs/utils.py b/python/zarrs/utils.py index 0881cc1e..f032f13b 100644 --- a/python/zarrs/utils.py +++ b/python/zarrs/utils.py @@ -12,7 +12,7 @@ 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 @@ -20,6 +20,10 @@ 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 diff --git a/tests/test_index_dtype_overflow.py b/tests/test_index_dtype_overflow.py new file mode 100644 index 00000000..43facbe1 --- /dev/null +++ b/tests/test_index_dtype_overflow.py @@ -0,0 +1,113 @@ +"""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], "uint8"))`` is +``[1]``, so the most extreme possible decrease reads as consecutive and the slice built +from it, ``slice(255, 1)``, is empty. uint64 is worse than wrong: it reaches us promoted +to float64 (zarr subtracts an int64 chunk offset) 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 + +# 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. +STRICT = { + "codec_pipeline.path": "zarrs.ZarrsCodecPipeline", + "codec_pipeline.strict": True, +} +# uint8 arrives as int64; uint64 alone arrives as float64. uint16/uint32 are uint8's path. +UNSIGNED = ["uint8", "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) + zarr.create_array( + path, shape=values.shape, dtype="float32", chunks=(4, 5), shards=(16, 20) + )[:] = 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(STRICT): + 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_unsigned_descending_rows_are_refused(dtype: str, sharded) -> None: + """Rows 27 and 3 land in different shards, so each arrives alone and looks orderable. + + What refuses them is the negative bound: zarr makes 3 chunk-relative against shard 1 + and hands over [-13]. Signed dtypes are unaffected, which is why this is dtype-specific. + """ + path, _ = sharded + with zarr.config.set(STRICT), 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: `slice(-13, -12)` is an empty + subset near the end of the chunk, not the row the caller asked for.""" + 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(STRICT): + array = zarr.open_array(path, mode="r") + for _ in range(50): + rows = np.sort(rng.choice(32, size=rng.integers(1, 8), replace=False)) + np.testing.assert_array_equal(array[rows, :], values[rows, :]) diff --git a/tests/test_sorted_fancy_indexing.py b/tests/test_sorted_fancy_indexing.py new file mode 100644 index 00000000..f8f41813 --- /dev/null +++ b/tests/test_sorted_fancy_indexing.py @@ -0,0 +1,189 @@ +"""Sorted integer-array reads reach the zarrs pipeline instead of falling back. + +`strict` makes these categorical rather than merely correct: with no fallback, a selection +zarrs cannot serve raises instead of quietly returning the right answer via zarr-python's +pipeline. It must be set before the array is opened -- that is when the pipeline decides +whether it has a fallback -- so every test opens its own handle via `open_strict`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import zarr + +from zarrs.utils import DiscontiguousArrayError, UnsupportedVIndexingError + +if TYPE_CHECKING: + from pathlib import Path + +SHAPE = (32, 24) +CHUNKS = (8, 6) +SHARDS = (16, 12) +ZARRS = {"codec_pipeline.path": "zarrs.ZarrsCodecPipeline"} + + +def _write(path, values, chunks, shards) -> Path: + zarr.create_array( + path, dtype=values.dtype, shape=values.shape, chunks=chunks, shards=shards + )[:] = values + return path + + +@pytest.fixture +def sharded(tmp_path: Path) -> tuple[Path, np.ndarray]: + values = np.arange(np.prod(SHAPE), dtype=np.float64).reshape(SHAPE) + return _write(tmp_path / "2d.zarr", values, CHUNKS, SHARDS), values + + +@pytest.fixture +def sharded_1d(tmp_path: Path) -> tuple[Path, np.ndarray]: + values = np.arange(64, dtype=np.float64) + return _write(tmp_path / "1d.zarr", values, (8,), (16,)), values + + +def open_strict(path: Path) -> zarr.Array: + """Open with no fallback, so an unsupported selection raises instead of being rerouted.""" + with zarr.config.set({**ZARRS, "codec_pipeline.strict": True}): + return zarr.open_array(path, mode="r+") + + +# Each case spans chunk *and* shard boundaries, and mixes runs of length 1 with longer ones. +@pytest.mark.parametrize( + "index", + [ + pytest.param(np.array([0, 3, 4, 5, 17, 30]), id="rows"), + pytest.param((np.array([2, 3, 20]), slice(4, 18)), id="rows-slice"), + pytest.param((slice(None), np.array([0, 1, 7, 23])), id="cols"), + # Do not drop: on main this raises `RuntimeError: incompatible offset`, which the + # fallback does not catch, so it is the one case here fixing a hard crash. + pytest.param((slice(6, 9), np.array([5, 6, 13])), id="slice-cols"), + pytest.param((np.array([4, 5, 6, 29]), 7), id="rows-int"), + pytest.param(np.array([0, 31]), id="rows-endpoints"), + pytest.param(np.array([11]), id="single-row"), + pytest.param(np.arange(0, 32), id="every-row"), + # A repeat ends its run early and reads that index again into the next output slot. + pytest.param(np.array([3, 3, 3]), id="all-repeats"), + pytest.param(np.array([0, 0, 1, 2, 2]), id="repeats-either-side-of-a-run"), + ], +) +def test_sorted_integer_array_read( + sharded: tuple[Path, np.ndarray], index: object +) -> None: + path, expected = sharded + np.testing.assert_array_equal(open_strict(path)[index], expected[index]) + + +def test_sorted_vindex_1d(sharded_1d: tuple[Path, np.ndarray]) -> None: + path, expected = sharded_1d + index = np.array([0, 1, 5, 12, 13, 14, 63]) + z = open_strict(path) + np.testing.assert_array_equal(z.vindex[index], expected[index]) + np.testing.assert_array_equal(z[index], expected[index]) + + +# Indices within one shard, so a single chunk item really does get several of them -- spread +# across shards each item gets one index, which is a box and was always supported. +@pytest.mark.parametrize( + "index", + [ + # Unsorted: zarr-python reorders the output, so a run's position in the selection is + # not its position in the output. + pytest.param(np.array([9, 2]), id="unsorted-rows"), + # Two array axes: outer and coordinate indexing disagree on what this means. + pytest.param((np.array([1, 3]), np.array([0, 2])), id="two-array-axes"), + pytest.param((slice(None), slice(None, None, 2)), id="strided"), + ], +) +def test_unsupported_raises_strictly_but_falls_back_correctly( + sharded: tuple[Path, np.ndarray], index: object +) -> None: + path, expected = sharded + with pytest.raises((DiscontiguousArrayError, UnsupportedVIndexingError)): + open_strict(path)[index] + with zarr.config.set(ZARRS): + z = zarr.open_array(path, mode="r") + np.testing.assert_array_equal(z[index], expected[index]) + + +def test_writes_are_not_split( + sharded: tuple[Path, np.ndarray], monkeypatch: pytest.MonkeyPatch +) -> None: + """A split write would make several read-modify-writes of one chunk race. + + The values alone prove nothing: without strict mode the write falls back to zarr-python + whatever happens, so this passed identically when writes WERE split. What it asserts is + that the splitter is never reached. Rows 1, 3 and 4 share a chunk -- the losing case. + """ + path, expected = sharded + index = np.array([1, 3, 4]) + monkeypatch.setattr( + "zarrs.utils.split_selection_runs", + lambda *_: pytest.fail("a write reached split_selection_runs"), + ) + with zarr.config.set(ZARRS): + z = zarr.open_array(path, mode="r+") + z[index, :] = np.full((len(index), SHAPE[1]), -1.0) + + # Undone before reading back: a READ is meant to reach the splitter. + monkeypatch.undo() + expected[index, :] = -1.0 + np.testing.assert_array_equal(z[...], expected) + + +def test_contiguous_output_does_not_imply_sorted_input( + sharded_1d: tuple[Path, np.ndarray], +) -> None: + """A rectangular output side is not evidence the input was ordered. + + `CoordinateIndexer` sorts only when the chunk-raveled order is wrong, and 7 and 3 both + live in chunk 0, so `out_selection` comes back `slice(0, 2)` while the indices descend. + Building runs from that would give the inverted box `slice(7, 4)`. + """ + path, expected = sharded_1d + index = np.array([7, 3]) + with pytest.raises((DiscontiguousArrayError, UnsupportedVIndexingError)): + open_strict(path).vindex[index] + with zarr.config.set(ZARRS): + got = zarr.open_array(path, mode="r").vindex[index] + np.testing.assert_array_equal(got, expected[index]) + + +def test_a_split_read_is_rejected_without_the_splitter( + sharded: tuple[Path, np.ndarray], monkeypatch: pytest.MonkeyPatch +) -> None: + """The tests above must be exercising the splitter, not passing for some other reason.""" + monkeypatch.setattr( + "zarrs.utils.split_selection_runs", + lambda chunk_sel, out_sel, chunk_shape=None: ((chunk_sel, out_sel),), + ) + path, _ = sharded + with pytest.raises(DiscontiguousArrayError): + open_strict(path)[np.array([0, 3, 4, 5, 17, 30])] + + +@pytest.mark.parametrize( + "mask", + [ + pytest.param(np.arange(SHAPE[0]) < 16, id="aligned-to-chunks"), + pytest.param(np.ones(SHAPE[0], dtype=bool), id="every-row"), + pytest.param(np.isin(np.arange(SHAPE[0]), [3, 17, 30]), id="scattered"), + pytest.param(np.zeros(SHAPE[0], dtype=bool), id="no-rows"), + ], +) +def test_boolean_mask_reads_the_positions_it_marks( + sharded: tuple[Path, np.ndarray], mask: np.ndarray +) -> None: + """A mask is not an index array, and casting one is silently wrong. + + `BoolArrayDimIndexer` hands over a BOOLEAN chunk selection with a slice out-selection. + Cast to int64 it becomes [1, 1, ...] -- non-decreasing, and exactly as long as the + output slice -- so it passes every eligibility test and reads element 1 once per True. + Values only, no exception. Masks aligned to chunk boundaries hit it every time. + """ + path, expected = sharded + with zarr.config.set(ZARRS): + got = zarr.open_array(path, mode="r")[mask, :] + np.testing.assert_array_equal(got, expected[mask, :])