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/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..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, ) @@ -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, 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/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..1ee6de2b 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,33 +46,32 @@ 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, 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 { 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 +83,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 +95,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 +109,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<()> { @@ -219,6 +213,7 @@ impl CodecPipelineImpl { num_threads=None, direct_io=false, file_handle_cache_size=0, + store_is_read_only=false, ))] #[new] fn new( @@ -230,6 +225,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); @@ -240,8 +236,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 +269,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, @@ -280,6 +282,7 @@ impl CodecPipelineImpl { num_threads, fill_value, data_type, + store_is_read_only, }) } @@ -310,18 +313,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 +360,6 @@ impl CodecPipelineImpl { self.codec_chain.decode_into( Cow::Owned(chunk_encoded), &item.shape, - &self.data_type, - &self.fill_value, target, &codec_options, ) @@ -399,6 +395,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_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_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) 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, :])