From faa5ef502a364809a32656c6942d8d501091610c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 14 Apr 2026 14:27:49 +0200 Subject: [PATCH 01/43] feat: add IndexTransform library for composable, lazy coordinate mappings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new `src/zarr/core/transforms/` package implementing TensorStore-inspired index transforms. The core idea: every indexing operation (slicing, fancy indexing, etc.) produces a coordinate mapping from user space to storage space. These mappings compose lazily — no I/O until explicitly resolved. Key types: - `IndexDomain` — rectangular region in N-dimensional integer space - `ConstantMap`, `DimensionMap`, `ArrayMap` — three representations of a set of storage coordinates (singleton, arithmetic progression, explicit enumeration) - `IndexTransform` — pairs an input domain with output maps (one per storage dim) - `compose(outer, inner)` — chain two transforms Key operations on IndexTransform: - `__getitem__`, `.oindex[]`, `.vindex[]` — indexing produces new transforms - `.intersect(domain)` — restrict to coordinates within a region (chunk resolution) - `.translate(shift)` — shift coordinates (make chunk-local) The transform library is standalone with no dependency on Array. Includes comprehensive test suite (143 tests covering all types, operations, composition, chunk resolution, and edge cases). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/zarr/core/array.py | 601 +++++++++-- src/zarr/core/transforms/__init__.py | 30 + src/zarr/core/transforms/chunk_resolution.py | 207 ++++ src/zarr/core/transforms/composition.py | 113 +++ src/zarr/core/transforms/domain.py | 178 ++++ src/zarr/core/transforms/output_map.py | 83 ++ src/zarr/core/transforms/transform.py | 932 ++++++++++++++++++ tests/test_array.py | 3 +- tests/test_lazy_indexing.py | 164 +++ tests/test_transforms/__init__.py | 0 .../test_transforms/test_chunk_resolution.py | 178 ++++ tests/test_transforms/test_composition.py | 166 ++++ tests/test_transforms/test_domain.py | 202 ++++ tests/test_transforms/test_output_map.py | 56 ++ tests/test_transforms/test_transform.py | 516 ++++++++++ 15 files changed, 3369 insertions(+), 60 deletions(-) create mode 100644 src/zarr/core/transforms/__init__.py create mode 100644 src/zarr/core/transforms/chunk_resolution.py create mode 100644 src/zarr/core/transforms/composition.py create mode 100644 src/zarr/core/transforms/domain.py create mode 100644 src/zarr/core/transforms/output_map.py create mode 100644 src/zarr/core/transforms/transform.py create mode 100644 tests/test_lazy_indexing.py create mode 100644 tests/test_transforms/__init__.py create mode 100644 tests/test_transforms/test_chunk_resolution.py create mode 100644 tests/test_transforms/test_composition.py create mode 100644 tests/test_transforms/test_domain.py create mode 100644 tests/test_transforms/test_output_map.py create mode 100644 tests/test_transforms/test_transform.py diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4736805b9d..1d087a7945 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -96,12 +96,18 @@ VIndex, _iter_grid, _iter_regions, + boundscheck_indices, check_fields, check_no_multi_fields, + ensure_tuple, + is_basic_selection, + is_coordinate_selection, is_pure_fancy_indexing, is_pure_orthogonal_indexing, is_scalar, pop_fields, + replace_lists, + wraparound_indices, ) from zarr.core.metadata import ( ArrayMetadata, @@ -126,6 +132,13 @@ resolve_chunks, ) from zarr.core.sync import sync +from zarr.core.transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr.core.transforms.output_map import ArrayMap +from zarr.core.transforms.transform import ( + IndexTransform, + _normalize_negative_indices, + selection_to_transform, +) from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, @@ -329,6 +342,7 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: store_path: StorePath codec_pipeline: CodecPipeline = field(init=False) _chunk_grid: ChunkGrid = field(init=False) + _transform: IndexTransform = field(init=False) config: ArrayConfig @overload @@ -365,6 +379,7 @@ def __init__( "codec_pipeline", create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), ) + object.__setattr__(self, "_transform", IndexTransform.from_shape(metadata_parsed.shape)) # this overload defines the function signature when zarr_format is 2 @overload @@ -1040,6 +1055,17 @@ async def example(): _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) return cls(store_path=store_path, metadata=_metadata_dict) + def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetadata]: + """Return a new AsyncArray sharing storage but with a different transform.""" + new = object.__new__(type(self)) + object.__setattr__(new, "metadata", self.metadata) + object.__setattr__(new, "store_path", self.store_path) + object.__setattr__(new, "config", self.config) + object.__setattr__(new, "_chunk_grid", self._chunk_grid) + object.__setattr__(new, "codec_pipeline", self.codec_pipeline) + object.__setattr__(new, "_transform", transform) + return new + @property def store(self) -> Store: return self.store_path.store @@ -1058,7 +1084,7 @@ def ndim(self) -> int: int The number of dimensions in the Array. """ - return len(self.metadata.shape) + return len(self.shape) @property def shape(self) -> tuple[int, ...]: @@ -1069,6 +1095,11 @@ def shape(self) -> tuple[int, ...]: tuple The shape of the Array. """ + return self._transform.domain.shape + + @property + def storage_shape(self) -> tuple[int, ...]: + """The shape of the underlying storage array (ignoring any view transform).""" return self.metadata.shape @property @@ -1828,6 +1859,40 @@ async def _set_selection( fields=fields, ) + async def _get_selection_t( + self, + transform: IndexTransform, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + ) -> NDArrayLikeOrScalar: + return await _get_selection_via_transform( + self.store_path, + self.metadata, + self.config, + transform, + self.codec_pipeline, + prototype=prototype, + out=out, + ) + + async def _set_selection_t( + self, + transform: IndexTransform, + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + ) -> None: + return await _set_selection_via_transform( + self.store_path, + self.metadata, + self.config, + transform, + value, + self.codec_pipeline, + prototype=prototype, + ) + async def setitem( self, selection: BasicSelection, @@ -2086,6 +2151,11 @@ def _chunk_grid(self) -> ChunkGrid: """The chunk grid for this array, bound to the array's shape.""" return self.async_array._chunk_grid + def _with_transform(self, transform: IndexTransform) -> Array[T_ArrayMetadata]: + """Return a new Array sharing storage but with a different transform.""" + new_async = self._async_array._with_transform(transform) + return type(self)(new_async) + @classmethod @deprecated("Use zarr.create_array instead.", category=ZarrDeprecationWarning) def create( @@ -3225,14 +3295,19 @@ def get_basic_selection( if prototype is None: prototype = default_buffer_prototype() - return sync( - self.async_array._get_selection( - BasicIndexer(selection, self.shape, self._chunk_grid), - out=out, - fields=fields, - prototype=prototype, + if fields is not None: + # Fall back to legacy path for structured dtype field selection + return sync( + self.async_array._get_selection( + BasicIndexer(selection, self.shape, self._chunk_grid), + out=out, + fields=fields, + prototype=prototype, + ) ) - ) + selection = _normalize_negative_indices(selection, self.shape) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_basic_selection( self, @@ -3334,8 +3409,16 @@ def set_basic_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + if fields is not None: + # Fall back to legacy path for structured dtype field selection + indexer = BasicIndexer(selection, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + selection = _normalize_negative_indices(selection, self.shape) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_orthogonal_selection( self, @@ -3462,12 +3545,17 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) - return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + if fields is not None or not is_basic_selection(selection): + # Fall back to legacy path for structured dtypes or advanced selections + indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) + return sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) - ) + selection = _normalize_negative_indices(selection, self.shape) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_orthogonal_selection( self, @@ -3581,10 +3669,16 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) - return sync( - self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) - ) + if fields is not None or not is_basic_selection(selection): + # Fall back to legacy path for structured dtypes or advanced selections + indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + selection = _normalize_negative_indices(selection, self.shape) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_mask_selection( self, @@ -3669,12 +3763,28 @@ def get_mask_selection( if prototype is None: prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, self.shape, self._chunk_grid) - return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + if fields is not None: + indexer = MaskIndexer(mask, self.shape, self._chunk_grid) + return sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) - ) + # Unwrap if VIndex passed a tuple + if isinstance(mask, tuple) and len(mask) == 1: + mask = mask[0] + # Validate mask + mask_arr = np.asarray(mask) + if mask_arr.dtype != np.bool_: + raise IndexError("invalid mask selection; expected Boolean array") + if mask_arr.shape != self.shape: + raise IndexError( + f"invalid mask selection; expected Boolean array with shape {self.shape}, " + f"got {mask_arr.shape}" + ) + selection = (mask_arr,) + transform = selection_to_transform(selection, self._async_array._transform, "vectorized") + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_mask_selection( self, @@ -3759,8 +3869,25 @@ def set_mask_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + if fields is not None: + indexer = MaskIndexer(mask, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + if isinstance(mask, tuple) and len(mask) == 1: + mask = mask[0] + mask_arr = np.asarray(mask) + if mask_arr.dtype != np.bool_: + raise IndexError("invalid mask selection; expected Boolean array") + if mask_arr.shape != self.shape: + raise IndexError( + f"invalid mask selection; expected Boolean array with shape {self.shape}, " + f"got {mask_arr.shape}" + ) + selection = (mask_arr,) + transform = selection_to_transform(selection, self._async_array._transform, "vectorized") + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_coordinate_selection( self, @@ -3847,16 +3974,45 @@ def get_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) - out_array = sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + if fields is not None: + indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) + out_array = sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) + if hasattr(out_array, "shape"): + out_array = np.array(out_array).reshape(indexer.sel_shape) + return out_array + # Validate and normalize as coordinate selection + sel_normalized = ensure_tuple(selection) + sel_normalized = tuple( + np.asarray([s], dtype=np.intp) if isinstance(s, (int, np.integer)) else s + for s in sel_normalized ) - - if hasattr(out_array, "shape"): - # restore shape - out_array = np.array(out_array).reshape(indexer.sel_shape) + sel_normalized = replace_lists(sel_normalized) + if not is_coordinate_selection(sel_normalized, self.shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + # Handle wraparound and bounds checking + for dim_sel, dim_len in zip(sel_normalized, self.shape, strict=True): + wraparound_indices(dim_sel, dim_len) + boundscheck_indices(dim_sel, dim_len) + transform = selection_to_transform( + sel_normalized, self._async_array._transform, "vectorized" + ) + out_array = sync( + self._async_array._get_selection_t(transform, out=out, prototype=prototype) + ) + # Reshape to the broadcast shape of the coordinate arrays + sel_tuple = sel_normalized + sel_arrays = [np.asarray(s) for s in sel_tuple] + sel_shape = np.broadcast_shapes(*(s.shape for s in sel_arrays)) + if hasattr(out_array, "shape") and sel_shape != (): + out_array = np.array(out_array).reshape(sel_shape) return out_array def set_coordinate_selection( @@ -3939,30 +4095,46 @@ def set_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - # setup indexer - indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) - - # handle value - need ndarray-like flatten value - if not is_scalar(value, self.dtype): - try: - from numcodecs.compat import ensure_ndarray_like - - value = ensure_ndarray_like(value) # TODO replace with agnostic - except TypeError: - # Handle types like `list` or `tuple` - value = np.array(value) # TODO replace with agnostic - if hasattr(value, "shape") and len(value.shape) > 1: - value = np.array(value).reshape(-1) - - if not is_scalar(value, self.dtype) and ( - isinstance(value, NDArrayLike) and indexer.shape != value.shape - ): - raise ValueError( - f"Attempting to set a selection of {indexer.sel_shape[0]} " - f"elements with an array of {value.shape[0]} elements." + # Normalize empty fields list to None + if not fields: + fields = None + if fields is not None: + indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) + if not is_scalar(value, self.dtype): + try: + from numcodecs.compat import ensure_ndarray_like + + value = ensure_ndarray_like(value) + except TypeError: + value = np.array(value) + if hasattr(value, "shape") and len(value.shape) > 1: + value = np.array(value).reshape(-1) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) - - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + return + sel_normalized = ensure_tuple(selection) + sel_normalized = tuple( + np.asarray([s], dtype=np.intp) if isinstance(s, (int, np.integer)) else s + for s in sel_normalized + ) + sel_normalized = replace_lists(sel_normalized) + if not is_coordinate_selection(sel_normalized, self.shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + for dim_sel, dim_len in zip(sel_normalized, self.shape, strict=True): + wraparound_indices(dim_sel, dim_len) + boundscheck_indices(dim_sel, dim_len) + transform = selection_to_transform( + sel_normalized, self._async_array._transform, "vectorized" + ) + # Flatten value for coordinate selection + if not is_scalar(value, self.dtype) and hasattr(value, "shape") and len(value.shape) > 1: + value = np.asarray(value).reshape(-1) + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_block_selection( self, @@ -4192,6 +4364,18 @@ def blocks(self) -> BlockIndex: examples.""" return BlockIndex(self) + @property + def z(self) -> _LazyIndexAccessor: + """Lazy indexing accessor. Returns a new Array with composed transform, no I/O.""" + return _LazyIndexAccessor(self) + + def resolve(self, prototype: BufferPrototype | None = None) -> NDArrayLikeOrScalar: + """Read and return the data for this array view. + + Equivalent to ``self[...]`` but more explicit for lazy views. + """ + return self[...] + def resize(self, new_shape: ShapeLike) -> None: """ Change the shape of the array by growing or shrinking one or more @@ -4295,7 +4479,11 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: return type(self)(new_array) def __repr__(self) -> str: - return f"" + t = self._async_array._transform + return ( + f"" + ) @property def info(self) -> Any: @@ -4405,6 +4593,65 @@ async def _shards_initialized( type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"] +class _LazyOIndex: + """Lazy orthogonal indexing via ``array.z.oindex[...]``.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Any) -> Array[Any]: + new_t = selection_to_transform(selection, self._array._async_array._transform, "orthogonal") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: + new_t = selection_to_transform(selection, self._array._async_array._transform, "orthogonal") + self._array._with_transform(new_t)[...] = value + + +class _LazyVIndex: + """Lazy vectorized indexing via ``array.z.vindex[...]``.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Any) -> Array[Any]: + new_t = selection_to_transform(selection, self._array._async_array._transform, "vectorized") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: + new_t = selection_to_transform(selection, self._array._async_array._transform, "vectorized") + self._array._with_transform(new_t)[...] = value + + +class _LazyIndexAccessor: + """Provides lazy indexing via ``array.z[...]``.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Selection) -> Array[Any]: + new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None: + new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") + self._array._with_transform(new_t)[...] = value + + @property + def oindex(self) -> _LazyOIndex: + return _LazyOIndex(self._array) + + @property + def vindex(self) -> _LazyVIndex: + return _LazyVIndex(self._array) + + class ShardsConfigParam(TypedDict): shape: tuple[int, ...] index_location: ShardingCodecIndexLocation | None @@ -5778,6 +6025,241 @@ def _get_chunk_spec( ) +def _is_complete_chunk( + sub_transform: IndexTransform, chunk_grid: ChunkGrid, chunk_coords: tuple[int, ...] +) -> bool: + """Check if a sub-transform covers an entire chunk.""" + from zarr.core.transforms.output_map import ConstantMap, DimensionMap + + spec = chunk_grid[chunk_coords] + if spec is None: + return False + for out_dim, m in enumerate(sub_transform.output): + if isinstance(m, ConstantMap): + # A ConstantMap means a single element is selected along this output dimension, + # so the write does not cover the full chunk along this dimension. + chunk_dim_size = spec.shape[out_dim] + if chunk_dim_size > 1: + return False + continue # chunk dim size is 1, so selecting the single element is complete + if isinstance(m, DimensionMap): + chunk_dim_size = spec.shape[out_dim] + # Compute actual storage range: storage = offset + stride * input_coord + dim_lo = sub_transform.domain.inclusive_min[m.input_dimension] + dim_hi = sub_transform.domain.exclusive_max[m.input_dimension] + if m.stride == 1: + storage_start = m.offset + dim_lo + storage_stop = m.offset + dim_hi + if storage_start != 0 or storage_stop != chunk_dim_size: + return False + else: + return False # strided access is never a complete chunk + else: + return False # ArrayMap is never a "complete chunk" + return True + + +async def _get_selection_via_transform( + store_path: StorePath, + metadata: ArrayMetadata, + config: ArrayConfig, + transform: IndexTransform, + codec_pipeline: CodecPipeline, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, +) -> NDArrayLikeOrScalar: + """Read data using an IndexTransform.""" + chunk_grid = ChunkGrid.from_metadata(metadata) + + # Get dtype (same logic as existing _get_selection) + if metadata.zarr_format == 2: + zdtype = metadata.dtype + order = metadata.order + else: + zdtype = metadata.data_type + order = config.order + dtype = zdtype.to_native_dtype() + + out_shape = transform.domain.shape + + # When the transform has ArrayMap outputs, chunk resolution produces + # flat scatter indices (out_indices). The output buffer must be 1D + # during the read, then reshaped to out_shape afterwards. + # For vectorized indexing (all outputs are ArrayMaps), chunk resolution + # produces flat scatter indices. The buffer must be 1D during the read. + # For orthogonal indexing (mixed ArrayMap + DimensionMap), the buffer + # stays multi-dimensional — each dim gets its own out_sel entry. + needs_flat_buffer = all(isinstance(m, ArrayMap) for m in transform.output) + buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape + + # Setup output buffer + if out is not None: + if not isinstance(out, NDBuffer): + raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") + if out.shape != out_shape: + raise ValueError( + f"shape of out argument doesn't match. Expected {out_shape}, got {out.shape}" + ) + out_buffer = out + else: + out_buffer = prototype.nd_buffer.empty(shape=buffer_shape, dtype=dtype, order=order) + + if product(out_shape) > 0: + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # Build batch_info using transforms + batch_info = [] + drop_axes: tuple[int, ...] = () + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms( + transform, chunk_grid + ): + chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) + drop_axes = da # same for all chunks + is_complete = _is_complete_chunk(sub_transform, chunk_grid, chunk_coords) + batch_info.append( + ( + store_path / metadata.encode_chunk_key(chunk_coords), + _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), + chunk_sel, + out_sel, + is_complete, + ) + ) + + results = await codec_pipeline.read(batch_info, out_buffer, drop_axes=drop_axes) + + # Handle read_missing_chunks + if _config.read_missing_chunks is False: + missing_info = [] + for i, result in enumerate(results): + if result["status"] == "missing": + coords_path = batch_info[i][0] + missing_info.append(f" chunk at '{coords_path}'") + if missing_info: + chunks_str = "\n".join(missing_info) + raise ChunkNotFoundError( + f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" + f"Set the 'array.read_missing_chunks' config to True to fill " + f"missing chunks with the fill value.\n" + f"Missing chunks:\n{chunks_str}" + ) + + # Return scalar for 0-d results + if out_shape == (): + return out_buffer.as_scalar() + out_result = out_buffer.as_ndarray_like() + # Reshape if we flattened for array indexing + if needs_flat_buffer and hasattr(out_result, "reshape"): + out_result = np.array(out_result).reshape(out_shape) + return out_result + + +async def _set_selection_via_transform( + store_path: StorePath, + metadata: ArrayMetadata, + config: ArrayConfig, + transform: IndexTransform, + value: npt.ArrayLike, + codec_pipeline: CodecPipeline, + *, + prototype: BufferPrototype, +) -> None: + """Write data using an IndexTransform.""" + chunk_grid = ChunkGrid.from_metadata(metadata) + + # Get dtype from metadata + if metadata.zarr_format == 2: + zdtype = metadata.dtype + else: + zdtype = metadata.data_type + dtype = zdtype.to_native_dtype() + + # check value shape + if np.isscalar(value): + array_like = prototype.buffer.create_zero_length().as_array_like() + if isinstance(array_like, np._typing._SupportsArrayFunc): + array_like_ = cast("np._typing._SupportsArrayFunc", array_like) + value = np.asanyarray(value, dtype=dtype, like=array_like_) + else: + if not hasattr(value, "shape"): + value = np.asarray(value, dtype) + if not hasattr(value, "dtype") or value.dtype.name != dtype.name: + if hasattr(value, "astype"): + value = value.astype(dtype=dtype, order="A") + else: + value = np.array(value, dtype=dtype, order="A") + value = cast("NDArrayLike", value) + + # Validate value shape against selection shape + sel_shape = transform.domain.shape + needs_flat_buffer = all(isinstance(m, ArrayMap) for m in transform.output) + if hasattr(value, "shape") and value.shape != () and value.shape != sel_shape: + if needs_flat_buffer: + # For ArrayMap (coordinate/vindex), values are flattened so check total size + sel_size = product(sel_shape) + val_size = product(value.shape) + if val_size != sel_size and val_size != 1: + raise ValueError( + f"Attempting to set a selection with a value of incompatible shape. " + f"The selection has shape {sel_shape}, but the value has shape {value.shape}." + ) + else: + # Check if value is broadcastable to sel_shape + try: + np.broadcast_shapes(value.shape, sel_shape) + except ValueError: + raise ValueError( + f"Attempting to set a selection with a value of incompatible shape. " + f"The selection has shape {sel_shape}, but the value has shape {value.shape}." + ) from None + + # When the transform has ArrayMap outputs, chunk resolution produces + # flat scatter indices (out_indices). The value buffer must be 1D + # during the write, matching the flat index layout. + if ( + needs_flat_buffer + and hasattr(value, "reshape") + and not np.isscalar(value) + and np.ndim(value) > 0 + ): + value = np.asarray(value).reshape(-1) + + # Convert to NDBuffer + value_buffer = prototype.nd_buffer.from_ndarray_like(value) + + # Determine memory order + if metadata.zarr_format == 2: + order = metadata.order + else: + order = config.order + + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # Build batch_info using transforms + batch_info = [] + drop_axes: tuple[int, ...] = () + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): + chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) + drop_axes = da # same for all chunks + is_complete = _is_complete_chunk(sub_transform, chunk_grid, chunk_coords) + batch_info.append( + ( + store_path / metadata.encode_chunk_key(chunk_coords), + _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), + chunk_sel, + out_sel, + is_complete, + ) + ) + + await codec_pipeline.write(batch_info, value_buffer, drop_axes=drop_axes) + + async def _get_selection( store_path: StorePath, metadata: ArrayMetadata, @@ -6315,9 +6797,10 @@ async def _delete_key(key: str) -> None: # Write new metadata await save_metadata(array.store_path, new_metadata) - # Update metadata and chunk_grid (in place) + # Update metadata, chunk_grid, and transform (in place) object.__setattr__(array, "metadata", new_metadata) object.__setattr__(array, "_chunk_grid", new_chunk_grid) + object.__setattr__(array, "_transform", IndexTransform.from_shape(new_shape)) async def _append( diff --git a/src/zarr/core/transforms/__init__.py b/src/zarr/core/transforms/__init__.py new file mode 100644 index 0000000000..530dd39cea --- /dev/null +++ b/src/zarr/core/transforms/__init__.py @@ -0,0 +1,30 @@ +"""Composable, lazy coordinate transforms for zarr array indexing. + +This package implements TensorStore-inspired index transforms. The core idea: +every indexing operation (slicing, fancy indexing, etc.) produces a coordinate +mapping from user space to storage space. These mappings compose lazily — no +I/O until you explicitly read or write. + +Key types: + +- ``IndexDomain`` — a rectangular region of integer coordinates +- ``IndexTransform`` — maps input coordinates to storage coordinates +- ``ConstantMap``, ``DimensionMap``, ``ArrayMap`` — the three ways a single + output dimension can depend on the input (see ``output_map.py``) +- ``compose`` — chain two transforms into one +""" + +from zarr.core.transforms.composition import compose +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.transform import IndexTransform + +__all__ = [ + "ArrayMap", + "ConstantMap", + "DimensionMap", + "IndexDomain", + "IndexTransform", + "OutputIndexMap", + "compose", +] diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py new file mode 100644 index 0000000000..db066a2525 --- /dev/null +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -0,0 +1,207 @@ +"""Chunk resolution — mapping transforms to chunk-level I/O. + +Given an ``IndexTransform`` (which coordinates a user wants to access) and a +``ChunkGrid`` (how storage is divided into chunks), chunk resolution answers: + + For each chunk, which storage coordinates does this transform touch, + and where do those values land in the output buffer? + +The algorithm is: + +1. **Enumerate candidate chunks** — determine which chunks could possibly + be touched by the transform's output coordinate ranges. + +2. **Intersect** — for each candidate chunk, call + ``transform.intersect(chunk_domain)`` to restrict the transform to + coordinates within that chunk. If the intersection is empty, skip it. + +3. **Translate** — shift the restricted transform to chunk-local coordinates + via ``transform.translate(-chunk_origin)``. + +4. **Yield** — produce ``(chunk_coords, local_transform, surviving_indices)`` + triples that the codec pipeline consumes. + +``sub_transform_to_selections`` bridges from the transform representation +back to the raw ``(chunk_selection, out_selection, drop_axes)`` tuples that +the current codec pipeline expects. This bridge will go away when the codec +pipeline accepts transforms natively. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr.core.chunk_grids import ChunkGrid + +ChunkTransformResult = tuple[ + tuple[int, ...], + IndexTransform, + np.ndarray[Any, np.dtype[np.intp]] | None, +] + + +def iter_chunk_transforms( + transform: IndexTransform, + chunk_grid: ChunkGrid, +) -> Iterator[ChunkTransformResult]: + """Resolve a composed IndexTransform against a ChunkGrid. + + Yields ``(chunk_coords, sub_transform, out_indices)`` triples: + + - ``chunk_coords``: which chunk to access. + - ``sub_transform``: maps output buffer coords to chunk-local coords. + - ``out_indices``: for vectorized/array indexing, the output scatter + indices (integer array). ``None`` for basic/slice indexing. + """ + dim_grids = chunk_grid._dimensions + + # Enumerate all possible chunks via cartesian product of per-dim chunk ranges + # For each candidate chunk, intersect the transform with the chunk domain. + # The transform.intersect method handles both orthogonal and vectorized cases. + chunk_ranges: list[range] = [] + for out_dim, m in enumerate(transform.output): + dg = dim_grids[out_dim] + if isinstance(m, ConstantMap): + # Single chunk + c = dg.index_to_chunk(m.offset) + chunk_ranges.append(range(c, c + 1)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + dim_lo = transform.domain.inclusive_min[d] + dim_hi = transform.domain.exclusive_max[d] + if dim_lo >= dim_hi: + return # empty domain + if m.stride > 0: + s_min = m.offset + m.stride * dim_lo + s_max = m.offset + m.stride * (dim_hi - 1) + else: + s_min = m.offset + m.stride * (dim_hi - 1) + s_max = m.offset + m.stride * dim_lo + first = dg.index_to_chunk(s_min) + last = dg.index_to_chunk(s_max) + chunk_ranges.append(range(first, last + 1)) + elif isinstance(m, ArrayMap): + storage = m.offset + m.stride * m.index_array + flat = storage.ravel().astype(np.intp) + chunk_ids = dg.indices_to_chunks(flat) + first = int(chunk_ids.min()) + last = int(chunk_ids.max()) + chunk_ranges.append(range(first, last + 1)) + + import itertools + + for chunk_coords_tuple in itertools.product(*chunk_ranges): + chunk_coords = tuple(int(c) for c in chunk_coords_tuple) + + # Build the chunk domain in storage space + chunk_min: list[int] = [] + chunk_max: list[int] = [] + chunk_shift: list[int] = [] + for out_dim, c in enumerate(chunk_coords): + dg = dim_grids[out_dim] + c_start = dg.chunk_offset(c) + c_size = dg.chunk_size(c) + chunk_min.append(c_start) + chunk_max.append(c_start + c_size) + chunk_shift.append(-c_start) + + chunk_domain = IndexDomain( + inclusive_min=tuple(chunk_min), + exclusive_max=tuple(chunk_max), + ) + + # Intersect transform with chunk domain + result = transform.intersect(chunk_domain) + if result is None: + continue + + restricted, surviving = result + + # Translate to chunk-local coordinates + local = restricted.translate(tuple(chunk_shift)) + + yield (chunk_coords, local, surviving) + + +def sub_transform_to_selections( + sub_transform: IndexTransform, + out_indices: np.ndarray[Any, np.dtype[np.intp]] | None = None, +) -> tuple[ + tuple[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[int, ...], +]: + """Convert a chunk-local sub-transform to raw selections for the codec pipeline. + + Parameters + ---------- + sub_transform + A chunk-local IndexTransform (output maps already translated to + chunk-local coordinates). + out_indices + For vectorized indexing: the output scatter indices for this chunk. + None for orthogonal/basic indexing. + + Returns + ------- + tuple + ``(chunk_selection, out_selection, drop_axes)`` + """ + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + drop_axes: list[int] = [] + + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + dim_lo = sub_transform.domain.inclusive_min[m.input_dimension] + dim_hi = sub_transform.domain.exclusive_max[m.input_dimension] + start = m.offset + m.stride * dim_lo + stop = m.offset + m.stride * dim_hi + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + elif isinstance(m, ArrayMap): + if m.offset == 0 and m.stride == 1: + chunk_sel.append(m.index_array) + else: + storage_coords = m.offset + m.stride * m.index_array + chunk_sel.append(storage_coords.astype(np.intp)) + + # Build out_sel: one entry per non-dropped output dim. + out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + + # Vectorized: multiple correlated ArrayMaps share one scatter index + is_vectorized = ( + out_indices is not None + and sum(1 for m in sub_transform.output if isinstance(m, ArrayMap)) >= 2 + ) + + if is_vectorized: + assert out_indices is not None + out_sel.append(out_indices) + else: + for m in sub_transform.output: + if isinstance(m, ConstantMap): + continue + if isinstance(m, DimensionMap): + lo = sub_transform.domain.inclusive_min[m.input_dimension] + hi = sub_transform.domain.exclusive_max[m.input_dimension] + out_sel.append(slice(lo, hi)) + elif isinstance(m, ArrayMap): + if out_indices is not None: + # Orthogonal ArrayMap: out_indices has the surviving positions + out_sel.append(out_indices) + else: + out_sel.append(slice(0, len(m.index_array))) + + return tuple(chunk_sel), tuple(out_sel), tuple(drop_axes) diff --git a/src/zarr/core/transforms/composition.py b/src/zarr/core/transforms/composition.py new file mode 100644 index 0000000000..9d07bd3324 --- /dev/null +++ b/src/zarr/core/transforms/composition.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.transform import IndexTransform + + +def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: + """Compose two IndexTransforms. + + ``outer`` maps user coords (rank m) to intermediate coords (rank n). + ``inner`` maps intermediate coords (rank n) to storage coords (rank p). + The result maps user coords (rank m) to storage coords (rank p). + + Precondition: ``outer.output_rank == inner.domain.ndim``. + """ + if outer.output_rank != inner.domain.ndim: + raise ValueError( + f"outer output rank ({outer.output_rank}) must match inner input rank " + f"({inner.domain.ndim})" + ) + + result_output = [_compose_single(outer, inner_map) for inner_map in inner.output] + + return IndexTransform(domain=outer.domain, output=tuple(result_output)) + + +def _compose_single(outer: IndexTransform, inner_map: OutputIndexMap) -> OutputIndexMap: + """Compose a single inner output map with the full outer transform.""" + if isinstance(inner_map, ConstantMap): + return ConstantMap(offset=inner_map.offset) + + if isinstance(inner_map, DimensionMap): + return _compose_dimension(outer, inner_map) + + if isinstance(inner_map, ArrayMap): + return _compose_array(outer, inner_map) + + raise TypeError(f"Unknown output map type: {type(inner_map)}") # pragma: no cover + + +def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> OutputIndexMap: + """Compose when inner is a DimensionMap. + + storage = offset_i + stride_i * intermediate[dim_i] + where intermediate[dim_i] = outer.output[dim_i](user_input) + """ + dim_i = inner_map.input_dimension + offset_i = inner_map.offset + stride_i = inner_map.stride + outer_map = outer.output[dim_i] + + if isinstance(outer_map, ConstantMap): + return ConstantMap(offset=offset_i + stride_i * outer_map.offset) + + if isinstance(outer_map, DimensionMap): + return DimensionMap( + input_dimension=outer_map.input_dimension, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + ) + + if isinstance(outer_map, ArrayMap): + return ArrayMap( + index_array=outer_map.index_array, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + ) + + raise TypeError(f"Unknown output map type: {type(outer_map)}") # pragma: no cover + + +def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap: + """Compose when inner is an ArrayMap. + + storage = offset_i + stride_i * arr_i[intermediate] + We need to evaluate arr_i at the intermediate coordinates produced by outer. + """ + arr_i = inner_map.index_array + offset_i = inner_map.offset + stride_i = inner_map.stride + + # Check if all outer outputs are constant + all_constant = all(isinstance(m, ConstantMap) for m in outer.output) + + if all_constant: + # Evaluate arr_i at the single constant point + idx = tuple(m.offset for m in outer.output if isinstance(m, ConstantMap)) + value = int(arr_i[idx]) + return ConstantMap(offset=offset_i + stride_i * value) + + # For 1D inner array with a single outer output (simple case) + if arr_i.ndim == 1 and len(outer.output) == 1: + outer_map = outer.output[0] + + if isinstance(outer_map, DimensionMap): + dim_size = outer.domain.shape[outer_map.input_dimension] + user_indices = np.arange(dim_size, dtype=np.intp) + intermediate_vals = outer_map.offset + outer_map.stride * user_indices + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + if isinstance(outer_map, ArrayMap): + intermediate_vals = outer_map.offset + outer_map.stride * outer_map.index_array + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + # General multi-dim case: not yet implemented + raise NotImplementedError( + "Composing a multi-dimensional inner array map with non-constant outer maps " + "is not yet supported." + ) diff --git a/src/zarr/core/transforms/domain.py b/src/zarr/core/transforms/domain.py new file mode 100644 index 0000000000..90bcc08ace --- /dev/null +++ b/src/zarr/core/transforms/domain.py @@ -0,0 +1,178 @@ +"""Index domains — rectangular regions in N-dimensional integer space. + +An ``IndexDomain`` represents the set of valid coordinates for an array or +array view. It is the cartesian product of per-dimension integer ranges:: + + IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + # represents {(i, j) : 2 <= i < 10, 5 <= j < 20} + +Unlike NumPy, domains can have **non-zero origins**. After slicing +``arr[5:10]``, the result has origin 5 and shape 5 — coordinates 5 through +9 are valid. This follows the TensorStore convention. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True, slots=True) +class IndexDomain: + """A rectangular region in N-dimensional index space. + + The valid coordinates are the integers in + ``[inclusive_min[d], exclusive_max[d])`` for each dimension ``d``. + """ + + inclusive_min: tuple[int, ...] + exclusive_max: tuple[int, ...] + labels: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + if len(self.inclusive_min) != len(self.exclusive_max): + raise ValueError( + f"inclusive_min and exclusive_max must have the same length. " + f"Got {len(self.inclusive_min)} and {len(self.exclusive_max)}." + ) + for i, (lo, hi) in enumerate(zip(self.inclusive_min, self.exclusive_max, strict=True)): + if lo > hi: + raise ValueError( + f"inclusive_min must be <= exclusive_max for all dimensions. " + f"Dimension {i}: {lo} > {hi}" + ) + if self.labels is not None and len(self.labels) != len(self.inclusive_min): + raise ValueError( + f"labels must have the same length as dimensions. " + f"Got {len(self.labels)} labels for {len(self.inclusive_min)} dimensions." + ) + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> IndexDomain: + """Create a domain with origin at zero.""" + return cls( + inclusive_min=(0,) * len(shape), + exclusive_max=shape, + ) + + @property + def ndim(self) -> int: + return len(self.inclusive_min) + + @property + def origin(self) -> tuple[int, ...]: + return self.inclusive_min + + @property + def shape(self) -> tuple[int, ...]: + return tuple(hi - lo for lo, hi in zip(self.inclusive_min, self.exclusive_max, strict=True)) + + def contains(self, index: tuple[int, ...]) -> bool: + if len(index) != self.ndim: + return False + return all( + lo <= idx < hi + for lo, hi, idx in zip(self.inclusive_min, self.exclusive_max, index, strict=True) + ) + + def contains_domain(self, other: IndexDomain) -> bool: + if other.ndim != self.ndim: + return False + return all( + self_lo <= other_lo and other_hi <= self_hi + for self_lo, self_hi, other_lo, other_hi in zip( + self.inclusive_min, + self.exclusive_max, + other.inclusive_min, + other.exclusive_max, + strict=True, + ) + ) + + def intersect(self, other: IndexDomain) -> IndexDomain | None: + if other.ndim != self.ndim: + raise ValueError( + f"Cannot intersect domains with different ranks: {self.ndim} vs {other.ndim}" + ) + new_min = tuple( + max(a, b) for a, b in zip(self.inclusive_min, other.inclusive_min, strict=True) + ) + new_max = tuple( + min(a, b) for a, b in zip(self.exclusive_max, other.exclusive_max, strict=True) + ) + if any(lo >= hi for lo, hi in zip(new_min, new_max, strict=True)): + return None + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def translate(self, offset: tuple[int, ...]) -> IndexDomain: + if len(offset) != self.ndim: + raise ValueError( + f"Offset must have same length as domain dimensions. " + f"Domain has {self.ndim} dimensions, offset has {len(offset)}." + ) + new_min = tuple(lo + off for lo, off in zip(self.inclusive_min, offset, strict=True)) + new_max = tuple(hi + off for hi, off in zip(self.exclusive_max, offset, strict=True)) + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def narrow(self, selection: Any) -> IndexDomain: + """Apply a basic selection and return a narrowed domain. + Indices are absolute coordinates. Integer indices produce length-1 extent. + Strided slices are not supported — use IndexTransform for strides. + """ + normalized = _normalize_selection(selection, self.ndim) + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + for dim_idx, (sel, dim_lo, dim_hi) in enumerate( + zip(normalized, self.inclusive_min, self.exclusive_max, strict=True) + ): + if isinstance(sel, int): + if sel < dim_lo or sel >= dim_hi: + raise IndexError( + f"index {sel} is out of bounds for dimension {dim_idx} " + f"with domain [{dim_lo}, {dim_hi})" + ) + new_inclusive_min.append(sel) + new_exclusive_max.append(sel + 1) + else: + start, stop, step = sel.start, sel.stop, sel.step + if step is not None and step != 1: + raise IndexError( + "IndexDomain.narrow only supports step=1 slices. " + f"Got step={step}. Use IndexTransform for strided access." + ) + abs_start = dim_lo if start is None else start + abs_stop = dim_hi if stop is None else stop + abs_start = max(abs_start, dim_lo) + abs_stop = min(abs_stop, dim_hi) + abs_stop = max(abs_stop, abs_start) + new_inclusive_min.append(abs_start) + new_exclusive_max.append(abs_stop) + return IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + +def _normalize_selection(selection: Any, ndim: int) -> tuple[int | slice, ...]: + """Normalize a basic selection to a tuple of ints/slices with length ndim.""" + if not isinstance(selection, tuple): + selection = (selection,) + result: list[int | slice] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - (len(selection) - 1) + result.extend([slice(None)] * num_missing) + else: + result.append(sel) + while len(result) < ndim: + result.append(slice(None)) + if len(result) > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, " + f"but {len(result)} were indexed" + ) + return tuple(result) diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py new file mode 100644 index 0000000000..5e17a0ae82 --- /dev/null +++ b/src/zarr/core/transforms/output_map.py @@ -0,0 +1,83 @@ +"""Output index maps — three representations of a set of integer coordinates. + +An output index map describes, for one dimension of storage, which coordinates +an array access will touch. Conceptually it is a **set of integers**. Three +representations cover the cases that arise in practice: + +- ``ConstantMap(offset=5)`` — a singleton set: ``{5}`` +- ``DimensionMap(input_dimension=0, offset=3, stride=2)`` over input ``[0, 5)`` + — an arithmetic progression: ``{3, 5, 7, 9, 11}`` +- ``ArrayMap(index_array=[1, 5, 9])`` — an explicit enumeration: ``{1, 5, 9}`` + +Every output map supports two set-theoretic operations (defined on +``IndexTransform``, which provides the input domain context these maps lack): + +- **intersect** — restrict to coordinates within a range (e.g., a chunk). + ``{3, 5, 7, 9, 11} ∩ [4, 8) = {5, 7}`` +- **translate** — shift every coordinate by a constant (e.g., make chunk-local). + ``{5, 7} - 4 = {1, 3}`` + +These two operations are the foundation of chunk resolution: for each chunk, +intersect the map with the chunk's range, then translate to chunk-local +coordinates. + +The three types exist because they trade off generality for efficiency: + +- ``ConstantMap``: O(1) storage, O(1) intersection +- ``DimensionMap``: O(1) storage, O(1) intersection (analytical) +- ``ArrayMap``: O(n) storage, O(n) intersection (must scan the array) + +Collapsing everything to ``ArrayMap`` would be correct but wasteful — a +billion-element slice would materialize a billion coordinates just to group +them by chunk, when ``DimensionMap`` does it with three integers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + +@dataclass(frozen=True, slots=True) +class ConstantMap: + """A singleton set: one storage coordinate. + + Represents ``{offset}``. Arises from integer indexing (e.g., ``arr[5]`` + fixes one dimension to coordinate 5). + """ + + offset: int = 0 + + +@dataclass(frozen=True, slots=True) +class DimensionMap: + """An arithmetic progression of storage coordinates. + + Represents ``{offset + stride * i : i in input_range}``, where the input + range comes from the enclosing ``IndexTransform``'s domain. Arises from + slice indexing (e.g., ``arr[2:10:3]`` gives offset=2, stride=3). + """ + + input_dimension: int + offset: int = 0 + stride: int = 1 + + +@dataclass(frozen=True, slots=True) +class ArrayMap: + """An explicit enumeration of storage coordinates. + + Represents ``{offset + stride * index_array[i] : i in input_range}``. + Arises from fancy indexing (e.g., ``arr[[1, 5, 9]]`` or boolean masks). + """ + + index_array: npt.NDArray[np.intp] + offset: int = 0 + stride: int = 1 + + +OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py new file mode 100644 index 0000000000..ee2f3ce4c2 --- /dev/null +++ b/src/zarr/core/transforms/transform.py @@ -0,0 +1,932 @@ +"""Index transforms — composable, lazy coordinate mappings. + +An ``IndexTransform`` pairs an **input domain** (the coordinates a user sees) +with a tuple of **output maps** (the storage coordinates those inputs map to). +One output map per storage dimension. See ``output_map.py`` for the three +output map types. + +Key operations: + +- **Indexing** (``transform[2:8]``, ``.oindex[idx]``, ``.vindex[idx]``) — + produces a new transform with a narrower input domain and adjusted output + maps. No I/O occurs. This is how lazy slicing works. + +- **intersect(output_domain)** — restrict to storage coordinates within a + region. This is chunk resolution: "which of my coordinates fall in this + chunk?" + +- **translate(shift)** — shift all output coordinates. This makes coordinates + chunk-local: "express my coordinates relative to the chunk origin." + +- **compose(outer, inner)** — chain two transforms. See ``composition.py``. + +The transform is the atomic unit that connects user-facing indexing to +chunk-level I/O. Every ``Array`` holds a transform (identity by default). +``Array.z[...]`` composes a new transform lazily. Reading resolves the +transform against the chunk grid via intersect + translate. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap + + +@dataclass(frozen=True, slots=True) +class IndexTransform: + """A composable mapping from input coordinates to storage coordinates. + + An ``IndexTransform`` has: + + - ``domain``: an ``IndexDomain`` describing the valid input coordinates + (the user-facing shape, possibly with non-zero origin). + - ``output``: a tuple of output maps (one per storage dimension), each + describing which storage coordinates the inputs touch. + + For a freshly opened array, the transform is the identity: input + coordinate ``i`` maps to storage coordinate ``i``. Indexing operations + compose new transforms without I/O. + """ + + domain: IndexDomain + output: tuple[OutputIndexMap, ...] + + def __post_init__(self) -> None: + for i, m in enumerate(self.output): + if isinstance(m, DimensionMap): + if m.input_dimension < 0 or m.input_dimension >= self.domain.ndim: + raise ValueError( + f"output[{i}].input_dimension = {m.input_dimension} " + f"is out of range for input rank {self.domain.ndim}" + ) + elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim: + raise ValueError( + f"output[{i}].index_array has {m.index_array.ndim} dims " + f"but input domain has {self.domain.ndim} dims" + ) + + @property + def input_rank(self) -> int: + return self.domain.ndim + + @property + def output_rank(self) -> int: + return len(self.output) + + @classmethod + def identity(cls, domain: IndexDomain) -> IndexTransform: + output = tuple(DimensionMap(input_dimension=i) for i in range(domain.ndim)) + return cls(domain=domain, output=output) + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> IndexTransform: + return cls.identity(IndexDomain.from_shape(shape)) + + @property + def selection_repr(self) -> str: + """Compact domain string, e.g. ``'{ [2, 8), [0, 10) }'``. + + Follows TensorStore's IndexDomain notation: each dimension shown + as ``[inclusive_min, exclusive_max)`` with stride annotation if not 1. + Constant (integer-indexed) dimensions show as a single value. + Array-indexed dimensions show the set of selected coordinates. + """ + parts: list[str] = [] + for m in self.output: + if isinstance(m, ConstantMap): + parts.append(str(m.offset)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + lo = self.domain.inclusive_min[d] + hi = self.domain.exclusive_max[d] + start = m.offset + m.stride * lo + stop = m.offset + m.stride * hi + if m.stride == 1: + parts.append(f"[{start}, {stop})") + else: + parts.append(f"[{start}, {stop}) step {m.stride}") + elif isinstance(m, ArrayMap): + storage = m.offset + m.stride * m.index_array + n = len(storage) + if n <= 5: + vals = ", ".join(str(int(v)) for v in storage.ravel()) + parts.append("{" + vals + "}") + else: + parts.append("{" + f"array({n})" + "}") + return "{ " + ", ".join(parts) + " }" + + def __repr__(self) -> str: + maps: list[str] = [] + for i, m in enumerate(self.output): + if isinstance(m, ConstantMap): + maps.append(f"out[{i}] = {m.offset}") + elif isinstance(m, DimensionMap): + maps.append(f"out[{i}] = {m.offset} + {m.stride} * in[{m.input_dimension}]") + elif isinstance(m, ArrayMap): + maps.append(f"out[{i}] = {m.offset} + {m.stride} * arr{m.index_array.shape}[in]") + maps_str = ", ".join(maps) + return f"IndexTransform(domain={self.domain}, {maps_str})" + + def intersect( + self, output_domain: IndexDomain + ) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: + """Restrict this transform to storage coordinates within output_domain. + + Returns ``(restricted_transform, surviving_indices)`` or None if empty. + + ``surviving_indices`` is an integer array of which input positions + survived the intersection (for ArrayMap dimensions), or None if all + positions survived (ConstantMap/DimensionMap only). + """ + return _intersect(self, output_domain) + + def translate(self, shift: tuple[int, ...]) -> IndexTransform: + """Shift all output coordinates by ``shift``.""" + if len(shift) != self.output_rank: + raise ValueError(f"shift must have length {self.output_rank}, got {len(shift)}") + new_output: list[OutputIndexMap] = [] + for m, s in zip(self.output, shift, strict=True): + if isinstance(m, ConstantMap): + new_output.append(ConstantMap(offset=m.offset + s)) + elif isinstance(m, DimensionMap): + new_output.append( + DimensionMap( + input_dimension=m.input_dimension, + offset=m.offset + s, + stride=m.stride, + ) + ) + elif isinstance(m, ArrayMap): + new_output.append( + ArrayMap( + index_array=m.index_array, + offset=m.offset + s, + stride=m.stride, + ) + ) + return IndexTransform(domain=self.domain, output=tuple(new_output)) + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_basic_indexing(self, selection) + + @property + def oindex(self) -> _OIndexHelper: + return _OIndexHelper(self) + + @property + def vindex(self) -> _VIndexHelper: + return _VIndexHelper(self) + + +def _intersect( + transform: IndexTransform, output_domain: IndexDomain +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: + """Intersect a transform with an output domain (e.g., a chunk's bounds). + + For each output dimension, restrict to storage coordinates within + [output_domain.inclusive_min[d], output_domain.exclusive_max[d]). + + For orthogonal transforms (ConstantMap, DimensionMap, independent ArrayMaps), + each dimension is intersected independently and the input domain is narrowed. + + For vectorized transforms (correlated ArrayMaps), all array dimensions + must be checked simultaneously — a point survives only if ALL its + coordinates fall within the output domain. + + Returns None if the intersection is empty. + """ + if output_domain.ndim != transform.output_rank: + raise ValueError( + f"output_domain rank ({output_domain.ndim}) != " + f"transform output rank ({transform.output_rank})" + ) + + # Check if we have correlated ArrayMaps (vectorized) + array_dims = [i for i, m in enumerate(transform.output) if isinstance(m, ArrayMap)] + if len(array_dims) >= 2: + return _intersect_vectorized(transform, output_domain, array_dims) + + # Orthogonal: intersect each output dimension independently + new_min = list(transform.domain.inclusive_min) + new_max = list(transform.domain.exclusive_max) + new_output: list[OutputIndexMap] = [] + surviving_indices: np.ndarray[Any, np.dtype[np.intp]] | None = None + + for out_dim, m in enumerate(transform.output): + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + + if isinstance(m, ConstantMap): + if lo <= m.offset < hi: + new_output.append(m) + else: + return None + + elif isinstance(m, DimensionMap): + d = m.input_dimension + input_lo = new_min[d] + input_hi = new_max[d] + if input_lo >= input_hi: + return None + + # Find input range that produces storage coords in [lo, hi) + if m.stride > 0: + new_input_lo = max(input_lo, math.ceil((lo - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((hi - m.offset) / m.stride)) + elif m.stride < 0: + new_input_lo = max(input_lo, math.ceil((hi - 1 - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((lo - 1 - m.offset) / m.stride)) + else: + if lo <= m.offset < hi: + new_input_lo, new_input_hi = input_lo, input_hi + else: + return None + + if new_input_lo >= new_input_hi: + return None + + new_min[d] = new_input_lo + new_max[d] = new_input_hi + new_output.append(m) + + elif isinstance(m, ArrayMap): + storage = m.offset + m.stride * m.index_array + mask = (storage >= lo) & (storage < hi) + if not np.any(mask): + return None + surviving_indices = np.nonzero(mask.ravel())[0].astype(np.intp) + filtered = m.index_array.ravel()[surviving_indices] + new_output.append( + ArrayMap( + index_array=filtered, + offset=m.offset, + stride=m.stride, + ) + ) + + new_domain = IndexDomain( + inclusive_min=tuple(new_min), + exclusive_max=tuple(new_max), + ) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + return (result, surviving_indices) + + +def _intersect_vectorized( + transform: IndexTransform, + output_domain: IndexDomain, + array_dims: list[int], +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: + """Intersect a vectorized transform with an output domain. + + All ArrayMap outputs are correlated — a point survives only if ALL its + storage coordinates fall within the output domain. + """ + # Compute storage coords per array dim and check bounds simultaneously + n_points: int | None = None + masks: list[np.ndarray[Any, np.dtype[np.bool_]]] = [] + + for out_dim in array_dims: + m = transform.output[out_dim] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + masks.append((storage >= lo) & (storage < hi)) + if n_points is None: + n_points = storage.size + + # A point survives only if it's in-bounds on ALL array dims + combined_mask = masks[0] + for mask in masks[1:]: + combined_mask = combined_mask & mask + + if not np.any(combined_mask): + return None + + surviving = np.nonzero(combined_mask.ravel())[0].astype(np.intp) + + # Build new output maps + new_output: list[OutputIndexMap] = [] + for out_dim, m in enumerate(transform.output): + if isinstance(m, ArrayMap): + filtered = m.index_array.ravel()[surviving] + new_output.append( + ArrayMap( + index_array=filtered, + offset=m.offset, + stride=m.stride, + ) + ) + elif isinstance(m, ConstantMap): + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + if lo <= m.offset < hi: + new_output.append(m) + else: + return None + elif isinstance(m, DimensionMap): + new_output.append(m) + + new_domain = IndexDomain.from_shape((len(surviving),)) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + return (result, surviving) + + +def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | None, ...]: + """Normalize a selection to a tuple of int, slice, or None (newaxis), + expanding ellipsis and padding with slice(None) as needed. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Count non-newaxis, non-ellipsis entries to determine how many real dims are addressed + n_newaxis = sum(1 for s in selection if s is None) + has_ellipsis = any(s is Ellipsis for s in selection) + n_real = len(selection) - n_newaxis - (1 if has_ellipsis else 0) + + if n_real > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {n_real} were indexed" + ) + + result: list[int | slice | None] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, (int, np.integer)): + result.append(int(sel)) + elif isinstance(sel, slice) or sel is None: + result.append(sel) + else: + raise IndexError(f"unsupported selection type for basic indexing: {type(sel)!r}") + + # Pad remaining dimensions with slice(None) + while sum(1 for s in result if s is not None) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _reindex_array( + arr: np.ndarray[Any, np.dtype[np.intp]], + normalized: tuple[int | slice | None, ...], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply basic indexing operations to an ArrayMap's index_array. + + The array's axes correspond to the transform's input dimensions (0-indexed + over the domain shape). When input dimensions are dropped (int), sliced, + or inserted (newaxis), the array must be updated accordingly. + """ + # Build a numpy indexing tuple: one entry per old input dimension + idx: list[Any] = [] + old_dim = 0 + newaxis_positions: list[int] = [] + result_axis = 0 + + for sel in normalized: + if sel is None: + newaxis_positions.append(result_axis) + result_axis += 1 + elif isinstance(sel, int): + if old_dim < arr.ndim: + # Convert absolute domain coordinate to 0-based array index + array_idx = sel - domain.inclusive_min[old_dim] + idx.append(array_idx) + old_dim += 1 + elif isinstance(sel, slice): + if old_dim < arr.ndim: + dim_size = domain.shape[old_dim] + # sel.indices gives 0-based start/stop/step for the array axis + start, stop, step = sel.indices(dim_size) + idx.append(slice(start, stop, step)) + old_dim += 1 + result_axis += 1 + + result = arr[tuple(idx)] if idx else arr + + for pos in newaxis_positions: + result = np.expand_dims(result, axis=pos) + + return np.asarray(result, dtype=np.intp) + + +def _reindex_array_oindex( + arr: np.ndarray[Any, np.dtype[np.intp]], + normalized: tuple[Any, ...] | list[Any], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply oindex/vindex selection to an existing ArrayMap's index_array. + + Each old input dimension gets either an array (fancy index that axis) + or a slice applied to the corresponding array axis. + """ + idx: list[Any] = [] + for old_dim, sel in enumerate(normalized): + if old_dim >= arr.ndim: + break + if isinstance(sel, np.ndarray): + idx.append(sel) + elif isinstance(sel, slice): + dim_size = domain.shape[old_dim] + start, stop, step = sel.indices(dim_size) + idx.append(slice(start, stop, step)) + else: + idx.append(slice(None)) + + result = arr[tuple(idx)] if idx else arr + return np.asarray(result, dtype=np.intp) + + +def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply basic indexing (int, slice, ellipsis, newaxis) to an IndexTransform.""" + normalized = _normalize_basic_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + old_dim = 0 + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + dropped_dims: set[int] = set() + + # Per old-dim: the slice parameters (for computing new output maps) + dim_slice_params: dict[int, tuple[int, int, int]] = {} # old_dim -> (start, stop, step) + dim_int_val: dict[int, int] = {} # old_dim -> integer index value + + for sel in normalized: + if sel is None: + # newaxis: add a size-1 dimension + new_inclusive_min.append(0) + new_exclusive_max.append(1) + new_dim_idx += 1 + elif isinstance(sel, int): + # Integer index: drop this input dimension. + # Negative indices are literal coordinates (TensorStore convention), + # NOT "from the end" like NumPy. The Array layer handles conversion. + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + idx = sel + if idx < lo or idx >= hi: + raise IndexError( + f"index {sel} is out of bounds for dimension {old_dim} with domain [{lo}, {hi})" + ) + dropped_dims.add(old_dim) + dim_int_val[old_dim] = idx + old_dim += 1 + elif isinstance(sel, slice): + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + dim_size = hi - lo + + # Resolve slice relative to the current domain (origin-based) + start, stop, step = sel.indices(dim_size) + # start, stop, step are now relative to a 0-based range of size dim_size + + if step <= 0: + raise IndexError("slice step must be positive") + + new_size = max(0, math.ceil((stop - start) / step)) + new_inclusive_min.append(0) + new_exclusive_max.append(new_size) + + # Absolute start in the original domain coordinates + abs_start = lo + start + dim_slice_params[old_dim] = (abs_start, stop, step) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + old_dim += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Now update output maps + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dropped_dims: + # Integer index: this output becomes constant + new_offset = m.offset + m.stride * dim_int_val[d] + new_output.append(ConstantMap(offset=new_offset)) + elif d in old_to_new_dim: + # Slice: update offset and stride + abs_start, _, step = dim_slice_params[d] + new_offset = m.offset + m.stride * abs_start + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + elif isinstance(m, ArrayMap): + new_arr = _reindex_array(m.index_array, normalized, transform.domain) + new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +class _OIndexHelper: + """Helper that provides orthogonal (outer) indexing via ``transform.oindex[...]``.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_oindex(self._transform, selection) + + +def _normalize_oindex_selection( + selection: Any, ndim: int +) -> tuple[np.ndarray[Any, np.dtype[np.intp]] | slice, ...]: + """Normalize an oindex selection: arrays, slices, booleans, integers.""" + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis + has_ellipsis = any(s is Ellipsis for s in selection) + n_ellipsis = 1 if has_ellipsis else 0 + n_real = len(selection) - n_ellipsis + + result: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + # Boolean array -> integer indices + (indices,) = np.nonzero(sel) + result.append(indices.astype(np.intp)) + elif isinstance(sel, np.ndarray): + result.append(sel.astype(np.intp)) + elif isinstance(sel, slice): + result.append(sel) + elif isinstance(sel, (int, np.integer)): + # Convert integer scalars to 1-element arrays for orthogonal indexing + result.append(np.array([int(sel)], dtype=np.intp)) + elif isinstance(sel, (list, tuple)): + result.append(np.asarray(sel, dtype=np.intp)) + else: + result.append(sel) + + # Pad with slice(None) + while len(result) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply orthogonal indexing to an IndexTransform. + + Each index array is applied independently per dimension (outer product). + """ + normalized = _normalize_oindex_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + + # Info per old dim + dim_array: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + dim_slice_params: dict[int, tuple[int, int, int]] = {} + + for old_dim, sel in enumerate(normalized): + if isinstance(sel, np.ndarray): + dim_array[old_dim] = sel + new_inclusive_min.append(0) + new_exclusive_max.append(len(sel)) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + elif isinstance(sel, slice): + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + dim_size = hi - lo + start, stop, step = sel.indices(dim_size) + if step <= 0: + raise IndexError("slice step must be positive") + new_size = max(0, math.ceil((stop - start) / step)) + new_inclusive_min.append(0) + new_exclusive_max.append(new_size) + abs_start = lo + start + dim_slice_params[old_dim] = (abs_start, stop, step) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dim_array: + new_output.append( + ArrayMap( + index_array=dim_array[d], + offset=m.offset, + stride=m.stride, + ) + ) + elif d in dim_slice_params: + abs_start, _, step = dim_slice_params[d] + new_offset = m.offset + m.stride * abs_start + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + elif isinstance(m, ArrayMap): + new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) + new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +class _VIndexHelper: + """Helper that provides vectorized (fancy) indexing via ``transform.vindex[...]``.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_vindex(self._transform, selection) + + +def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply vectorized indexing to an IndexTransform. + + All array indices are broadcast together. Broadcast dimensions are prepended, + followed by non-array (slice) dimensions. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis and count consumed dimensions + # Boolean arrays with ndim > 1 consume ndim dims + n_consumed = 0 + for s in selection: + if s is Ellipsis: + continue + if isinstance(s, np.ndarray) and s.dtype == np.bool_ and s.ndim > 1: + n_consumed += s.ndim + else: + n_consumed += 1 + ndim = transform.domain.ndim + + expanded: list[Any] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_consumed + expanded.extend([slice(None)] * num_missing) + else: + expanded.append(sel) + # Count dimensions already consumed by expanded entries + n_expanded_dims = 0 + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_ and sel.ndim > 1: + n_expanded_dims += sel.ndim + else: + n_expanded_dims += 1 + while n_expanded_dims < ndim: + expanded.append(slice(None)) + n_expanded_dims += 1 + + # Convert booleans, lists, ints to integer arrays + processed: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + indices_tuple = np.nonzero(sel) + processed.extend(indices.astype(np.intp) for indices in indices_tuple) + elif isinstance(sel, np.ndarray): + processed.append(sel.astype(np.intp)) + elif isinstance(sel, (list, tuple)): + processed.append(np.asarray(sel, dtype=np.intp)) + elif isinstance(sel, (int, np.integer)): + processed.append(np.array([int(sel)], dtype=np.intp)) + else: + processed.append(sel) + + # Separate array dims and slice dims + array_dims: list[int] = [] + slice_dims: list[int] = [] + arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + + for i, sel in enumerate(processed): + if isinstance(sel, np.ndarray): + array_dims.append(i) + arrays.append(sel) + else: + slice_dims.append(i) + + # Broadcast all arrays together + broadcast_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] + if arrays: + broadcast_arrays = list(np.broadcast_arrays(*arrays)) + broadcast_shape = broadcast_arrays[0].shape + else: + broadcast_arrays = [] + broadcast_shape = () + + # Build new domain: broadcast dims first, then slice dims + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + + # Broadcast dimensions + for s in broadcast_shape: + new_inclusive_min.append(0) + new_exclusive_max.append(s) + + # Slice dimensions + slice_dim_params: dict[int, tuple[int, int, int]] = {} + for old_dim in slice_dims: + sel = processed[old_dim] + assert isinstance(sel, slice) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + dim_size = hi - lo + start, stop, step = sel.indices(dim_size) + if step <= 0: + raise IndexError("slice step must be positive") + new_size = max(0, math.ceil((stop - start) / step)) + new_inclusive_min.append(0) + new_exclusive_max.append(new_size) + abs_start = lo + start + slice_dim_params[old_dim] = (abs_start, stop, step) + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Build output maps + array_dim_to_broadcast: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for i, d in enumerate(array_dims): + array_dim_to_broadcast[d] = broadcast_arrays[i] + + # New dim index for slice dims starts after broadcast dims + n_broadcast_dims = len(broadcast_shape) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in array_dim_to_broadcast: + new_output.append( + ArrayMap( + index_array=array_dim_to_broadcast[d], + offset=m.offset, + stride=m.stride, + ) + ) + else: + # Slice dim + abs_start, _, step = slice_dim_params[d] + new_offset = m.offset + m.stride * abs_start + new_stride = m.stride * step + new_input_dim = n_broadcast_dims + slice_dims.index(d) + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + elif isinstance(m, ArrayMap): + new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) + new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: + """Convert negative indices to positive ones using the array shape. + + Only normalizes integer and array-like index components; leaves + slices, Ellipsis, None, etc. untouched. + """ + if not isinstance(selection, tuple): + selection_tuple: tuple[Any, ...] = (selection,) + else: + selection_tuple = selection + + # Count real dimensions (non-None, non-Ellipsis) to map each entry to a shape dim + has_ellipsis = any(s is Ellipsis for s in selection_tuple) + n_non_newaxis = sum(1 for s in selection_tuple if s is not None and s is not Ellipsis) + n_ellipsis_dims = len(shape) - n_non_newaxis + (1 if has_ellipsis else 0) + + result: list[Any] = [] + dim = 0 + + for sel in selection_tuple: + if sel is Ellipsis: + result.append(sel) + dim += max(0, n_ellipsis_dims) + elif sel is None: + result.append(sel) + elif isinstance(sel, (int, np.integer)) and not isinstance(sel, bool): + idx = int(sel) + if idx < 0 and dim < len(shape): + idx = idx + shape[dim] + result.append(idx) + dim += 1 + elif isinstance(sel, np.ndarray) and sel.dtype != np.bool_: + arr = sel.copy() + if dim < len(shape): + arr = np.where(arr < 0, arr + shape[dim], arr) + result.append(arr) + dim += 1 + elif isinstance(sel, list): + # Convert lists to arrays with negative index normalization + arr = np.asarray(sel, dtype=np.intp) + if dim < len(shape): + arr = np.where(arr < 0, arr + shape[dim], arr) + result.append(arr) + dim += 1 + else: + # slice, bool array, or anything else: pass through + result.append(sel) + if sel is not None and sel is not Ellipsis: + dim += 1 + + if not isinstance(selection, tuple) and len(result) == 1: + return result[0] + return tuple(result) + + +def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: + """Validate array-based selections (orthogonal, vectorized). + + Rejects types that are not valid for coordinate/vectorized indexing. + Does not check bounds — the transform operations handle that. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for sel in items: + if sel is Ellipsis or isinstance(sel, (int, np.integer, slice)): + continue + if isinstance(sel, (list, np.ndarray)): + continue + raise IndexError(f"unsupported selection type for {mode} indexing: {type(sel)!r}") + + +def _validate_basic_selection(selection: Any) -> None: + """Validate that a selection only contains basic indexing types (int, slice, Ellipsis). + + Rejects None (newaxis), arrays, lists, floats, strings, etc. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for s in items: + if s is Ellipsis or isinstance(s, (int, np.integer, slice)): + continue + raise IndexError(f"unsupported selection type for basic indexing: {type(s)!r}") + + +def selection_to_transform( + selection: Any, + transform: IndexTransform, + mode: Literal["basic", "orthogonal", "vectorized"], +) -> IndexTransform: + """Convert a user selection into a composed IndexTransform. + + Negative indices are treated as literal coordinates (TensorStore convention). + The caller (Array layer) is responsible for converting numpy-style negative + indices before calling this function. + """ + if mode == "basic": + _validate_basic_selection(selection) + return transform[selection] + elif mode == "orthogonal": + _validate_array_selection(selection, transform.domain.shape, mode) + return transform.oindex[selection] + elif mode == "vectorized": + _validate_array_selection(selection, transform.domain.shape, mode) + return transform.vindex[selection] + else: + raise ValueError(f"Unknown mode: {mode!r}") diff --git a/tests/test_array.py b/tests/test_array.py index f7f564f30e..396c09d28a 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1956,7 +1956,8 @@ def test_array_repr(store: Store) -> None: shape = (2, 3, 4) dtype = "uint8" arr = zarr.create_array(store, shape=shape, dtype=dtype) - assert str(arr) == f"" + domain = "{ [0, 2), [0, 3), [0, 4) }" + assert str(arr) == f"" class UnknownObjectDtype(UTF8Base[np.dtypes.ObjectDType]): diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py new file mode 100644 index 0000000000..f27a242109 --- /dev/null +++ b/tests/test_lazy_indexing.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +import zarr +from zarr.storage import MemoryStore + + +@pytest.fixture +def arr() -> zarr.Array[Any]: + """Create a 2D array with known data.""" + store = MemoryStore() + a = zarr.create(shape=(20, 30), chunks=(5, 10), dtype="i4", store=store) + data = np.arange(600, dtype="i4").reshape(20, 30) + a[...] = data + return a + + +@pytest.fixture +def data() -> np.ndarray[Any, Any]: + return np.arange(600, dtype="i4").reshape(20, 30) + + +class TestEagerRead: + def test_basic_slice(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[2:8, 5:15] + np.testing.assert_array_equal(result, data[2:8, 5:15]) + + def test_basic_int(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[3] + np.testing.assert_array_equal(result, data[3]) + + def test_basic_int_scalar(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[3, 5] + assert result == data[3, 5] + + def test_ellipsis(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[...] + np.testing.assert_array_equal(result, data) + + def test_strided_slice(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[::2, ::3] + np.testing.assert_array_equal(result, data[::2, ::3]) + + def test_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + idx = np.array([1, 5, 10], dtype=np.intp) + result = arr.oindex[idx, :] + np.testing.assert_array_equal(result, data[idx, :]) + + def test_vindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + idx0 = np.array([1, 5, 10], dtype=np.intp) + idx1 = np.array([2, 8, 15], dtype=np.intp) + result = arr.vindex[idx0, idx1] + np.testing.assert_array_equal(result, data[idx0, idx1]) + + def test_slice_across_chunks(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + """Slice that spans multiple chunks.""" + result = arr[3:17, 8:22] + np.testing.assert_array_equal(result, data[3:17, 8:22]) + + def test_single_element(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[0:1, 0:1] + np.testing.assert_array_equal(result, data[0:1, 0:1]) + + def test_full_read(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + result = arr[:] + np.testing.assert_array_equal(result, data) + + +class TestEagerWrite: + def test_write_slice(self, arr: zarr.Array[Any]) -> None: + arr[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 + result = arr[2:5, 10:20] + np.testing.assert_array_equal(result, np.ones((3, 10), dtype="i4") * 99) + + def test_write_scalar(self, arr: zarr.Array[Any]) -> None: + arr[0, 0] = 42 + assert arr[0, 0] == 42 + + def test_roundtrip(self, arr: zarr.Array[Any]) -> None: + new_data = np.random.randint(0, 100, size=(20, 30), dtype="i4") + arr[...] = new_data + np.testing.assert_array_equal(arr[...], new_data) + + def test_write_across_chunks(self, arr: zarr.Array[Any]) -> None: + """Write spanning multiple chunks.""" + val = np.ones((14, 14), dtype="i4") * 77 + arr[3:17, 8:22] = val + result = arr[3:17, 8:22] + np.testing.assert_array_equal(result, val) + + +class TestLazyRead: + def test_lazy_shape(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.z[2:8, 5:15] + assert isinstance(v, zarr.Array) + assert v.shape == (6, 10) + + def test_lazy_resolve(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.z[2:8, 5:15] + result = v[...] + np.testing.assert_array_equal(result, data[2:8, 5:15]) + + def test_lazy_np_asarray(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.z[2:8] + result = np.asarray(v) + np.testing.assert_array_equal(result, data[2:8]) + + def test_lazy_composition(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.z[2:12].z[3:8] + assert v.shape == (5, 30) + result = v[...] + np.testing.assert_array_equal(result, data[5:10]) + + def test_lazy_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + idx = np.array([1, 5, 10], dtype=np.intp) + v = arr.z.oindex[idx, :] + assert isinstance(v, zarr.Array) + assert v.shape == (3, 30) + result = v[...] + np.testing.assert_array_equal(result, data[idx, :]) + + def test_lazy_vindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + idx0 = np.array([1, 5, 10], dtype=np.intp) + idx1 = np.array([2, 8, 15], dtype=np.intp) + v = arr.z.vindex[idx0, idx1] + assert isinstance(v, zarr.Array) + assert v.shape == (3,) + result = v[...] + np.testing.assert_array_equal(result, data[idx0, idx1]) + + def test_lazy_resolve_method(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.z[2:8] + result = v.resolve() + np.testing.assert_array_equal(result, data[2:8]) + + def test_lazy_across_chunks(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + """Lazy slice spanning multiple chunks resolves correctly.""" + v = arr.z[3:17, 8:22] + result = v[...] + np.testing.assert_array_equal(result, data[3:17, 8:22]) + + +class TestLazyWrite: + def test_lazy_write(self, arr: zarr.Array[Any]) -> None: + arr.z[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 + result = arr[2:5, 10:20] + np.testing.assert_array_equal(result, np.ones((3, 10), dtype="i4") * 99) + + def test_lazy_oindex_write(self, arr: zarr.Array[Any]) -> None: + idx = np.array([0, 5, 10], dtype=np.intp) + arr.z.oindex[idx, :] = np.zeros((3, 30), dtype="i4") + result = arr.oindex[idx, :] + np.testing.assert_array_equal(result, np.zeros((3, 30), dtype="i4")) + + def test_lazy_vindex_write(self, arr: zarr.Array[Any]) -> None: + idx0 = np.array([0, 5, 10], dtype=np.intp) + idx1 = np.array([0, 5, 10], dtype=np.intp) + arr.z.vindex[idx0, idx1] = np.array([77, 88, 99], dtype="i4") + result = arr.vindex[idx0, idx1] + np.testing.assert_array_equal(result, np.array([77, 88, 99], dtype="i4")) diff --git a/tests/test_transforms/__init__.py b/tests/test_transforms/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_transforms/test_chunk_resolution.py b/tests/test_transforms/test_chunk_resolution.py new file mode 100644 index 0000000000..63246a9caf --- /dev/null +++ b/tests/test_transforms/test_chunk_resolution.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.chunk_grids import ChunkGrid, FixedDimension +from zarr.core.transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + + +class TestChunkResolutionIdentity: + def test_single_chunk(self) -> None: + """Array fits in one chunk.""" + t = IndexTransform.from_shape((10,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=10),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert sub_t.domain.shape == (10,) + + def test_multiple_chunks_1d(self) -> None: + """1D array spanning 3 chunks.""" + t = IndexTransform.from_shape((30,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 3 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + def test_multiple_chunks_2d(self) -> None: + """2D array spanning 2x3 chunks.""" + t = IndexTransform.from_shape((20, 30)) + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=20), + FixedDimension(size=10, extent=30), + ) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 6 + coords_list = [r[0] for r in results] + assert (0, 0) in coords_list + assert (1, 2) in coords_list + + +class TestChunkResolutionSliced: + def test_slice_within_chunk(self) -> None: + """Slice that falls within a single chunk.""" + t = IndexTransform.from_shape((100,))[5:8] + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert isinstance(sub_t.output[0], DimensionMap) + assert sub_t.output[0].offset == 5 + + def test_slice_across_chunks(self) -> None: + """Slice that spans two chunks.""" + t = IndexTransform.from_shape((100,))[8:15] + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 2 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + + +class TestChunkResolutionConstant: + def test_integer_index(self) -> None: + """Integer index produces constant map — single chunk per constant dim.""" + t = IndexTransform.from_shape((100, 100))[25, :] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=100), + FixedDimension(size=10, extent=100), + ) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 10 + for coords, _, _ in results: + assert coords[0] == 2 + + +class TestChunkResolutionArray: + def test_array_index(self) -> None: + """Array index map — chunks determined by array values.""" + idx = np.array([5, 15, 25], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=idx),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid)) + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + +class TestSubTransformToSelections: + def test_constant_map(self) -> None: + """ConstantMap produces int selection + drop axis.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (5,) + assert out_sel == () + assert drop_axes == () + + def test_dimension_map_stride_1(self) -> None: + """DimensionMap with stride=1 produces contiguous slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=3, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(3, 13, 1),) + assert out_sel == (slice(0, 10),) + assert drop_axes == () + + def test_dimension_map_strided(self) -> None: + """DimensionMap with stride>1 produces strided slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=2, stride=3),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(2, 17, 3),) + assert out_sel == (slice(0, 5),) + assert drop_axes == () + + def test_array_map(self) -> None: + """ArrayMap produces integer array selection.""" + arr = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=0, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], arr) + # Without chunk_mask, out_sel falls back to domain-based slices + assert out_sel == (slice(0, 3),) + assert drop_axes == () + + def test_array_map_with_offset_stride(self) -> None: + """ArrayMap with offset and stride computes storage coords.""" + arr = np.array([0, 1, 2], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=10, stride=5),), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], np.array([10, 15, 20])) + assert drop_axes == () + + def test_mixed_maps_2d(self) -> None: + """Mix of ConstantMap and DimensionMap.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel[0] == 5 + assert chunk_sel[1] == slice(0, 10, 1) + # drop_axes is empty — integer in chunk_sel naturally drops the dim via numpy + assert drop_axes == () diff --git a/tests/test_transforms/test_composition.py b/tests/test_transforms/test_composition.py new file mode 100644 index 0000000000..e96616933d --- /dev/null +++ b/tests/test_transforms/test_composition.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.core.transforms.composition import compose +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + + +class TestComposeConstantInner: + """Inner = constant. Result is always constant.""" + + def test_constant_inner_any_outer(self) -> None: + outer = IndexTransform.from_shape((5,)) + inner = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=42),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 42 + + +class TestComposeDimensionInner: + """Inner = DimensionMap.""" + + def test_dimension_inner_constant_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 25 + + def test_dimension_inner_dimension_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + assert result.output[0].input_dimension == 0 + + def test_dimension_inner_array_outer(self) -> None: + arr = np.array([0, 2, 4], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + +class TestComposeArrayInner: + """Inner = ArrayMap.""" + + def test_array_inner_constant_outer(self) -> None: + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 20 + + def test_array_inner_array_outer(self) -> None: + outer_arr = np.array([0, 2, 1], dtype=np.intp) + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=outer_arr, offset=0, stride=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + expected = np.array([10, 30, 20], dtype=np.intp) + np.testing.assert_array_equal(result.output[0].index_array, expected) + + +class TestComposeMultiDim: + def test_2d_identity_compose(self) -> None: + a = IndexTransform.from_shape((10, 20)) + b = IndexTransform.from_shape((10, 20)) + result = compose(a, b) + assert result.domain.shape == (10, 20) + for i in range(2): + m = result.output[i] + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_mixed_map_types(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10, 10)), + output=( + DimensionMap(input_dimension=0, offset=2, stride=3), + DimensionMap(input_dimension=1, offset=0, stride=1), + ), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 17 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + assert result.output[1].offset == 0 + assert result.output[1].stride == 1 + + def test_rank_mismatch_raises(self) -> None: + outer = IndexTransform.from_shape((10,)) + inner = IndexTransform.from_shape((10, 20)) + with pytest.raises(ValueError, match="rank"): + compose(outer, inner) + + +class TestComposeChain: + def test_three_transforms(self) -> None: + a = IndexTransform.from_shape((100,)) + b = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=1),), + ) + c = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + bc = compose(b, c) + abc = compose(a, bc) + assert isinstance(abc.output[0], DimensionMap) + assert abc.output[0].offset == 25 + assert abc.output[0].stride == 2 diff --git a/tests/test_transforms/test_domain.py b/tests/test_transforms/test_domain.py new file mode 100644 index 0000000000..5a222a548c --- /dev/null +++ b/tests/test_transforms/test_domain.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import pytest + +from zarr.core.transforms.domain import IndexDomain + + +class TestIndexDomainConstruction: + def test_from_shape(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.inclusive_min == (0, 0) + assert d.exclusive_max == (10, 20) + assert d.ndim == 2 + assert d.origin == (0, 0) + assert d.shape == (10, 20) + + def test_from_shape_0d(self) -> None: + d = IndexDomain.from_shape(()) + assert d.ndim == 0 + assert d.shape == () + + def test_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5, 10), exclusive_max=(15, 30)) + assert d.origin == (5, 10) + assert d.shape == (10, 20) + assert d.ndim == 2 + + def test_validation_mismatched_lengths(self) -> None: + with pytest.raises(ValueError, match="same length"): + IndexDomain(inclusive_min=(0,), exclusive_max=(10, 20)) + + def test_validation_min_greater_than_max(self) -> None: + with pytest.raises(ValueError, match="inclusive_min must be <="): + IndexDomain(inclusive_min=(10,), exclusive_max=(5,)) + + def test_empty_domain(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(5,)) + assert d.shape == (0,) + + def test_labels(self) -> None: + d = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + assert d.labels == ("x", "y") + + def test_labels_none(self) -> None: + d = IndexDomain.from_shape((10,)) + assert d.labels is None + + +class TestIndexDomainContains: + def test_contains_inside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((0, 0)) is True + assert d.contains((9, 19)) is True + assert d.contains((5, 10)) is True + + def test_contains_outside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((10, 0)) is False + assert d.contains((-1, 0)) is False + assert d.contains((0, 20)) is False + + def test_contains_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert d.contains((5,)) is True + assert d.contains((9,)) is True + assert d.contains((4,)) is False + assert d.contains((10,)) is False + + def test_contains_wrong_ndim(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((5,)) is False + + def test_contains_domain_inside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(8, 15)) + assert outer.contains_domain(inner) is True + + def test_contains_domain_outside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(11, 15)) + assert outer.contains_domain(inner) is False + + def test_contains_domain_wrong_ndim(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain.from_shape((5,)) + assert outer.contains_domain(inner) is False + + +class TestIndexDomainIntersect: + def test_overlapping(self) -> None: + a = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 10)) + b = IndexDomain(inclusive_min=(5, 5), exclusive_max=(15, 15)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5, 5) + assert result.exclusive_max == (10, 10) + + def test_disjoint(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(10,), exclusive_max=(15,)) + assert a.intersect(b) is None + + def test_touching_boundary(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert a.intersect(b) is None + + def test_contained(self) -> None: + a = IndexDomain.from_shape((20,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5,) + assert result.exclusive_max == (10,) + + def test_wrong_ndim(self) -> None: + a = IndexDomain.from_shape((10,)) + b = IndexDomain.from_shape((10, 20)) + with pytest.raises(ValueError, match="different ranks"): + a.intersect(b) + + +class TestIndexDomainTranslate: + def test_translate_positive(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.translate((5, 10)) + assert result.inclusive_min == (5, 10) + assert result.exclusive_max == (15, 30) + + def test_translate_negative(self) -> None: + d = IndexDomain(inclusive_min=(10, 20), exclusive_max=(30, 40)) + result = d.translate((-10, -20)) + assert result.inclusive_min == (0, 0) + assert result.exclusive_max == (20, 20) + + def test_translate_wrong_length(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(ValueError, match="same length"): + d.translate((1, 2)) + + +class TestIndexDomainNarrow: + def test_narrow_slice(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((slice(2, 8), slice(5, 15))) + assert result.inclusive_min == (2, 5) + assert result.exclusive_max == (8, 15) + + def test_narrow_int(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((3, slice(None))) + assert result.inclusive_min == (3, 0) + assert result.exclusive_max == (4, 20) + + def test_narrow_ellipsis(self) -> None: + d = IndexDomain.from_shape((10, 20, 30)) + result = d.narrow((slice(1, 5), ...)) + assert result.inclusive_min == (1, 0, 0) + assert result.exclusive_max == (5, 20, 30) + + def test_narrow_slice_none(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(None),)) + assert result == d + + def test_narrow_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(10,), exclusive_max=(20,)) + result = d.narrow((slice(12, 18),)) + assert result.inclusive_min == (12,) + assert result.exclusive_max == (18,) + + def test_narrow_int_out_of_bounds(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((10,)) + + def test_narrow_int_below_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((4,)) + + def test_narrow_clamps_to_domain(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(-5, 100),)) + assert result.inclusive_min == (0,) + assert result.exclusive_max == (10,) + + def test_narrow_bare_slice(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow(slice(2, 8)) + assert result.inclusive_min == (2,) + assert result.exclusive_max == (8,) + + def test_narrow_too_many_indices(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="too many indices"): + d.narrow((1, 2)) + + def test_narrow_step_not_one(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="step=1"): + d.narrow((slice(0, 10, 2),)) diff --git a/tests/test_transforms/test_output_map.py b/tests/test_transforms/test_output_map.py new file mode 100644 index 0000000000..57dbebaf44 --- /dev/null +++ b/tests/test_transforms/test_output_map.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap + + +class TestConstantMap: + def test_construction(self) -> None: + m = ConstantMap(offset=42) + assert m.offset == 42 + + def test_default_offset(self) -> None: + m = ConstantMap() + assert m.offset == 0 + + def test_frozen(self) -> None: + m = ConstantMap(offset=5) + assert isinstance(m, ConstantMap) + + +class TestDimensionMap: + def test_construction(self) -> None: + m = DimensionMap(input_dimension=3, offset=5, stride=2) + assert m.input_dimension == 3 + assert m.offset == 5 + assert m.stride == 2 + + def test_defaults(self) -> None: + m = DimensionMap(input_dimension=0) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + m = DimensionMap(input_dimension=0) + assert isinstance(m, DimensionMap) + + +class TestArrayMap: + def test_construction(self) -> None: + arr = np.array([1, 3, 5], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=10, stride=2) + assert m.offset == 10 + assert m.stride == 2 + np.testing.assert_array_equal(m.index_array, arr) + + def test_defaults(self) -> None: + arr = np.array([0, 1], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + arr = np.array([0], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert isinstance(m, ArrayMap) diff --git a/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py new file mode 100644 index 0000000000..f531714045 --- /dev/null +++ b/tests/test_transforms/test_transform.py @@ -0,0 +1,516 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform, selection_to_transform + + +class TestIndexTransformConstruction: + def test_from_shape(self) -> None: + t = IndexTransform.from_shape((10, 20)) + assert t.input_rank == 2 + assert t.output_rank == 2 + assert t.domain.shape == (10, 20) + assert t.domain.origin == (0, 0) + for i, m in enumerate(t.output): + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_identity(self) -> None: + domain = IndexDomain(inclusive_min=(5,), exclusive_max=(15,)) + t = IndexTransform.identity(domain) + assert t.input_rank == 1 + assert t.output_rank == 1 + assert t.domain == domain + assert isinstance(t.output[0], DimensionMap) + assert t.output[0].input_dimension == 0 + + def test_from_shape_0d(self) -> None: + t = IndexTransform.from_shape(()) + assert t.input_rank == 0 + assert t.output_rank == 0 + assert t.domain.shape == () + + def test_custom_output_maps(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (ConstantMap(offset=42), DimensionMap(input_dimension=0, offset=5, stride=2)) + t = IndexTransform(domain=domain, output=maps) + assert t.input_rank == 1 + assert t.output_rank == 2 + + def test_validation_input_dimension_out_of_range(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (DimensionMap(input_dimension=5),) + with pytest.raises(ValueError, match="input_dimension"): + IndexTransform(domain=domain, output=maps) + + +class TestIndexTransformBasicIndexing: + def test_slice_identity(self) -> None: + """slice(None) on identity transform is a no-op.""" + t = IndexTransform.from_shape((10, 20)) + result = t[slice(None), slice(None)] + assert result.domain.shape == (10, 20) + assert result.input_rank == 2 + assert result.output_rank == 2 + + def test_slice_narrows(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8, 5:15] + assert result.domain.shape == (6, 10) + assert result.domain.origin == (0, 0) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 2 + assert result.output[0].stride == 1 + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 5 + assert result.output[1].input_dimension == 1 + + def test_strided_slice(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[::2] + assert result.domain.shape == (5,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 2 + + def test_strided_slice_with_start(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[1:9:3] + # indices: 1, 4, 7 -> 3 elements + assert result.domain.shape == (3,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 1 + assert result.output[0].stride == 3 + + def test_int_drops_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + assert result.output_rank == 2 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + + def test_int_middle_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[:, 5, :] + assert result.input_rank == 2 + assert result.output_rank == 3 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], ConstantMap) + assert result.output[1].offset == 5 + assert isinstance(result.output[2], DimensionMap) + assert result.output[2].input_dimension == 1 + + def test_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[2:8, ...] + assert result.input_rank == 3 + assert result.domain.shape == (6, 20, 30) + + def test_newaxis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[np.newaxis, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (1, 10, 20) + assert result.output_rank == 2 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 2 + + def test_int_out_of_bounds(self) -> None: + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[10] + + def test_negative_int_is_literal(self) -> None: + """Negative indices are literal coordinates (TensorStore convention), + not 'from the end' like NumPy.""" + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[-1] # -1 is out of bounds for domain [0, 10) + + def test_negative_int_valid_with_negative_origin(self) -> None: + """Negative index is valid if the domain includes negative coordinates.""" + domain = IndexDomain(inclusive_min=(-5,), exclusive_max=(5,)) + t = IndexTransform.identity(domain) + result = t[-3] + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == -3 + + def test_composition_of_slices(self) -> None: + """Slicing a sliced transform should compose offsets.""" + t = IndexTransform.from_shape((100,)) + result = t[10:50][5:20] + assert result.domain.shape == (15,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 15 + assert result.output[0].stride == 1 + + def test_composition_of_strides(self) -> None: + t = IndexTransform.from_shape((100,)) + result = t[::2][::3] + # t[::2] -> shape (50,), offset=0, stride=2 + # [::3] -> shape ceil(50/3)=17, offset=0, stride=2*3=6 + assert result.domain.shape == (17,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].stride == 6 + + def test_bare_int(self) -> None: + """Non-tuple selection.""" + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + + def test_bare_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8] + assert result.domain.shape == (6, 20) + + +class TestBasicIndexingOnArrayMaps: + """When a transform already has ArrayMap outputs, basic indexing must + apply the corresponding operation to the index_array's axes.""" + + def test_int_on_array_map_drops_axis(self) -> None: + """Integer index on a dimension referenced by an ArrayMap should + index into the array on that axis.""" + arr = np.array([[10, 20], [30, 40], [50, 60]], dtype=np.intp) + # 2D input domain (3, 2), one ArrayMap output + t = IndexTransform( + domain=IndexDomain.from_shape((3, 2)), + output=(ArrayMap(index_array=arr),), + ) + # Index with int on dim 0 -> pick row 1 -> arr[1, :] = [30, 40] + result = t[1] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([30, 40])) + + def test_slice_on_array_map(self) -> None: + """Slice on a dimension referenced by an ArrayMap should slice the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[1:4] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30, 40])) + + def test_strided_slice_on_array_map(self) -> None: + """Strided slice on ArrayMap should stride the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[::2] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([10, 30, 50])) + + def test_newaxis_on_array_map(self) -> None: + """Newaxis should insert an axis in the index_array.""" + arr = np.array([10, 20, 30], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[np.newaxis, :] + assert result.input_rank == 2 + assert result.domain.shape == (1, 3) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].index_array.shape == (1, 3) + np.testing.assert_array_equal(result.output[0].index_array, np.array([[10, 20, 30]])) + + def test_int_drops_one_of_two_array_dims(self) -> None: + """2D array map, int on dim 0, slice on dim 1.""" + arr = np.array([[10, 20, 30], [40, 50, 60]], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2, 3)), + output=(ArrayMap(index_array=arr),), + ) + result = t[0, 1:3] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + # arr[0, 1:3] = [20, 30] + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30])) + + +class TestIndexTransformOindex: + def test_oindex_int_array(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.oindex[idx, :] + assert result.input_rank == 2 + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 1 + + def test_oindex_bool_array(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.oindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal( + result.output[0].index_array, np.array([0, 2, 4], dtype=np.intp) + ) + + def test_oindex_mixed(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([2, 4], dtype=np.intp) + result = t.oindex[idx, 5:15] + assert result.input_rank == 2 + assert result.domain.shape == (2, 10) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 5 + + def test_oindex_multiple_arrays(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 10, 15], dtype=np.intp) + result = t.oindex[idx0, :, idx1] + assert result.input_rank == 3 + assert result.domain.shape == (2, 20, 3) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert isinstance(result.output[2], ArrayMap) + + +class TestIndexTransformVindex: + def test_vindex_single_array(self) -> None: + t = IndexTransform.from_shape((10,)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx] + assert result.input_rank == 1 + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx) + + def test_vindex_broadcast(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([[1, 2], [3, 4]], dtype=np.intp) + idx1 = np.array([[10, 11], [12, 13]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 2) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx0) + np.testing.assert_array_equal(result.output[1].index_array, idx1) + + def test_vindex_with_slice(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (3, 20, 30) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_bool_mask(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.vindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_broadcast_different_shapes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 2, 3], dtype=np.intp) + idx1 = np.array([[10], [11]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 3) + + +class TestSelectionToTransform: + def test_basic_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((slice(2, 8), slice(5, 15)), t, "basic") + assert result.domain.shape == (6, 10) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 2 + + def test_basic_int(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((3, slice(None)), t, "basic") + assert result.input_rank == 1 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + + def test_basic_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform(Ellipsis, t, "basic") + assert result.domain.shape == (10, 20) + + def test_orthogonal(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = selection_to_transform((idx, slice(None)), t, "orthogonal") + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + + def test_vectorized(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 7], dtype=np.intp) + result = selection_to_transform((idx0, idx1), t, "vectorized") + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + + def test_composition_with_non_identity(self) -> None: + """Indexing a sliced transform composes offsets.""" + t = IndexTransform.from_shape((100,))[10:50] + result = selection_to_transform(slice(5, 20), t, "basic") + assert result.domain.shape == (15,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 15 + + +class TestIndexTransformIntersect: + def test_constant_inside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(0,), exclusive_max=(10,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert surviving is None + + def test_constant_outside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) + assert result is None + + def test_dimension_partial(self) -> None: + """DimensionMap over [0,10) intersected with [5,15) narrows input to [5,10).""" + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(15,))) + assert result is not None + restricted, surviving = result + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + assert surviving is None + + def test_dimension_no_overlap(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(20,), exclusive_max=(30,))) + assert result is None + + def test_dimension_strided(self) -> None: + """stride=2, offset=1 over [0,5): storage 1,3,5,7,9. Chunk [4,8).""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=1, stride=2),), + ) + result = t.intersect(IndexDomain(inclusive_min=(4,), exclusive_max=(8,))) + assert result is not None + restricted, _surviving = result + # input 2->5, input 3->7. Both in [4,8). + assert restricted.domain.inclusive_min == (2,) + assert restricted.domain.exclusive_max == (4,) + + def test_array_partial(self) -> None: + arr = np.array([3, 8, 15, 22], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(ArrayMap(index_array=arr),), + ) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(20,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ArrayMap) + np.testing.assert_array_equal(restricted.output[0].index_array, np.array([8, 15])) + assert surviving is not None + np.testing.assert_array_equal(surviving, np.array([1, 2])) + + def test_array_none_inside(self) -> None: + arr = np.array([1, 2, 3], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + assert t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) is None + + def test_2d_mixed(self) -> None: + """2D: ConstantMap on dim 0, DimensionMap on dim 1.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk = IndexDomain(inclusive_min=(0, 5), exclusive_max=(10, 15)) + result = t.intersect(chunk) + assert result is not None + restricted, _ = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert isinstance(restricted.output[1], DimensionMap) + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + + +class TestIndexTransformTranslate: + def test_translate_constant(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.translate((-5,)) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 0 + + def test_translate_dimension(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.translate((-3,)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -3 + assert result.output[0].stride == 1 + + def test_translate_array(self) -> None: + arr = np.array([5, 10], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2,)), + output=(ArrayMap(index_array=arr, offset=3),), + ) + result = t.translate((-3,)) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 0 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + def test_translate_2d(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.translate((-5, -10)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -5 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == -10 From 79cd9c80a0b743f350f0683fa7bd75597342314e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 14 Apr 2026 15:08:22 +0200 Subject: [PATCH 02/43] feat: add JSON serialization for IndexTransform (TensorStore-compatible) Add TypedDict definitions and conversion functions for serializing IndexDomain, OutputIndexMap, and IndexTransform to/from JSON. The JSON format follows TensorStore's conventions for interoperability: - IndexDomain: input_inclusive_min, input_exclusive_max, input_labels - OutputIndexMap: offset + optional stride/input_dimension/index_array - IndexTransform: domain fields + output array TypedDicts: IndexDomainJSON, OutputIndexMapJSON, IndexTransformJSON Functions: index_domain_to_json, index_domain_from_json, index_transform_to_json, index_transform_from_json Co-Authored-By: Claude Opus 4.6 (1M context) --- src/zarr/core/transforms/__init__.py | 16 +++ src/zarr/core/transforms/json.py | 163 ++++++++++++++++++++++ tests/test_transforms/test_json.py | 199 +++++++++++++++++++++++++++ 3 files changed, 378 insertions(+) create mode 100644 src/zarr/core/transforms/json.py create mode 100644 tests/test_transforms/test_json.py diff --git a/src/zarr/core/transforms/__init__.py b/src/zarr/core/transforms/__init__.py index 530dd39cea..ec98e02d4d 100644 --- a/src/zarr/core/transforms/__init__.py +++ b/src/zarr/core/transforms/__init__.py @@ -16,6 +16,15 @@ from zarr.core.transforms.composition import compose from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.json import ( + IndexDomainJSON, + IndexTransformJSON, + OutputIndexMapJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, +) from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr.core.transforms.transform import IndexTransform @@ -24,7 +33,14 @@ "ConstantMap", "DimensionMap", "IndexDomain", + "IndexDomainJSON", "IndexTransform", + "IndexTransformJSON", "OutputIndexMap", + "OutputIndexMapJSON", "compose", + "index_domain_from_json", + "index_domain_to_json", + "index_transform_from_json", + "index_transform_to_json", ] diff --git a/src/zarr/core/transforms/json.py b/src/zarr/core/transforms/json.py new file mode 100644 index 0000000000..59a81cc1e8 --- /dev/null +++ b/src/zarr/core/transforms/json.py @@ -0,0 +1,163 @@ +"""JSON serialization for index transforms. + +Defines TypedDict types matching TensorStore's JSON representation of +IndexTransform and IndexDomain, plus conversion functions. + +The JSON format follows TensorStore's conventions for interoperability:: + + { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [100, 200], + "input_labels": ["x", "y"], + "output": [ + {"offset": 5}, + {"offset": 10, "stride": 2, "input_dimension": 1}, + {"offset": 0, "stride": 1, "index_array": [[1, 2, 0]]} + ] + } +""" + +from __future__ import annotations + +from typing import Required, TypedDict + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.transform import IndexTransform + +# --------------------------------------------------------------------------- +# TypedDict definitions (JSON shapes) +# --------------------------------------------------------------------------- + + +class IndexDomainJSON(TypedDict, total=False): + """JSON representation of an IndexDomain.""" + + input_inclusive_min: Required[list[int]] + input_exclusive_max: Required[list[int]] + input_labels: list[str] + + +class OutputIndexMapJSON(TypedDict, total=False): + """JSON representation of a single output index map. + + Exactly one of three forms: + - ``{"offset": 5}`` — constant + - ``{"offset": 0, "stride": 1, "input_dimension": 0}`` — dimension + - ``{"offset": 0, "stride": 1, "index_array": [...]}`` — array + """ + + offset: int + stride: int + input_dimension: int + index_array: list[int] | list[list[int]] + + +class IndexTransformJSON(TypedDict, total=False): + """JSON representation of an IndexTransform.""" + + input_inclusive_min: Required[list[int]] + input_exclusive_max: Required[list[int]] + input_labels: list[str] + output: Required[list[OutputIndexMapJSON]] + + +# --------------------------------------------------------------------------- +# IndexDomain serialization +# --------------------------------------------------------------------------- + + +def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON: + """Convert an IndexDomain to its JSON representation.""" + result: IndexDomainJSON = { + "input_inclusive_min": list(domain.inclusive_min), + "input_exclusive_max": list(domain.exclusive_max), + } + if domain.labels is not None: + result["input_labels"] = list(domain.labels) + return result + + +def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: + """Construct an IndexDomain from its JSON representation.""" + return IndexDomain( + inclusive_min=tuple(data["input_inclusive_min"]), + exclusive_max=tuple(data["input_exclusive_max"]), + labels=tuple(data["input_labels"]) if "input_labels" in data else None, + ) + + +# --------------------------------------------------------------------------- +# OutputIndexMap serialization +# --------------------------------------------------------------------------- + + +def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: + """Convert an output index map to its JSON representation.""" + if isinstance(m, ConstantMap): + result: OutputIndexMapJSON = {"offset": m.offset} + return result + + if isinstance(m, DimensionMap): + result = {"offset": m.offset, "input_dimension": m.input_dimension} + if m.stride != 1: + result["stride"] = m.stride + return result + + if isinstance(m, ArrayMap): + result = {"offset": m.offset, "index_array": m.index_array.tolist()} + if m.stride != 1: + result["stride"] = m.stride + return result + + raise TypeError(f"Unknown output map type: {type(m)}") + + +def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: + """Construct an output index map from its JSON representation.""" + if "index_array" in data: + return ArrayMap( + index_array=np.asarray(data["index_array"], dtype=np.intp), + offset=data.get("offset", 0), + stride=data.get("stride", 1), + ) + + if "input_dimension" in data: + return DimensionMap( + input_dimension=data["input_dimension"], + offset=data.get("offset", 0), + stride=data.get("stride", 1), + ) + + # Constant map: only offset present + return ConstantMap(offset=data.get("offset", 0)) + + +# --------------------------------------------------------------------------- +# IndexTransform serialization +# --------------------------------------------------------------------------- + + +def index_transform_to_json(transform: IndexTransform) -> IndexTransformJSON: + """Convert an IndexTransform to its JSON representation.""" + result: IndexTransformJSON = { + "input_inclusive_min": list(transform.domain.inclusive_min), + "input_exclusive_max": list(transform.domain.exclusive_max), + "output": [output_index_map_to_json(m) for m in transform.output], + } + if transform.domain.labels is not None: + result["input_labels"] = list(transform.domain.labels) + return result + + +def index_transform_from_json(data: IndexTransformJSON) -> IndexTransform: + """Construct an IndexTransform from its JSON representation.""" + domain = IndexDomain( + inclusive_min=tuple(data["input_inclusive_min"]), + exclusive_max=tuple(data["input_exclusive_max"]), + labels=tuple(data["input_labels"]) if "input_labels" in data else None, + ) + output = tuple(output_index_map_from_json(m) for m in data["output"]) + return IndexTransform(domain=domain, output=output) diff --git a/tests/test_transforms/test_json.py b/tests/test_transforms/test_json.py new file mode 100644 index 0000000000..db582a9561 --- /dev/null +++ b/tests/test_transforms/test_json.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.json import ( + IndexTransformJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, + output_index_map_from_json, + output_index_map_to_json, +) +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + + +class TestIndexDomainJSON: + def test_roundtrip(self) -> None: + domain = IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + json = index_domain_to_json(domain) + assert json == {"input_inclusive_min": [2, 5], "input_exclusive_max": [10, 20]} + restored = index_domain_from_json(json) + assert restored == domain + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + json = index_domain_to_json(domain) + assert json["input_labels"] == ["x", "y"] + restored = index_domain_from_json(json) + assert restored.labels == ("x", "y") + + def test_without_labels(self) -> None: + domain = IndexDomain.from_shape((5,)) + json = index_domain_to_json(domain) + assert "input_labels" not in json + restored = index_domain_from_json(json) + assert restored.labels is None + + def test_zero_origin(self) -> None: + domain = IndexDomain.from_shape((10, 20, 30)) + json = index_domain_to_json(domain) + assert json == { + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [10, 20, 30], + } + assert index_domain_from_json(json) == domain + + +class TestOutputIndexMapJSON: + def test_constant(self) -> None: + m = ConstantMap(offset=42) + json = output_index_map_to_json(m) + assert json == {"offset": 42} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 42 + + def test_constant_zero(self) -> None: + m = ConstantMap(offset=0) + json = output_index_map_to_json(m) + assert json == {"offset": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 0 + + def test_dimension(self) -> None: + m = DimensionMap(input_dimension=1, offset=10, stride=3) + json = output_index_map_to_json(m) + assert json == {"offset": 10, "stride": 3, "input_dimension": 1} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.input_dimension == 1 + assert restored.offset == 10 + assert restored.stride == 3 + + def test_dimension_stride_1_omitted(self) -> None: + """stride=1 is the default and should be omitted from JSON.""" + m = DimensionMap(input_dimension=0) + json = output_index_map_to_json(m) + assert "stride" not in json + assert json == {"offset": 0, "input_dimension": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.stride == 1 + + def test_array(self) -> None: + arr = np.array([1, 5, 9], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=2, stride=3) + json = output_index_map_to_json(m) + assert json == {"offset": 2, "stride": 3, "index_array": [1, 5, 9]} + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + assert restored.offset == 2 + assert restored.stride == 3 + + def test_array_stride_1_omitted(self) -> None: + arr = np.array([0, 1, 2], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert "stride" not in json + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + assert restored.stride == 1 + + def test_array_2d(self) -> None: + arr = np.array([[1, 2], [3, 4]], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert json["index_array"] == [[1, 2], [3, 4]] + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + + +class TestIndexTransformJSON: + def test_identity(self) -> None: + t = IndexTransform.from_shape((10, 20)) + json = index_transform_to_json(t) + assert json == { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [10, 20], + "output": [ + {"offset": 0, "input_dimension": 0}, + {"offset": 0, "input_dimension": 1}, + ], + } + restored = index_transform_from_json(json) + assert restored.domain == t.domain + assert len(restored.output) == 2 + for orig, rest in zip(t.output, restored.output, strict=True): + assert type(orig) is type(rest) + + def test_sliced(self) -> None: + t = IndexTransform.from_shape((100,))[10:50:2] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert restored.domain.shape == t.domain.shape + assert isinstance(restored.output[0], DimensionMap) + orig = t.output[0] + assert isinstance(orig, DimensionMap) + assert restored.output[0].offset == orig.offset + assert restored.output[0].stride == orig.stride + + def test_with_constant(self) -> None: + t = IndexTransform.from_shape((10, 20))[3] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ConstantMap) + assert restored.output[0].offset == 3 + assert isinstance(restored.output[1], DimensionMap) + + def test_with_array(self) -> None: + idx = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform.from_shape((10, 20)).oindex[idx, :] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ArrayMap) + np.testing.assert_array_equal(restored.output[0].index_array, idx) + assert isinstance(restored.output[1], DimensionMap) + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + t = IndexTransform.identity(domain) + json = index_transform_to_json(t) + assert json["input_labels"] == ["x", "y"] + restored = index_transform_from_json(json) + assert restored.domain.labels == ("x", "y") + + def test_tensorstore_compatible_format(self) -> None: + """Verify the JSON matches TensorStore's format exactly.""" + json: IndexTransformJSON = { + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [100, 200, 3], + "input_labels": ["x", "y", "channel"], + "output": [ + {"offset": 5}, + {"offset": 10, "stride": 2, "input_dimension": 1}, + {"offset": 0, "stride": 1, "index_array": [1, 2, 0]}, + ], + } + t = index_transform_from_json(json) + assert t.domain.shape == (100, 200, 3) + assert t.domain.labels == ("x", "y", "channel") + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == 5 + assert isinstance(t.output[1], DimensionMap) + assert t.output[1].offset == 10 + assert t.output[1].stride == 2 + assert t.output[1].input_dimension == 1 + assert isinstance(t.output[2], ArrayMap) + np.testing.assert_array_equal(t.output[2].index_array, [1, 2, 0]) + + # Roundtrip + json_rt = index_transform_to_json(t) + t_rt = index_transform_from_json(json_rt) + assert t_rt.domain == t.domain From ce78f489d992ca4382572ef09ed33ac41c358238 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 19 May 2026 11:03:00 +0200 Subject: [PATCH 03/43] rename accessor to lazy, and use result instead of resolve --- src/zarr/core/array.py | 11 ++++++----- src/zarr/core/transforms/transform.py | 2 +- tests/test_lazy_indexing.py | 26 +++++++++++++------------- 3 files changed, 20 insertions(+), 19 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index ce5a8ece9c..cd7df1b4ff 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3961,14 +3961,15 @@ def blocks(self) -> BlockIndex: return BlockIndex(self) @property - def z(self) -> _LazyIndexAccessor: + def lazy(self) -> _LazyIndexAccessor: """Lazy indexing accessor. Returns a new Array with composed transform, no I/O.""" return _LazyIndexAccessor(self) - def resolve(self, prototype: BufferPrototype | None = None) -> NDArrayLikeOrScalar: + def result(self, prototype: BufferPrototype | None = None) -> NDArrayLikeOrScalar: """Read and return the data for this array view. Equivalent to ``self[...]`` but more explicit for lazy views. + Named after `tensorstore.Future.result`. """ return self[...] @@ -4190,7 +4191,7 @@ async def _shards_initialized( class _LazyOIndex: - """Lazy orthogonal indexing via ``array.z.oindex[...]``.""" + """Lazy orthogonal indexing via ``array.lazy.oindex[...]``.""" __slots__ = ("_array",) @@ -4207,7 +4208,7 @@ def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: class _LazyVIndex: - """Lazy vectorized indexing via ``array.z.vindex[...]``.""" + """Lazy vectorized indexing via ``array.lazy.vindex[...]``.""" __slots__ = ("_array",) @@ -4224,7 +4225,7 @@ def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: class _LazyIndexAccessor: - """Provides lazy indexing via ``array.z[...]``.""" + """Provides lazy indexing via ``array.lazy[...]``.""" __slots__ = ("_array",) diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index ee2f3ce4c2..ab6a24464f 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -22,7 +22,7 @@ The transform is the atomic unit that connects user-facing indexing to chunk-level I/O. Every ``Array`` holds a transform (identity by default). -``Array.z[...]`` composes a new transform lazily. Reading resolves the +``Array.lazy[...]`` composes a new transform lazily. Reading resolves the transform against the chunk grid via intersect + translate. """ diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index f27a242109..fa0d218bd3 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -95,29 +95,29 @@ def test_write_across_chunks(self, arr: zarr.Array[Any]) -> None: class TestLazyRead: def test_lazy_shape(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.z[2:8, 5:15] + v = arr.lazy[2:8, 5:15] assert isinstance(v, zarr.Array) assert v.shape == (6, 10) def test_lazy_resolve(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.z[2:8, 5:15] + v = arr.lazy[2:8, 5:15] result = v[...] np.testing.assert_array_equal(result, data[2:8, 5:15]) def test_lazy_np_asarray(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.z[2:8] + v = arr.lazy[2:8] result = np.asarray(v) np.testing.assert_array_equal(result, data[2:8]) def test_lazy_composition(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.z[2:12].z[3:8] + v = arr.lazy[2:12].lazy[3:8] assert v.shape == (5, 30) result = v[...] np.testing.assert_array_equal(result, data[5:10]) def test_lazy_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: idx = np.array([1, 5, 10], dtype=np.intp) - v = arr.z.oindex[idx, :] + v = arr.lazy.oindex[idx, :] assert isinstance(v, zarr.Array) assert v.shape == (3, 30) result = v[...] @@ -126,39 +126,39 @@ def test_lazy_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> def test_lazy_vindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: idx0 = np.array([1, 5, 10], dtype=np.intp) idx1 = np.array([2, 8, 15], dtype=np.intp) - v = arr.z.vindex[idx0, idx1] + v = arr.lazy.vindex[idx0, idx1] assert isinstance(v, zarr.Array) assert v.shape == (3,) result = v[...] np.testing.assert_array_equal(result, data[idx0, idx1]) - def test_lazy_resolve_method(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.z[2:8] - result = v.resolve() + def test_lazy_result_method(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: + v = arr.lazy[2:8] + result = v.result() np.testing.assert_array_equal(result, data[2:8]) def test_lazy_across_chunks(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: """Lazy slice spanning multiple chunks resolves correctly.""" - v = arr.z[3:17, 8:22] + v = arr.lazy[3:17, 8:22] result = v[...] np.testing.assert_array_equal(result, data[3:17, 8:22]) class TestLazyWrite: def test_lazy_write(self, arr: zarr.Array[Any]) -> None: - arr.z[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 + arr.lazy[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 result = arr[2:5, 10:20] np.testing.assert_array_equal(result, np.ones((3, 10), dtype="i4") * 99) def test_lazy_oindex_write(self, arr: zarr.Array[Any]) -> None: idx = np.array([0, 5, 10], dtype=np.intp) - arr.z.oindex[idx, :] = np.zeros((3, 30), dtype="i4") + arr.lazy.oindex[idx, :] = np.zeros((3, 30), dtype="i4") result = arr.oindex[idx, :] np.testing.assert_array_equal(result, np.zeros((3, 30), dtype="i4")) def test_lazy_vindex_write(self, arr: zarr.Array[Any]) -> None: idx0 = np.array([0, 5, 10], dtype=np.intp) idx1 = np.array([0, 5, 10], dtype=np.intp) - arr.z.vindex[idx0, idx1] = np.array([77, 88, 99], dtype="i4") + arr.lazy.vindex[idx0, idx1] = np.array([77, 88, 99], dtype="i4") result = arr.vindex[idx0, idx1] np.testing.assert_array_equal(result, np.array([77, 88, 99], dtype="i4")) From 93e8c94557a12b20324693f697ca907419f17d8b Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 16:21:03 +0000 Subject: [PATCH 04/43] perf: cache Array shape/ndim and reuse chunk grid in transform resolvers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `AsyncArray.shape` recomputed `self._transform.domain.shape` (a zip+genexpr) on every access — ~10x slower than returning the stored tuple — and `ndim` went through it as well. Cache `_shape` wherever the transform is set and read it directly. Also thread the array's already-built `_chunk_grid` into the transform read/write resolvers instead of rebuilding it via `ChunkGrid.from_metadata` on every call. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_011D1uGKtHP9s7E4WqiUSYG4 --- src/zarr/core/array.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index fe70434029..e39f46f366 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -383,6 +383,7 @@ def __init__( create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), ) object.__setattr__(self, "_transform", IndexTransform.from_shape(metadata_parsed.shape)) + object.__setattr__(self, "_shape", self._transform.domain.shape) @classmethod async def _create( @@ -805,6 +806,7 @@ def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetada object.__setattr__(new, "_chunk_grid", self._chunk_grid) object.__setattr__(new, "codec_pipeline", self.codec_pipeline) object.__setattr__(new, "_transform", transform) + object.__setattr__(new, "_shape", transform.domain.shape) return new @property @@ -825,7 +827,7 @@ def ndim(self) -> int: int The number of dimensions in the Array. """ - return len(self.shape) + return len(self._shape) @property def shape(self) -> tuple[int, ...]: @@ -836,7 +838,7 @@ def shape(self) -> tuple[int, ...]: tuple The shape of the Array. """ - return self._transform.domain.shape + return self._shape @property def storage_shape(self) -> tuple[int, ...]: @@ -1617,6 +1619,7 @@ async def _get_selection_t( self.codec_pipeline, prototype=prototype, out=out, + chunk_grid=self._chunk_grid, ) async def _set_selection_t( @@ -1634,6 +1637,7 @@ async def _set_selection_t( value, self.codec_pipeline, prototype=prototype, + chunk_grid=self._chunk_grid, ) async def setitem( @@ -5645,9 +5649,11 @@ async def _get_selection_via_transform( *, prototype: BufferPrototype, out: NDBuffer | None = None, + chunk_grid: ChunkGrid | None = None, ) -> NDArrayLikeOrScalar: """Read data using an IndexTransform.""" - chunk_grid = ChunkGrid.from_metadata(metadata) + if chunk_grid is None: + chunk_grid = ChunkGrid.from_metadata(metadata) # Get dtype (same logic as existing _get_selection) if metadata.zarr_format == 2: @@ -5743,9 +5749,11 @@ async def _set_selection_via_transform( codec_pipeline: CodecPipeline, *, prototype: BufferPrototype, + chunk_grid: ChunkGrid | None = None, ) -> None: """Write data using an IndexTransform.""" - chunk_grid = ChunkGrid.from_metadata(metadata) + if chunk_grid is None: + chunk_grid = ChunkGrid.from_metadata(metadata) # Get dtype from metadata if metadata.zarr_format == 2: @@ -6404,6 +6412,7 @@ async def _delete_key(key: str) -> None: object.__setattr__(array, "metadata", new_metadata) object.__setattr__(array, "_chunk_grid", new_chunk_grid) object.__setattr__(array, "_transform", IndexTransform.from_shape(new_shape)) + object.__setattr__(array, "_shape", array._transform.domain.shape) async def _append( From ab795ca764a9c2b65c563c17341eea0a96918f13 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 16:21:03 +0000 Subject: [PATCH 05/43] perf: skip the transform resolver for eager (identity-transform) indexing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A freshly-opened array carries the identity transform, so eager indexing yields exactly the coordinates the original indexers compute. Route such arrays straight to the legacy Basic/Orthogonal/Mask/Coordinate indexer path; only non-identity transforms (opt-in `.lazy[...]` views) go through the transform resolver. This removes the per-chunk transform-resolution overhead from the common eager path while preserving lazy semantics. Also restore the incompatible-shape ValueError in the coordinate-selection legacy branch — the transform path carried it but the fields-only branch had dropped it, and the eager fast-path now exercises that branch. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_011D1uGKtHP9s7E4WqiUSYG4 --- src/zarr/core/array.py | 72 +++++++++++++++++++++++++++++++++++------- 1 file changed, 60 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index e39f46f366..1b611bc34c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -384,6 +384,12 @@ def __init__( ) object.__setattr__(self, "_transform", IndexTransform.from_shape(metadata_parsed.shape)) object.__setattr__(self, "_shape", self._transform.domain.shape) + # A freshly-opened array has the identity transform: input coord i maps to + # storage coord i over the full storage domain. Eager indexing on such an + # array can use the original (legacy) indexers directly, avoiding the + # transform-resolution overhead. Lazy views (created via _with_transform) + # carry a non-identity transform and must go through the transform path. + object.__setattr__(self, "_is_identity", True) @classmethod async def _create( @@ -807,6 +813,9 @@ def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetada object.__setattr__(new, "codec_pipeline", self.codec_pipeline) object.__setattr__(new, "_transform", transform) object.__setattr__(new, "_shape", transform.domain.shape) + object.__setattr__( + new, "_is_identity", _transform_is_identity(transform, self.metadata.shape) + ) return new @property @@ -2904,8 +2913,9 @@ def get_basic_selection( if prototype is None: prototype = default_buffer_prototype() - if fields is not None: - # Fall back to legacy path for structured dtype field selection + if fields is not None or self._async_array._is_identity: + # Eager (identity-transform) arrays and structured-dtype field + # selection use the original indexer path directly. return sync( self.async_array._get_selection( BasicIndexer(selection, self.shape, self._chunk_grid), @@ -3018,8 +3028,9 @@ def set_basic_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None: - # Fall back to legacy path for structured dtype field selection + if fields is not None or self._async_array._is_identity: + # Eager (identity-transform) arrays and structured-dtype field + # selection use the original indexer path directly. indexer = BasicIndexer(selection, self.shape, self._chunk_grid) sync( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) @@ -3154,8 +3165,9 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None or not is_basic_selection(selection): - # Fall back to legacy path for structured dtypes or advanced selections + if fields is not None or self._async_array._is_identity or not is_basic_selection(selection): + # Eager (identity) arrays, structured dtypes, and advanced selections + # use the original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) return sync( self.async_array._get_selection( @@ -3277,8 +3289,9 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None or not is_basic_selection(selection): - # Fall back to legacy path for structured dtypes or advanced selections + if fields is not None or self._async_array._is_identity or not is_basic_selection(selection): + # Eager (identity) arrays, structured dtypes, and advanced selections + # use the original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) sync( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) @@ -3371,7 +3384,7 @@ def get_mask_selection( if prototype is None: prototype = default_buffer_prototype() - if fields is not None: + if fields is not None or self._async_array._is_identity: indexer = MaskIndexer(mask, self.shape, self._chunk_grid) return sync( self.async_array._get_selection( @@ -3476,7 +3489,7 @@ def set_mask_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None: + if fields is not None or self._async_array._is_identity: indexer = MaskIndexer(mask, self.shape, self._chunk_grid) sync( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) @@ -3581,7 +3594,7 @@ def get_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None: + if fields is not None or self._async_array._is_identity: indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) out_array = sync( self.async_array._get_selection( @@ -3704,7 +3717,7 @@ def set_coordinate_selection( # Normalize empty fields list to None if not fields: fields = None - if fields is not None: + if fields is not None or self._async_array._is_identity: indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) if not is_scalar(value, self.dtype): try: @@ -3715,6 +3728,13 @@ def set_coordinate_selection( value = np.array(value) if hasattr(value, "shape") and len(value.shape) > 1: value = np.array(value).reshape(-1) + if not is_scalar(value, self.dtype) and ( + isinstance(value, NDArrayLike) and indexer.shape != value.shape + ): + raise ValueError( + f"Attempting to set a selection of {indexer.sel_shape[0]} " + f"elements with an array of {value.shape[0]} elements." + ) sync( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) @@ -5606,6 +5626,33 @@ def _get_chunk_spec( ) +def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, ...]) -> bool: + """Return True if ``transform`` is the identity over the full storage domain. + + An identity transform maps input coordinate ``i`` to storage coordinate ``i`` + across the array's whole storage shape (origin 0, unit stride, dimensions in + order). Such an array is an ordinary eager array — indexing it produces the + same coordinates the legacy indexers compute, so the legacy fast path is + safe. Any narrowing, striding, reordering, or fancy selection (i.e. a lazy + view) yields a non-identity transform that must go through the transform + resolver. Cheap: O(ndim), no array work. + """ + from zarr.core.transforms.output_map import DimensionMap + + domain = transform.domain + ndim = len(storage_shape) + if domain.ndim != ndim or len(transform.output) != ndim: + return False + if domain.inclusive_min != (0,) * ndim or domain.exclusive_max != storage_shape: + return False + for i, m in enumerate(transform.output): + if not ( + type(m) is DimensionMap and m.input_dimension == i and m.offset == 0 and m.stride == 1 + ): + return False + return True + + def _is_complete_chunk( sub_transform: IndexTransform, chunk_grid: ChunkGrid, chunk_coords: tuple[int, ...] ) -> bool: @@ -6413,6 +6460,7 @@ async def _delete_key(key: str) -> None: object.__setattr__(array, "_chunk_grid", new_chunk_grid) object.__setattr__(array, "_transform", IndexTransform.from_shape(new_shape)) object.__setattr__(array, "_shape", array._transform.domain.shape) + object.__setattr__(array, "_is_identity", True) async def _append( From b43023a75865d994d0e9a7b57b5cde58de3b0669 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 16:21:03 +0000 Subject: [PATCH 06/43] perf(lazy): resolve chunk_grid[coords] once per chunk in transform read/write MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_is_complete_chunk` and `_get_chunk_spec` each looked up `chunk_grid[chunk_coords]` — the single most expensive line in the per-chunk loop. Split the ArraySpec builder out (`_array_spec_from_chunk_spec`), look the ChunkSpec up once in the read/write loops, and pass it to both. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_011D1uGKtHP9s7E4WqiUSYG4 --- src/zarr/core/array.py | 64 +++++++++++++++++++++++++++--------------- 1 file changed, 42 insertions(+), 22 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 1b611bc34c..745cbd48df 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -41,6 +41,7 @@ from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, + ChunkSpec, _is_rectilinear_chunks, as_regular_shape, guess_chunks, @@ -5606,17 +5607,18 @@ async def _nbytes_stored( return await store_path.store.getsize_prefix(store_path.path) -def _get_chunk_spec( +def _array_spec_from_chunk_spec( metadata: ArrayMetadata, - chunk_grid: ChunkGrid, - chunk_coords: tuple[int, ...], + spec: ChunkSpec, array_config: ArrayConfig, prototype: BufferPrototype, ) -> ArraySpec: - """Build an ArraySpec for a single chunk using the ChunkGrid.""" - spec = chunk_grid[chunk_coords] - if spec is None: - raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + """Build an ArraySpec from an already-resolved ChunkSpec. + + Split out from :func:`_get_chunk_spec` so the transform read/write path can + resolve ``chunk_grid[chunk_coords]`` once per chunk and feed the same + ``spec`` to both this and :func:`_is_complete_chunk`. + """ return ArraySpec( shape=spec.codec_shape, dtype=metadata.dtype, @@ -5626,6 +5628,20 @@ def _get_chunk_spec( ) +def _get_chunk_spec( + metadata: ArrayMetadata, + chunk_grid: ChunkGrid, + chunk_coords: tuple[int, ...], + array_config: ArrayConfig, + prototype: BufferPrototype, +) -> ArraySpec: + """Build an ArraySpec for a single chunk using the ChunkGrid.""" + spec = chunk_grid[chunk_coords] + if spec is None: + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + return _array_spec_from_chunk_spec(metadata, spec, array_config, prototype) + + def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, ...]) -> bool: """Return True if ``transform`` is the identity over the full storage domain. @@ -5653,25 +5669,25 @@ def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, return True -def _is_complete_chunk( - sub_transform: IndexTransform, chunk_grid: ChunkGrid, chunk_coords: tuple[int, ...] -) -> bool: - """Check if a sub-transform covers an entire chunk.""" +def _is_complete_chunk(sub_transform: IndexTransform, spec: ChunkSpec) -> bool: + """Check if a sub-transform covers an entire chunk. + + ``spec`` is the chunk's already-resolved :class:`ChunkSpec` (the caller looks + it up once and shares it with :func:`_array_spec_from_chunk_spec`). + """ from zarr.core.transforms.output_map import ConstantMap, DimensionMap - spec = chunk_grid[chunk_coords] - if spec is None: - return False + shape = spec.shape for out_dim, m in enumerate(sub_transform.output): if isinstance(m, ConstantMap): # A ConstantMap means a single element is selected along this output dimension, # so the write does not cover the full chunk along this dimension. - chunk_dim_size = spec.shape[out_dim] + chunk_dim_size = shape[out_dim] if chunk_dim_size > 1: return False continue # chunk dim size is 1, so selecting the single element is complete if isinstance(m, DimensionMap): - chunk_dim_size = spec.shape[out_dim] + chunk_dim_size = shape[out_dim] # Compute actual storage range: storage = offset + stride * input_coord dim_lo = sub_transform.domain.inclusive_min[m.input_dimension] dim_hi = sub_transform.domain.exclusive_max[m.input_dimension] @@ -5746,16 +5762,18 @@ async def _get_selection_via_transform( for chunk_coords, sub_transform, out_indices in iter_chunk_transforms( transform, chunk_grid ): + chunk_spec = chunk_grid[chunk_coords] + if chunk_spec is None: + continue chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) drop_axes = da # same for all chunks - is_complete = _is_complete_chunk(sub_transform, chunk_grid, chunk_coords) batch_info.append( ( store_path / metadata.encode_chunk_key(chunk_coords), - _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), + _array_spec_from_chunk_spec(metadata, chunk_spec, _config, prototype), chunk_sel, out_sel, - is_complete, + _is_complete_chunk(sub_transform, chunk_spec), ) ) @@ -5876,16 +5894,18 @@ async def _set_selection_via_transform( batch_info = [] drop_axes: tuple[int, ...] = () for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): + chunk_spec = chunk_grid[chunk_coords] + if chunk_spec is None: + continue chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) drop_axes = da # same for all chunks - is_complete = _is_complete_chunk(sub_transform, chunk_grid, chunk_coords) batch_info.append( ( store_path / metadata.encode_chunk_key(chunk_coords), - _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), + _array_spec_from_chunk_spec(metadata, chunk_spec, _config, prototype), chunk_sel, out_sel, - is_complete, + _is_complete_chunk(sub_transform, chunk_spec), ) ) From 32cd3afedec2cd9fbe6963ae464e8c0803fe25e2 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 16:21:03 +0000 Subject: [PATCH 07/43] perf(lazy): single-pass sub_transform_to_selections with hoisted domain bounds Build chunk_sel and out_sel in a single pass over the output maps instead of two, with the domain min/max tuples hoisted to locals. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_011D1uGKtHP9s7E4WqiUSYG4 --- src/zarr/core/transforms/chunk_resolution.py | 59 ++++++++------------ 1 file changed, 24 insertions(+), 35 deletions(-) diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py index db066a2525..e3997a3071 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -157,51 +157,40 @@ def sub_transform_to_selections( ``(chunk_selection, out_selection, drop_axes)`` """ chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] - drop_axes: list[int] = [] + out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + + # Hoist the per-dimension domain bounds out of the loop (attribute-chain + + # tuple indexing per output dim otherwise). Build chunk_sel and out_sel in a + # single pass; ConstantMap dims are dropped (no out_sel entry). + inclusive_min = sub_transform.domain.inclusive_min + exclusive_max = sub_transform.domain.exclusive_max + n_array_maps = 0 for m in sub_transform.output: - if isinstance(m, ConstantMap): + t = type(m) + if t is ConstantMap: chunk_sel.append(m.offset) - elif isinstance(m, DimensionMap): - dim_lo = sub_transform.domain.inclusive_min[m.input_dimension] - dim_hi = sub_transform.domain.exclusive_max[m.input_dimension] + elif t is DimensionMap: + d = m.input_dimension + dim_lo = inclusive_min[d] + dim_hi = exclusive_max[d] start = m.offset + m.stride * dim_lo stop = m.offset + m.stride * dim_hi if m.stride < 0: start, stop = stop + 1, start + 1 chunk_sel.append(slice(start, stop, m.stride)) - elif isinstance(m, ArrayMap): + out_sel.append(slice(dim_lo, dim_hi)) + else: # ArrayMap + n_array_maps += 1 if m.offset == 0 and m.stride == 1: chunk_sel.append(m.index_array) else: - storage_coords = m.offset + m.stride * m.index_array - chunk_sel.append(storage_coords.astype(np.intp)) + chunk_sel.append((m.offset + m.stride * m.index_array).astype(np.intp)) + # Orthogonal ArrayMap: out_indices holds the surviving positions. + out_sel.append(out_indices if out_indices is not None else slice(0, len(m.index_array))) - # Build out_sel: one entry per non-dropped output dim. - out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + # Vectorized: ≥2 correlated ArrayMaps scatter through a single shared index. + if out_indices is not None and n_array_maps >= 2: + out_sel = [out_indices] - # Vectorized: multiple correlated ArrayMaps share one scatter index - is_vectorized = ( - out_indices is not None - and sum(1 for m in sub_transform.output if isinstance(m, ArrayMap)) >= 2 - ) - - if is_vectorized: - assert out_indices is not None - out_sel.append(out_indices) - else: - for m in sub_transform.output: - if isinstance(m, ConstantMap): - continue - if isinstance(m, DimensionMap): - lo = sub_transform.domain.inclusive_min[m.input_dimension] - hi = sub_transform.domain.exclusive_max[m.input_dimension] - out_sel.append(slice(lo, hi)) - elif isinstance(m, ArrayMap): - if out_indices is not None: - # Orthogonal ArrayMap: out_indices has the surviving positions - out_sel.append(out_indices) - else: - out_sel.append(slice(0, len(m.index_array))) - - return tuple(chunk_sel), tuple(out_sel), tuple(drop_axes) + return tuple(chunk_sel), tuple(out_sel), () From f16ce13740a07ae7e0577d9e6136a36751c9c8e8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 19:14:00 +0200 Subject: [PATCH 08/43] fix(indexing): restore type narrowing and memoize shape on IndexDomain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lazy-indexing perf commits broke mypy: `type(m) is X` in sub_transform_to_selections defeats union narrowing, and the array-level `_shape`/`_is_identity` caches were set via object.__setattr__ but never declared on AsyncArray. Restore `isinstance` narrowing (just as fast, and it narrows). Move shape caching onto the immutable, slotted IndexDomain as a lazy slot-backed memo (excluded from init/repr/eq/hash) instead of denormalizing it onto the array. AsyncArray.shape/ndim/_is_identity become plain derived properties again — no duplicated state, no _resize cache invalidation, and the memo lives on the immutable object so it can never go stale. Also cast the pre-existing coordinate-selection reshape result to NDArrayLikeOrScalar to clear a latent mypy assignment error. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 45 ++++++++++++-------- src/zarr/core/transforms/chunk_resolution.py | 5 +-- src/zarr/core/transforms/domain.py | 15 ++++++- 3 files changed, 42 insertions(+), 23 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 745cbd48df..e274933994 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -384,13 +384,6 @@ def __init__( create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), ) object.__setattr__(self, "_transform", IndexTransform.from_shape(metadata_parsed.shape)) - object.__setattr__(self, "_shape", self._transform.domain.shape) - # A freshly-opened array has the identity transform: input coord i maps to - # storage coord i over the full storage domain. Eager indexing on such an - # array can use the original (legacy) indexers directly, avoiding the - # transform-resolution overhead. Lazy views (created via _with_transform) - # carry a non-identity transform and must go through the transform path. - object.__setattr__(self, "_is_identity", True) @classmethod async def _create( @@ -813,10 +806,6 @@ def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetada object.__setattr__(new, "_chunk_grid", self._chunk_grid) object.__setattr__(new, "codec_pipeline", self.codec_pipeline) object.__setattr__(new, "_transform", transform) - object.__setattr__(new, "_shape", transform.domain.shape) - object.__setattr__( - new, "_is_identity", _transform_is_identity(transform, self.metadata.shape) - ) return new @property @@ -837,7 +826,7 @@ def ndim(self) -> int: int The number of dimensions in the Array. """ - return len(self._shape) + return len(self.shape) @property def shape(self) -> tuple[int, ...]: @@ -848,7 +837,21 @@ def shape(self) -> tuple[int, ...]: tuple The shape of the Array. """ - return self._shape + return self._transform.domain.shape + + @property + def _is_identity(self) -> bool: + """Whether this array's transform is the identity over the full storage domain. + + A freshly-opened or resized array has the identity transform: input coord + ``i`` maps to storage coord ``i`` over the whole storage shape. Eager + indexing on such an array produces the same coordinates the legacy + indexers compute, so it can take the legacy fast path and skip + transform resolution. Lazy views (created via :meth:`_with_transform`) + carry a non-identity transform and must go through the transform path. + Cheap (O(ndim)); the domain's shape lookup it relies on is memoized. + """ + return _transform_is_identity(self._transform, self.metadata.shape) @property def storage_shape(self) -> tuple[int, ...]: @@ -3166,7 +3169,11 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None or self._async_array._is_identity or not is_basic_selection(selection): + if ( + fields is not None + or self._async_array._is_identity + or not is_basic_selection(selection) + ): # Eager (identity) arrays, structured dtypes, and advanced selections # use the original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) @@ -3290,7 +3297,11 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if fields is not None or self._async_array._is_identity or not is_basic_selection(selection): + if ( + fields is not None + or self._async_array._is_identity + or not is_basic_selection(selection) + ): # Eager (identity) arrays, structured dtypes, and advanced selections # use the original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) @@ -3633,7 +3644,7 @@ def get_coordinate_selection( sel_arrays = [np.asarray(s) for s in sel_tuple] sel_shape = np.broadcast_shapes(*(s.shape for s in sel_arrays)) if hasattr(out_array, "shape") and sel_shape != (): - out_array = np.array(out_array).reshape(sel_shape) + out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(sel_shape)) return out_array def set_coordinate_selection( @@ -6479,8 +6490,6 @@ async def _delete_key(key: str) -> None: object.__setattr__(array, "metadata", new_metadata) object.__setattr__(array, "_chunk_grid", new_chunk_grid) object.__setattr__(array, "_transform", IndexTransform.from_shape(new_shape)) - object.__setattr__(array, "_shape", array._transform.domain.shape) - object.__setattr__(array, "_is_identity", True) async def _append( diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py index e3997a3071..b713dd06fe 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -167,10 +167,9 @@ def sub_transform_to_selections( n_array_maps = 0 for m in sub_transform.output: - t = type(m) - if t is ConstantMap: + if isinstance(m, ConstantMap): chunk_sel.append(m.offset) - elif t is DimensionMap: + elif isinstance(m, DimensionMap): d = m.input_dimension dim_lo = inclusive_min[d] dim_hi = exclusive_max[d] diff --git a/src/zarr/core/transforms/domain.py b/src/zarr/core/transforms/domain.py index 90bcc08ace..bca6488ad1 100644 --- a/src/zarr/core/transforms/domain.py +++ b/src/zarr/core/transforms/domain.py @@ -13,7 +13,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any @@ -28,6 +28,11 @@ class IndexDomain: inclusive_min: tuple[int, ...] exclusive_max: tuple[int, ...] labels: tuple[str, ...] | None = None + # Lazily-memoized shape. Excluded from init/repr/eq/hash: it is derived + # state, not part of the domain's identity. The domain is frozen, so the + # value is computed at most once (see ``shape``). ``None`` is the unset + # sentinel; an empty shape caches as ``()``. + _shape: tuple[int, ...] | None = field(default=None, init=False, repr=False, compare=False) def __post_init__(self) -> None: if len(self.inclusive_min) != len(self.exclusive_max): @@ -65,7 +70,13 @@ def origin(self) -> tuple[int, ...]: @property def shape(self) -> tuple[int, ...]: - return tuple(hi - lo for lo, hi in zip(self.inclusive_min, self.exclusive_max, strict=True)) + cached = self._shape + if cached is None: + cached = tuple( + hi - lo for lo, hi in zip(self.inclusive_min, self.exclusive_max, strict=True) + ) + object.__setattr__(self, "_shape", cached) + return cached def contains(self, index: tuple[int, ...]) -> bool: if len(index) != self.ndim: From cd217c703ef76bf8f2e9a31007c29d5fc9bbe460 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 29 Jun 2026 19:14:00 +0200 Subject: [PATCH 09/43] docs: update Array repr doctests for domain= suffix The IndexTransform work made Array.__repr__ emit a `domain={ ... }` suffix, but several docstring examples (from_array, Group.__setitem__, arrays, array_values) were never updated, failing the doctest run. Update them. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/api/synchronous.py | 8 ++++---- src/zarr/core/group.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 8386427b3f..d6775b9104 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -1138,20 +1138,20 @@ def from_array( ... ) >>> arr2 = zarr.from_array(store, data=arr, overwrite=True) >>> arr2 - + >>> asyncio.run(store.clear()) # Remove files generated by test Create an array from an existing NumPy array: >>> import numpy as np >>> zarr.from_array({}, data=np.arange(10000, dtype="i4").reshape(100, 100)) - + Create an array from any array-like object: >>> arr3 = zarr.from_array({}, data=[[1, 2], [3, 4]]) >>> arr3 - + >>> arr3[...] array([[1, 2], [3, 4]]) @@ -1160,7 +1160,7 @@ def from_array( >>> arr4 = zarr.from_array({}, data=[[1, 2], [3, 4]]) >>> arr5 = zarr.from_array({}, data=arr4, write_data=False) >>> arr5 - + >>> arr5[...] array([[0, 0], [0, 0]]) """ diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index de8c8e9a68..766e0755d6 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1946,7 +1946,7 @@ def __setitem__(self, key: str, value: Any) -> None: >>> group = zarr.group() >>> group["foo"] = np.array(zarr.zeros((10,))) >>> group["foo"] - + """ self._sync(self._async_group.setitem(key, value)) @@ -2266,7 +2266,7 @@ def arrays(self) -> Generator[tuple[str, AnyArray], None]: >>> group = zarr.group() >>> subarray = group.create_array("subarray", dtype="i1", shape=(10,), chunks=(10,)) >>> list(group.arrays()) - [('subarray', )] + [('subarray', )] """ for name, async_array in self._sync_iter(self._async_group.arrays()): yield name, Array(async_array) @@ -2295,7 +2295,7 @@ def array_values(self) -> Generator[AnyArray, None]: >>> group = zarr.group() >>> subarray = group.create_array("subarray", dtype="i1", shape=(10,), chunks=(10,)) >>> list(group.array_values()) - [] + [] """ for _, array in self.arrays(): yield array From a2091a3cada18a390cf0c72aa0334e06cf98c54d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 07:02:57 +0200 Subject: [PATCH 10/43] test(lazy-indexing): dense parametrized matrix over 0-D / sharded / unsharded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the single-fixture lazy-indexing tests with a matrix parametrized over array geometry — 0-D, and 1/2/3-D in both unsharded and sharded form — crossed with basic/strided/across-boundary selections plus oindex, vindex, write-through and composition, each checked against a NumPy reference. Sharded configs exercise the inner-chunk read-modify-write path of the transform resolver. Surfaced two issues: - negative slice steps are unsupported (eager raises NegativeStepError, lazy IndexError); added a focused error test rather than a happy-path case. - lazy oindex across >=2 array axes returns a pointwise scatter instead of an outer product (eager oindex is correct). Captured as a strict xfail so it flags when fixed. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_lazy_indexing.py | 392 ++++++++++++++++++++++-------------- 1 file changed, 239 insertions(+), 153 deletions(-) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index fa0d218bd3..3c80ae5960 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -1,164 +1,250 @@ +"""End-to-end tests for the public lazy-indexing API (``Array.lazy``). + +These exercise lazy views — ``arr.lazy[sel]``, ``arr.lazy.oindex[sel]``, +``arr.lazy.vindex[sel]`` — and their write-through and composition behaviour, +in every case comparing against an equivalent NumPy reference. + +Coverage is parametrized over a matrix of array geometries so the same dense +set of selections runs against 0-D arrays, and 1/2/3-D arrays in both +**unsharded** and **sharded** form. Sharded arrays route through inner-chunk +geometry (partial-shard read-modify-write), so running the full matrix against +them guards the transform read/write path where it is most likely to break. +""" + from __future__ import annotations +from dataclasses import dataclass from typing import Any import numpy as np +import numpy.typing as npt import pytest import zarr from zarr.storage import MemoryStore -@pytest.fixture -def arr() -> zarr.Array[Any]: - """Create a 2D array with known data.""" - store = MemoryStore() - a = zarr.create(shape=(20, 30), chunks=(5, 10), dtype="i4", store=store) - data = np.arange(600, dtype="i4").reshape(20, 30) - a[...] = data - return a - - -@pytest.fixture -def data() -> np.ndarray[Any, Any]: - return np.arange(600, dtype="i4").reshape(20, 30) - - -class TestEagerRead: - def test_basic_slice(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[2:8, 5:15] - np.testing.assert_array_equal(result, data[2:8, 5:15]) - - def test_basic_int(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[3] - np.testing.assert_array_equal(result, data[3]) - - def test_basic_int_scalar(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[3, 5] - assert result == data[3, 5] - - def test_ellipsis(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[...] - np.testing.assert_array_equal(result, data) - - def test_strided_slice(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[::2, ::3] - np.testing.assert_array_equal(result, data[::2, ::3]) - - def test_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - idx = np.array([1, 5, 10], dtype=np.intp) - result = arr.oindex[idx, :] - np.testing.assert_array_equal(result, data[idx, :]) - - def test_vindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - idx0 = np.array([1, 5, 10], dtype=np.intp) - idx1 = np.array([2, 8, 15], dtype=np.intp) - result = arr.vindex[idx0, idx1] - np.testing.assert_array_equal(result, data[idx0, idx1]) - - def test_slice_across_chunks(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - """Slice that spans multiple chunks.""" - result = arr[3:17, 8:22] - np.testing.assert_array_equal(result, data[3:17, 8:22]) - - def test_single_element(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[0:1, 0:1] - np.testing.assert_array_equal(result, data[0:1, 0:1]) - - def test_full_read(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - result = arr[:] - np.testing.assert_array_equal(result, data) - - -class TestEagerWrite: - def test_write_slice(self, arr: zarr.Array[Any]) -> None: - arr[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 - result = arr[2:5, 10:20] - np.testing.assert_array_equal(result, np.ones((3, 10), dtype="i4") * 99) - - def test_write_scalar(self, arr: zarr.Array[Any]) -> None: - arr[0, 0] = 42 - assert arr[0, 0] == 42 - - def test_roundtrip(self, arr: zarr.Array[Any]) -> None: - new_data = np.random.randint(0, 100, size=(20, 30), dtype="i4") - arr[...] = new_data - np.testing.assert_array_equal(arr[...], new_data) - - def test_write_across_chunks(self, arr: zarr.Array[Any]) -> None: - """Write spanning multiple chunks.""" - val = np.ones((14, 14), dtype="i4") * 77 - arr[3:17, 8:22] = val - result = arr[3:17, 8:22] - np.testing.assert_array_equal(result, val) - - -class TestLazyRead: - def test_lazy_shape(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.lazy[2:8, 5:15] - assert isinstance(v, zarr.Array) - assert v.shape == (6, 10) - - def test_lazy_resolve(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.lazy[2:8, 5:15] - result = v[...] - np.testing.assert_array_equal(result, data[2:8, 5:15]) - - def test_lazy_np_asarray(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.lazy[2:8] - result = np.asarray(v) - np.testing.assert_array_equal(result, data[2:8]) - - def test_lazy_composition(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.lazy[2:12].lazy[3:8] - assert v.shape == (5, 30) - result = v[...] - np.testing.assert_array_equal(result, data[5:10]) - - def test_lazy_oindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - idx = np.array([1, 5, 10], dtype=np.intp) - v = arr.lazy.oindex[idx, :] - assert isinstance(v, zarr.Array) - assert v.shape == (3, 30) - result = v[...] - np.testing.assert_array_equal(result, data[idx, :]) - - def test_lazy_vindex(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - idx0 = np.array([1, 5, 10], dtype=np.intp) - idx1 = np.array([2, 8, 15], dtype=np.intp) - v = arr.lazy.vindex[idx0, idx1] - assert isinstance(v, zarr.Array) - assert v.shape == (3,) - result = v[...] - np.testing.assert_array_equal(result, data[idx0, idx1]) - - def test_lazy_result_method(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - v = arr.lazy[2:8] - result = v.result() - np.testing.assert_array_equal(result, data[2:8]) - - def test_lazy_across_chunks(self, arr: zarr.Array[Any], data: np.ndarray[Any, Any]) -> None: - """Lazy slice spanning multiple chunks resolves correctly.""" - v = arr.lazy[3:17, 8:22] - result = v[...] - np.testing.assert_array_equal(result, data[3:17, 8:22]) - - -class TestLazyWrite: - def test_lazy_write(self, arr: zarr.Array[Any]) -> None: - arr.lazy[2:5, 10:20] = np.ones((3, 10), dtype="i4") * 99 - result = arr[2:5, 10:20] - np.testing.assert_array_equal(result, np.ones((3, 10), dtype="i4") * 99) - - def test_lazy_oindex_write(self, arr: zarr.Array[Any]) -> None: - idx = np.array([0, 5, 10], dtype=np.intp) - arr.lazy.oindex[idx, :] = np.zeros((3, 30), dtype="i4") - result = arr.oindex[idx, :] - np.testing.assert_array_equal(result, np.zeros((3, 30), dtype="i4")) - - def test_lazy_vindex_write(self, arr: zarr.Array[Any]) -> None: - idx0 = np.array([0, 5, 10], dtype=np.intp) - idx1 = np.array([0, 5, 10], dtype=np.intp) - arr.lazy.vindex[idx0, idx1] = np.array([77, 88, 99], dtype="i4") - result = arr.vindex[idx0, idx1] - np.testing.assert_array_equal(result, np.array([77, 88, 99], dtype="i4")) +@dataclass(frozen=True) +class Config: + """A single array geometry: shape, inner chunk shape, optional shard shape.""" + + id: str + shape: tuple[int, ...] + chunks: tuple[int, ...] + shards: tuple[int, ...] | None + + +CONFIGS = [ + Config("0d", (), (), None), + Config("1d-unsharded", (24,), (4,), None), + Config("1d-sharded", (24,), (4,), (12,)), + Config("2d-unsharded", (20, 30), (5, 10), None), + Config("2d-sharded", (20, 30), (5, 10), (10, 30)), + Config("3d-unsharded", (8, 6, 10), (2, 3, 5), None), + Config("3d-sharded", (8, 6, 10), (2, 3, 5), (4, 6, 10)), +] + +# Basic (slice/int/ellipsis) selections valid for each rank. Includes +# across-boundary and positive-strided cases. Negative steps are unsupported +# by zarr (see TestLazyErrors.test_negative_step_raises), so they are excluded. +BASIC_SELECTIONS: dict[int, list[Any]] = { + 0: [Ellipsis], + 1: [slice(3, 20), 7, Ellipsis, slice(None, None, 2), slice(5, 22)], + 2: [ + (slice(2, 8), slice(5, 15)), + 3, + (3, 5), + Ellipsis, + (slice(None, None, 2), slice(None, None, 3)), + (slice(3, 17), slice(8, 22)), + (slice(None), 5), + ], + 3: [ + (slice(1, 7), slice(0, 4), slice(2, 9)), + 2, + (1, 2, 3), + Ellipsis, + (slice(None, None, 2), slice(None, None, 2), slice(None, None, 2)), + (slice(2, 7), slice(1, 5), slice(3, 9)), + ], +} + +# Per-rank integer index arrays for orthogonal / vectorized fancy indexing. +FANCY_INDICES: dict[int, tuple[npt.NDArray[np.intp], ...]] = { + 1: (np.array([1, 3, 7, 20], dtype=np.intp),), + 2: (np.array([1, 5, 10, 18], dtype=np.intp), np.array([0, 3, 9, 29], dtype=np.intp)), + 3: ( + np.array([1, 5, 7], dtype=np.intp), + np.array([0, 3, 5], dtype=np.intp), + np.array([2, 7, 9], dtype=np.intp), + ), +} + +ND_CONFIGS = [c for c in CONFIGS if len(c.shape) >= 1] + + +def _make(cfg: Config) -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + """Build a zarr array for ``cfg`` and an identical NumPy reference.""" + a = zarr.create_array( + MemoryStore(), shape=cfg.shape, chunks=cfg.chunks, shards=cfg.shards, dtype="i4" + ) + n = int(np.prod(cfg.shape, dtype=int)) # prod(()) == 1 -> a single 0-D element + ref = np.arange(n, dtype="i4").reshape(cfg.shape) + a[...] = ref + return a, ref + + +def _value_like(expected: npt.NDArray[Any]) -> Any: + """A distinctly-valued array (or scalar) shaped like ``expected`` for writes.""" + shape = np.shape(expected) + val = (np.arange(int(np.prod(shape, dtype=int)), dtype="i4").reshape(shape) + 1) * 7 + 1 + return val[()] if shape == () else val # 0-D -> python/np scalar + + +def _fmt(sel: Any) -> str: + if sel is Ellipsis: + return "..." + if isinstance(sel, tuple): + return ",".join(_fmt(s) for s in sel) + if isinstance(sel, slice): + start = "" if sel.start is None else str(sel.start) + stop = "" if sel.stop is None else str(sel.stop) + step = f":{sel.step}" if sel.step is not None else "" + return f"{start}:{stop}{step}" + if isinstance(sel, np.ndarray): + return "x".join(map(str, sel.shape)) + return str(sel) + + +BASIC_CASES = [ + pytest.param(cfg, sel, id=f"{cfg.id}:{_fmt(sel)}") + for cfg in CONFIGS + for sel in BASIC_SELECTIONS[len(cfg.shape)] +] +ND_CASES = [pytest.param(cfg, id=cfg.id) for cfg in ND_CONFIGS] +# Orthogonal indexing across >=2 array axes is currently broken (see +# TestLazyOIndex.test_multi_axis_read_xfail); these are the configs that have +# enough dimensions to exercise it. +MULTI_AXIS_CASES = [pytest.param(cfg, id=cfg.id) for cfg in ND_CONFIGS if len(cfg.shape) >= 2] + + +def _oindex_one_axis(cfg: Config) -> tuple[Any, ...]: + """An orthogonal selection with a single fancy axis (axis 0) and slices elsewhere.""" + idx = FANCY_INDICES[len(cfg.shape)][0] + return (idx, *([slice(None)] * (len(cfg.shape) - 1))) + + +class TestLazyBasicRead: + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_matches_numpy(self, cfg: Config, sel: Any) -> None: + """A lazy basic view reads identically to NumPy (and to eager indexing).""" + a, ref = _make(cfg) + expected = ref[sel] + view = a.lazy[sel] + assert tuple(view.shape) == np.shape(expected) + np.testing.assert_array_equal(view[...], expected) + np.testing.assert_array_equal(a[sel], expected) # eager parity, all geometries + + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_result_and_asarray(self, cfg: Config, sel: Any) -> None: + """``view.result()`` and ``np.asarray(view)`` agree with direct resolution.""" + a, ref = _make(cfg) + expected = ref[sel] + view = a.lazy[sel] + np.testing.assert_array_equal(view.result(), expected) + np.testing.assert_array_equal(np.asarray(view), np.asarray(expected)) + + +class TestLazyBasicWrite: + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_write_through(self, cfg: Config, sel: Any) -> None: + """Assigning through a lazy basic view mutates exactly the selected region.""" + a, ref = _make(cfg) + expected = ref.copy() + val = _value_like(ref[sel]) + expected[sel] = val + a.lazy[sel] = val + np.testing.assert_array_equal(a[...], expected) + + +class TestLazyOIndex: + @pytest.mark.parametrize("cfg", ND_CASES) + def test_read_single_axis(self, cfg: Config) -> None: + """Lazy orthogonal indexing on a single fancy axis matches NumPy.""" + a, ref = _make(cfg) + sel = _oindex_one_axis(cfg) + expected = ref[sel] + view = a.lazy.oindex[sel] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_write_single_axis(self, cfg: Config) -> None: + """Write-through lazy orthogonal indexing on a single fancy axis updates the region.""" + a, ref = _make(cfg) + sel = _oindex_one_axis(cfg) + expected = ref.copy() + val = _value_like(ref[sel]) + expected[sel] = val + a.lazy.oindex[sel] = val + np.testing.assert_array_equal(a[...], expected) + + @pytest.mark.xfail( + strict=True, + reason="lazy oindex with >=2 array axes returns a pointwise scatter, not an outer product", + ) + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_multi_axis_read_xfail(self, cfg: Config) -> None: + """Lazy orthogonal indexing across >=2 array axes should equal ``np.ix_``. + + Currently broken: it collapses to the vectorized (pointwise) scatter and + fills the rest with the fill value. Eager ``oindex`` is correct, so this + is a lazy-path bug; the strict xfail will flag when it is fixed. + """ + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref[np.ix_(*idx)] + view = a.lazy.oindex[idx] + np.testing.assert_array_equal(view[...], expected) + + +class TestLazyVIndex: + @pytest.mark.parametrize("cfg", ND_CASES) + def test_read(self, cfg: Config) -> None: + """Lazy vectorized indexing matches NumPy's point (coordinate) selection.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref[idx] + view = a.lazy.vindex[idx] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_write(self, cfg: Config) -> None: + """Write-through lazy vectorized indexing updates the selected points.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref.copy() + val = _value_like(ref[idx]) + expected[idx] = val + a.lazy.vindex[idx] = val + np.testing.assert_array_equal(a[...], expected) + + +class TestLazyComposition: + @pytest.mark.parametrize("cfg", ND_CASES) + def test_chained_views_compose(self, cfg: Config) -> None: + """Composing two lazy slices equals applying them in sequence on NumPy.""" + a, ref = _make(cfg) + view = a.lazy[1:-1].lazy[1:-1] + expected = ref[1:-1][1:-1] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + +class TestLazyErrors: + def test_negative_step_raises(self) -> None: + """A negative slice step is unsupported and raises on the lazy path.""" + a, _ = _make(CONFIGS[1]) # 1d-unsharded + with pytest.raises(IndexError, match="step must be positive"): + a.lazy[::-1] From 24aaaaf9bb984802bb06ca806c5dc40896d38474 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 07:48:32 +0200 Subject: [PATCH 11/43] test(lazy-indexing): model-based randomized round-trips over sharded/unsharded Add NumPy-oracle round-trip tests (TensorStore driver_testutil pattern): apply random selections to both the zarr array and a NumPy reference and assert they stay equal, parametrized over sharded and unsharded 2-D/3-D grids with one body so the same selections exercise chunk- and shard-boundary read-modify-write. Covers random strided basic writes, coordinate (vindex) writes with distinct points, and single-axis orthogonal writes. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_lazy_indexing.py | 75 +++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 3c80ae5960..3683e565c1 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -134,6 +134,28 @@ def _oindex_one_axis(cfg: Config) -> tuple[Any, ...]: return (idx, *([slice(None)] * (len(cfg.shape) - 1))) +def _rand_slice(rng: np.random.Generator, size: int) -> slice: + """A random non-empty positive-step slice within ``[0, size)``.""" + start = int(rng.integers(0, size)) + stop = int(rng.integers(start + 1, size + 1)) + step = int(rng.integers(1, 4)) + return slice(start, stop, step) + + +def _unique_coords( + rng: np.random.Generator, shape: tuple[int, ...], n: int +) -> tuple[npt.NDArray[np.intp], ...]: + """``n`` distinct coordinate tuples (no duplicate points → write order irrelevant).""" + total = int(np.prod(shape, dtype=int)) + flat = rng.choice(total, size=min(n, total), replace=False) + return tuple(c.astype(np.intp) for c in np.unravel_index(flat, shape)) + + +# Multi-dim geometries used for the randomized model-based round-trips (the 0-D / +# 1-D configs have too small a selection space to be interesting here). +RANDOM_CASES = [pytest.param(cfg, id=cfg.id) for cfg in CONFIGS if len(cfg.shape) >= 2] + + class TestLazyBasicRead: @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) def test_matches_numpy(self, cfg: Config, sel: Any) -> None: @@ -248,3 +270,56 @@ def test_negative_step_raises(self) -> None: a, _ = _make(CONFIGS[1]) # 1d-unsharded with pytest.raises(IndexError, match="step must be positive"): a.lazy[::-1] + + +class TestLazyRandomizedRoundtrip: + """Model-based round-trips: apply random selections to both the zarr array and a + NumPy reference, then assert they stay equal. Parametrized over sharded and + unsharded grids with one body, so the same selections exercise chunk- and + shard-boundary read-modify-write (the pattern TensorStore's driver_testutil uses). + """ + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_basic_writes_track_numpy(self, cfg: Config) -> None: + """Random strided basic writes/reads match a NumPy model across boundaries.""" + rng = np.random.default_rng(0) + a, ref = _make(cfg) + ref = ref.copy() + for _ in range(25): + sel = tuple(_rand_slice(rng, s) for s in cfg.shape) + val = rng.integers(-9999, 9999, size=ref[sel].shape, dtype="i4") + ref[sel] = val + a.lazy[sel] = val + np.testing.assert_array_equal(a.lazy[sel][...], ref[sel]) + np.testing.assert_array_equal(a[...], ref) + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_vindex_writes_track_numpy(self, cfg: Config) -> None: + """Random coordinate (vindex) writes/reads match a NumPy model across boundaries.""" + rng = np.random.default_rng(1) + a, ref = _make(cfg) + ref = ref.copy() + for _ in range(25): + idx = _unique_coords(rng, cfg.shape, int(rng.integers(1, 7))) + val = rng.integers(-9999, 9999, size=idx[0].shape, dtype="i4") + ref[idx] = val + a.lazy.vindex[idx] = val + np.testing.assert_array_equal(a.lazy.vindex[idx][...], ref[idx]) + np.testing.assert_array_equal(a[...], ref) + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_oindex_single_axis_writes_track_numpy(self, cfg: Config) -> None: + """Random single-fancy-axis orthogonal writes/reads match a NumPy model.""" + rng = np.random.default_rng(2) + a, ref = _make(cfg) + ref = ref.copy() + size0 = cfg.shape[0] + for _ in range(25): + k = int(rng.integers(1, size0 + 1)) + idx0 = rng.choice(size0, size=k, replace=False).astype(np.intp) + sel: Any = (idx0, *(_rand_slice(rng, s) for s in cfg.shape[1:])) + val = rng.integers(-9999, 9999, size=ref[sel].shape, dtype="i4") + ref[sel] = val + a.lazy.oindex[sel] = val + np.testing.assert_array_equal(a.lazy.oindex[sel][...], ref[sel]) + np.testing.assert_array_equal(a[...], ref) From c51ef58c903817d1bd7415f9d452b344f03d5d3e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 08:17:20 +0200 Subject: [PATCH 12/43] fix(lazy): correct orthogonal (oindex) indexing across multiple array axes `arr.lazy.oindex[(i0, i1)]` with >= 2 integer-array axes returned a pointwise diagonal scattered into the first row instead of the outer product, because `ArrayMap` carried no input-dimension binding: `_intersect` could not tell orthogonal (oindex, distinct input dims) from vectorized (vindex, shared input dim) and routed any >= 2 ArrayMaps through the vectorized scatter. Give `ArrayMap` an `input_dimension` (set by oindex, None for vindex) and thread it through translate/intersect/reindex/json. `_intersect` now treats >= 2 ArrayMaps as vectorized only when all are unbound; orthogonal arrays are intersected per axis and return a per-output-dim dict of surviving positions, which `sub_transform_to_selections` turns into np.ix_-style selections (as the legacy OrthogonalIndexer does). `_is_vectorized_transform` replaces the `all(isinstance ArrayMap)` flat-buffer test so oindex-all-axes keeps a multi-dimensional buffer. Lazy oindex now matches eager/numpy for read (sharded + unsharded, 0-D..3-D) and unsharded write. Sharded orthogonal multi-array writes remain unsupported in the sharding partial-write codec (eager oindex raises identically); pinned by a strict xfail. Flips the previous strict-xfail regression test to passing. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 30 +++-- src/zarr/core/transforms/chunk_resolution.py | 42 +++++-- src/zarr/core/transforms/json.py | 4 + src/zarr/core/transforms/output_map.py | 15 +++ src/zarr/core/transforms/transform.py | 114 +++++++++++++++---- tests/test_lazy_indexing.py | 52 ++++++--- 6 files changed, 206 insertions(+), 51 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b0ca8fc6c7..76d1fcef87 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5648,6 +5648,19 @@ def _get_chunk_spec( return _array_spec_from_chunk_spec(metadata, spec, array_config, prototype) +def _is_vectorized_transform(transform: IndexTransform) -> bool: + """Return True for a vectorized (vindex) transform: >= 2 correlated ArrayMaps. + + Correlated ArrayMaps (``input_dimension is None``) are jointly indexed and + scatter through a single flat index, so the output buffer is flattened during + read/write. Orthogonal ArrayMaps (oindex), each bound to a distinct input + dimension, form an outer product and keep a multi-dimensional buffer — even + when every output happens to be an ArrayMap (e.g. ``oindex[i0, i1]``). + """ + array_maps = [m for m in transform.output if isinstance(m, ArrayMap)] + return len(array_maps) >= 2 and all(m.input_dimension is None for m in array_maps) + + def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, ...]) -> bool: """Return True if ``transform`` is the identity over the full storage domain. @@ -5735,14 +5748,13 @@ async def _get_selection_via_transform( out_shape = transform.domain.shape - # When the transform has ArrayMap outputs, chunk resolution produces - # flat scatter indices (out_indices). The output buffer must be 1D - # during the read, then reshaped to out_shape afterwards. - # For vectorized indexing (all outputs are ArrayMaps), chunk resolution - # produces flat scatter indices. The buffer must be 1D during the read. - # For orthogonal indexing (mixed ArrayMap + DimensionMap), the buffer - # stays multi-dimensional — each dim gets its own out_sel entry. - needs_flat_buffer = all(isinstance(m, ArrayMap) for m in transform.output) + # Only vectorized indexing (>= 2 correlated ArrayMaps, input_dimension None) + # scatters through a single flat index, so the buffer must be 1D during the + # read and reshaped to out_shape afterwards. Orthogonal indexing — including + # the case where every output is an ArrayMap bound to a distinct input + # dimension (oindex[i0, i1]) — keeps a multi-dimensional buffer, one out_sel + # entry per dim. + needs_flat_buffer = _is_vectorized_transform(transform) buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape # Setup output buffer @@ -5851,7 +5863,7 @@ async def _set_selection_via_transform( # Validate value shape against selection shape sel_shape = transform.domain.shape - needs_flat_buffer = all(isinstance(m, ArrayMap) for m in transform.output) + needs_flat_buffer = _is_vectorized_transform(transform) if hasattr(value, "shape") and value.shape != () and value.shape != sel_shape: if needs_flat_buffer: # For ArrayMap (coordinate/vindex), values are flattened so check total size diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py index b713dd06fe..e57d71d408 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -42,10 +42,14 @@ from zarr.core.chunk_grids import ChunkGrid +OutIndices = ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None +) + ChunkTransformResult = tuple[ tuple[int, ...], IndexTransform, - np.ndarray[Any, np.dtype[np.intp]] | None, + OutIndices, ] @@ -134,7 +138,7 @@ def iter_chunk_transforms( def sub_transform_to_selections( sub_transform: IndexTransform, - out_indices: np.ndarray[Any, np.dtype[np.intp]] | None = None, + out_indices: OutIndices = None, ) -> tuple[ tuple[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], tuple[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], @@ -156,14 +160,38 @@ def sub_transform_to_selections( tuple ``(chunk_selection, out_selection, drop_axes)`` """ + inclusive_min = sub_transform.domain.inclusive_min + exclusive_max = sub_transform.domain.exclusive_max + + # Orthogonal outer product: >= 2 ArrayMaps each bound to a distinct input + # dimension. out_indices is a per-output-dim dict of surviving positions. The + # codec applies chunk_array[chunk_sel] / out[out_sel] with NumPy semantics, so + # build np.ix_-style selections (mirroring the legacy OrthogonalIndexer): one + # 1-D selector per dimension, expanded to an open mesh. ConstantMap dims are + # size-1 in chunk space and squeezed out via drop_axes. + if isinstance(out_indices, dict): + chunk_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + out_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + drop_axes: list[int] = [] + for out_dim, m in enumerate(sub_transform.output): + if isinstance(m, ConstantMap): + chunk_arrays.append(np.array([m.offset], dtype=np.intp)) + drop_axes.append(out_dim) + elif isinstance(m, DimensionMap): + rng = np.arange(inclusive_min[m.input_dimension], exclusive_max[m.input_dimension]) + chunk_arrays.append((m.offset + m.stride * rng).astype(np.intp)) + out_arrays.append(rng.astype(np.intp)) + else: # ArrayMap + idx = m.index_array.ravel() + chunk_arrays.append((m.offset + m.stride * idx).astype(np.intp)) + out_arrays.append(out_indices[out_dim]) + return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] - # Hoist the per-dimension domain bounds out of the loop (attribute-chain + - # tuple indexing per output dim otherwise). Build chunk_sel and out_sel in a - # single pass; ConstantMap dims are dropped (no out_sel entry). - inclusive_min = sub_transform.domain.inclusive_min - exclusive_max = sub_transform.domain.exclusive_max + # Single-pass build for the basic / single-array / vectorized cases. + # ConstantMap dims are dropped (no out_sel entry). n_array_maps = 0 for m in sub_transform.output: diff --git a/src/zarr/core/transforms/json.py b/src/zarr/core/transforms/json.py index 59a81cc1e8..79fe6e58b8 100644 --- a/src/zarr/core/transforms/json.py +++ b/src/zarr/core/transforms/json.py @@ -110,6 +110,9 @@ def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: result = {"offset": m.offset, "index_array": m.index_array.tolist()} if m.stride != 1: result["stride"] = m.stride + # Present only for orthogonal (oindex) arrays; omitted for vectorized. + if m.input_dimension is not None: + result["input_dimension"] = m.input_dimension return result raise TypeError(f"Unknown output map type: {type(m)}") @@ -122,6 +125,7 @@ def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: index_array=np.asarray(data["index_array"], dtype=np.intp), offset=data.get("offset", 0), stride=data.get("stride", 1), + input_dimension=data.get("input_dimension"), ) if "input_dimension" in data: diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py index 5e17a0ae82..c53402be4d 100644 --- a/src/zarr/core/transforms/output_map.py +++ b/src/zarr/core/transforms/output_map.py @@ -73,11 +73,26 @@ class ArrayMap: Represents ``{offset + stride * index_array[i] : i in input_range}``. Arises from fancy indexing (e.g., ``arr[[1, 5, 9]]`` or boolean masks). + + ``input_dimension`` records which single input dimension indexes the array, + binding it the way :class:`DimensionMap` is bound. It distinguishes the two + flavours of multi-array fancy indexing: + + - **orthogonal** (``oindex``): each array indexes a *distinct* input + dimension, so ``input_dimension`` is set; the result is their outer + product. + - **vectorized** (``vindex``): the arrays are correlated and jointly + indexed by the same (possibly multi-dimensional) input range, so + ``input_dimension`` is ``None``. + + This binding is what lets chunk resolution tell an outer product from a + pointwise scatter when more than one ``ArrayMap`` is present. """ index_array: npt.NDArray[np.intp] offset: int = 0 stride: int = 1 + input_dimension: int | None = None OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index ab6a24464f..1080c67b89 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -30,7 +30,7 @@ import math from dataclasses import dataclass -from typing import Any, Literal +from typing import Any, Literal, cast import numpy as np @@ -135,14 +135,23 @@ def __repr__(self) -> str: def intersect( self, output_domain: IndexDomain - ) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: + ) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] + | np.ndarray[Any, np.dtype[np.intp]] + | None, + ] + | None + ): """Restrict this transform to storage coordinates within output_domain. - Returns ``(restricted_transform, surviving_indices)`` or None if empty. + Returns ``(restricted_transform, out_indices)`` or None if empty. - ``surviving_indices`` is an integer array of which input positions - survived the intersection (for ArrayMap dimensions), or None if all - positions survived (ConstantMap/DimensionMap only). + ``out_indices`` carries the surviving output positions: ``None`` when all + positions survive (ConstantMap/DimensionMap only), a single integer array + for one ArrayMap (or correlated/vectorized ArrayMaps), or a dict keyed by + output dimension for >= 2 orthogonal ArrayMaps (an outer product). """ return _intersect(self, output_domain) @@ -168,6 +177,7 @@ def translate(self, shift: tuple[int, ...]) -> IndexTransform: index_array=m.index_array, offset=m.offset + s, stride=m.stride, + input_dimension=m.input_dimension, ) ) return IndexTransform(domain=self.domain, output=tuple(new_output)) @@ -186,7 +196,13 @@ def vindex(self) -> _VIndexHelper: def _intersect( transform: IndexTransform, output_domain: IndexDomain -) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): """Intersect a transform with an output domain (e.g., a chunk's bounds). For each output dimension, restrict to storage coordinates within @@ -207,16 +223,23 @@ def _intersect( f"transform output rank ({transform.output_rank})" ) - # Check if we have correlated ArrayMaps (vectorized) + # Correlated ArrayMaps (vindex) are jointly indexed (input_dimension is None); + # >= 2 of them means a vectorized intersection. Orthogonal ArrayMaps (oindex) + # each bind a distinct input dimension and are handled per-dimension below. array_dims = [i for i, m in enumerate(transform.output) if isinstance(m, ArrayMap)] - if len(array_dims) >= 2: - return _intersect_vectorized(transform, output_domain, array_dims) - - # Orthogonal: intersect each output dimension independently + vectorized_dims = [ + i for i in array_dims if cast("ArrayMap", transform.output[i]).input_dimension is None + ] + if len(vectorized_dims) >= 2: + return _intersect_vectorized(transform, output_domain, vectorized_dims) + + # Orthogonal: intersect each output dimension independently. Multiple + # ArrayMaps bound to distinct input dimensions form an outer product, so each + # array dimension's surviving *output* positions are tracked separately. new_min = list(transform.domain.inclusive_min) new_max = list(transform.domain.exclusive_max) new_output: list[OutputIndexMap] = [] - surviving_indices: np.ndarray[Any, np.dtype[np.intp]] | None = None + out_positions: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} for out_dim, m in enumerate(transform.output): lo = output_domain.inclusive_min[out_dim] @@ -258,24 +281,43 @@ def _intersect( elif isinstance(m, ArrayMap): storage = m.offset + m.stride * m.index_array mask = (storage >= lo) & (storage < hi) - if not np.any(mask): + survivors = np.nonzero(mask.ravel())[0].astype(np.intp) + if survivors.size == 0: return None - surviving_indices = np.nonzero(mask.ravel())[0].astype(np.intp) - filtered = m.index_array.ravel()[surviving_indices] + filtered = m.index_array.ravel()[survivors] new_output.append( ArrayMap( index_array=filtered, offset=m.offset, stride=m.stride, + input_dimension=m.input_dimension, ) ) + # Narrow this array's own input dimension to the surviving count. + if m.input_dimension is not None: + new_min[m.input_dimension] = 0 + new_max[m.input_dimension] = int(survivors.size) + out_positions[out_dim] = survivors new_domain = IndexDomain( inclusive_min=tuple(new_min), exclusive_max=tuple(new_max), ) result = IndexTransform(domain=new_domain, output=tuple(new_output)) - return (result, surviving_indices) + + # Hand back the surviving output positions in the shape the bridge expects: + # None (no arrays), a single vector (one array), or a per-output-dim dict + # (>= 2 orthogonal arrays → outer product). + out_indices: ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None + ) + if len(out_positions) == 0: + out_indices = None + elif len(out_positions) == 1: + out_indices = next(iter(out_positions.values())) + else: + out_indices = out_positions + return (result, out_indices) def _intersect_vectorized( @@ -322,6 +364,7 @@ def _intersect_vectorized( index_array=filtered, offset=m.offset, stride=m.stride, + input_dimension=m.input_dimension, ) ) elif isinstance(m, ConstantMap): @@ -539,7 +582,17 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra raise RuntimeError(f"unexpected: dimension {d} not handled") elif isinstance(m, ArrayMap): new_arr = _reindex_array(m.index_array, normalized, transform.domain) - new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) return IndexTransform(domain=new_domain, output=tuple(new_output)) @@ -649,6 +702,10 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: index_array=dim_array[d], offset=m.offset, stride=m.stride, + # Each oindex array indexes its own input dimension; this + # binding is what marks the selection orthogonal (outer + # product) rather than vectorized. + input_dimension=old_to_new_dim[d], ) ) elif d in dim_slice_params: @@ -665,7 +722,17 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: raise RuntimeError(f"unexpected: dimension {d} not handled") elif isinstance(m, ArrayMap): new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) - new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) return IndexTransform(domain=new_domain, output=tuple(new_output)) @@ -821,7 +888,14 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: ) elif isinstance(m, ArrayMap): new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) - new_output.append(ArrayMap(index_array=new_arr, offset=m.offset, stride=m.stride)) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) return IndexTransform(domain=new_domain, output=tuple(new_output)) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 3683e565c1..82539108dc 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -122,10 +122,16 @@ def _fmt(sel: Any) -> str: for sel in BASIC_SELECTIONS[len(cfg.shape)] ] ND_CASES = [pytest.param(cfg, id=cfg.id) for cfg in ND_CONFIGS] -# Orthogonal indexing across >=2 array axes is currently broken (see -# TestLazyOIndex.test_multi_axis_read_xfail); these are the configs that have -# enough dimensions to exercise it. -MULTI_AXIS_CASES = [pytest.param(cfg, id=cfg.id) for cfg in ND_CONFIGS if len(cfg.shape) >= 2] +# Configs with >= 2 dimensions, i.e. enough to exercise orthogonal indexing +# across multiple array axes (the outer-product case). +MULTI_AXIS = [cfg for cfg in ND_CONFIGS if len(cfg.shape) >= 2] +MULTI_AXIS_CASES = [pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS] +MULTI_AXIS_UNSHARDED_CASES = [ + pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS if cfg.shards is None +] +MULTI_AXIS_SHARDED_CASES = [ + pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS if cfg.shards is not None +] def _oindex_one_axis(cfg: Config) -> tuple[Any, ...]: @@ -211,24 +217,40 @@ def test_write_single_axis(self, cfg: Config) -> None: a.lazy.oindex[sel] = val np.testing.assert_array_equal(a[...], expected) - @pytest.mark.xfail( - strict=True, - reason="lazy oindex with >=2 array axes returns a pointwise scatter, not an outer product", - ) @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) - def test_multi_axis_read_xfail(self, cfg: Config) -> None: - """Lazy orthogonal indexing across >=2 array axes should equal ``np.ix_``. - - Currently broken: it collapses to the vectorized (pointwise) scatter and - fills the rest with the fill value. Eager ``oindex`` is correct, so this - is a lazy-path bug; the strict xfail will flag when it is fixed. - """ + def test_multi_axis_read(self, cfg: Config) -> None: + """Lazy orthogonal indexing across >=2 array axes is the outer product (np.ix_).""" a, ref = _make(cfg) idx = FANCY_INDICES[len(cfg.shape)] expected = ref[np.ix_(*idx)] view = a.lazy.oindex[idx] + assert tuple(view.shape) == expected.shape np.testing.assert_array_equal(view[...], expected) + @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) + def test_multi_axis_write(self, cfg: Config) -> None: + """Write-through lazy orthogonal indexing across >=2 array axes (unsharded).""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref.copy() + val = _value_like(ref[np.ix_(*idx)]) + expected[np.ix_(*idx)] = val + a.lazy.oindex[idx] = val + np.testing.assert_array_equal(a[...], expected) + + @pytest.mark.xfail( + strict=True, + reason="orthogonal multi-array writes are unsupported by the sharding " + "partial-write codec (eager oindex raises the same way); a strict xfail " + "flags if sharded support is added", + ) + @pytest.mark.parametrize("cfg", MULTI_AXIS_SHARDED_CASES) + def test_multi_axis_write_sharded_unsupported(self, cfg: Config) -> None: + """Sharded orthogonal multi-array write — a pre-existing codec limitation.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + a.lazy.oindex[idx] = _value_like(ref[np.ix_(*idx)]) + class TestLazyVIndex: @pytest.mark.parametrize("cfg", ND_CASES) From 6f1a829f72f75ef561b5de4b2aae46d26f3f32e0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 10:00:28 +0200 Subject: [PATCH 13/43] fix(indexing): honor the transform for fancy/field-less selections on lazy views MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit roborev review of the branch surfaced four correctness gaps, all rooted in how indexing on a non-identity (lazy view) Array was dispatched: 1. pop_fields returned [] (not None) for field-less *tuple* selections, while the non-tuple branch returned None. The get/set_*_selection methods route on `fields is not None`, so any view-method call with a tuple selection (`v[i, j]`, `v.oindex[...]`, `v.vindex[...]`) was treated as a field request and fell back to the legacy storage-grid indexer, silently ignoring the view's transform. Make pop_fields return None for no fields (consistent with the non-tuple branch); this fixes the whole family at the root. 2. get/set_orthogonal_selection on a view sent fancy selections to the legacy indexer (built from the view shape against the storage grid) — silently wrong data. Route non-identity views through the transform, choosing basic mode for plain int/slice selections (so integer axes drop) and orthogonal mode for fancy ones. 3. _normalize_negative_indices over-counted ellipsis dims by one, so a negative index after an ellipsis on a view (`v[..., -1]`) was left unnormalized. 4. vindex is coordinate-only; mixing a slice now raises VindexInvalidSelectionError (parity with eager) instead of crashing deep in chunk resolution. Array.result(prototype=...) now forwards its prototype instead of dropping it. Adds TestLazyViewMethods covering view-object indexing (basic tuple, int-drop, oindex, vindex, negative-after-ellipsis, write-through) — the surface the existing matrix missed because it only exercised the arr.lazy.* accessor. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 42 ++++++++------ src/zarr/core/indexing.py | 12 +++- src/zarr/core/transforms/transform.py | 21 +++++-- tests/test_lazy_indexing.py | 83 +++++++++++++++++++++++++++ 4 files changed, 133 insertions(+), 25 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 76d1fcef87..b0ae2f7f51 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3166,21 +3166,23 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if ( - fields is not None - or self._async_array._is_identity - or not is_basic_selection(selection) - ): - # Eager (identity) arrays, structured dtypes, and advanced selections - # use the original indexer path directly. + if fields is not None or self._async_array._is_identity: + # Eager (identity) arrays and structured-dtype field selection use the + # original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) return sync( self.async_array._get_selection( indexer=indexer, out=out, fields=fields, prototype=prototype ) ) + # Lazy view (non-identity transform): route through the transform so the + # view's offset/stride are honored. Use basic mode for plain int/slice + # selections (so integer axes drop), orthogonal mode for fancy selections. selection = _normalize_negative_indices(selection, self.shape) - transform = selection_to_transform(selection, self._async_array._transform, "basic") + mode: Literal["basic", "orthogonal"] = ( + "basic" if is_basic_selection(selection) else "orthogonal" + ) + transform = selection_to_transform(selection, self._async_array._transform, mode) return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_orthogonal_selection( @@ -3294,20 +3296,22 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - if ( - fields is not None - or self._async_array._is_identity - or not is_basic_selection(selection) - ): - # Eager (identity) arrays, structured dtypes, and advanced selections - # use the original indexer path directly. + if fields is not None or self._async_array._is_identity: + # Eager (identity) arrays and structured-dtype field selection use the + # original indexer path directly. indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) sync( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) return + # Lazy view (non-identity transform): route through the transform so the + # view's offset/stride are honored. Use basic mode for plain int/slice + # selections (so integer axes drop), orthogonal mode for fancy selections. selection = _normalize_negative_indices(selection, self.shape) - transform = selection_to_transform(selection, self._async_array._transform, "basic") + mode: Literal["basic", "orthogonal"] = ( + "basic" if is_basic_selection(selection) else "orthogonal" + ) + transform = selection_to_transform(selection, self._async_array._transform, mode) sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_mask_selection( @@ -4005,10 +4009,10 @@ def lazy(self) -> _LazyIndexAccessor: def result(self, prototype: BufferPrototype | None = None) -> NDArrayLikeOrScalar: """Read and return the data for this array view. - Equivalent to ``self[...]`` but more explicit for lazy views. - Named after `tensorstore.Future.result`. + Equivalent to ``self[...]`` but more explicit for lazy views, and forwards + ``prototype``. Named after `tensorstore.Future.result`. """ - return self[...] + return self.get_basic_selection(Ellipsis, prototype=prototype) def resize(self, new_shape: ShapeLike) -> None: """ diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f6eb495cd9..26eb06c996 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1442,8 +1442,16 @@ def pop_fields(selection: SelectionWithFields) -> tuple[Fields | None, Selection return None, cast("Selection", selection) else: # multiple items, split fields from selection items - fields: Fields = [f for f in selection if isinstance(f, str)] - fields = fields[0] if len(fields) == 1 else fields + field_names = [f for f in selection if isinstance(f, str)] + # No fields -> None (consistent with the non-tuple branch above), so callers + # can use `fields is not None` to mean "a field was requested". + fields: Fields | None + if len(field_names) == 0: + fields = None + elif len(field_names) == 1: + fields = field_names[0] + else: + fields = field_names selection_tuple = tuple(s for s in selection if not isinstance(s, str)) selection = cast( "Selection", selection_tuple[0] if len(selection_tuple) == 1 else selection_tuple diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 1080c67b89..861915b6ab 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -36,6 +36,7 @@ from zarr.core.transforms.domain import IndexDomain from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.errors import VindexInvalidSelectionError @dataclass(frozen=True, slots=True) @@ -911,10 +912,12 @@ def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: else: selection_tuple = selection - # Count real dimensions (non-None, non-Ellipsis) to map each entry to a shape dim - has_ellipsis = any(s is Ellipsis for s in selection_tuple) + # Count real dimensions (non-None, non-Ellipsis) to map each entry to a shape + # dim. The ellipsis covers the dims not consumed by those entries; since + # n_non_newaxis already excludes the ellipsis itself, that count is exactly + # len(shape) - n_non_newaxis (no +1). n_non_newaxis = sum(1 for s in selection_tuple if s is not None and s is not Ellipsis) - n_ellipsis_dims = len(shape) - n_non_newaxis + (1 if has_ellipsis else 0) + n_ellipsis_dims = len(shape) - n_non_newaxis result: list[Any] = [] dim = 0 @@ -963,7 +966,17 @@ def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) """ items = selection if isinstance(selection, tuple) else (selection,) for sel in items: - if sel is Ellipsis or isinstance(sel, (int, np.integer, slice)): + if isinstance(sel, slice): + # vindex is coordinate-only (matches eager zarr): every axis needs an + # integer/boolean array, never a slice. Orthogonal (oindex) allows slices. + if mode == "vectorized": + raise VindexInvalidSelectionError( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {selection!r}" + ) + continue + if sel is Ellipsis or isinstance(sel, (int, np.integer)): continue if isinstance(sel, (list, np.ndarray)): continue diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 82539108dc..8906f5c61d 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -15,12 +15,14 @@ from dataclasses import dataclass from typing import Any +from unittest import mock import numpy as np import numpy.typing as npt import pytest import zarr +from zarr.core.buffer import default_buffer_prototype from zarr.storage import MemoryStore @@ -286,6 +288,72 @@ def test_chained_views_compose(self, cfg: Config) -> None: np.testing.assert_array_equal(view[...], expected) +class TestLazyViewMethods: + """Indexing methods called on a lazy *view* object (``v = arr.lazy[...]``) must + honor the view's transform, not silently fall back to the storage grid. + + The accessor (``arr.lazy.oindex[...]``) was covered elsewhere, but the methods + on the *returned* non-identity Array (``v.oindex[...]``, ``v[..., -1]``) route + through ``Array``'s own dispatch, which had correctness gaps. + """ + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_view_oindex_respects_transform(self, cfg: Config) -> None: + """``v.oindex[idx]`` on a sub-view reads from the view, not the base array.""" + a, ref = _make(cfg) + n0 = cfg.shape[0] + cut = n0 // 2 + vslice = (slice(cut, n0), *([slice(None)] * (len(cfg.shape) - 1))) + vref = ref[vslice] + v = a.lazy[vslice] + idx = np.array([0, vref.shape[0] - 1], dtype=np.intp) + osel: Any = (idx, *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v.oindex[osel], vref[osel]) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_negative_index_after_ellipsis(self, cfg: Config) -> None: + """``v[..., -1]`` on a view selects the last element of the last axis.""" + a, ref = _make(cfg) + v = a.lazy[1 : cfg.shape[0]] + vref = ref[1 : cfg.shape[0]] + np.testing.assert_array_equal(v[..., -1], vref[..., -1]) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_basic_tuple_and_int(self, cfg: Config) -> None: + """Basic tuple selections on a view (incl. integer axes that drop) honor the view.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + vref = ref[cut:] + full: Any = (slice(0, 2), *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[full], vref[full]) + intsel: Any = (1, *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[intsel], vref[intsel]) # int axis drops + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_vindex(self, cfg: Config) -> None: + """``v.vindex[...]`` (coordinate selection) on a view honors the view.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + vref = ref[cut:] + idx = tuple(np.array([0, 1, 2], dtype=np.intp) for _ in cfg.shape) + np.testing.assert_array_equal(v.vindex[idx], vref[idx]) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) + def test_view_write_through_tuple(self, cfg: Config) -> None: + """Writing through a view with a basic tuple selection lands in view coords.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + expected = ref.copy() + wsel: Any = (slice(0, 2), *([slice(None)] * (len(cfg.shape) - 1))) + val = _value_like(ref[cut:][wsel]) + expected[cut:][wsel] = val + v[wsel] = val + np.testing.assert_array_equal(a[...], expected) + + class TestLazyErrors: def test_negative_step_raises(self) -> None: """A negative slice step is unsupported and raises on the lazy path.""" @@ -293,6 +361,21 @@ def test_negative_step_raises(self) -> None: with pytest.raises(IndexError, match="step must be positive"): a.lazy[::-1] + def test_vindex_with_slice_rejected(self) -> None: + """vindex is coordinate-only; mixing a slice raises (parity with eager).""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded + with pytest.raises(IndexError): + a.lazy.vindex[(np.array([0, 1], dtype=np.intp), slice(None))] + + def test_result_threads_prototype(self) -> None: + """``result(prototype=...)`` forwards the prototype rather than dropping it.""" + a, ref = _make(CONFIGS[3]) # 2d-unsharded + proto = default_buffer_prototype() + with mock.patch.object(type(a), "get_basic_selection", autospec=True) as gbs: + gbs.return_value = ref + a.result(prototype=proto) + assert gbs.call_args.kwargs.get("prototype") is proto + class TestLazyRandomizedRoundtrip: """Model-based round-trips: apply random selections to both the zarr array and a From d38205cee56250b9705beff5a87f99593618fc15 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 10:07:45 +0200 Subject: [PATCH 14/43] test(lazy-indexing): hypothesis property test for the lazy-view surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a NumPy-oracle property test (hypothesis.extra.numpy, already a dep) that builds a lazy view from a random non-negative slice window and asserts that the view's *methods* — v[...] (basic, incl. negative indices) and v.oindex[...] (orthogonal) — match NumPy on the viewed data. This is the durable guard for the class of view-method dispatch bugs the deterministic matrix missed; it extends the existing oracle pattern in test_properties.py (simple_arrays + basic_indices + orthogonal_indices). Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 47 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..d2988bd843 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -448,3 +448,50 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert serialized_complex_float_is_valid(asdict_dict["fill_value"]) elif dtype_native.kind in ("M", "m") and np.isnat(meta.fill_value): assert asdict_dict["fill_value"] == -9223372036854775808 + + +@st.composite +def _window(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[slice, ...]: + """A tuple of non-negative, full-rank slice windows (one per axis). + + Used to build a lazy view via the accessor, which treats negative indices as + literal coordinates (TensorStore convention) — so the view-defining selection + must stay non-negative. The view *methods* tested below normalize negatives. + """ + out: list[slice] = [] + for size in shape: + if size == 0: + out.append(slice(0, 0)) + continue + start = draw(st.integers(min_value=0, max_value=size - 1)) + stop = draw(st.integers(min_value=start + 1, max_value=size)) + out.append(slice(start, stop)) + return tuple(out) + + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_lazy_view_indexing_matches_numpy(data: st.DataObject) -> None: + """Indexing the *methods* of a lazy view must match NumPy on the viewed data. + + Regression guard for the class of bugs where view-object indexing + (``v[...]``, ``v.oindex[...]``, ``v.vindex[...]``) bypassed the view's + transform. The view is built from a non-negative slice window; the second + selection is drawn freely (basic, orthogonal, or coordinate). + """ + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) + nparray = zarray[:] + base = data.draw(_window(nparray.shape)) + view = zarray.lazy[base] + vref = nparray[base] + assume(vref.ndim >= 1) + + # basic indexing on the view (slices/ints/ellipsis, incl. negative indices) + sub = data.draw(basic_indices(shape=vref.shape)) + assert_array_equal(np.asarray(view[sub]), np.asarray(vref[sub])) + + # orthogonal indexing on the view + if all(s > 0 for s in vref.shape): + zidx, npidx = data.draw(orthogonal_indices(shape=vref.shape)) + assert_array_equal(np.asarray(view.oindex[zidx]), np.asarray(vref[npidx])) From d9aad9f643a111e8c1cd5e1c7fbe5733b82af1e3 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 10:39:35 +0200 Subject: [PATCH 15/43] test(indexing): comprehensive cross-path indexing parity harness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add test_indexing_parity: a NumPy-oracle property test over the cartesian product {basic, oindex, vindex, mask} × {eager array, lazy view} × {out=None, out=buffer}. Every lazy/view indexing bug found in review was "the lazy/view path diverges from NumPy/eager for some (method, parameter) combination"; enumerating that surface catches the whole class instead of one case per review round. Doubles as a guard for the existing eager indexing paths (incl. out=), independent of lazy indexing. The harness immediately surfaced two fixes: - _get_selection_via_transform validated a provided out= buffer against the multi-dimensional out_shape, but vectorized reads scatter through a flat index; validate against the (flat) buffer_shape instead, matching the eager contract. - _is_vectorized_transform required >= 2 coordinate ArrayMaps, missing a single coordinate map with a multi-dimensional index array (vindex[idx_2d] on a 1-D view), which also scatters flat. Flag a flat buffer whenever any coordinate (input_dimension is None) ArrayMap is present. Also adds a deterministic view-method out= regression test. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 29 +++++++----- tests/test_lazy_indexing.py | 20 ++++++++ tests/test_properties.py | 93 +++++++++++++++++++++++++++++-------- 3 files changed, 110 insertions(+), 32 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b0ae2f7f51..e383d9eda0 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5653,16 +5653,19 @@ def _get_chunk_spec( def _is_vectorized_transform(transform: IndexTransform) -> bool: - """Return True for a vectorized (vindex) transform: >= 2 correlated ArrayMaps. - - Correlated ArrayMaps (``input_dimension is None``) are jointly indexed and - scatter through a single flat index, so the output buffer is flattened during - read/write. Orthogonal ArrayMaps (oindex), each bound to a distinct input - dimension, form an outer product and keep a multi-dimensional buffer — even - when every output happens to be an ArrayMap (e.g. ``oindex[i0, i1]``). + """Return True for a vectorized (vindex/coordinate) transform. + + Coordinate ArrayMaps (``input_dimension is None``) are jointly indexed over + the broadcast input domain and scatter through a single flat index, so the + output buffer is flattened during read/write and reshaped afterwards. This + holds whenever *any* such map is present — including a single coordinate map + with a multi-dimensional index array (``vindex[idx_2d]`` on a 1-D array), + whose flat result still needs a flat buffer. Orthogonal ArrayMaps (oindex), + each bound to a distinct input dimension, form an outer product and keep a + multi-dimensional buffer — even when every output is an ArrayMap + (e.g. ``oindex[i0, i1]``). """ - array_maps = [m for m in transform.output if isinstance(m, ArrayMap)] - return len(array_maps) >= 2 and all(m.input_dimension is None for m in array_maps) + return any(isinstance(m, ArrayMap) and m.input_dimension is None for m in transform.output) def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, ...]) -> bool: @@ -5761,13 +5764,15 @@ async def _get_selection_via_transform( needs_flat_buffer = _is_vectorized_transform(transform) buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape - # Setup output buffer + # Setup output buffer. Vectorized reads scatter through a flat index, so the + # caller must supply a flat out buffer (matching the eager path); otherwise it + # is the full multi-dimensional out_shape. if out is not None: if not isinstance(out, NDBuffer): raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") - if out.shape != out_shape: + if out.shape != buffer_shape: raise ValueError( - f"shape of out argument doesn't match. Expected {out_shape}, got {out.shape}" + f"shape of out argument doesn't match. Expected {buffer_shape}, got {out.shape}" ) out_buffer = out else: diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 8906f5c61d..676382faba 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -340,6 +340,26 @@ def test_view_vindex(self, cfg: Config) -> None: idx = tuple(np.array([0, 1, 2], dtype=np.intp) for _ in cfg.shape) np.testing.assert_array_equal(v.vindex[idx], vref[idx]) + def test_view_vindex_with_flat_out_buffer(self) -> None: + """vindex with a multi-dim result and out= on a view uses a flat out buffer. + + Vectorized indexing scatters through a single flat index, so (as in the + eager path) the out buffer must be flat (shape = number of points). + """ + a, ref = _make(CONFIGS[3]) # 2d-unsharded + v = a.lazy[2:18] + vref = ref[2:18] + i0 = np.array([[0, 1], [2, 3]], dtype=np.intp) + i1 = np.array([[0, 5], [10, 15]], dtype=np.intp) + expected = vref[i0, i1] + buf = default_buffer_prototype().nd_buffer.empty( + shape=(expected.size,), dtype=np.dtype("i4") + ) + v.get_coordinate_selection((i0, i1), out=buf) + np.testing.assert_array_equal( + np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected + ) + @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) def test_view_write_through_tuple(self, cfg: Config) -> None: """Writing through a view with a basic tuple selection lands in view coords.""" diff --git a/tests/test_properties.py b/tests/test_properties.py index d2988bd843..1129558f0a 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -469,29 +469,82 @@ def _window(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[slice, ...]: return tuple(out) +# The indexing modes and which Array method implements each. vindex/mask are +# "vectorized" — they scatter through a single flat index, so an out= buffer must +# be flat (number of selected points) rather than the multi-dimensional result. +_INDEX_MODES = ("basic", "oindex", "vindex", "mask") +_VECTORIZED_MODES = frozenset({"vindex", "mask"}) + + +def _draw_indexer(data: st.DataObject, mode: str, shape: tuple[int, ...]) -> tuple[Any, Any]: + """Draw a (zarr_selection, numpy_selection) pair valid for ``mode`` on ``shape``.""" + if mode == "basic": + sel = data.draw(basic_indices(shape=shape)) + return sel, sel + if mode == "oindex": + return data.draw(orthogonal_indices(shape=shape)) + if mode == "vindex": + idx = data.draw( + npst.integer_array_indices( + shape=shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) + ) + ) + return idx, idx + if mode == "mask": + m = data.draw(npst.arrays(dtype=np.bool_, shape=st.just(shape))) + return m, m + raise AssertionError(mode) + + +def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: + """Read ``zsel`` from ``target`` via the get-method for ``mode``.""" + if mode == "basic": + return target.get_basic_selection(zsel, out=out) + if mode == "oindex": + return target.get_orthogonal_selection(zsel, out=out) + if mode == "vindex": + return target.get_coordinate_selection(zsel, out=out) + if mode == "mask": + return target.get_mask_selection(zsel, out=out) + raise AssertionError(mode) + + @settings(deadline=None) @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) -async def test_lazy_view_indexing_matches_numpy(data: st.DataObject) -> None: - """Indexing the *methods* of a lazy view must match NumPy on the viewed data. - - Regression guard for the class of bugs where view-object indexing - (``v[...]``, ``v.oindex[...]``, ``v.vindex[...]``) bypassed the view's - transform. The view is built from a non-negative slice window; the second - selection is drawn freely (basic, orthogonal, or coordinate). +async def test_indexing_parity(data: st.DataObject) -> None: + """Every indexing method matches NumPy on both an eager array and a lazy view. + + Comprehensive cross-path oracle covering the cartesian product of + {basic, oindex, vindex, mask} × {eager array, lazy view} × {out=None, out=}. + All of the lazy-view indexing bugs found in review were "the lazy/view path + diverges from NumPy/eager for some (method, parameter) combination"; this + enumerates that surface so the whole class is caught, not one case at a time. + The view is built from a non-negative slice window (the accessor treats + negative indices as literal); view methods normalize negatives like NumPy. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) nparray = zarray[:] - base = data.draw(_window(nparray.shape)) - view = zarray.lazy[base] - vref = nparray[base] - assume(vref.ndim >= 1) - - # basic indexing on the view (slices/ints/ellipsis, incl. negative indices) - sub = data.draw(basic_indices(shape=vref.shape)) - assert_array_equal(np.asarray(view[sub]), np.asarray(vref[sub])) - - # orthogonal indexing on the view - if all(s > 0 for s in vref.shape): - zidx, npidx = data.draw(orthogonal_indices(shape=vref.shape)) - assert_array_equal(np.asarray(view.oindex[zidx]), np.asarray(vref[npidx])) + window = data.draw(_window(nparray.shape)) + view = zarray.lazy[window] + vref = nparray[window] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + + for target, ref in ((zarray, nparray), (view, vref)): + if ref.ndim == 0: + continue + # integer_array_indices / orthogonal strategies can't handle 0-size dims. + if mode != "basic" and not all(s > 0 for s in ref.shape): + continue + zsel, npsel = _draw_indexer(data, mode, ref.shape) + expected = np.asarray(ref[npsel]) + + # 1) plain read + assert_array_equal(np.asarray(_get(target, mode, zsel)), expected) + + # 2) read into an out= buffer (flat for vectorized modes, full otherwise) + if expected.ndim >= 1 and expected.size > 0: + buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else expected.shape + buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=nparray.dtype) + _get(target, mode, zsel, out=buf) + assert_array_equal(np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected) From 3734357ee4cb49a2a5b1257048bd1c07557ec2b4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 10:52:56 +0200 Subject: [PATCH 16/43] test(strategies): reusable indexing-selection fixtures (indexers, windows) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promote the inlined test-data generators from the indexing parity test into zarr.testing.strategies as documented, reusable strategies: - indexers(mode, shape): one strategy returning a (zarr_selection, numpy_selection) pair for any of basic / oindex / vindex / mask, so indexing tests are written once and parametrized over mode instead of re-deriving selection setup. - windows(shape): non-negative, full-rank slice windows for building lazy views. test_indexing_parity now draws from these shared strategies. The reusable data-generation building blocks — not any single test — are the durable asset: new indexing tests (and downstream users of zarr.testing) compose them rather than reinventing setup. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/testing/strategies.py | 59 ++++++++++++++++++++++++++++++++++ tests/test_properties.py | 47 +++------------------------ 2 files changed, 64 insertions(+), 42 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index f2c83677cc..74de984440 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -598,6 +598,65 @@ def orthogonal_indices( return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) +IndexMode = Literal["basic", "oindex", "vindex", "mask"] + + +@st.composite +def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: + """A non-negative, full-rank tuple of slice windows — one per axis. + + Useful for building a lazy view (``arr.lazy[window]``) whose rank matches the + source array. The lazy accessor treats negative indices as literal + coordinates (TensorStore convention), so a view-defining selection must stay + non-negative. Each axis gets ``start:stop`` with ``0 <= start < stop <= size`` + (an empty ``0:0`` slice for a zero-length axis). + """ + out: list[slice] = [] + for size in shape: + if size == 0: + out.append(slice(0, 0)) + continue + start = draw(st.integers(min_value=0, max_value=size - 1)) + stop = draw(st.integers(min_value=start + 1, max_value=size)) + out.append(slice(start, stop)) + return tuple(out) + + +@st.composite +def indexers(draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...]) -> tuple[Any, Any]: + """A ``(zarr_selection, numpy_selection)`` pair for ``mode`` on ``shape``. + + One strategy covering every indexing mode, so a test can be written once and + parametrized over mode instead of re-deriving selection setup per test: + + - ``"basic"`` — slices / ints / ellipsis (no newaxis, no negative slices) + - ``"oindex"`` — per-axis integer arrays or slices (orthogonal / outer product) + - ``"vindex"`` — broadcast integer coordinate arrays (vectorized) + - ``"mask"`` — a boolean array of ``shape`` + + The two returned selections differ only for ``oindex`` (zarr's per-axis + spelling vs the ``np.ix_``-style spelling numpy needs); for the other modes + the same object indexes both a zarr array and its numpy reference. Fancy modes + (``oindex``/``vindex``) require non-zero-size axes; ``mask``/``basic`` do not. + """ + if mode == "basic": + sel = draw(basic_indices(shape=shape)) + return sel, sel + if mode == "oindex": + return draw(orthogonal_indices(shape=shape)) + if mode == "vindex": + idx = draw( + npst.integer_array_indices( + shape=shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) + ) + ) + return idx, idx + if mode == "mask": + m = draw(npst.arrays(dtype=np.bool_, shape=st.just(shape))) + return m, m + raise ValueError(f"unknown indexing mode: {mode!r}") + + @st.composite def block_indices( draw: st.DrawFn, *, chunk_sizes: tuple[tuple[int, ...], ...] diff --git a/tests/test_properties.py b/tests/test_properties.py index 1129558f0a..663c3d575b 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -28,11 +28,13 @@ block_indices, block_test_arrays, complex_rectilinear_arrays, + indexers, numpy_arrays, orthogonal_indices, rectilinear_arrays, simple_arrays, stores, + windows, zarr_formats, ) @@ -450,25 +452,6 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert asdict_dict["fill_value"] == -9223372036854775808 -@st.composite -def _window(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[slice, ...]: - """A tuple of non-negative, full-rank slice windows (one per axis). - - Used to build a lazy view via the accessor, which treats negative indices as - literal coordinates (TensorStore convention) — so the view-defining selection - must stay non-negative. The view *methods* tested below normalize negatives. - """ - out: list[slice] = [] - for size in shape: - if size == 0: - out.append(slice(0, 0)) - continue - start = draw(st.integers(min_value=0, max_value=size - 1)) - stop = draw(st.integers(min_value=start + 1, max_value=size)) - out.append(slice(start, stop)) - return tuple(out) - - # The indexing modes and which Array method implements each. vindex/mask are # "vectorized" — they scatter through a single flat index, so an out= buffer must # be flat (number of selected points) rather than the multi-dimensional result. @@ -476,26 +459,6 @@ def _window(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[slice, ...]: _VECTORIZED_MODES = frozenset({"vindex", "mask"}) -def _draw_indexer(data: st.DataObject, mode: str, shape: tuple[int, ...]) -> tuple[Any, Any]: - """Draw a (zarr_selection, numpy_selection) pair valid for ``mode`` on ``shape``.""" - if mode == "basic": - sel = data.draw(basic_indices(shape=shape)) - return sel, sel - if mode == "oindex": - return data.draw(orthogonal_indices(shape=shape)) - if mode == "vindex": - idx = data.draw( - npst.integer_array_indices( - shape=shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) - ) - ) - return idx, idx - if mode == "mask": - m = data.draw(npst.arrays(dtype=np.bool_, shape=st.just(shape))) - return m, m - raise AssertionError(mode) - - def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: """Read ``zsel`` from ``target`` via the get-method for ``mode``.""" if mode == "basic": @@ -516,7 +479,7 @@ async def test_indexing_parity(data: st.DataObject) -> None: """Every indexing method matches NumPy on both an eager array and a lazy view. Comprehensive cross-path oracle covering the cartesian product of - {basic, oindex, vindex, mask} × {eager array, lazy view} × {out=None, out=}. + {basic, oindex, vindex, mask} by {eager array, lazy view} by {out=None, out=}. All of the lazy-view indexing bugs found in review were "the lazy/view path diverges from NumPy/eager for some (method, parameter) combination"; this enumerates that surface so the whole class is caught, not one case at a time. @@ -525,7 +488,7 @@ async def test_indexing_parity(data: st.DataObject) -> None: """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) nparray = zarray[:] - window = data.draw(_window(nparray.shape)) + window = data.draw(windows(shape=nparray.shape)) view = zarray.lazy[window] vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) @@ -536,7 +499,7 @@ async def test_indexing_parity(data: st.DataObject) -> None: # integer_array_indices / orthogonal strategies can't handle 0-size dims. if mode != "basic" and not all(s > 0 for s in ref.shape): continue - zsel, npsel = _draw_indexer(data, mode, ref.shape) + zsel, npsel = data.draw(indexers(mode=mode, shape=ref.shape)) expected = np.asarray(ref[npsel]) # 1) plain read From 889787a25550da577c2d3c0a60ba37b771d7ee25 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 10:52:56 +0200 Subject: [PATCH 17/43] fix(indexing): reject block selection on a lazy view instead of corrupting get_block_selection / set_block_selection were not updated for non-identity (lazy view) arrays: BlockIndexer used the view shape against the storage chunk grid, so v.blocks[i] silently read/wrote raw storage blocks, ignoring the view's offset/stride. Block selection addresses the storage chunk grid and is ill-defined on a view, so raise NotImplementedError (pointing users to materialize the view) rather than return/write the wrong region. Element-level lazy indexing (basic/oindex/vindex/mask) is unaffected. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 15 +++++++++++++++ tests/test_lazy_indexing.py | 9 +++++++++ 2 files changed, 24 insertions(+) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index e383d9eda0..e502188913 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3872,6 +3872,14 @@ def get_block_selection( """ if prototype is None: prototype = default_buffer_prototype() + if not self._async_array._is_identity: + # Block selection addresses the storage chunk grid; it is ill-defined + # on a lazy view (offset/stride) and the legacy indexer would read raw + # storage blocks, ignoring the transform. Reject rather than corrupt. + raise NotImplementedError( + "block selection is not supported on a lazy view; materialize the " + "view first (e.g. zarr.array(view[...]))" + ) indexer = BlockIndexer(selection, self.shape, self._chunk_grid) return sync( self.async_array._get_selection( @@ -3972,6 +3980,13 @@ def set_block_selection( """ if prototype is None: prototype = default_buffer_prototype() + if not self._async_array._is_identity: + # See get_block_selection: block selection is ill-defined on a lazy + # view, so reject rather than write to the wrong storage region. + raise NotImplementedError( + "block selection is not supported on a lazy view; materialize the " + "view first (e.g. zarr.array(view[...]))" + ) indexer = BlockIndexer(selection, self.shape, self._chunk_grid) sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 676382faba..621e86f7a9 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -381,6 +381,15 @@ def test_negative_step_raises(self) -> None: with pytest.raises(IndexError, match="step must be positive"): a.lazy[::-1] + def test_block_selection_on_view_rejected(self) -> None: + """Block selection is ill-defined on a lazy view and must raise, not corrupt.""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded + v = a.lazy[5:15] + with pytest.raises(NotImplementedError, match="block selection"): + _ = v.blocks[0] + with pytest.raises(NotImplementedError, match="block selection"): + v.blocks[0] = 0 + def test_vindex_with_slice_rejected(self) -> None: """vindex is coordinate-only; mixing a slice raises (parity with eager).""" a, _ = _make(CONFIGS[3]) # 2d-unsharded From 2027a776b711e1037b9dd72dfe14edaee14448c5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 11:00:48 +0200 Subject: [PATCH 18/43] test(indexing): make the parity harness consumer-agnostic (eager-mergeable) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split the cross-path indexing oracle so the test improvements can land ahead of lazy indexing: - assert_indexing_matches_numpy(target, ref, mode, data): a consumer-agnostic helper — `target` is any zarr.Array (eager array or lazy view), checked against a NumPy reference for a drawn selection, with and without out=. - test_indexing_parity: eager-only, no reference to Array.lazy. Comprehensively guards the existing eager indexing API ({basic,oindex,vindex,mask} × out=) and can merge on its own. - test_lazy_view_indexing_parity: reuses the same helper for a lazy view, and is skipped unless Array.lazy exists — so it's inert on a tree without the feature. The lazy view is "just another zarr.Array" flowing through the same oracle, so adding the lazy consumer is additive rather than a rewrite. windows() is reworded as a generic sub-region strategy (not lazy-specific). Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/testing/strategies.py | 10 ++-- tests/test_properties.py | 86 +++++++++++++++++++++++----------- 2 files changed, 63 insertions(+), 33 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 74de984440..41355c8e42 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -605,11 +605,11 @@ def orthogonal_indices( def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: """A non-negative, full-rank tuple of slice windows — one per axis. - Useful for building a lazy view (``arr.lazy[window]``) whose rank matches the - source array. The lazy accessor treats negative indices as literal - coordinates (TensorStore convention), so a view-defining selection must stay - non-negative. Each axis gets ``start:stop`` with ``0 <= start < stop <= size`` - (an empty ``0:0`` slice for a zero-length axis). + A rank-preserving sub-region selection: each axis gets ``start:stop`` with + ``0 <= start < stop <= size`` (an empty ``0:0`` slice for a zero-length axis). + Bounds stay non-negative so the window is valid for any consumer, including + those that treat negative indices as literal coordinates rather than + from-the-end (e.g. building a sub-array view). """ out: list[slice] = [] for size in shape: diff --git a/tests/test_properties.py b/tests/test_properties.py index 663c3d575b..86c2312b96 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -472,19 +472,67 @@ def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: raise AssertionError(mode) +def assert_indexing_matches_numpy( + target: zarr.Array, ref: np.ndarray[Any, Any], mode: str, data: st.DataObject +) -> None: + """Assert ``target``'s ``mode`` indexing matches NumPy on ``ref``. + + Consumer-agnostic: ``target`` is any ``zarr.Array`` (an ordinary eager array + or a lazy view) and ``ref`` is the NumPy array it should behave like. Draws a + selection for ``mode`` and checks both a plain read and a read into an ``out=`` + buffer (flat for the vectorized vindex/mask modes, full-shape otherwise). + """ + if ref.ndim == 0: + return + # integer_array_indices / orthogonal strategies can't handle 0-size dims. + if mode != "basic" and not all(s > 0 for s in ref.shape): + return + zsel, npsel = data.draw(indexers(mode=mode, shape=ref.shape)) + expected = np.asarray(ref[npsel]) + + # 1) plain read + assert_array_equal(np.asarray(_get(target, mode, zsel)), expected) + + # 2) read into an out= buffer (flat for vectorized modes, full otherwise) + if expected.ndim >= 1 and expected.size > 0: + buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else expected.shape + buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=expected.dtype) + _get(target, mode, zsel, out=buf) + assert_array_equal(np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected) + + @settings(deadline=None) @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) async def test_indexing_parity(data: st.DataObject) -> None: - """Every indexing method matches NumPy on both an eager array and a lazy view. - - Comprehensive cross-path oracle covering the cartesian product of - {basic, oindex, vindex, mask} by {eager array, lazy view} by {out=None, out=}. - All of the lazy-view indexing bugs found in review were "the lazy/view path - diverges from NumPy/eager for some (method, parameter) combination"; this - enumerates that surface so the whole class is caught, not one case at a time. - The view is built from a non-negative slice window (the accessor treats - negative indices as literal); view methods normalize negatives like NumPy. + """Every indexing method on an eager array matches NumPy (all modes, with/without out=). + + Comprehensive oracle over {basic, oindex, vindex, mask} by {out=None, out=}. + Deliberately independent of lazy indexing so it can guard the eager API on its + own; ``test_lazy_view_indexing_parity`` reuses the same oracle for the lazy + consumer (a view is just another ``zarr.Array``). + """ + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) + nparray = zarray[:] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + assert_indexing_matches_numpy(zarray, nparray, mode, data) + + +@pytest.mark.skipif( + not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" +) +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: + """The eager indexing oracle, applied to a lazy view (the lazy consumer). + + A view (built from a non-negative window) is just another ``zarr.Array``, so + it flows through the same ``assert_indexing_matches_numpy`` harness. The lazy + indexing bugs found in review were all "the view path diverges from NumPy for + some (method, parameter) combination" — enumerating that surface here catches + the class. This test is skipped until ``Array.lazy`` exists, so the harness + above can merge ahead of the lazy feature. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) nparray = zarray[:] @@ -492,22 +540,4 @@ async def test_indexing_parity(data: st.DataObject) -> None: view = zarray.lazy[window] vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) - - for target, ref in ((zarray, nparray), (view, vref)): - if ref.ndim == 0: - continue - # integer_array_indices / orthogonal strategies can't handle 0-size dims. - if mode != "basic" and not all(s > 0 for s in ref.shape): - continue - zsel, npsel = data.draw(indexers(mode=mode, shape=ref.shape)) - expected = np.asarray(ref[npsel]) - - # 1) plain read - assert_array_equal(np.asarray(_get(target, mode, zsel)), expected) - - # 2) read into an out= buffer (flat for vectorized modes, full otherwise) - if expected.ndim >= 1 and expected.size > 0: - buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else expected.shape - buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=nparray.dtype) - _get(target, mode, zsel, out=buf) - assert_array_equal(np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected) + assert_indexing_matches_numpy(view, vref, mode, data) From 8dc45e25d21e708baca388dd2f64ef88225cd93e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 11:13:16 +0200 Subject: [PATCH 19/43] test(indexing): fold per-mode eager tests into the single parity oracle Replace test_basic_indexing / test_oindex / test_vindex / test_mask_indexing with one enumerated oracle, test_indexing_parity, covering every mode against a NumPy reference: sync read, read into an out= buffer, async read, and write. The historical per-mode write guards are preserved in one place (vindex setitem is chunk-dependent for repeated coords; oindex/mask writes to sharded arrays are unsupported per GH2834; oindex repeated-index writes are unspecified). Selection generation comes from the shared indexers() strategy, and the read half is shared with the lazy-view test via assert_read_matches_numpy. Dedicated tests remain for inputs the oracle's strategies don't generate: complex rectilinear arrays and block indexing. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 246 +++++++++++++++------------------------ 1 file changed, 91 insertions(+), 155 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 86c2312b96..410aeba991 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -30,7 +30,6 @@ complex_rectilinear_arrays, indexers, numpy_arrays, - orthogonal_indices, rectilinear_arrays, simple_arrays, stores, @@ -118,33 +117,10 @@ def test_array_creates_implicit_groups(array): ) -# this decorator removes timeout; not ideal but it should avoid intermittent CI failures - - -@settings(deadline=None) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -@given(data=st.data()) -async def test_basic_indexing(data: st.DataObject) -> None: - zarray = data.draw(st.one_of(simple_arrays(), rectilinear_arrays())) - nparray = zarray[:] - indexer = data.draw(basic_indices(shape=nparray.shape)) - - # sync get - actual = zarray[indexer] - assert_array_equal(nparray[indexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.getitem(indexer) - assert_array_equal(nparray[indexer], actual) - - # sync set - new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) - zarray[indexer] = new_data - nparray[indexer] = new_data - assert_array_equal(nparray, zarray[:]) - - # TODO test async setitem? +# Eager basic/oindex/vindex/mask indexing is exercised comprehensively (read, +# out=, async read, and write) by the single oracle ``test_indexing_parity`` +# defined below. Special array shapes that the oracle's strategies don't cover +# keep dedicated tests (complex rectilinear arrays here; block indexing later). @settings(deadline=None) @@ -156,105 +132,6 @@ async def test_basic_indexing_complex_rectilinear(data: st.DataObject) -> None: assert_array_equal(nparray[indexer], zarray[indexer]) -@given(data=st.data()) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -async def test_oindex(data: st.DataObject) -> None: - # integer_array_indices can't handle 0-size dimensions. - zarray = data.draw( - st.one_of( - simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), - rectilinear_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1, max_side=20)), - ) - ) - nparray = zarray[:] - zindexer, npindexer = data.draw(orthogonal_indices(shape=nparray.shape)) - - # sync get - actual = zarray.oindex[zindexer] - assert_array_equal(nparray[npindexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.oindex.getitem(zindexer) - assert_array_equal(nparray[npindexer], actual) - - # sync get - assume(zarray.shards is None) # GH2834 - for idxr in npindexer: - if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr).size: - # behaviour of setitem with repeated indices is not guaranteed in practice - assume(False) - new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) - nparray[npindexer] = new_data - zarray.oindex[zindexer] = new_data - assert_array_equal(nparray, zarray[:]) - - # note: async oindex setitem not yet implemented - - -@given(data=st.data()) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -async def test_vindex(data: st.DataObject) -> None: - # integer_array_indices can't handle 0-size dimensions. - zarray = data.draw( - st.one_of( - simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), - rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), - ) - ) - nparray = zarray[:] - indexer = data.draw( - npst.integer_array_indices( - shape=nparray.shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) - ) - ) - - # sync get - actual = zarray.vindex[indexer] - assert_array_equal(nparray[indexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.vindex.getitem(indexer) - assert_array_equal(nparray[indexer], actual) - - # sync set - # FIXME! - # when the indexer is such that a value gets overwritten multiple times, - # I think the output depends on chunking. - # new_data = data.draw(npst.arrays(shape=st.just(actual.shape), dtype=nparray.dtype)) - # nparray[indexer] = new_data - # zarray.vindex[indexer] = new_data - # assert_array_equal(nparray, zarray[:]) - - # note: async vindex setitem not yet implemented - - -@settings(deadline=None) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -@given(data=st.data()) -def test_mask_indexing(data: st.DataObject) -> None: - zarray = data.draw(st.one_of(simple_arrays(), rectilinear_arrays())) - nparray = zarray[:] - mask = data.draw(npst.arrays(dtype=np.bool_, shape=st.just(nparray.shape))) - - expected = nparray[mask] - - # sync get, via both the dedicated method and the vindex interface - assert_array_equal(expected, zarray.get_mask_selection(mask)) - assert_array_equal(expected, zarray.vindex[mask]) - - # sync set, via both interfaces - assume(zarray.shards is None) # GH2834 - new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) - nparray[mask] = new_data - zarray.set_mask_selection(mask, new_data) - assert_array_equal(nparray, zarray[:]) - - zarray.vindex[mask] = new_data - assert_array_equal(nparray, zarray[:]) - - @settings(deadline=None) @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) @@ -472,28 +349,56 @@ def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: raise AssertionError(mode) -def assert_indexing_matches_numpy( - target: zarr.Array, ref: np.ndarray[Any, Any], mode: str, data: st.DataObject +def _async_get(async_array: Any, mode: str, zsel: Any) -> Any: + """The async read coroutine for ``mode`` (vindex/mask share the vectorized accessor).""" + if mode == "basic": + return async_array.getitem(zsel) + if mode == "oindex": + return async_array.oindex.getitem(zsel) + return async_array.vindex.getitem(zsel) + + +def _setitem(zarray: zarr.Array, mode: str, zsel: Any, value: Any) -> None: + """Write ``value`` at ``zsel`` via the set-method for ``mode``.""" + if mode == "basic": + zarray[zsel] = value + elif mode == "oindex": + zarray.oindex[zsel] = value + elif mode == "mask": + zarray.set_mask_selection(zsel, value) + else: + raise AssertionError(f"writes are not exercised for mode {mode!r}") + + +def _has_repeated_indices(npsel: Any) -> bool: + """True if any integer-array component selects a coordinate more than once.""" + sel = npsel if isinstance(npsel, tuple) else (npsel,) + return any(isinstance(i, np.ndarray) and i.size != np.unique(i).size for i in sel) + + +def _eligible(mode: str, shape: tuple[int, ...]) -> bool: + """Whether ``mode`` can be exercised on ``shape``. + + Rank-0 arrays have no interesting selections; the fancy modes + (oindex/vindex/mask via integer/boolean arrays) can't handle zero-size axes. + """ + if len(shape) == 0: + return False + return mode == "basic" or all(s > 0 for s in shape) + + +def assert_read_matches_numpy( + target: zarr.Array, ref: np.ndarray[Any, Any], mode: str, zsel: Any, npsel: Any ) -> None: - """Assert ``target``'s ``mode`` indexing matches NumPy on ``ref``. + """Assert ``target``'s read of ``zsel`` (mode) matches ``ref[npsel]``, with/without out=. - Consumer-agnostic: ``target`` is any ``zarr.Array`` (an ordinary eager array - or a lazy view) and ``ref`` is the NumPy array it should behave like. Draws a - selection for ``mode`` and checks both a plain read and a read into an ``out=`` - buffer (flat for the vectorized vindex/mask modes, full-shape otherwise). + Consumer-agnostic: ``target`` is any ``zarr.Array`` (an eager array or a lazy + view) and ``ref`` is the NumPy array it must behave like. The ``out=`` buffer + is flat for the vectorized vindex/mask modes (which scatter through a flat + index) and full-shape otherwise. """ - if ref.ndim == 0: - return - # integer_array_indices / orthogonal strategies can't handle 0-size dims. - if mode != "basic" and not all(s > 0 for s in ref.shape): - return - zsel, npsel = data.draw(indexers(mode=mode, shape=ref.shape)) expected = np.asarray(ref[npsel]) - - # 1) plain read assert_array_equal(np.asarray(_get(target, mode, zsel)), expected) - - # 2) read into an out= buffer (flat for vectorized modes, full otherwise) if expected.ndim >= 1 and expected.size > 0: buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else expected.shape buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=expected.dtype) @@ -505,17 +410,45 @@ def assert_indexing_matches_numpy( @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) async def test_indexing_parity(data: st.DataObject) -> None: - """Every indexing method on an eager array matches NumPy (all modes, with/without out=). + """Single eager-indexing oracle: every method matches NumPy. - Comprehensive oracle over {basic, oindex, vindex, mask} by {out=None, out=}. + Covers {basic, oindex, vindex, mask} against a NumPy reference: sync read, + read into an ``out=`` buffer, async read, and (where defined) write. Replaces + the per-mode test_basic/oindex/vindex/mask tests with one enumeration. Deliberately independent of lazy indexing so it can guard the eager API on its - own; ``test_lazy_view_indexing_parity`` reuses the same oracle for the lazy - consumer (a view is just another ``zarr.Array``). + own; ``test_lazy_view_indexing_parity`` reuses the read half for lazy views. """ - zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) + zarray = data.draw( + st.one_of( + simple_arrays(shapes=npst.array_shapes(max_dims=4)), + rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), + ) + ) nparray = zarray[:] mode = data.draw(st.sampled_from(_INDEX_MODES)) - assert_indexing_matches_numpy(zarray, nparray, mode, data) + if not _eligible(mode, nparray.shape): + return + zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) + + # read: sync + out= (shared with the lazy-view test) and async + assert_read_matches_numpy(zarray, nparray, mode, zsel, npsel) + expected = np.asarray(nparray[npsel]) + assert_array_equal(np.asarray(await _async_get(zarray._async_array, mode, zsel)), expected) + + # write, mirroring the historical per-mode guards: + # - vindex setitem is chunk-dependent for repeated coords (skipped) + # - oindex/mask writes to sharded arrays are unsupported (GH2834) + # - oindex writes with repeated indices are unspecified + if mode == "vindex": + return + if mode in ("oindex", "mask") and zarray.shards is not None: + return + if mode == "oindex" and _has_repeated_indices(npsel): + return + new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) + _setitem(zarray, mode, zsel, new_data) + nparray[npsel] = new_data + assert_array_equal(nparray, zarray[:]) @pytest.mark.skipif( @@ -525,14 +458,14 @@ async def test_indexing_parity(data: st.DataObject) -> None: @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: - """The eager indexing oracle, applied to a lazy view (the lazy consumer). + """The eager read oracle, applied to a lazy view (the lazy consumer). A view (built from a non-negative window) is just another ``zarr.Array``, so - it flows through the same ``assert_indexing_matches_numpy`` harness. The lazy + it flows through the same ``assert_read_matches_numpy`` harness. The lazy indexing bugs found in review were all "the view path diverges from NumPy for some (method, parameter) combination" — enumerating that surface here catches - the class. This test is skipped until ``Array.lazy`` exists, so the harness - above can merge ahead of the lazy feature. + the class. Skipped until ``Array.lazy`` exists, so the eager oracle can merge + ahead of the lazy feature. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) nparray = zarray[:] @@ -540,4 +473,7 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: view = zarray.lazy[window] vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) - assert_indexing_matches_numpy(view, vref, mode, data) + if not _eligible(mode, vref.shape): + return + zsel, npsel = data.draw(indexers(mode=mode, shape=vref.shape)) + assert_read_matches_numpy(view, vref, mode, zsel, npsel) From ab2e35293a67775e90965ab71f336add253e77e1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 12:12:37 +0200 Subject: [PATCH 20/43] refactor(indexing): address review nits (comments, dead code, test coverage) Low-severity follow-ups from the branch review (jobs 303/304), no behavior change: - array.py: correct the stale "needs_flat_buffer" comment (any coordinate ArrayMap triggers the flat buffer, not just >= 2). - array.py set_coordinate_selection: normalize empty fields with an explicit `== [] or == ()` check instead of a falsy `if not fields`. - transform._intersect_vectorized: drop the unused n_points variable. - test_properties: restore the vindex[mask] read+write equivalence coverage that the per-mode fold dropped (the two mask spellings must agree). - strategies.indexers: fix the docstring's zero-size-axis claim to match _eligible. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 17 +++++++++-------- src/zarr/core/transforms/transform.py | 3 --- src/zarr/testing/strategies.py | 5 +++-- tests/test_properties.py | 9 +++++++++ 4 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index e502188913..6e4b1301be 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3727,8 +3727,9 @@ def set_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - # Normalize empty fields list to None - if not fields: + # Normalize an empty fields list/tuple to None (pop_fields yields [] for + # no fields); keep the exact-emptiness check rather than a falsy one. + if fields == [] or fields == (): fields = None if fields is not None or self._async_array._is_identity: indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) @@ -5770,12 +5771,12 @@ async def _get_selection_via_transform( out_shape = transform.domain.shape - # Only vectorized indexing (>= 2 correlated ArrayMaps, input_dimension None) - # scatters through a single flat index, so the buffer must be 1D during the - # read and reshaped to out_shape afterwards. Orthogonal indexing — including - # the case where every output is an ArrayMap bound to a distinct input - # dimension (oindex[i0, i1]) — keeps a multi-dimensional buffer, one out_sel - # entry per dim. + # Vectorized indexing (any coordinate ArrayMap, input_dimension None — even a + # single one, e.g. vindex[idx_2d]) scatters through a single flat index, so + # the buffer must be 1D during the read and reshaped to out_shape afterwards. + # Orthogonal indexing — including the case where every output is an ArrayMap + # bound to a distinct input dimension (oindex[i0, i1]) — keeps a + # multi-dimensional buffer, one out_sel entry per dim. needs_flat_buffer = _is_vectorized_transform(transform) buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 861915b6ab..0847474da0 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -332,7 +332,6 @@ def _intersect_vectorized( storage coordinates fall within the output domain. """ # Compute storage coords per array dim and check bounds simultaneously - n_points: int | None = None masks: list[np.ndarray[Any, np.dtype[np.bool_]]] = [] for out_dim in array_dims: @@ -342,8 +341,6 @@ def _intersect_vectorized( lo = output_domain.inclusive_min[out_dim] hi = output_domain.exclusive_max[out_dim] masks.append((storage >= lo) & (storage < hi)) - if n_points is None: - n_points = storage.size # A point survives only if it's in-bounds on ALL array dims combined_mask = masks[0] diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 41355c8e42..247f8e3916 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -636,8 +636,9 @@ def indexers(draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...]) -> tup The two returned selections differ only for ``oindex`` (zarr's per-axis spelling vs the ``np.ix_``-style spelling numpy needs); for the other modes - the same object indexes both a zarr array and its numpy reference. Fancy modes - (``oindex``/``vindex``) require non-zero-size axes; ``mask``/``basic`` do not. + the same object indexes both a zarr array and its numpy reference. The + array-based modes (``oindex``/``vindex``/``mask``) need ``shape`` to have no + zero-length axis; ``basic`` has no such requirement. """ if mode == "basic": sel = draw(basic_indices(shape=shape)) diff --git a/tests/test_properties.py b/tests/test_properties.py index 410aeba991..04109a0f87 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -435,6 +435,10 @@ async def test_indexing_parity(data: st.DataObject) -> None: expected = np.asarray(nparray[npsel]) assert_array_equal(np.asarray(await _async_get(zarray._async_array, mode, zsel)), expected) + # a mask can also be spelled via vindex[...]; the two interfaces must agree + if mode == "mask": + assert_array_equal(np.asarray(zarray.vindex[zsel]), expected) + # write, mirroring the historical per-mode guards: # - vindex setitem is chunk-dependent for repeated coords (skipped) # - oindex/mask writes to sharded arrays are unsupported (GH2834) @@ -450,6 +454,11 @@ async def test_indexing_parity(data: st.DataObject) -> None: nparray[npsel] = new_data assert_array_equal(nparray, zarray[:]) + # the vindex[mask] = ... spelling must agree with set_mask_selection + if mode == "mask": + zarray.vindex[zsel] = new_data + assert_array_equal(nparray, zarray[:]) + @pytest.mark.skipif( not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" From 57f3c002bffae1306378e03da1927c5f1d8cad70 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 12:43:02 +0200 Subject: [PATCH 21/43] fix(indexing): bounds-check lazy oindex/vindex array values The lazy accessor copies TensorStore: index-array values are literal coordinates (origin 0), not NumPy from-the-end. Out-of-domain values (negative, or >= the axis length) previously resolved to silently-wrong data (e.g. lazy.oindex[[-1]] returned element 0); now they raise BoundsCheckError at transform construction, matching the eager bounds check. Applies to both the literal accessor and the NumPy-normalizing view methods (which normalize first, then bounds-check the result). Basic accessor negatives already raise as literal out-of-range. Keeps the intentional accessor-vs-method asymmetry (the accessor is the TensorStore-style transform builder; the Array methods normalize NumPy-style); documented and pinned by test_accessor_negative_index_is_literal. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/transforms/transform.py | 20 +++++++++++++++++++- tests/test_lazy_indexing.py | 17 +++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 0847474da0..68f32e7ced 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -36,7 +36,7 @@ from zarr.core.transforms.domain import IndexDomain from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap -from zarr.errors import VindexInvalidSelectionError +from zarr.errors import BoundsCheckError, VindexInvalidSelectionError @dataclass(frozen=True, slots=True) @@ -663,6 +663,9 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: for old_dim, sel in enumerate(normalized): if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + _check_array_in_bounds(sel, hi - lo) dim_array[old_dim] = sel new_inclusive_min.append(0) new_exclusive_max.append(len(sel)) @@ -806,6 +809,9 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: for i, sel in enumerate(processed): if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[i] + hi = transform.domain.exclusive_max[i] + _check_array_in_bounds(sel, hi - lo) array_dims.append(i) arrays.append(sel) else: @@ -955,6 +961,18 @@ def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: return tuple(result) +def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], dim_size: int) -> None: + """Reject index-array values outside ``[0, dim_size)``. + + Index-array values are literal coordinates (origin 0), so a negative value is + out of bounds rather than counting from the end — the Array layer normalizes + NumPy-style negatives before building the transform. Matches the eager + bounds-check error so out-of-range values raise instead of silently wrapping. + """ + if arr.size and (int(arr.min()) < 0 or int(arr.max()) >= dim_size): + raise BoundsCheckError(f"index out of bounds for dimension with length {dim_size}") + + def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: """Validate array-based selections (orthogonal, vectorized). diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 621e86f7a9..79e37b9570 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -381,6 +381,23 @@ def test_negative_step_raises(self) -> None: with pytest.raises(IndexError, match="step must be positive"): a.lazy[::-1] + def test_accessor_negative_index_is_literal(self) -> None: + """The lazy accessor copies TensorStore: indices are literal coordinates. + + Negative indices are absolute (origin 0), not from-the-end, so they fall + outside the ``[0, N)`` domain and raise — unlike the NumPy-normalizing + eager/view-method paths. Out-of-bounds array values raise cleanly rather + than silently wrapping. + """ + a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + with pytest.raises(IndexError): + _ = a.lazy[-1][...] + for sel in (np.array([-1], dtype=np.intp), np.array([24], dtype=np.intp)): + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.oindex[(sel,)][...] + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.vindex[(sel,)][...] + def test_block_selection_on_view_rejected(self) -> None: """Block selection is ill-defined on a lazy view and must raise, not corrupt.""" a, _ = _make(CONFIGS[3]) # 2d-unsharded From ebf9f8f90691800bdf3b115eb03731e36e94fc06 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 17:01:28 +0200 Subject: [PATCH 22/43] test: type mode params and indexers return; consistent mode collections Mirror the review fixes from the eager test-infra PR onto the lazy branch: - type `mode` as the `IndexMode` literal across the indexing test helpers - `_VECTORIZED_MODES` is a tuple like `_INDEX_MODES` - `indexers(...)` returns `tuple[Selection, Selection]` (real type) Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/testing/strategies.py | 5 ++++- tests/test_properties.py | 17 ++++++++++------- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 247f8e3916..b40f116baa 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -27,6 +27,7 @@ from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat from zarr.core.dtype import get_data_type_from_native_dtype +from zarr.core.indexing import Selection from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata from zarr.core.sync import sync @@ -623,7 +624,9 @@ def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: @st.composite -def indexers(draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...]) -> tuple[Any, Any]: +def indexers( + draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...] +) -> tuple[Selection, Selection]: """A ``(zarr_selection, numpy_selection)`` pair for ``mode`` on ``shape``. One strategy covering every indexing mode, so a test can be written once and diff --git a/tests/test_properties.py b/tests/test_properties.py index 04109a0f87..d388769a0a 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -22,6 +22,7 @@ from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.sync import sync from zarr.testing.strategies import ( + IndexMode, array_metadata, arrays, basic_indices, @@ -332,11 +333,13 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N # The indexing modes and which Array method implements each. vindex/mask are # "vectorized" — they scatter through a single flat index, so an out= buffer must # be flat (number of selected points) rather than the multi-dimensional result. -_INDEX_MODES = ("basic", "oindex", "vindex", "mask") -_VECTORIZED_MODES = frozenset({"vindex", "mask"}) +_INDEX_MODES: tuple[IndexMode, ...] = ("basic", "oindex", "vindex", "mask") +# Modes that scatter through a flat index (so an out= buffer must be flat). Kept a +# tuple like _INDEX_MODES; membership is checked against it below. +_VECTORIZED_MODES: tuple[IndexMode, ...] = ("vindex", "mask") -def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: +def _get(target: zarr.Array, mode: IndexMode, zsel: Any, *, out: Any = None) -> Any: """Read ``zsel`` from ``target`` via the get-method for ``mode``.""" if mode == "basic": return target.get_basic_selection(zsel, out=out) @@ -349,7 +352,7 @@ def _get(target: zarr.Array, mode: str, zsel: Any, *, out: Any = None) -> Any: raise AssertionError(mode) -def _async_get(async_array: Any, mode: str, zsel: Any) -> Any: +def _async_get(async_array: Any, mode: IndexMode, zsel: Any) -> Any: """The async read coroutine for ``mode`` (vindex/mask share the vectorized accessor).""" if mode == "basic": return async_array.getitem(zsel) @@ -358,7 +361,7 @@ def _async_get(async_array: Any, mode: str, zsel: Any) -> Any: return async_array.vindex.getitem(zsel) -def _setitem(zarray: zarr.Array, mode: str, zsel: Any, value: Any) -> None: +def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None: """Write ``value`` at ``zsel`` via the set-method for ``mode``.""" if mode == "basic": zarray[zsel] = value @@ -376,7 +379,7 @@ def _has_repeated_indices(npsel: Any) -> bool: return any(isinstance(i, np.ndarray) and i.size != np.unique(i).size for i in sel) -def _eligible(mode: str, shape: tuple[int, ...]) -> bool: +def _eligible(mode: IndexMode, shape: tuple[int, ...]) -> bool: """Whether ``mode`` can be exercised on ``shape``. Rank-0 arrays have no interesting selections; the fancy modes @@ -388,7 +391,7 @@ def _eligible(mode: str, shape: tuple[int, ...]) -> bool: def assert_read_matches_numpy( - target: zarr.Array, ref: np.ndarray[Any, Any], mode: str, zsel: Any, npsel: Any + target: zarr.Array, ref: np.ndarray[Any, Any], mode: IndexMode, zsel: Any, npsel: Any ) -> None: """Assert ``target``'s read of ``zsel`` (mode) matches ``ref[npsel]``, with/without out=. From 8aee92129352ef0294795824bf5513aa1633fb49 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 17:14:45 +0200 Subject: [PATCH 23/43] test(indexing): split read/write parity; xfail the one known-unsupported write MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the brittle silent-`return` write guards in the indexing oracle with a principled split: - test_indexing_read_parity: reads (sync/out=/async) for every mode, no guards — reads are well-defined for all selections. - test_indexing_write_parity: writes, filtering with `assume` only the cases that have no well-defined oracle — repeated coordinates (order-dependent, a fundamental limitation shared with NumPy, not a bug). This now actually exercises vindex writes, sharded single-axis oindex/mask writes, and unsharded multi-axis oindex writes, which the old blanket `assume(shards is None)` skipped. - test_oindex_sharded_multiaxis_write_xfail: pins the one genuinely-unsupported write (orthogonal across >=2 array axes on a sharded array, GH2834) as a strict xfail, so it is visible at the pytest level and flips to a failure when fixed. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 112 +++++++++++++++++++++++++++------------ 1 file changed, 77 insertions(+), 35 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index d388769a0a..2afeecef8d 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -367,10 +367,12 @@ def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None zarray[zsel] = value elif mode == "oindex": zarray.oindex[zsel] = value + elif mode == "vindex": + zarray.vindex[zsel] = value elif mode == "mask": zarray.set_mask_selection(zsel, value) else: - raise AssertionError(f"writes are not exercised for mode {mode!r}") + raise AssertionError(mode) def _has_repeated_indices(npsel: Any) -> bool: @@ -379,6 +381,12 @@ def _has_repeated_indices(npsel: Any) -> bool: return any(isinstance(i, np.ndarray) and i.size != np.unique(i).size for i in sel) +def _n_array_axes(npsel: Any) -> int: + """Number of integer-array (fancy) axes in a selection.""" + sel = npsel if isinstance(npsel, tuple) else (npsel,) + return sum(isinstance(s, np.ndarray) for s in sel) + + def _eligible(mode: IndexMode, shape: tuple[int, ...]) -> bool: """Whether ``mode`` can be exercised on ``shape``. @@ -409,60 +417,95 @@ def assert_read_matches_numpy( assert_array_equal(np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected) -@settings(deadline=None) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -@given(data=st.data()) -async def test_indexing_parity(data: st.DataObject) -> None: - """Single eager-indexing oracle: every method matches NumPy. - - Covers {basic, oindex, vindex, mask} against a NumPy reference: sync read, - read into an ``out=`` buffer, async read, and (where defined) write. Replaces - the per-mode test_basic/oindex/vindex/mask tests with one enumeration. - Deliberately independent of lazy indexing so it can guard the eager API on its - own; ``test_lazy_view_indexing_parity`` reuses the read half for lazy views. - """ - zarray = data.draw( +def _indexing_array(data: st.DataObject) -> zarr.Array: + """An eager array (regular or rectilinear) for the indexing-parity oracles.""" + return data.draw( st.one_of( simple_arrays(shapes=npst.array_shapes(max_dims=4)), rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), ) ) + + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_indexing_read_parity(data: st.DataObject) -> None: + """Every indexing *read* on an eager array matches NumPy (all modes). + + Covers {basic, oindex, vindex, mask}: sync read, read into an out= buffer, and + async read. No special-casing — reads are well-defined for every selection. + Consumer-agnostic via assert_read_matches_numpy, which test_lazy_view_indexing_parity + reuses for lazy views. + """ + zarray = _indexing_array(data) nparray = zarray[:] mode = data.draw(st.sampled_from(_INDEX_MODES)) - if not _eligible(mode, nparray.shape): - return + assume(_eligible(mode, nparray.shape)) zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) - # read: sync + out= (shared with the lazy-view test) and async assert_read_matches_numpy(zarray, nparray, mode, zsel, npsel) expected = np.asarray(nparray[npsel]) assert_array_equal(np.asarray(await _async_get(zarray._async_array, mode, zsel)), expected) - - # a mask can also be spelled via vindex[...]; the two interfaces must agree - if mode == "mask": + if mode == "mask": # a mask read may also be spelled vindex[mask]; they must agree assert_array_equal(np.asarray(zarray.vindex[zsel]), expected) - # write, mirroring the historical per-mode guards: - # - vindex setitem is chunk-dependent for repeated coords (skipped) - # - oindex/mask writes to sharded arrays are unsupported (GH2834) - # - oindex writes with repeated indices are unspecified - if mode == "vindex": - return - if mode in ("oindex", "mask") and zarray.shards is not None: - return - if mode == "oindex" and _has_repeated_indices(npsel): - return + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_indexing_write_parity(data: st.DataObject) -> None: + """Every *supported, well-defined* indexing write on an eager array matches NumPy. + + Two families of selections are filtered out with ``assume`` (they are not + failures to test here): + + - **Repeated coordinates** (vindex/oindex selecting a cell more than once): + the write order, and thus the final value, is undefined — NumPy has the same + ambiguity. This is a fundamental limitation, not a zarr bug. + - **Sharded orthogonal writes across >= 2 array axes**: unsupported by the + sharding partial-write codec; pinned at the pytest level by + ``test_oindex_sharded_multiaxis_write_xfail`` rather than hidden here. + """ + zarray = _indexing_array(data) + nparray = zarray[:] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + assume(_eligible(mode, nparray.shape)) + zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) + if mode in ("oindex", "vindex"): + assume(not _has_repeated_indices(npsel)) + if mode == "oindex" and zarray.shards is not None: + assume(_n_array_axes(npsel) < 2) + + expected = np.asarray(nparray[npsel]) new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) _setitem(zarray, mode, zsel, new_data) nparray[npsel] = new_data assert_array_equal(nparray, zarray[:]) - - # the vindex[mask] = ... spelling must agree with set_mask_selection - if mode == "mask": + if mode == "mask": # the vindex[mask] = ... spelling must agree with set_mask_selection zarray.vindex[zsel] = new_data assert_array_equal(nparray, zarray[:]) +@pytest.mark.xfail( + reason="GH2834: the sharding partial-write codec does not support orthogonal " + "writes spanning two or more array axes", + strict=True, +) +def test_oindex_sharded_multiaxis_write_xfail() -> None: + """Pin the one known-unsupported write: oindex across >= 2 array axes on a shard. + + Single-axis oindex, vindex, and mask writes to sharded arrays all work; only + the multi-array-axis orthogonal case is broken. Strict xfail so this flips to a + failure (prompting removal) if/when GH2834 is fixed. + """ + a = zarr.create_array({}, shape=(12, 12), chunks=(3, 3), shards=(6, 6), dtype="i4") + a[...] = np.arange(144, dtype="i4").reshape(12, 12) + i0, i1 = np.array([0, 5, 9]), np.array([1, 6, 10]) + a.set_orthogonal_selection((i0, i1), np.zeros((3, 3), dtype="i4")) + assert_array_equal(a.oindex[i0, i1], np.zeros((3, 3), dtype="i4")) + + @pytest.mark.skipif( not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" ) @@ -485,7 +528,6 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: view = zarray.lazy[window] vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) - if not _eligible(mode, vref.shape): - return + assume(_eligible(mode, vref.shape)) zsel, npsel = data.draw(indexers(mode=mode, shape=vref.shape)) assert_read_matches_numpy(view, vref, mode, zsel, npsel) From 734febd783a521b763b0aa704db7b47eb5f47e90 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 17:32:27 +0200 Subject: [PATCH 24/43] test(indexing): filter writes whose targets collide after negative-index normalization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI (slow-hypothesis, ci profile) found test_indexing_write_parity falsified by a vindex write with selection [2, -2, 0, 1] on length 4: -2 normalizes to 2, so cell 2 is written twice. That is the undefined repeated-target case (order- dependent, NumPy has the same ambiguity) — but the old filter checked *raw* index values and missed the post-normalization collision. Replace _has_repeated_indices with _write_is_unambiguous(mode, npsel, shape), which normalizes negative indices first and is mode-aware: per-axis uniqueness for oindex (independent axes), unique coordinate tuples for vindex (correlated axes). Not a zarr bug — a test-oracle gap. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 2afeecef8d..75d7a5f9c0 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -375,10 +375,36 @@ def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None raise AssertionError(mode) -def _has_repeated_indices(npsel: Any) -> bool: - """True if any integer-array component selects a coordinate more than once.""" +def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) -> bool: + """Whether a write to ``npsel`` targets each cell at most once. + + Negative indices are normalized first (``[2, -2]`` on length 4 both target cell + 2), then duplicate targets are detected. A repeated target makes the write + order-dependent — NumPy has the same ambiguity — so there is no well-defined + oracle and such an example is skipped (a fundamental limitation, not a bug). + """ sel = npsel if isinstance(npsel, tuple) else (npsel,) - return any(isinstance(i, np.ndarray) and i.size != np.unique(i).size for i in sel) + if mode == "oindex": + # Independent axes: each fancy axis must select each index at most once. + for axis, s in enumerate(sel): + if isinstance(s, np.ndarray): + norm = np.where(s < 0, s + shape[axis], s) + if norm.size != np.unique(norm).size: + return False + return True + if mode == "vindex": + # Correlated coordinate arrays: a target is the tuple of normalized coords, + # so the write is well-defined iff no coordinate tuple repeats. + norm = [ + np.where(s < 0, s + shape[axis], s) + for axis, s in enumerate(sel) + if isinstance(s, np.ndarray) + ] + if not norm: + return True + coords = np.stack([a.ravel() for a in np.broadcast_arrays(*norm)], axis=-1) + return len(coords) == len(np.unique(coords, axis=0)) + return True # basic / mask never target a cell twice def _n_array_axes(npsel: Any) -> int: @@ -473,7 +499,7 @@ async def test_indexing_write_parity(data: st.DataObject) -> None: assume(_eligible(mode, nparray.shape)) zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) if mode in ("oindex", "vindex"): - assume(not _has_repeated_indices(npsel)) + assume(_write_is_unambiguous(mode, npsel, nparray.shape)) if mode == "oindex" and zarray.shards is not None: assume(_n_array_axes(npsel) < 2) From dce5d4685ef6560af25ee4376584459b46925093 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 21:03:27 +0200 Subject: [PATCH 25/43] test(indexing): stateful model-lockstep harness + bounds error-parity + dtype-strict oracle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Borrow the strongest ideas from TensorStore's and ndindex's indexing test suites: - Stateful harness (TensorStore driver_testutil pattern, NumPy oracle): _IndexingStateMachine applies many random reads/writes (basic/oindex/vindex/mask) to one array while a NumPy model tracks it in lockstep, asserting equality after every step. Catches read-modify-write, chunk-boundary, and shard-merge bugs the single-shot parity tests cannot. Runs over regular/sharded/3D geometries, gated behind --run-slow-hypothesis like the store stateful tests. - Error parity (ndindex check_same): test_indexing_bounds_error_parity generates possibly-out-of-bounds integer selections and asserts zarr raises iff NumPy raises (zarr's bounds errors subclass IndexError). Catches the silent-wrong-data class — returning garbage where NumPy would raise. - Strict read oracle: assert dtype, not just values/shape, for ndim>=1 results. - Regression anchor: test_write_is_unambiguous pins the negative-index collision fix ([2,-2,0,1] -> cell 2 written twice). Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 147 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 146 insertions(+), 1 deletion(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 75d7a5f9c0..6afa1019ac 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -16,6 +16,7 @@ import hypothesis.extra.numpy as npst import hypothesis.strategies as st from hypothesis import assume, given, settings +from hypothesis.stateful import RuleBasedStateMachine, invariant, rule from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON @@ -435,7 +436,13 @@ def assert_read_matches_numpy( index) and full-shape otherwise. """ expected = np.asarray(ref[npsel]) - assert_array_equal(np.asarray(_get(target, mode, zsel)), expected) + actual = np.asarray(_get(target, mode, zsel)) + # dtype must match, not just values (catches silent dtype coercion); ndindex's + # assert_equal does the same. Guarded to ndim>=1 to avoid python-vs-numpy + # scalar dtype noise on 0-d results. + if expected.ndim >= 1: + assert actual.dtype == expected.dtype, (actual.dtype, expected.dtype) + assert_array_equal(actual, expected) if expected.ndim >= 1 and expected.size > 0: buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else expected.shape buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=expected.dtype) @@ -532,6 +539,144 @@ def test_oindex_sharded_multiaxis_write_xfail() -> None: assert_array_equal(a.oindex[i0, i1], np.zeros((3, 3), dtype="i4")) +@pytest.mark.parametrize( + ("mode", "npsel", "shape", "well_defined"), + [ + pytest.param("vindex", (np.array([2, -2, 0, 1]),), (4,), False, id="vindex-neg-collision"), + pytest.param("vindex", (np.array([0, 1, 2]),), (4,), True, id="vindex-distinct"), + pytest.param("oindex", (np.array([1, -3]),), (4,), False, id="oindex-neg-collision"), + pytest.param( + "oindex", (np.array([0, 2]), np.array([1, 1])), (4, 4), False, id="oindex-repeat" + ), + pytest.param( + "oindex", (np.array([0, 2]), np.array([1, 3])), (4, 4), True, id="oindex-distinct" + ), + pytest.param("basic", (slice(None),), (4,), True, id="basic-always"), + ], +) +def test_write_is_unambiguous( + mode: IndexMode, npsel: Any, shape: tuple[int, ...], well_defined: bool +) -> None: + """A write is well-defined iff each target cell is hit at most once *after* + negative indices are normalized. + + Regression anchor for the CI fix: vindex ``[2, -2, 0, 1]`` on length 4 maps + ``-2 -> 2``, so cell 2 is written twice (order-dependent) and must be rejected, + even though the raw values are all distinct. + """ + assert _write_is_unambiguous(mode, npsel, shape) is well_defined + + +@st.composite +def _bounds_selection(draw: st.DrawFn, mode: IndexMode, shape: tuple[int, ...]) -> Any: + """An integer selection whose components may fall outside ``[-size, size)``, to + probe out-of-bounds behavior. Returns ``(zarr_sel, numpy_sel)``.""" + wide = {s: st.integers(-2 * s - 1, 2 * s) for s in set(shape)} + if mode == "basic": + sel = tuple(draw(wide[s]) for s in shape) + return sel, sel + n = draw(st.integers(1, 4)) + arrs = tuple( + np.array(draw(st.lists(wide[s], min_size=n, max_size=n)), dtype=np.intp) for s in shape + ) + if mode == "oindex": + return arrs, np.ix_(*arrs) # numpy outer-product oracle + return arrs, arrs # vindex: pointwise + + +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +def test_indexing_bounds_error_parity(data: st.DataObject) -> None: + """Out-of-bounds integer indexing raises in zarr iff it raises in NumPy. + + zarr's bounds errors (BoundsCheckError etc.) subclass IndexError, so the eager + methods must agree with NumPy on *whether* an index is in bounds. This is the + ndindex ``check_same`` error-parity idea, and it catches the silent-wrong-data + class — returning garbage where NumPy would raise. + """ + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=6))) + nparray = zarray[:] + mode = data.draw(st.sampled_from(("basic", "oindex", "vindex"))) + zsel, npsel = data.draw(_bounds_selection(mode, nparray.shape)) + try: + expected = np.asarray(nparray[npsel]) + except IndexError: + with pytest.raises(IndexError): + _get(zarray, mode, zsel) + return + assert_array_equal(np.asarray(_get(zarray, mode, zsel)), expected) + + +class _IndexingStateMachine(RuleBasedStateMachine): + """Apply many random indexing writes to one array while a NumPy model tracks it + in lockstep; the model must match after every step. + + This is the TensorStore ``driver_testutil`` pattern with NumPy as the oracle: + a single array accumulates many writes, so it catches read-modify-write, + chunk-boundary, and shard-merge/persistence bugs that the single-shot + ``test_indexing_write_parity`` cannot see. Concrete geometries are defined by + subclasses; the rules are dtype-agnostic (fixed i4 — dtype breadth lives in the + property tests). + """ + + shape: tuple[int, ...] = () + chunks: tuple[int, ...] = () + shards: tuple[int, ...] | None = None + + def __init__(self) -> None: + super().__init__() + n = int(np.prod(self.shape, dtype=int)) + self.model = np.arange(n, dtype="i4").reshape(self.shape) + self.zarray = zarr.create_array( + {}, shape=self.shape, chunks=self.chunks, shards=self.shards, dtype="i4" + ) + self.zarray[...] = self.model + + @rule(data=st.data(), mode=st.sampled_from(_INDEX_MODES)) + def write(self, data: st.DataObject, mode: IndexMode) -> None: + if not _eligible(mode, self.shape): + return + zsel, npsel = data.draw(indexers(mode=mode, shape=self.shape)) + if mode in ("oindex", "vindex") and not _write_is_unambiguous(mode, npsel, self.shape): + return + if mode == "oindex" and self.shards is not None and _n_array_axes(npsel) >= 2: + return # GH2834, see test_oindex_sharded_multiaxis_write_xfail + expected = np.asarray(self.model[npsel]) + value = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=self.model.dtype)) + self.model[npsel] = value + _setitem(self.zarray, mode, zsel, value) + + @rule(data=st.data(), mode=st.sampled_from(_INDEX_MODES)) + def read(self, data: st.DataObject, mode: IndexMode) -> None: + if not _eligible(mode, self.shape): + return + zsel, npsel = data.draw(indexers(mode=mode, shape=self.shape)) + assert_array_equal(np.asarray(_get(self.zarray, mode, zsel)), np.asarray(self.model[npsel])) + + @invariant() + def matches_model(self) -> None: + assert_array_equal(self.zarray[:], self.model) + + +class _RegularStateMachine(_IndexingStateMachine): + shape, chunks, shards = (12,), (3,), None + + +class _Sharded2DStateMachine(_IndexingStateMachine): + shape, chunks, shards = (8, 9), (2, 3), (4, 9) + + +class _ThreeDStateMachine(_IndexingStateMachine): + shape, chunks, shards = (4, 5, 6), (2, 3, 3), None + + +# Run these under the slow-hypothesis CI job (like the store stateful tests): a +# stateful run is a sequence of steps, so it is heavier than the single-shot tests. +TestIndexingStateMachineRegular = pytest.mark.slow_hypothesis(_RegularStateMachine.TestCase) +TestIndexingStateMachineSharded = pytest.mark.slow_hypothesis(_Sharded2DStateMachine.TestCase) +TestIndexingStateMachine3D = pytest.mark.slow_hypothesis(_ThreeDStateMachine.TestCase) + + @pytest.mark.skipif( not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" ) From 93fb2d4fd3baaae9ea26728eff253d11b7c7af41 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 21:22:28 +0200 Subject: [PATCH 26/43] test(indexing): rename indexers -> numpy_array_indexers and clarify scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The strategy covers only the element-space indexing modes that have a direct NumPy-array oracle (basic/oindex/vindex/mask). Rename it so the name states that, and document that block indexing is deliberately excluded — it addresses the chunk grid (parametrized by the grid, not shape) and has no NumPy equivalent, so it lives in the separate block_indices strategy. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/testing/strategies.py | 14 +++++++++++--- tests/test_properties.py | 12 ++++++------ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index b40f116baa..4b15416b1d 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -624,13 +624,15 @@ def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: @st.composite -def indexers( +def numpy_array_indexers( draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...] ) -> tuple[Selection, Selection]: """A ``(zarr_selection, numpy_selection)`` pair for ``mode`` on ``shape``. - One strategy covering every indexing mode, so a test can be written once and - parametrized over mode instead of re-deriving selection setup per test: + Scope: the *element-space* indexing modes that have a direct NumPy-array + equivalent, so a NumPy array of ``shape`` can serve as the correctness oracle. + One strategy covers them all, so a test can be written once and parametrized + over mode instead of re-deriving selection setup per test: - ``"basic"`` — slices / ints / ellipsis (no newaxis, no negative slices) - ``"oindex"`` — per-axis integer arrays or slices (orthogonal / outer product) @@ -642,6 +644,12 @@ def indexers( the same object indexes both a zarr array and its numpy reference. The array-based modes (``oindex``/``vindex``/``mask``) need ``shape`` to have no zero-length axis; ``basic`` has no such requirement. + + Deliberately excluded is **block** indexing (``Array.blocks`` / + ``get_block_selection``): it addresses the *chunk grid*, not array elements, + so it is parametrized by the chunk grid rather than ``shape`` and has no + NumPy-array equivalent to compare against — its oracle is a coordinate + translation built by the separate `block_indices` strategy. """ if mode == "basic": sel = draw(basic_indices(shape=shape)) diff --git a/tests/test_properties.py b/tests/test_properties.py index 6afa1019ac..cad26ba7cd 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -30,7 +30,7 @@ block_indices, block_test_arrays, complex_rectilinear_arrays, - indexers, + numpy_array_indexers, numpy_arrays, rectilinear_arrays, simple_arrays, @@ -475,7 +475,7 @@ async def test_indexing_read_parity(data: st.DataObject) -> None: nparray = zarray[:] mode = data.draw(st.sampled_from(_INDEX_MODES)) assume(_eligible(mode, nparray.shape)) - zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=nparray.shape)) assert_read_matches_numpy(zarray, nparray, mode, zsel, npsel) expected = np.asarray(nparray[npsel]) @@ -504,7 +504,7 @@ async def test_indexing_write_parity(data: st.DataObject) -> None: nparray = zarray[:] mode = data.draw(st.sampled_from(_INDEX_MODES)) assume(_eligible(mode, nparray.shape)) - zsel, npsel = data.draw(indexers(mode=mode, shape=nparray.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=nparray.shape)) if mode in ("oindex", "vindex"): assume(_write_is_unambiguous(mode, npsel, nparray.shape)) if mode == "oindex" and zarray.shards is not None: @@ -636,7 +636,7 @@ def __init__(self) -> None: def write(self, data: st.DataObject, mode: IndexMode) -> None: if not _eligible(mode, self.shape): return - zsel, npsel = data.draw(indexers(mode=mode, shape=self.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=self.shape)) if mode in ("oindex", "vindex") and not _write_is_unambiguous(mode, npsel, self.shape): return if mode == "oindex" and self.shards is not None and _n_array_axes(npsel) >= 2: @@ -650,7 +650,7 @@ def write(self, data: st.DataObject, mode: IndexMode) -> None: def read(self, data: st.DataObject, mode: IndexMode) -> None: if not _eligible(mode, self.shape): return - zsel, npsel = data.draw(indexers(mode=mode, shape=self.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=self.shape)) assert_array_equal(np.asarray(_get(self.zarray, mode, zsel)), np.asarray(self.model[npsel])) @invariant() @@ -700,5 +700,5 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) assume(_eligible(mode, vref.shape)) - zsel, npsel = data.draw(indexers(mode=mode, shape=vref.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=vref.shape)) assert_read_matches_numpy(view, vref, mode, zsel, npsel) From 8153d4843b3ff8fc57f32cbf356c8a2e16d27170 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 21:28:58 +0200 Subject: [PATCH 27/43] docs: use single backticks in indexing test-infra docstrings mkdocs renders single backticks for inline code; the docstrings used RST-style double backticks. Normalize to single backticks across the indexing strategies and property tests. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/testing/strategies.py | 64 +++++++++++++++++----------------- tests/test_properties.py | 44 +++++++++++------------ 2 files changed, 54 insertions(+), 54 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 4b15416b1d..60c99e92a1 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -606,8 +606,8 @@ def orthogonal_indices( def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: """A non-negative, full-rank tuple of slice windows — one per axis. - A rank-preserving sub-region selection: each axis gets ``start:stop`` with - ``0 <= start < stop <= size`` (an empty ``0:0`` slice for a zero-length axis). + A rank-preserving sub-region selection: each axis gets `start:stop` with + `0 <= start < stop <= size` (an empty `0:0` slice for a zero-length axis). Bounds stay non-negative so the window is valid for any consumer, including those that treat negative indices as literal coordinates rather than from-the-end (e.g. building a sub-array view). @@ -627,27 +627,27 @@ def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: def numpy_array_indexers( draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...] ) -> tuple[Selection, Selection]: - """A ``(zarr_selection, numpy_selection)`` pair for ``mode`` on ``shape``. + """A `(zarr_selection, numpy_selection)` pair for `mode` on `shape`. Scope: the *element-space* indexing modes that have a direct NumPy-array - equivalent, so a NumPy array of ``shape`` can serve as the correctness oracle. + equivalent, so a NumPy array of `shape` can serve as the correctness oracle. One strategy covers them all, so a test can be written once and parametrized over mode instead of re-deriving selection setup per test: - - ``"basic"`` — slices / ints / ellipsis (no newaxis, no negative slices) - - ``"oindex"`` — per-axis integer arrays or slices (orthogonal / outer product) - - ``"vindex"`` — broadcast integer coordinate arrays (vectorized) - - ``"mask"`` — a boolean array of ``shape`` + - `"basic"` — slices / ints / ellipsis (no newaxis, no negative slices) + - `"oindex"` — per-axis integer arrays or slices (orthogonal / outer product) + - `"vindex"` — broadcast integer coordinate arrays (vectorized) + - `"mask"` — a boolean array of `shape` - The two returned selections differ only for ``oindex`` (zarr's per-axis - spelling vs the ``np.ix_``-style spelling numpy needs); for the other modes + The two returned selections differ only for `oindex` (zarr's per-axis + spelling vs the `np.ix_`-style spelling numpy needs); for the other modes the same object indexes both a zarr array and its numpy reference. The - array-based modes (``oindex``/``vindex``/``mask``) need ``shape`` to have no - zero-length axis; ``basic`` has no such requirement. + array-based modes (`oindex`/`vindex`/`mask`) need `shape` to have no + zero-length axis; `basic` has no such requirement. - Deliberately excluded is **block** indexing (``Array.blocks`` / - ``get_block_selection``): it addresses the *chunk grid*, not array elements, - so it is parametrized by the chunk grid rather than ``shape`` and has no + Deliberately excluded is **block** indexing (`Array.blocks` / + `get_block_selection`): it addresses the *chunk grid*, not array elements, + so it is parametrized by the chunk grid rather than `shape` and has no NumPy-array equivalent to compare against — its oracle is a coordinate translation built by the separate `block_indices` strategy. """ @@ -677,17 +677,17 @@ def block_indices( Strategy for block-selection indexers over a chunk grid. Block indexing is basic indexing applied to the block grid (the grid of - chunks), so each axis is drawn with ``basic_indices`` over that axis's chunk - count, mirroring how ``orthogonal_indices`` reuses ``basic_indices`` per - axis. ``chunk_sizes`` gives the per-chunk data sizes of the array's *outer* - (block) grid for every axis — i.e. ``Array.write_chunk_sizes``, the grid that - ``Array.blocks`` addresses (the shard grid when sharding is used). For - example ``(3, 3, 3, 1)`` for a length-10 axis with a regular chunk size of 3, - or the explicit edges of a rectilinear axis; ``nchunks`` for an axis is - ``len(chunk_sizes[axis])``. + chunks), so each axis is drawn with `basic_indices` over that axis's chunk + count, mirroring how `orthogonal_indices` reuses `basic_indices` per + axis. `chunk_sizes` gives the per-chunk data sizes of the array's *outer* + (block) grid for every axis — i.e. `Array.write_chunk_sizes`, the grid that + `Array.blocks` addresses (the shard grid when sharding is used). For + example `(3, 3, 3, 1)` for a length-10 axis with a regular chunk size of 3, + or the explicit edges of a rectilinear axis; `nchunks` for an axis is + `len(chunk_sizes[axis])`. The array-space translation uses the cumulative sum of those sizes, matching - ``BlockIndexer``'s use of ``dim_grid.chunk_offset``. Because the sizes are + `BlockIndexer`'s use of `dim_grid.chunk_offset`. Because the sizes are clipped to the array extent, the final offset equals the extent and the translation is exact for regular (uniform), rectilinear, and sharded grids alike. @@ -700,7 +700,7 @@ def block_indices( ------- block_indexer A per-axis tuple of ints / step-1 slices addressing whole chunks, - suitable for ``Array.blocks`` / ``get_block_selection`` / ``set_block_selection``. + suitable for `Array.blocks` / `get_block_selection` / `set_block_selection`. array_indexer The equivalent array-space selection (a tuple of slices) for indexing the corresponding numpy array, used as the comparison oracle. @@ -736,7 +736,7 @@ def predicate(value: tuple[Any, ...]) -> bool: # basic_indices draws slices far more often than bare integers, so the # integer (single-block) branch below would only be hit on rare draws. # Union in an explicit integer so it is reliably exercised — keeping - # coverage deterministic under the derandomized ``ci`` Hypothesis profile. + # coverage deterministic under the derandomized `ci` Hypothesis profile. (dim_sel,) = draw( dim_strategy | st.integers(min_value=0, max_value=nchunks - 1).map(lambda i: (i,)) ) @@ -761,9 +761,9 @@ def block_test_arrays( - **regular**: a regular chunk grid, optionally wrapped in sharding. - **rectilinear**: a variable (rectilinear) chunk grid, always unsharded. - Returns ``(zarray, nparray)``. The per-axis block sizes the oracle needs are - ``zarray.write_chunk_sizes`` — the array's *outer* (block / shard) grid, which - is exactly the grid ``Array.blocks`` addresses; the caller reads it directly. + Returns `(zarray, nparray)`. The per-axis block sizes the oracle needs are + `zarray.write_chunk_sizes` — the array's *outer* (block / shard) grid, which + is exactly the grid `Array.blocks` addresses; the caller reads it directly. """ chunks: tuple[int, ...] | list[list[int]] if draw(st.booleans()): @@ -809,10 +809,10 @@ def key_ranges( [(key, byte_request), (key, byte_request),...] - where ``byte_request`` is ``None`` or any of the concrete ``ByteRequest`` + where `byte_request` is `None` or any of the concrete `ByteRequest` subtypes. The bounds are drawn independently of each value's length, so the offsets/suffixes routinely exceed the data and exercise the clamping logic - in ``_normalize_byte_range_index``. + in `_normalize_byte_range_index`. """ def make_range(start: int, length: int) -> RangeByteRequest: @@ -841,7 +841,7 @@ def complex_rectilinear_arrays( """Generate a rectilinear array with many small chunks. The shape is derived from the chunk edges (5-10 chunks per dim, - sizes 1-5), exercising higher chunk counts than ``rectilinear_arrays``. + sizes 1-5), exercising higher chunk counts than `rectilinear_arrays`. """ ndim = draw(st.integers(min_value=1, max_value=3)) nchunks = draw(st.integers(min_value=5, max_value=10)) diff --git a/tests/test_properties.py b/tests/test_properties.py index cad26ba7cd..05960f318c 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -120,7 +120,7 @@ def test_array_creates_implicit_groups(array): # Eager basic/oindex/vindex/mask indexing is exercised comprehensively (read, -# out=, async read, and write) by the single oracle ``test_indexing_parity`` +# out=, async read, and write) by the single oracle `test_indexing_parity` # defined below. Special array shapes that the oracle's strategies don't cover # keep dedicated tests (complex rectilinear arrays here; block indexing later). @@ -341,7 +341,7 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N def _get(target: zarr.Array, mode: IndexMode, zsel: Any, *, out: Any = None) -> Any: - """Read ``zsel`` from ``target`` via the get-method for ``mode``.""" + """Read `zsel` from `target` via the get-method for `mode`.""" if mode == "basic": return target.get_basic_selection(zsel, out=out) if mode == "oindex": @@ -354,7 +354,7 @@ def _get(target: zarr.Array, mode: IndexMode, zsel: Any, *, out: Any = None) -> def _async_get(async_array: Any, mode: IndexMode, zsel: Any) -> Any: - """The async read coroutine for ``mode`` (vindex/mask share the vectorized accessor).""" + """The async read coroutine for `mode` (vindex/mask share the vectorized accessor).""" if mode == "basic": return async_array.getitem(zsel) if mode == "oindex": @@ -363,7 +363,7 @@ def _async_get(async_array: Any, mode: IndexMode, zsel: Any) -> Any: def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None: - """Write ``value`` at ``zsel`` via the set-method for ``mode``.""" + """Write `value` at `zsel` via the set-method for `mode`.""" if mode == "basic": zarray[zsel] = value elif mode == "oindex": @@ -377,9 +377,9 @@ def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) -> bool: - """Whether a write to ``npsel`` targets each cell at most once. + """Whether a write to `npsel` targets each cell at most once. - Negative indices are normalized first (``[2, -2]`` on length 4 both target cell + Negative indices are normalized first (`[2, -2]` on length 4 both target cell 2), then duplicate targets are detected. A repeated target makes the write order-dependent — NumPy has the same ambiguity — so there is no well-defined oracle and such an example is skipped (a fundamental limitation, not a bug). @@ -415,7 +415,7 @@ def _n_array_axes(npsel: Any) -> int: def _eligible(mode: IndexMode, shape: tuple[int, ...]) -> bool: - """Whether ``mode`` can be exercised on ``shape``. + """Whether `mode` can be exercised on `shape`. Rank-0 arrays have no interesting selections; the fancy modes (oindex/vindex/mask via integer/boolean arrays) can't handle zero-size axes. @@ -428,10 +428,10 @@ def _eligible(mode: IndexMode, shape: tuple[int, ...]) -> bool: def assert_read_matches_numpy( target: zarr.Array, ref: np.ndarray[Any, Any], mode: IndexMode, zsel: Any, npsel: Any ) -> None: - """Assert ``target``'s read of ``zsel`` (mode) matches ``ref[npsel]``, with/without out=. + """Assert `target`'s read of `zsel` (mode) matches `ref[npsel]`, with/without out=. - Consumer-agnostic: ``target`` is any ``zarr.Array`` (an eager array or a lazy - view) and ``ref`` is the NumPy array it must behave like. The ``out=`` buffer + Consumer-agnostic: `target` is any `zarr.Array` (an eager array or a lazy + view) and `ref` is the NumPy array it must behave like. The `out=` buffer is flat for the vectorized vindex/mask modes (which scatter through a flat index) and full-shape otherwise. """ @@ -490,7 +490,7 @@ async def test_indexing_read_parity(data: st.DataObject) -> None: async def test_indexing_write_parity(data: st.DataObject) -> None: """Every *supported, well-defined* indexing write on an eager array matches NumPy. - Two families of selections are filtered out with ``assume`` (they are not + Two families of selections are filtered out with `assume` (they are not failures to test here): - **Repeated coordinates** (vindex/oindex selecting a cell more than once): @@ -498,7 +498,7 @@ async def test_indexing_write_parity(data: st.DataObject) -> None: ambiguity. This is a fundamental limitation, not a zarr bug. - **Sharded orthogonal writes across >= 2 array axes**: unsupported by the sharding partial-write codec; pinned at the pytest level by - ``test_oindex_sharded_multiaxis_write_xfail`` rather than hidden here. + `test_oindex_sharded_multiaxis_write_xfail` rather than hidden here. """ zarray = _indexing_array(data) nparray = zarray[:] @@ -560,8 +560,8 @@ def test_write_is_unambiguous( """A write is well-defined iff each target cell is hit at most once *after* negative indices are normalized. - Regression anchor for the CI fix: vindex ``[2, -2, 0, 1]`` on length 4 maps - ``-2 -> 2``, so cell 2 is written twice (order-dependent) and must be rejected, + Regression anchor for the CI fix: vindex `[2, -2, 0, 1]` on length 4 maps + `-2 -> 2`, so cell 2 is written twice (order-dependent) and must be rejected, even though the raw values are all distinct. """ assert _write_is_unambiguous(mode, npsel, shape) is well_defined @@ -569,8 +569,8 @@ def test_write_is_unambiguous( @st.composite def _bounds_selection(draw: st.DrawFn, mode: IndexMode, shape: tuple[int, ...]) -> Any: - """An integer selection whose components may fall outside ``[-size, size)``, to - probe out-of-bounds behavior. Returns ``(zarr_sel, numpy_sel)``.""" + """An integer selection whose components may fall outside `[-size, size)`, to + probe out-of-bounds behavior. Returns `(zarr_sel, numpy_sel)`.""" wide = {s: st.integers(-2 * s - 1, 2 * s) for s in set(shape)} if mode == "basic": sel = tuple(draw(wide[s]) for s in shape) @@ -591,7 +591,7 @@ def test_indexing_bounds_error_parity(data: st.DataObject) -> None: zarr's bounds errors (BoundsCheckError etc.) subclass IndexError, so the eager methods must agree with NumPy on *whether* an index is in bounds. This is the - ndindex ``check_same`` error-parity idea, and it catches the silent-wrong-data + ndindex `check_same` error-parity idea, and it catches the silent-wrong-data class — returning garbage where NumPy would raise. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=6))) @@ -611,10 +611,10 @@ class _IndexingStateMachine(RuleBasedStateMachine): """Apply many random indexing writes to one array while a NumPy model tracks it in lockstep; the model must match after every step. - This is the TensorStore ``driver_testutil`` pattern with NumPy as the oracle: + This is the TensorStore `driver_testutil` pattern with NumPy as the oracle: a single array accumulates many writes, so it catches read-modify-write, chunk-boundary, and shard-merge/persistence bugs that the single-shot - ``test_indexing_write_parity`` cannot see. Concrete geometries are defined by + `test_indexing_write_parity` cannot see. Concrete geometries are defined by subclasses; the rules are dtype-agnostic (fixed i4 — dtype breadth lives in the property tests). """ @@ -686,11 +686,11 @@ class _ThreeDStateMachine(_IndexingStateMachine): async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: """The eager read oracle, applied to a lazy view (the lazy consumer). - A view (built from a non-negative window) is just another ``zarr.Array``, so - it flows through the same ``assert_read_matches_numpy`` harness. The lazy + A view (built from a non-negative window) is just another `zarr.Array`, so + it flows through the same `assert_read_matches_numpy` harness. The lazy indexing bugs found in review were all "the view path diverges from NumPy for some (method, parameter) combination" — enumerating that surface here catches - the class. Skipped until ``Array.lazy`` exists, so the eager oracle can merge + the class. Skipped until `Array.lazy` exists, so the eager oracle can merge ahead of the lazy feature. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) From 04dab0aad702cca25861b8bce9fc271f9e002f56 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 21:35:34 +0200 Subject: [PATCH 28/43] docs: clarify why repeated-target writes are rejected in _write_is_unambiguous MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NumPy's deterministic last-write-wins is an implementation detail of serial, C-order index iteration, not a guarantee. zarr writes scalars into chunks and cannot in general promise a write order, so it cannot be expected to match NumPy's choice — hence no well-defined oracle for duplicate-target writes. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 05960f318c..8e62642965 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -380,9 +380,13 @@ def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) - """Whether a write to `npsel` targets each cell at most once. Negative indices are normalized first (`[2, -2]` on length 4 both target cell - 2), then duplicate targets are detected. A repeated target makes the write - order-dependent — NumPy has the same ambiguity — so there is no well-defined - oracle and such an example is skipped (a fundamental limitation, not a bug). + 2), then duplicate targets are detected. A repeated target makes the surviving + value depend on the order the writes are applied. NumPy happens to be + deterministic here (last-write-wins), but only as an implementation detail of + iterating the index serially in C order — it is not a guarantee. zarr writes + scalars into chunks and cannot in general promise a write order, so it cannot + be expected to match NumPy's choice. There is therefore no well-defined oracle + for such a write, and we reject it (a fundamental limitation, not a bug). """ sel = npsel if isinstance(npsel, tuple) else (npsel,) if mode == "oindex": From a192ca4773d05253b8d0ce076f70e369ff29be39 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 30 Jun 2026 21:41:53 +0200 Subject: [PATCH 29/43] test(indexing): use the Expect helper for test_write_is_unambiguous cases Replace raw pytest.param tuples with the conftest Expect dataclass (input / output / id), matching the established idiom in test_chunk_grids etc. Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 35 ++++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 8e62642965..55fb349fa4 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -18,6 +18,7 @@ from hypothesis import assume, given, settings from hypothesis.stateful import RuleBasedStateMachine, invariant, rule +from tests.conftest import Expect from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata @@ -544,22 +545,33 @@ def test_oindex_sharded_multiaxis_write_xfail() -> None: @pytest.mark.parametrize( - ("mode", "npsel", "shape", "well_defined"), + "case", [ - pytest.param("vindex", (np.array([2, -2, 0, 1]),), (4,), False, id="vindex-neg-collision"), - pytest.param("vindex", (np.array([0, 1, 2]),), (4,), True, id="vindex-distinct"), - pytest.param("oindex", (np.array([1, -3]),), (4,), False, id="oindex-neg-collision"), - pytest.param( - "oindex", (np.array([0, 2]), np.array([1, 1])), (4, 4), False, id="oindex-repeat" + Expect( + input=("vindex", (np.array([2, -2, 0, 1]),), (4,)), + output=False, + id="vindex-neg-collision", ), - pytest.param( - "oindex", (np.array([0, 2]), np.array([1, 3])), (4, 4), True, id="oindex-distinct" + Expect(input=("vindex", (np.array([0, 1, 2]),), (4,)), output=True, id="vindex-distinct"), + Expect( + input=("oindex", (np.array([1, -3]),), (4,)), output=False, id="oindex-neg-collision" ), - pytest.param("basic", (slice(None),), (4,), True, id="basic-always"), + Expect( + input=("oindex", (np.array([0, 2]), np.array([1, 1])), (4, 4)), + output=False, + id="oindex-repeat", + ), + Expect( + input=("oindex", (np.array([0, 2]), np.array([1, 3])), (4, 4)), + output=True, + id="oindex-distinct", + ), + Expect(input=("basic", (slice(None),), (4,)), output=True, id="basic-always"), ], + ids=lambda c: c.id, ) def test_write_is_unambiguous( - mode: IndexMode, npsel: Any, shape: tuple[int, ...], well_defined: bool + case: Expect[tuple[IndexMode, Any, tuple[int, ...]], bool], ) -> None: """A write is well-defined iff each target cell is hit at most once *after* negative indices are normalized. @@ -568,7 +580,8 @@ def test_write_is_unambiguous( `-2 -> 2`, so cell 2 is written twice (order-dependent) and must be rejected, even though the raw values are all distinct. """ - assert _write_is_unambiguous(mode, npsel, shape) is well_defined + mode, npsel, shape = case.input + assert _write_is_unambiguous(mode, npsel, shape) is case.output @st.composite From 34d3b1bcf14d4cf24552d2f2ecc9bad83544bbcb Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 1 Jul 2026 15:27:33 +0200 Subject: [PATCH 30/43] feat(array): guard grid-describing members on lazy views (LazyViewError) Layer C of the lazy-view chunk-layout design. Members that assume the array fills its chunk grid now raise LazyViewError (a NotImplementedError subclass) on a non-identity lazy view instead of silently describing the backing grid: chunks, shards, read_chunk_sizes, write_chunk_sizes, cdata_shape, nchunks, nchunks_initialized, nbytes_stored, resize, append, info, info_complete. Normal (identity) arrays are unaffected, and views are new, so nothing breaks. Guards go through a single AsyncArray._require_identity helper (sync members delegate). Also fix size/nbytes, which are logical "keep" members: they now derive from the view's shape rather than the backing metadata shape (a.lazy[2:10].size == 8, not 12). cdata_shape's sync accessor is redirected to the guarded public property. Design: docs/superpowers/specs/2026-07-01-lazy-view-chunk-layout-design.md (gist: https://gist.github.com/d-v-b/f50c93bc5f365e6986dadba6e4aa1902). Next: Layer A (ChunkLayout, #4040) and Layer B (chunk_projections partition API). Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 34 ++++++++++++++++++++-- src/zarr/errors.py | 13 +++++++++ tests/test_lazy_indexing.py | 56 +++++++++++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 3 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 6e4b1301be..354ec8366f 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -147,6 +147,7 @@ from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, + LazyViewError, MetadataValidationError, ZarrDeprecationWarning, ZarrUserWarning, @@ -850,6 +851,20 @@ def _is_identity(self) -> bool: """ return _transform_is_identity(self._transform, self.metadata.shape) + def _require_identity(self, name: str) -> None: + """Raise `LazyViewError` if this array is a non-identity lazy view. + + Grid-describing/mutating members assume the array fills its chunk grid, + which a view (a sliced/indexed array) generally does not. See + `LazyViewError`; `name` is the member being guarded, for the message. + """ + if not self._is_identity: + raise LazyViewError( + f"`{name}` is not defined for a lazy view (a sliced or indexed array that " + f"does not fill its chunk grid). Use `chunk_layout` for the backing array's " + f"structure, or `chunk_projections` for this view's chunk granularity." + ) + @property def storage_shape(self) -> tuple[int, ...]: """The shape of the underlying storage array (ignoring any view transform).""" @@ -868,6 +883,7 @@ def chunks(self) -> tuple[int, ...]: tuple[int, ...]: The chunk shape of the Array. """ + self._require_identity("chunks") # TODO: move sharding awareness out of metadata return self.metadata.chunks @@ -894,6 +910,7 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: >>> arr.read_chunk_sizes ((30, 30, 30, 10), (40, 40)) """ + self._require_identity("read_chunk_sizes") from zarr.codecs.sharding import ShardingCodec @@ -925,7 +942,7 @@ def write_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: >>> arr.write_chunk_sizes ((30, 30, 30, 10), (40, 40)) """ - + self._require_identity("write_chunk_sizes") return self._chunk_grid.chunk_sizes @property @@ -941,6 +958,7 @@ def shards(self) -> tuple[int, ...] | None: tuple[int, ...]: The shard shape of the Array. """ + self._require_identity("shards") return self.metadata.shards @property @@ -952,7 +970,9 @@ def size(self) -> int: int Total number of elements in the array """ - return math.prod(self.metadata.shape) + # `self.shape` is the view's logical shape (== metadata.shape for a + # non-view); using it keeps `size`/`nbytes` correct for lazy views. + return math.prod(self.shape) @property def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]: @@ -1118,6 +1138,7 @@ def cdata_shape(self) -> tuple[int, ...]: tuple[int, ...] The number of chunks along each dimension. """ + self._require_identity("cdata_shape") return self._chunk_grid_shape @property @@ -1174,6 +1195,7 @@ def nchunks(self) -> int: int The total number of chunks in the array. """ + self._require_identity("nchunks") return product(self._chunk_grid_shape) @property @@ -1258,6 +1280,7 @@ async def example(): result = asyncio.run(example()) ``` """ + self._require_identity("nchunks_initialized") return await _nchunks_initialized(self) async def _nshards_initialized(self) -> int: @@ -1299,6 +1322,7 @@ async def example(): return await _nshards_initialized(self) async def nbytes_stored(self) -> int: + self._require_identity("nbytes_stored") return await _nbytes_stored(self.store_path) def _iter_chunk_coords( @@ -1740,6 +1764,7 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) ----- - This method is asynchronous and should be awaited. """ + self._require_identity("resize") return await _resize(self, new_shape, delete_outside_chunks) async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: @@ -1761,6 +1786,7 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: The size of all dimensions other than `axis` must match between this array and `data`. """ + self._require_identity("append") return await _append(self, data, axis) async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: @@ -1832,6 +1858,7 @@ def info(self) -> Any: Compressors : (ZstdCodec(level=0, checksum=False),) No. bytes : 480 """ + self._require_identity("info") return self._info() async def info_complete(self) -> Any: @@ -1851,6 +1878,7 @@ async def info_complete(self) -> Any: ------- [zarr.AsyncArray.info][] - A property giving just the statically known information about an array. """ + self._require_identity("info_complete") return await _info_complete(self) def _info( @@ -2253,7 +2281,7 @@ def cdata_shape(self) -> tuple[int, ...]: When sharding is used, this counts inner chunks (not shards) per dimension. """ - return self.async_array._chunk_grid_shape + return self.async_array.cdata_shape @property def _chunk_grid_shape(self) -> tuple[int, ...]: diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 781bebe534..e59c55da35 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -9,6 +9,7 @@ "ContainsGroupError", "DataTypeValidationError", "GroupNotFoundError", + "LazyViewError", "MetadataValidationError", "NegativeStepError", "NodeTypeValidationError", @@ -151,6 +152,18 @@ class BoundsCheckError(IndexError): ... class ArrayIndexError(IndexError): ... +class LazyViewError(NotImplementedError): + """Raised when an operation that assumes an array fills its chunk grid is used + on a non-identity lazy view (created via ``Array.lazy[...]``). + + Grid-describing members (``chunks``, ``shards``, ``nchunks``, ``read_chunk_sizes``, + ...) and grid-mutating ones (``resize``, ``append``) have no well-defined answer + for a view onto a subset of the backing grid. Use ``chunk_layout`` for the backing + structure and ``chunk_projections`` for the view's granularity. Subclasses + ``NotImplementedError`` so existing consumers that catch it keep working. + """ + + class ChunkNotFoundError(BaseZarrError): """ Raised when a chunk that was expected to exist in storage was not retrieved successfully. diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 79e37b9570..28e9d6cbb4 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -23,6 +23,7 @@ import zarr from zarr.core.buffer import default_buffer_prototype +from zarr.errors import LazyViewError from zarr.storage import MemoryStore @@ -474,3 +475,58 @@ def test_oindex_single_axis_writes_track_numpy(self, cfg: Config) -> None: a.lazy.oindex[sel] = val np.testing.assert_array_equal(a.lazy.oindex[sel][...], ref[sel]) np.testing.assert_array_equal(a[...], ref) + + +_VIEW_GUARDED_PROPERTIES = ( + "chunks", + "shards", + "read_chunk_sizes", + "write_chunk_sizes", + "cdata_shape", + "nchunks", + "nchunks_initialized", + "info", +) + +# method name -> call arguments +_VIEW_GUARDED_METHODS: dict[str, tuple[Any, ...]] = { + "nbytes_stored": (), + "info_complete": (), + "resize": ((24,),), + "append": (np.zeros((3,), dtype="i4"),), +} + + +class TestLazyViewGridGuards: + """Grid-describing members assume the array fills its chunk grid, so on a + non-identity lazy view they must raise instead of silently describing the + backing grid (a footgun for consumers that size reads off ``.chunks``).""" + + @staticmethod + def _array() -> zarr.Array[Any]: + a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + return a + + @pytest.mark.parametrize("name", _VIEW_GUARDED_PROPERTIES) + def test_property_raises_on_view(self, name: str) -> None: + """A grid-describing property works on the backing array but raises LazyViewError on a view.""" + a = self._array() + getattr(a, name) # backing (identity) array: no raise + with pytest.raises(LazyViewError): + getattr(a.lazy[2:10], name) + + @pytest.mark.parametrize("name", list(_VIEW_GUARDED_METHODS)) + def test_method_raises_on_view(self, name: str) -> None: + """A grid-describing/mutating method raises LazyViewError on a view.""" + a = self._array() + with pytest.raises(LazyViewError): + getattr(a.lazy[2:10], name)(*_VIEW_GUARDED_METHODS[name]) + + def test_size_and_nbytes_reflect_the_view(self) -> None: + """size / nbytes are logical members: they describe the view's extent, not the backing array.""" + a = self._array() # shape (12,), i4 (4 bytes) + view = a.lazy[2:10] # logical shape (8,) + assert view.shape == (8,) + assert view.size == 8 + assert view.nbytes == 8 * 4 From 16e96dc2fb1ef9067cb73f43a052744a65b78743 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 1 Jul 2026 15:43:53 +0200 Subject: [PATCH 31/43] feat(array): add view-aware chunk_projections partition API (Layer B) Layer B of the lazy-view chunk-layout design: expose the L4 partition that a lazy view already resolves. `Array.chunk_projections(unit="read"|"write")` yields a `ChunkProjection(coord, key, shape, chunk_selection, array_selection, is_partial)` per stored chunk the array/view projects onto, reusing the I/O path's `iter_chunk_transforms` + `sub_transform_to_selections`. `is_chunk_aligned()` is the thin L3 wrapper. - Identity arrays: projections tile the whole domain (none partial), recovering the chunk grid. Views: only touched chunks, with partial-coverage flagged (RMW). - Arbitrary selections compose through the lazy accessor: `array.lazy[sel].chunk_projections()` (no `selection=` param). - `unit="write"` is the store-object (shard, else chunk) grid; `unit="read"` equals it unless sharded. Inner-chunk read partitioning of sharded arrays is deferred (raises NotImplementedError, pointing at unit="write"). - New public `zarr.ChunkProjection`; `shape` is extent-clipped via the grid's chunk_sizes so boundary chunks are not misflagged partial. Layer A (the static ChunkLayout, #4040) is intentionally left to maxrjones's PR. Design: gist f50c93bc5f365e6986dadba6e4aa1902. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/__init__.py | 2 + src/zarr/core/array.py | 70 +++++++++++++++++++++++ src/zarr/core/chunk_partition.py | 98 ++++++++++++++++++++++++++++++++ tests/test_lazy_indexing.py | 64 +++++++++++++++++++++ 4 files changed, 234 insertions(+) create mode 100644 src/zarr/core/chunk_partition.py diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index cdf3840c3b..cbda7ad0ae 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -35,6 +35,7 @@ zeros_like, ) from zarr.core.array import Array, AsyncArray +from zarr.core.chunk_partition import ChunkProjection from zarr.core.config import config from zarr.core.group import AsyncGroup, Group @@ -146,6 +147,7 @@ def set_format(log_format: str) -> None: "Array", "AsyncArray", "AsyncGroup", + "ChunkProjection", "Group", "__version__", "array", diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 354ec8366f..07a1dd1286 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -56,6 +56,7 @@ V2ChunkKeyEncoding, parse_chunk_key_encoding, ) +from zarr.core.chunk_partition import ChunkProjection, iter_chunk_projections from zarr.core.common import ( JSON, ZARR_JSON, @@ -865,6 +866,54 @@ def _require_identity(self, name: str) -> None: f"structure, or `chunk_projections` for this view's chunk granularity." ) + @property + def _is_sharded(self) -> bool: + """Whether the array stores inner chunks inside shards (a sharding codec).""" + from zarr.codecs.sharding import ShardingCodec + + codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) + return len(codecs) == 1 and isinstance(codecs[0], ShardingCodec) + + def chunk_projections( + self, *, unit: Literal["read", "write"] = "read" + ) -> Iterator[ChunkProjection]: + """Enumerate the stored chunks this array (or lazy view) projects onto. + + Yields a `ChunkProjection` per stored chunk: its coordinate, store key, and + (extent-clipped) shape; the region of the chunk this array covers; the region + of this array it maps to; and whether the coverage is partial (a partial + write is a read-modify-write). For an identity array every chunk is fully + covered and the projections tile the whole domain; for a view only the touched + chunks appear. + + `unit` selects the granularity: `"write"` is the store-object grid (the shard + when sharded, else the chunk); `"read"` is the chunk grid. They coincide + unless the array is sharded. Read-unit (inner-chunk) partitioning of a sharded + array is not yet implemented; use `unit="write"` there. + + To partition an arbitrary selection, compose through the lazy accessor: + `array.lazy[sel].chunk_projections()`. + """ + if unit not in ("read", "write"): + raise ValueError(f"unit must be 'read' or 'write', got {unit!r}") + if unit == "read" and self._is_sharded: + raise NotImplementedError( + "read-unit (inner-chunk) `chunk_projections` for sharded arrays is not yet " + "implemented; use `unit='write'` for shard-granularity projections." + ) + return iter_chunk_projections( + self._transform, self._chunk_grid, self.metadata.encode_chunk_key + ) + + def is_chunk_aligned(self) -> bool: + """Whether this array/view aligns to write-unit (store-object) boundaries. + + True iff no stored write unit is only partially covered — i.e. every unit can + be written without a read-modify-write. A cheap wrapper over + `chunk_projections`. + """ + return all(not p.is_partial for p in self.chunk_projections(unit="write")) + @property def storage_shape(self) -> tuple[int, ...]: """The shape of the underlying storage array (ignoring any view transform).""" @@ -2359,6 +2408,27 @@ def nbytes(self) -> int: """ return self.async_array.nbytes + def chunk_projections( + self, *, unit: Literal["read", "write"] = "read" + ) -> Iterator[ChunkProjection]: + """Enumerate the stored chunks this array (or lazy view) projects onto. + + See [zarr.AsyncArray.chunk_projections][] for the full description. Each + `ChunkProjection` reports a stored chunk's coordinate/key/shape, the region of + it this array covers, the region of this array it maps to, and whether the + coverage is partial. Compose through `lazy` to partition an arbitrary + selection: `array.lazy[sel].chunk_projections()`. + """ + return self.async_array.chunk_projections(unit=unit) + + def is_chunk_aligned(self) -> bool: + """Whether this array/view aligns to write-unit (store-object) boundaries. + + True iff no stored write unit is only partially covered. See + [zarr.AsyncArray.is_chunk_aligned][]. + """ + return self.async_array.is_chunk_aligned() + @property def nchunks_initialized(self) -> int: """ diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py new file mode 100644 index 0000000000..7edc586183 --- /dev/null +++ b/src/zarr/core/chunk_partition.py @@ -0,0 +1,98 @@ +"""View-aware chunk partitioning (Layer B of the lazy-view chunk-layout design). + +A lazy view is an ``IndexTransform`` applied to a backing array — i.e. a stored +selection already resolved. ``chunk_projections`` enumerates the stored chunks that +selection (or a whole identity array) projects onto, reusing the same resolution +machinery the read/write I/O path uses (``iter_chunk_transforms`` + +``sub_transform_to_selections``). Each projection reports the stored chunk, the +region of it this array covers, the region of this array it maps to, and whether the +coverage is partial (a partial write is a read-modify-write). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from zarr.core.transforms.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from zarr.core.chunk_grids import ChunkGrid + from zarr.core.transforms.transform import IndexTransform + +__all__ = ["ChunkProjection", "iter_chunk_projections"] + + +@dataclass(frozen=True) +class ChunkProjection: + """How one stored chunk contributes to an array (or lazy view). + + Attributes + ---------- + coord + Coordinate of the stored chunk in the chunk grid. + key + Store key of the stored chunk (which file). + shape + Stored size of the chunk, clipped to the array extent. + chunk_selection + The region *within* the stored chunk that this array covers. + array_selection + The region of this array (or view) that the chunk maps to. + is_partial + Whether the chunk is only partially covered — a write to a partial chunk + is a read-modify-write. + """ + + coord: tuple[int, ...] + key: str + shape: tuple[int, ...] + chunk_selection: tuple[Any, ...] + array_selection: tuple[Any, ...] + is_partial: bool + + +def _covers_full_chunk(chunk_selection: tuple[Any, ...], shape: tuple[int, ...]) -> bool: + """Whether ``chunk_selection`` selects every element of a chunk of ``shape``. + + Only a per-dimension full ``0:size:1`` slice is a full cover; an integer or + array selection touches a strict subset, so the chunk is partial. + """ + if len(chunk_selection) != len(shape): + return False + for sel, size in zip(chunk_selection, shape, strict=True): + if not isinstance(sel, slice): + return False + start, stop, step = sel.indices(size) + if not (start == 0 and stop == size and step == 1): + return False + return True + + +def iter_chunk_projections( + transform: IndexTransform, + chunk_grid: ChunkGrid, + encode_key: Callable[[tuple[int, ...]], str], +) -> Iterator[ChunkProjection]: + """Yield a `ChunkProjection` for each stored chunk ``transform`` projects onto.""" + chunk_sizes = chunk_grid.chunk_sizes # per-dimension, extent-clipped + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): + if chunk_grid[chunk_coords] is None: + continue + chunk_selection, array_selection, _drop_axes = sub_transform_to_selections( + sub_transform, out_indices + ) + shape = tuple(chunk_sizes[dim][c] for dim, c in enumerate(chunk_coords)) + yield ChunkProjection( + coord=chunk_coords, + key=encode_key(chunk_coords), + shape=shape, + chunk_selection=chunk_selection, + array_selection=array_selection, + is_partial=not _covers_full_chunk(chunk_selection, shape), + ) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 28e9d6cbb4..339b9969fd 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -530,3 +530,67 @@ def test_size_and_nbytes_reflect_the_view(self) -> None: assert view.shape == (8,) assert view.size == 8 assert view.nbytes == 8 * 4 + + +class TestChunkProjections: + """`chunk_projections` enumerates the stored chunks an (identity or view) array + projects onto, giving each one's stored region, the region of this array it maps + to, and whether it is partially covered.""" + + @staticmethod + def _array(shape: tuple[int, ...], chunks: tuple[int, ...]) -> zarr.Array[Any]: + a = zarr.create_array({}, shape=shape, chunks=chunks, dtype="i4") + a[...] = np.arange(int(np.prod(shape)), dtype="i4").reshape(shape) + return a + + def test_identity_tiles_the_whole_domain(self) -> None: + """On a full (identity) array, projections tile the domain exactly: one per chunk, none partial.""" + a = self._array((10,), (3,)) # chunks (3,3,3,1) + projs = list(a.chunk_projections()) + assert len(projs) == 4 + assert all(not p.is_partial for p in projs) + covered = np.zeros(10, dtype=bool) + for p in projs: + arr_sel: Any = p.array_selection + covered[arr_sel] = True + assert covered.all() # complete, and (bool assignment) each cell once + + def test_view_round_trip(self) -> None: + """Placing each projection's stored-chunk slice at its array_selection reconstructs the view.""" + a = self._array((10,), (3,)) + backing = np.asarray(a[...]) + view = a.lazy[2:9] # (7,) + out = np.empty(view.shape, dtype="i4") + for p in view.chunk_projections(): + offset = p.coord[0] * 3 # regular grid: chunk c starts at c*chunk_size + stored_chunk = backing[offset : offset + p.shape[0]] + arr_sel: Any = p.array_selection + chunk_sel: Any = p.chunk_selection + out[arr_sel] = stored_chunk[chunk_sel] + np.testing.assert_array_equal(out, backing[2:9]) + + def test_partial_flag_and_alignment(self) -> None: + """Boundary chunks a view only partially covers are flagged; a chunk-aligned view is not.""" + a = self._array((12,), (4,)) # chunks at [0,4) [4,8) [8,12) + partial = a.lazy[2:12] # first chunk partially covered + assert any(p.is_partial for p in partial.chunk_projections()) + assert not partial.is_chunk_aligned() + aligned = a.lazy[4:12] # exactly chunks 1 and 2 + assert all(not p.is_partial for p in aligned.chunk_projections()) + assert aligned.is_chunk_aligned() + + def test_sharded_write_unit(self) -> None: + """For a sharded array, unit='write' enumerates shards; unit='read' (inner chunks) is deferred.""" + a = zarr.create_array({}, shape=(12,), chunks=(2,), shards=(6,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + shards = list(a.chunk_projections(unit="write")) + assert len(shards) == 2 # two shards of size 6 + assert all(p.shape == (6,) and not p.is_partial for p in shards) + with pytest.raises(NotImplementedError): + a.chunk_projections(unit="read") + + def test_invalid_unit_rejected(self) -> None: + """An unknown granularity is rejected eagerly.""" + a = self._array((6,), (2,)) + with pytest.raises(ValueError, match="unit"): + a.chunk_projections(unit="bogus") # type: ignore[arg-type] From 8f51eade75e18422d73a87780e8b23748cb7df1a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 1 Jul 2026 18:18:39 +0200 Subject: [PATCH 32/43] test(indexing): fix the write filter to inspect the raw zarr selection (roborev #341) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit roborev flagged that _write_is_unambiguous / _n_array_axes were fed the broadcast np.ix_-style `npsel` (what orthogonal_indices returns for NumPy), not the raw per-axis `zsel`. For oindex the broadcast form tiles every component across all axes, so the per-axis uniqueness check tripped on *tiling* not duplicates: it rejected essentially every multi-axis orthogonal write. Effect: test_indexing_write_parity `assume`d them away and the state machine's write rule returned early for them, so the stateful harness performed no multi-axis writes at all — silently gutting the coverage it exists for. (The feature itself is fine; covered by TestLazyOIndex.test_multi_axis_write.) Fix: pass the raw `zsel` to both helpers (correct on per-axis arrays; slices are skipped, so _n_array_axes counts genuine fancy axes for the sharded guard). Clarify the helper docstrings, rename the misleading `npsel` params, and add test_write_filter_rejects_broadcast_oindex pinning the raw-vs-broadcast invariant (the old regression test masked the bug by feeding raw arrays callers never sent). Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_properties.py | 59 ++++++++++++++++++++++++++++++---------- 1 file changed, 45 insertions(+), 14 deletions(-) diff --git a/tests/test_properties.py b/tests/test_properties.py index 55fb349fa4..d292f7ce55 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -377,8 +377,13 @@ def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None raise AssertionError(mode) -def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) -> bool: - """Whether a write to `npsel` targets each cell at most once. +def _write_is_unambiguous(mode: IndexMode, zsel: Any, shape: tuple[int, ...]) -> bool: + """Whether a write to `zsel` targets each cell at most once. + + `zsel` must be the *raw zarr selection* (per-axis arrays/slices), NOT the + broadcast np.ix_-style form NumPy needs: for oindex the broadcast form tiles + each component across every axis, so the per-axis uniqueness check below would + trip on tiling instead of genuine duplicates and reject well-defined writes. Negative indices are normalized first (`[2, -2]` on length 4 both target cell 2), then duplicate targets are detected. A repeated target makes the surviving @@ -389,7 +394,7 @@ def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) - be expected to match NumPy's choice. There is therefore no well-defined oracle for such a write, and we reject it (a fundamental limitation, not a bug). """ - sel = npsel if isinstance(npsel, tuple) else (npsel,) + sel = zsel if isinstance(zsel, tuple) else (zsel,) if mode == "oindex": # Independent axes: each fancy axis must select each index at most once. for axis, s in enumerate(sel): @@ -413,9 +418,13 @@ def _write_is_unambiguous(mode: IndexMode, npsel: Any, shape: tuple[int, ...]) - return True # basic / mask never target a cell twice -def _n_array_axes(npsel: Any) -> int: - """Number of integer-array (fancy) axes in a selection.""" - sel = npsel if isinstance(npsel, tuple) else (npsel,) +def _n_array_axes(zsel: Any) -> int: + """Number of integer-array (fancy) axes in a *raw* zarr selection (slices excluded). + + Must be the raw per-axis selection, not the broadcast numpy form: the latter + turns every axis into an array, so this would always return ``ndim``. + """ + sel = zsel if isinstance(zsel, tuple) else (zsel,) return sum(isinstance(s, np.ndarray) for s in sel) @@ -511,9 +520,9 @@ async def test_indexing_write_parity(data: st.DataObject) -> None: assume(_eligible(mode, nparray.shape)) zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=nparray.shape)) if mode in ("oindex", "vindex"): - assume(_write_is_unambiguous(mode, npsel, nparray.shape)) + assume(_write_is_unambiguous(mode, zsel, nparray.shape)) if mode == "oindex" and zarray.shards is not None: - assume(_n_array_axes(npsel) < 2) + assume(_n_array_axes(zsel) < 2) expected = np.asarray(nparray[npsel]) new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) @@ -576,12 +585,34 @@ def test_write_is_unambiguous( """A write is well-defined iff each target cell is hit at most once *after* negative indices are normalized. - Regression anchor for the CI fix: vindex `[2, -2, 0, 1]` on length 4 maps - `-2 -> 2`, so cell 2 is written twice (order-dependent) and must be rejected, + Cases use the *raw zarr selection* (per-axis arrays) — the form real callers + pass — so `oindex-distinct` ([0,2] x [1,3]) is well-defined. See + `test_write_filter_rejects_broadcast_oindex` for why the broadcast form must not + be passed. Regression anchor for the CI fix: vindex `[2, -2, 0, 1]` on length 4 + maps `-2 -> 2`, so cell 2 is written twice (order-dependent) and must be rejected even though the raw values are all distinct. """ - mode, npsel, shape = case.input - assert _write_is_unambiguous(mode, npsel, shape) is case.output + mode, zsel, shape = case.input + assert _write_is_unambiguous(mode, zsel, shape) is case.output + + +def test_write_filter_rejects_broadcast_oindex() -> None: + """The write filter must be fed the raw per-axis selection, not the broadcast form. + + `numpy_array_indexers(mode="oindex")` returns a raw per-axis `zsel` and a broadcast + np.ix_-style `npsel`. A well-defined multi-axis write ([0,2] x [1,3]) is admitted + on the raw form but — because the broadcast form tiles each component across all + axes — is wrongly rejected on `npsel`. This pins the review finding that passing + `npsel` silently skipped essentially every multi-axis orthogonal write. + """ + zsel = (np.array([0, 2]), np.array([1, 3])) # raw per-axis, distinct -> well-defined + assert _write_is_unambiguous("oindex", zsel, (4, 4)) is True + # the fully-tiled form orthogonal_indices produces (np.broadcast_arrays over np.ix_) + broadcast = tuple(np.broadcast_arrays(*np.ix_(*zsel))) + assert _write_is_unambiguous("oindex", broadcast, (4, 4)) is False # tiling; do not pass + # _n_array_axes counts only genuine fancy axes on the raw form (slices excluded), + # so a single-fancy-axis oindex is correctly allowed on a sharded array. + assert _n_array_axes((slice(None), np.array([1, 3]))) == 1 @st.composite @@ -654,9 +685,9 @@ def write(self, data: st.DataObject, mode: IndexMode) -> None: if not _eligible(mode, self.shape): return zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=self.shape)) - if mode in ("oindex", "vindex") and not _write_is_unambiguous(mode, npsel, self.shape): + if mode in ("oindex", "vindex") and not _write_is_unambiguous(mode, zsel, self.shape): return - if mode == "oindex" and self.shards is not None and _n_array_axes(npsel) >= 2: + if mode == "oindex" and self.shards is not None and _n_array_axes(zsel) >= 2: return # GH2834, see test_oindex_sharded_multiaxis_write_xfail expected = np.asarray(self.model[npsel]) value = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=self.model.dtype)) From 97b0947e649232826d23cccedd318758b32e1a71 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 08:34:56 +0200 Subject: [PATCH 33/43] =?UTF-8?q?fix(lazy):=20make=20negative=20slice=20bo?= =?UTF-8?q?unds=20literal=20=E2=80=94=20one=20consistent=20coordinate=20ru?= =?UTF-8?q?le?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User-probe feedback showed the lazy path's negative-index handling split three ways: integers literal (plain IndexError, "domain" phrasing), index arrays literal (BoundsCheckError, "length" phrasing), but slice bounds silently wrapped from-the-end. Special-casing slice bounds is not consistent; the rule is now one sentence: a lazy selection is a declaration in literal coordinates, domains start at 0, so a negative value is never in bounds regardless of syntactic form. - Transform layer: reject negative slice start/stop (_check_slice_bounds) at all three lowering sites (basic/oindex/vindex). Positive overflow still clamps — a slice denotes a range and ranges intersect the domain; clamping can only shorten, never silently select different data the way wraparound does. - Eager dialect preserved: _normalize_negative_indices now wraps negative slice bounds NumPy-style before lowering, so eager methods on views keep full NumPy semantics (v[-3:] still works); only the .lazy accessor is literal. - Unified error type and phrasing: all literal-coordinate rejections raise BoundsCheckError (an IndexError) saying "valid indices [lo, hi)" — dropping the repr/error "domain" wording collision — plus a hint that negatives are not from-the-end and what to use instead. - Fix selection_repr crash on 0-d index arrays (len -> size), so integer-indexed fancy views can at least be displayed. - LazyViewError message no longer recommends the not-yet-existing chunk_layout; points at chunk_projections + metadata/chunk_grid. - Tests: literal slice bounds (reads/writes/views/oindex), positive-clamp retention, eager-dialect retention, unified error type, guard message, repr regression; strict-xfail pins for the known int-on-fancy-picked-dim defect (read crash; vindex scalar shape) to be fixed separately. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 5 +- src/zarr/core/transforms/transform.py | 67 +++++++++++++-- src/zarr/errors.py | 5 +- tests/test_lazy_indexing.py | 112 +++++++++++++++++++++++++- 4 files changed, 175 insertions(+), 14 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 07a1dd1286..5653d9b26a 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -862,8 +862,9 @@ def _require_identity(self, name: str) -> None: if not self._is_identity: raise LazyViewError( f"`{name}` is not defined for a lazy view (a sliced or indexed array that " - f"does not fill its chunk grid). Use `chunk_layout` for the backing array's " - f"structure, or `chunk_projections` for this view's chunk granularity." + f"does not fill its chunk grid). Use `chunk_projections` for this view's " + f"chunk granularity; the backing array's stored structure is available via " + f"`metadata` / `chunk_grid`." ) @property diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 68f32e7ced..e786df8fe2 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -114,7 +114,7 @@ def selection_repr(self) -> str: parts.append(f"[{start}, {stop}) step {m.stride}") elif isinstance(m, ArrayMap): storage = m.offset + m.stride * m.index_array - n = len(storage) + n = int(storage.size) # .size, not len(): index_array may be 0-d if n <= 5: vals = ", ".join(str(int(v)) for v in storage.ravel()) parts.append("{" + vals + "}") @@ -520,8 +520,10 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra hi = transform.domain.exclusive_max[old_dim] idx = sel if idx < lo or idx >= hi: - raise IndexError( - f"index {sel} is out of bounds for dimension {old_dim} with domain [{lo}, {hi})" + hint = _LITERAL_HINT if sel < 0 else "" + raise BoundsCheckError( + f"index {sel} is out of bounds for dimension {old_dim} " + f"(valid indices [{lo}, {hi})){hint}" ) dropped_dims.add(old_dim) dim_int_val[old_dim] = idx @@ -531,6 +533,9 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo + # Literal coordinates: negative bounds are out of the domain, not + # from-the-end (the eager dialect wraps them before reaching here). + _check_slice_bounds(sel, old_dim, lo, hi) # Resolve slice relative to the current domain (origin-based) start, stop, step = sel.indices(dim_size) # start, stop, step are now relative to a 0-based range of size dim_size @@ -675,6 +680,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo + _check_slice_bounds(sel, old_dim, lo, hi) start, stop, step = sel.indices(dim_size) if step <= 0: raise IndexError("slice step must be positive") @@ -843,6 +849,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo + _check_slice_bounds(sel, old_dim, lo, hi) start, stop, step = sel.indices(dim_size) if step <= 0: raise IndexError("slice step must be positive") @@ -950,8 +957,22 @@ def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: arr = np.where(arr < 0, arr + shape[dim], arr) result.append(arr) dim += 1 + elif isinstance(sel, slice): + # Wrap negative slice bounds NumPy-style here (the eager dialect); + # the transform layer treats every coordinate literally and rejects + # negatives, so the wrap must happen before lowering. Negative steps + # are left untouched (rejected later with their own error); a + # wrapped bound still below 0 clamps to 0, matching NumPy. + if (sel.step is None or sel.step > 0) and dim < len(shape): + n = shape[dim] + start = sel.start if sel.start is None or sel.start >= 0 else max(0, sel.start + n) + stop = sel.stop if sel.stop is None or sel.stop >= 0 else max(0, sel.stop + n) + result.append(slice(start, stop, sel.step)) + else: + result.append(sel) + dim += 1 else: - # slice, bool array, or anything else: pass through + # bool array, or anything else: pass through result.append(sel) if sel is not None and sel is not Ellipsis: dim += 1 @@ -961,6 +982,30 @@ def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: return tuple(result) +_LITERAL_HINT = ( + "; negative indices are literal coordinates in lazy indexing, not from-the-end " + "(use `shape[dim] - k`, or the eager methods, for NumPy semantics)" +) + + +def _check_slice_bounds(sel: slice, dim: int, lo: int, hi: int) -> None: + """Reject negative slice bounds: lazy-path coordinates are literal. + + Consistent with integer and index-array selections — a lazy selection is a + declaration in literal coordinates, and domains start at 0, so a negative + bound is never in the domain regardless of syntactic form. Positive + out-of-range bounds are NOT rejected: a slice denotes a range, and ranges + intersect the domain (Python/NumPy clamping), which can only shorten the + result, never silently select different data the way wraparound would. + """ + for name, bound in (("start", sel.start), ("stop", sel.stop)): + if bound is not None and bound < 0: + raise BoundsCheckError( + f"slice {name} {bound} is out of bounds for dimension {dim} " + f"(valid indices [{lo}, {hi})){_LITERAL_HINT}" + ) + + def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], dim_size: int) -> None: """Reject index-array values outside ``[0, dim_size)``. @@ -969,8 +1014,18 @@ def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], dim_size: in NumPy-style negatives before building the transform. Matches the eager bounds-check error so out-of-range values raise instead of silently wrapping. """ - if arr.size and (int(arr.min()) < 0 or int(arr.max()) >= dim_size): - raise BoundsCheckError(f"index out of bounds for dimension with length {dim_size}") + if arr.size == 0: + return + lo_val, hi_val = int(arr.min()), int(arr.max()) + if lo_val < 0: + raise BoundsCheckError( + f"index {lo_val} is out of bounds for dimension with length {dim_size}{_LITERAL_HINT}" + ) + if hi_val >= dim_size: + raise BoundsCheckError( + f"index {hi_val} is out of bounds for dimension with length {dim_size} " + f"(valid indices [0, {dim_size}))" + ) def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: diff --git a/src/zarr/errors.py b/src/zarr/errors.py index e59c55da35..d3bbb5ff74 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -158,8 +158,9 @@ class LazyViewError(NotImplementedError): Grid-describing members (``chunks``, ``shards``, ``nchunks``, ``read_chunk_sizes``, ...) and grid-mutating ones (``resize``, ``append``) have no well-defined answer - for a view onto a subset of the backing grid. Use ``chunk_layout`` for the backing - structure and ``chunk_projections`` for the view's granularity. Subclasses + for a view onto a subset of the backing grid. Use `chunk_projections` for the + view's granularity; the backing array's stored structure is available via + `metadata` / `chunk_grid`. Subclasses ``NotImplementedError`` so existing consumers that catch it keep working. """ diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 339b9969fd..df4d1cbdc2 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -23,7 +23,7 @@ import zarr from zarr.core.buffer import default_buffer_prototype -from zarr.errors import LazyViewError +from zarr.errors import BoundsCheckError, LazyViewError from zarr.storage import MemoryStore @@ -281,10 +281,16 @@ def test_write(self, cfg: Config) -> None: class TestLazyComposition: @pytest.mark.parametrize("cfg", ND_CASES) def test_chained_views_compose(self, cfg: Config) -> None: - """Composing two lazy slices equals applying them in sequence on NumPy.""" + """Composing two lazy slices equals applying them in sequence on NumPy. + + Bounds are spelled literally: lazy coordinates are literal, so the + from-the-end `1:-1` spelling is rejected on the lazy path (see + TestLazyErrors.test_negative_slice_bounds_are_literal). + """ a, ref = _make(cfg) - view = a.lazy[1:-1].lazy[1:-1] - expected = ref[1:-1][1:-1] + n = cfg.shape[0] + view = a.lazy[1 : n - 1].lazy[1 : n - 3] + expected = ref[1 : n - 1][1 : n - 3] assert tuple(view.shape) == expected.shape np.testing.assert_array_equal(view[...], expected) @@ -399,6 +405,104 @@ def test_accessor_negative_index_is_literal(self) -> None: with pytest.raises(IndexError, match="out of bounds"): _ = a.lazy.vindex[(sel,)][...] + def test_negative_slice_bounds_are_literal(self) -> None: + """Negative slice bounds are literal coordinates on the lazy path, like + every other index form — never from-the-end — so they raise. + + Consistency rule: a lazy selection is a declaration in literal + coordinates, and view domains start at 0, so a negative value is never + in bounds regardless of syntactic form (integer, slice bound, or index + array). + """ + a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + for sel in (slice(-3, None), slice(None, -1), slice(1, -1), slice(-24, None)): + with pytest.raises(BoundsCheckError, match="out of bounds"): + a.lazy[sel] + with pytest.raises(BoundsCheckError, match="out of bounds"): + a.lazy.oindex[(sel,)] + # ... on views too + v = a.lazy[2:10] + with pytest.raises(BoundsCheckError, match="out of bounds"): + v.lazy[-2:] + # ... and for writes + with pytest.raises(BoundsCheckError, match="out of bounds"): + a.lazy[-3:] = 0 + + def test_positive_slice_overflow_still_clamps(self) -> None: + """Positive out-of-range slice bounds intersect the domain (Python range + semantics): a long stop shortens the range, it cannot select wrong data.""" + a, ref = _make(CONFIGS[1]) # shape (24,) + np.testing.assert_array_equal(a.lazy[5:100].result(), ref[5:100]) + assert a.lazy[100:200].shape == (0,) + + def test_eager_view_methods_keep_numpy_negative_slices(self) -> None: + """The eager dialect on a view still wraps negative slice bounds like NumPy.""" + a, ref = _make(CONFIGS[1]) + v = a.lazy[2:10] + np.testing.assert_array_equal(v[-3:], ref[2:10][-3:]) + np.testing.assert_array_equal(v[1:-1], ref[2:10][1:-1]) + + def test_lazy_bounds_errors_share_one_type(self) -> None: + """All three literal-coordinate rejections raise BoundsCheckError (an + IndexError), with one message shape naming the valid range.""" + a, _ = _make(CONFIGS[1]) + for trigger in ( + lambda: a.lazy[-1], + lambda: a.lazy[-3:], + lambda: a.lazy.oindex[(np.array([-1], dtype=np.intp),)], + ): + with pytest.raises(BoundsCheckError, match="out of bounds"): + trigger() + + def test_guard_message_names_real_apis(self) -> None: + """LazyViewError points at APIs that exist (chunk_projections/metadata), + not the not-yet-added chunk_layout.""" + a, _ = _make(CONFIGS[1]) + with pytest.raises(LazyViewError) as ei: + _ = a.lazy[2:10].chunks + assert "chunk_projections" in str(ei.value) + assert "chunk_layout" not in str(ei.value) + + def test_fancy_view_repr_does_not_crash(self) -> None: + """repr of an integer-indexed fancy view must not raise (0-d index array + in selection_repr).""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + ov = b.lazy.oindex[np.array([1, 3]), slice(None)] + assert isinstance(repr(ov.lazy[0]), str) + + +class TestKnownFancyIntBugs: + """Strict-xfail pins for the int-on-fancy-picked-dim defect (see review notes): + integer indexing a dimension that an oindex/vindex selection created is + mis-lowered, crashing reads or mis-shaping results. These flip to failures + when the bug is fixed, forcing the pins to be removed.""" + + @pytest.mark.xfail( + reason="integer indexing on an oindex-picked dimension crashes in the " + "codec pipeline (ArrayMap not collapsed to ConstantMap)", + strict=True, + ) + def test_int_read_on_oindex_view(self) -> None: + """`rows[0]` on an oindex-created view must equal the picked row.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + ref = np.arange(48, dtype="i4").reshape(6, 8) + b[...] = ref + rows = b.lazy.oindex[np.array([1, 3]), slice(None)] + np.testing.assert_array_equal(rows[0], ref[[1, 3]][0]) + + @pytest.mark.xfail( + reason="vindex-created views return shape-(1,) data for integer reads " + "where NumPy returns a scalar", + strict=True, + ) + def test_int_read_on_vindex_view_is_scalar(self) -> None: + """`pts[0]` on a vindex-created view must be a scalar, as in NumPy.""" + a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + pts = a.lazy.vindex[np.array([1, 3, 5])] + assert np.shape(pts[0]) == () + def test_block_selection_on_view_rejected(self) -> None: """Block selection is ill-defined on a lazy view and must raise, not corrupt.""" a, _ = _make(CONFIGS[3]) # 2d-unsharded From 02a5c674ae1245b069127efb167300082a493689 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 08:35:35 +0200 Subject: [PATCH 34/43] test(lazy): type the bounds-error trigger list for mypy Assisted-by: ClaudeCode:claude-opus-4.8 --- tests/test_lazy_indexing.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index df4d1cbdc2..5e50a80b73 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -13,6 +13,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from typing import Any from unittest import mock @@ -446,11 +447,12 @@ def test_lazy_bounds_errors_share_one_type(self) -> None: """All three literal-coordinate rejections raise BoundsCheckError (an IndexError), with one message shape naming the valid range.""" a, _ = _make(CONFIGS[1]) - for trigger in ( + triggers: list[Callable[[], Any]] = [ lambda: a.lazy[-1], lambda: a.lazy[-3:], lambda: a.lazy.oindex[(np.array([-1], dtype=np.intp),)], - ): + ] + for trigger in triggers: with pytest.raises(BoundsCheckError, match="out of bounds"): trigger() From 767df8bec001ce1d8d8fe4b9331b23fa54d823d0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 08:38:38 +0200 Subject: [PATCH 35/43] docs: add a lazy-indexing user guide (theory + executed patterns) New docs/user-guide/lazy_indexing.md, registered in the nav after arrays.md. All examples are live (markdown-exec) so the docs build fails if behavior drifts. Structure: - Theory: indexing as a *declaration* (an index transform) vs an action; views are zero-based windows like NumPy views; the repr's domain is provenance, not the indexing coordinate system. - The literal-coordinate rule: one sentence, consistent across integers, slice bounds, and index arrays (negative is never in bounds), demonstrated with the unified BoundsCheckError; positive slice overflow clamps (range intersection). - The two dialects on one object: lazy = declare (literal), eager = access (NumPy semantics, negatives wrap), with a table. - Patterns, all executed: crop/compose, write-through (region, strided, composed), oindex/vindex/mask selection, materialization, chunk-aware processing via chunk_projections/is_chunk_aligned, LazyViewError guards. - "Coming from NumPy" checklist and current limitations (fancy-int defect, sharded read-unit projections, no newaxis / negative steps). Content grounded in two executed probe reports (semantics probe + naive-user journey), which also drove the preceding consistency fix. Assisted-by: ClaudeCode:claude-opus-4.8 --- docs/user-guide/lazy_indexing.md | 228 +++++++++++++++++++++++++++++++ mkdocs.yml | 1 + 2 files changed, 229 insertions(+) create mode 100644 docs/user-guide/lazy_indexing.md diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md new file mode 100644 index 0000000000..8a43626c3b --- /dev/null +++ b/docs/user-guide/lazy_indexing.md @@ -0,0 +1,228 @@ +# Lazy indexing + +Zarr arrays support *lazy* indexing through the `Array.lazy` accessor. Where +ordinary indexing reads or writes data immediately, `a.lazy[selection]` returns a +lightweight **view** — itself a `zarr.Array` — without touching storage. Views +compose, support orthogonal and coordinate selection, write through to the +backing array, and materialize on demand. + +```python exec="true" session="lazy" +import numpy as np +import zarr +``` + +```python exec="true" session="lazy" source="above" result="ansi" +a = zarr.create_array(store="memory://lazy-demo", shape=(12,), chunks=(3,), dtype="int32") +a[...] = np.arange(12) + +view = a.lazy[2:10] # no I/O happens here +print(view) +print(view.result()) # I/O happens here +``` + +## Theory: indexing as a declaration + +Eager indexing is an *action*: `a[2:10]` performs I/O now and hands back the +bytes. Lazy indexing is a *declaration*: `a.lazy[2:10]` records **which cells +you mean** — an index transform mapping the view's coordinates to storage +coordinates — and defers the I/O. Because the selection is data rather than an +action, zarr can: + +- **compose** it with further selections without reading anything, +- **write through** it (the same transform routes values back to storage), +- **plan** with it (`chunk_projections` enumerates exactly the stored chunks + the declaration touches). + +A view is a *window*: its shape is the selection's shape, and indexing a view is +always relative to the view itself, starting at zero — exactly like indexing a +NumPy view: + +```python exec="true" session="lazy" source="above" result="ansi" +v = a.lazy[2:10] # window onto cells 2..9 +print(v.shape) # the window's shape +print(v.lazy[0].result()) # the window's first element -> base cell 2 +print(v.lazy[3:5].result()) # window cells 3,4 -> base cells 5,6 +``` + +The `domain={ [2, 10) }` shown in a view's repr is **provenance** — the region +of the backing store the view maps onto. It is *not* the coordinate system you +index with; that is always `[0, shape)`. + +### Literal coordinates: `-1` is just another index + +A lazy selection is a declaration in **literal coordinates**. NumPy's negative +indexing ("count from the end") is sugar applied at *evaluation* time; in a +deferred, composable setting it would silently re-bind meaning as views compose. +Zarr therefore treats every coordinate in a lazy selection literally, and since +view domains start at 0, **a negative value is never in bounds — no matter the +syntactic form** (integer, slice bound, or index array): + +```python exec="true" session="lazy" source="above" result="ansi" +for make in (lambda: a.lazy[-1], lambda: a.lazy[-3:], lambda: a.lazy.oindex[[-1]]): + try: + make() + except IndexError as e: + print(type(e).__name__, "-", e) +``` + +To select from the end, say what you mean literally: + +```python exec="true" session="lazy" source="above" result="ansi" +n = a.shape[0] +print(a.lazy[n - 3 :].result()) # the last three elements +``` + +Positive out-of-range *slice* bounds are still fine — a slice denotes a range, +and ranges intersect the domain (as in Python and NumPy). Clamping can only +shorten a result; it can never silently select different data the way +wraparound can: + +```python exec="true" session="lazy" source="above" result="ansi" +print(a.lazy[5:100].result()) # clamps to [5, 12) +``` + +### Two dialects on one object + +A view is an ordinary `zarr.Array`, so it also has the ordinary *eager* methods +— and those keep full NumPy semantics, negatives included. The two dialects +split cleanly by intent: + +| You are... | Spelling | Negative indices | +| ----------- | --------------------------- | ----------------------- | +| declaring | `v.lazy[...]` (a new view) | literal — raise | +| accessing | `v[...]`, `v.oindex[...]` | NumPy — from the end | + +```python exec="true" session="lazy" source="above" result="ansi" +print(v[-1]) # eager access: NumPy dialect -> base cell 9 +print(v.lazy[v.shape[0] - 1].result()) # lazy declaration: same cell, said literally +``` + +## Common patterns + +### Crop, analyze, crop again + +```python exec="true" session="lazy" source="above" result="ansi" +img = zarr.create_array(store="memory://lazy-img", shape=(100, 100), chunks=(10, 10), dtype="float64") +img[...] = np.arange(100 * 100).reshape(100, 100) + +crop = img.lazy[25:75, 25:75] # 50x50 window, no I/O +inner = crop.lazy[10:40, 10:40] # 30x30 window of the window +print(crop.shape, inner.shape) +print(float(np.mean(inner))) # I/O happens here, for the inner crop only +``` + +### Write through a view + +Assignment through the accessor, or through a view, routes values back to +storage — including strided and composed selections: + +```python exec="true" session="lazy" source="above" result="ansi" +img.lazy[30:50, 40:60] = 0.0 # region write, no read-back needed here +tile = img.lazy[30:50, 40:60] +tile[0:5, 0:5] = 7.0 # eager write through the view +img.lazy[::2, ::2] = -1.0 # strided write, NumPy-equivalent +print(img[29:33, 39:43]) +``` + +### Orthogonal and coordinate selection + +`lazy.oindex` selects an outer product per axis; `lazy.vindex` selects points. +Both return views that compose: + +```python exec="true" session="lazy" source="above" result="ansi" +rows = img.lazy.oindex[[3, 17, 42], :] # three full rows -> shape (3, 100) +sub = rows.lazy[:, 10:20] # then a column window of those rows +print(rows.shape, sub.shape) + +pts = img.lazy.vindex[[3, 17], [5, 9]] # two points -> shape (2,) +print(pts.result()) +``` + +Boolean masks are array selections, so they go through `oindex`/`vindex`: + +```python exec="true" session="lazy" source="above" result="ansi" +mask = np.zeros(12, dtype=bool) +mask[[1, 4, 6]] = True +print(a.lazy.oindex[mask].result()) +``` + +### Materializing + +`view.result()`, `view[...]`, and `np.asarray(view)` are equivalent reads of the +whole view; views also work directly with NumPy reductions: + +```python exec="true" session="lazy" source="above" result="ansi" +w = a.lazy[3:9] +print(w.result(), np.asarray(w), float(np.mean(w))) +``` + +### Chunk-aware processing + +`chunk_projections` enumerates the stored chunks a view touches: which store +object (`coord`, `key`), its stored `shape`, the region *within the chunk* the +view covers (`chunk_selection`), the region *of the view* it maps to +(`array_selection`), and whether the chunk is only partially covered +(`is_partial` — a partial write requires a read-modify-write): + +```python exec="true" session="lazy" source="above" result="ansi" +for p in a.lazy[2:10].chunk_projections(): + print(p.key, p.chunk_selection, p.array_selection, p.is_partial) +print(a.lazy[2:10].is_chunk_aligned()) +print(a.lazy[3:9].is_chunk_aligned()) # starts and ends on chunk boundaries +``` + +This is the supported way to partition any selection for parallel or +chunk-at-a-time work — compose the selection through `.lazy`, then project: + +```python exec="true" session="lazy" source="above" result="ansi" +total = 0.0 +crop = img.lazy[25:75, 25:75] +for p in crop.chunk_projections(): + total += float(np.sum(crop[p.array_selection])) +print(total == float(np.sum(crop))) +``` + +For sharded arrays, pass `unit="write"` to enumerate at shard (write-unit) +granularity; read-unit projections for sharded arrays are not yet implemented. + +### What a view will not tell you + +Members that describe the chunk grid assume the array *fills* its grid, which a +view generally does not. On a view they raise `zarr.errors.LazyViewError` +instead of silently describing the backing array: + +```python exec="true" session="lazy" source="above" result="ansi" +try: + v.chunks +except zarr.errors.LazyViewError as e: + print(e) +``` + +Logical members (`shape`, `size`, `nbytes`, `dtype`, `attrs`, ...) reflect the +view; `metadata` and `chunk_grid` remain available and describe the *backing* +array. + +## Coming from NumPy + +- **Negative indices raise on the lazy path** — in every form (integer, slice + bound, index array). Use `shape[dim] - k`, or the eager methods, which keep + NumPy semantics. +- **No negative steps**: `a.lazy[::-1]` raises; reversal is not supported by + zarr indexing generally. +- **No `newaxis`**: `a.lazy[None]` raises; insert axes on the materialized + result instead. +- **The basic accessor takes basic selections only** (integers, slices, + ellipsis). Lists, arrays, and boolean masks go through `lazy.oindex` / + `lazy.vindex`. +- **Scalar reads return NumPy scalars**: `a.lazy[3].result()` is `np.int32(3)`, + matching `a[3]`'s value with scalar type. + +## Current limitations + +- Integer indexing a dimension *created by* an `oindex`/`vindex` selection + (e.g. `rows.lazy[0]` after `rows = a.lazy.oindex[[3, 17, 42], :]`) is not yet + supported reliably; slice the view instead (`rows.lazy[0:1]`). +- `chunk_projections(unit="read")` on sharded arrays (inner-chunk granularity) + is not yet implemented; use `unit="write"`. +- Views cannot be resized or appended to, and block selection is not defined + for views. diff --git a/mkdocs.yml b/mkdocs.yml index e4e757e630..10f47b4ff5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -15,6 +15,7 @@ nav: - user-guide/index.md - user-guide/installation.md - user-guide/arrays.md + - user-guide/lazy_indexing.md - user-guide/groups.md - user-guide/attributes.md - user-guide/storage.md From 76799cc90c9b5dd8897698d35124d59209d7b97d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 09:45:07 +0200 Subject: [PATCH 36/43] fix(transforms): slice bounds are literal coordinates at any domain origin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Probing "can a domain contain a negative number?" (IndexDomain.translate makes one in two lines) exposed that the literal-coordinate rule was only implemented for zero-origin domains: on domain [-10, 2), integers were already literal (t[-5] valid, t[-11] rejected with the true bounds) but slice bounds were interpreted as POSITIONS within the domain (t[0:2] silently selected storage [-10, -8)), and the new bounds check hard-coded `< 0` instead of `< lo` — the same ints-vs-slices inconsistency just removed at origin 0, one level down. Complete the rule uniformly: slice bounds name coordinates, shifted into the domain's 0-based range (_resolve_slice_literal; an identity for lo == 0, so every public view is byte-for-byte unchanged) after _check_slice_bounds rejects bounds below inclusive_min — the any-origin generalization of "negative raises". Bounds past exclusive_max still clamp (ranges intersect the domain). Not publicly reachable today: all .lazy lowering re-zeroes domains, now pinned by test_public_view_domains_are_zero_origin so a future translate-style API cannot silently invalidate the documented origin-0 behavior. New TestNegativeOriginDomains covers int/slice literal semantics on [-10, 2). Index-array paths on non-zero-origin domains remain to be audited when a translate API is exposed. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/transforms/transform.py | 44 ++++++++++----- tests/test_transforms/test_transform.py | 75 +++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 13 deletions(-) diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index e786df8fe2..9bac9babe9 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -533,11 +533,10 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo - # Literal coordinates: negative bounds are out of the domain, not - # from-the-end (the eager dialect wraps them before reaching here). + # Literal coordinates: bounds below the origin are out of the + # domain, not from-the-end (the eager dialect wraps them first). _check_slice_bounds(sel, old_dim, lo, hi) - # Resolve slice relative to the current domain (origin-based) - start, stop, step = sel.indices(dim_size) + start, stop, step = _resolve_slice_literal(sel, lo, dim_size) # start, stop, step are now relative to a 0-based range of size dim_size if step <= 0: @@ -681,7 +680,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo _check_slice_bounds(sel, old_dim, lo, hi) - start, stop, step = sel.indices(dim_size) + start, stop, step = _resolve_slice_literal(sel, lo, dim_size) if step <= 0: raise IndexError("slice step must be positive") new_size = max(0, math.ceil((stop - start) / step)) @@ -850,7 +849,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: hi = transform.domain.exclusive_max[old_dim] dim_size = hi - lo _check_slice_bounds(sel, old_dim, lo, hi) - start, stop, step = sel.indices(dim_size) + start, stop, step = _resolve_slice_literal(sel, lo, dim_size) if step <= 0: raise IndexError("slice step must be positive") new_size = max(0, math.ceil((stop - start) / step)) @@ -989,23 +988,42 @@ def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: def _check_slice_bounds(sel: slice, dim: int, lo: int, hi: int) -> None: - """Reject negative slice bounds: lazy-path coordinates are literal. + """Reject slice bounds below the domain origin: coordinates are literal. Consistent with integer and index-array selections — a lazy selection is a - declaration in literal coordinates, and domains start at 0, so a negative - bound is never in the domain regardless of syntactic form. Positive - out-of-range bounds are NOT rejected: a slice denotes a range, and ranges - intersect the domain (Python/NumPy clamping), which can only shorten the - result, never silently select different data the way wraparound would. + declaration in literal coordinates, so a bound below ``inclusive_min`` is + never in the domain regardless of syntactic form. For the zero-origin + domains all public views have, this is exactly "negative bounds raise". + Bounds past ``exclusive_max`` are NOT rejected: a slice denotes a range, + and ranges intersect the domain (Python/NumPy clamping), which can only + shorten the result, never silently select different data the way + wraparound would. """ for name, bound in (("start", sel.start), ("stop", sel.stop)): - if bound is not None and bound < 0: + if bound is not None and bound < lo: raise BoundsCheckError( f"slice {name} {bound} is out of bounds for dimension {dim} " f"(valid indices [{lo}, {hi})){_LITERAL_HINT}" ) +def _resolve_slice_literal(sel: slice, lo: int, dim_size: int) -> tuple[int, int, int]: + """Resolve slice bounds as literal domain coordinates. + + Bounds name coordinates, not positions: they are shifted into the domain's + 0-based range (an identity when ``lo == 0``, i.e. for every public view) + and then clamped by ``slice.indices`` — safe from wraparound because + ``_check_slice_bounds`` has already rejected anything below ``lo``. + Returns 0-based ``(start, stop, step)`` relative to the domain origin. + """ + rel = slice( + None if sel.start is None else sel.start - lo, + None if sel.stop is None else sel.stop - lo, + sel.step, + ) + return rel.indices(dim_size) + + def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], dim_size: int) -> None: """Reject index-array values outside ``[0, dim_size)``. diff --git a/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py index f531714045..92bb770c45 100644 --- a/tests/test_transforms/test_transform.py +++ b/tests/test_transforms/test_transform.py @@ -514,3 +514,78 @@ def test_translate_2d(self) -> None: assert result.output[0].offset == -5 assert isinstance(result.output[1], DimensionMap) assert result.output[1].offset == -10 + + +class TestNegativeOriginDomains: + """Literal-coordinate semantics must hold at ANY domain origin, not just 0. + + Public views always re-zero their domains, but the transform library admits + non-zero origins (`IndexDomain.translate`), and the literal rule — an index + names a coordinate, never a position or an offset from the end — must apply + uniformly to integers and slice bounds alike. On domain [-10, 2): coordinate + -5 is *in bounds*; slice bounds are coordinates (t[0:2] means storage + [0, 2), the last two cells, NOT the first two positions); bounds below + inclusive_min raise (the any-origin generalization of "negative raises at + origin 0"); bounds past exclusive_max clamp (range intersection). + """ + + def _t(self) -> IndexTransform: + return IndexTransform.identity(IndexDomain.from_shape((12,)).translate((-10,))) + + def test_integer_literal_in_bounds(self) -> None: + """An in-domain negative coordinate is a valid literal integer index.""" + t = self._t()[-5] + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == -5 + + def test_integer_below_origin_raises(self) -> None: + """A coordinate below inclusive_min is out of bounds.""" + with pytest.raises(IndexError, match=r"valid indices \[-10, 2\)"): + self._t()[-11] + + def test_slice_bounds_are_literal_coordinates(self) -> None: + """Slice bounds are coordinates, not positions: t[0:2] is storage [0, 2).""" + t = self._t()[0:2] + assert t.domain.shape == (2,) + assert isinstance(t.output[0], DimensionMap) + # first element of the result must map to storage coordinate 0, not -10 + assert t.output[0].offset == 0 + + def test_slice_with_negative_literal_bound(self) -> None: + """t[-5:] on domain [-10, 2) is the literal interval [-5, 2).""" + t = self._t()[-5:] + assert t.domain.shape == (7,) + assert isinstance(t.output[0], DimensionMap) + assert t.output[0].offset == -5 + + def test_slice_bound_below_origin_raises(self) -> None: + """A slice bound below inclusive_min raises, mirroring the origin-0 rule.""" + with pytest.raises(IndexError, match=r"valid indices \[-10, 2\)"): + self._t()[-12:] + + def test_slice_overflow_clamps(self) -> None: + """Bounds past exclusive_max clamp: ranges intersect the domain.""" + t = self._t()[-5:100] + assert t.domain.shape == (7,) + + +def test_public_view_domains_are_zero_origin() -> None: + """Every view the public `.lazy` API can produce has a zero-origin domain. + + The user-guide rule "a negative value is never in bounds" is the origin-0 + corollary of literal coordinates; this pins the premise so a future + translate-style API cannot silently invalidate the documented behavior. + """ + import zarr + + a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + views = ( + a.lazy[2:10], + a.lazy[2:10].lazy[1:5], + a.lazy[::2], + a.lazy.oindex[np.array([1, 5, 9], dtype=np.intp)], + a.lazy.vindex[np.array([1, 5], dtype=np.intp)], + ) + for v in views: + assert all(m == 0 for m in v._async_array._transform.domain.inclusive_min) From e70056dd5db38baa44c30efcc2116aac990cb8f1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 11:19:26 +0200 Subject: [PATCH 37/43] =?UTF-8?q?feat(lazy)!:=20TensorStore-parity=20domai?= =?UTF-8?q?n=20semantics=20=E2=80=94=20preservation,=20one=20dialect,=20tr?= =?UTF-8?q?anslate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Match TensorStore's indexing model exactly, grounded in an executed behavioral oracle (tensorstore 0.1.84; every rule below verified empirically and ported to tests/test_transforms/test_tensorstore_parity.py): - Domain preservation: a step-1 slice keeps its literal coordinates as the new domain (a.lazy[2:10] has domain [2,10); v[3] is coordinate 3 = base cell 3; v[0] is out of bounds). Nothing re-zeros implicitly. - Strided-domain rule: origin = trunc(start/step) (toward zero), size = ceil((stop-start)/step), coordinate origin+k -> base start + k*step. - Strict containment: non-empty slice intervals must lie within the domain (no clamping, no negative wrapping); empty intervals are valid anywhere; reversed bounds are an invalid-interval error, not an empty result. - One dialect: views use domain coordinates on EVERY path (lazy accessor and eager methods alike); the legacy wraparound/boundscheck pre-validation on the view branches is removed, as is _normalize_negative_indices (identity arrays still take the unchanged NumPy-dialect legacy fast path). - translate_domain_by/translate_domain_to on IndexTransform, exposed as Array.translate_by / Array.translate_to (TensorStore's translate_*): shift a domain, or re-zero a view explicitly, preserving which cells are addressed. - Index-array values are literal domain coordinates; fancy dims keep fresh [0, n) zero-origin domains (already matching TensorStore). - Views are not iterable (TypeError, like TensorStore): the getitem-from-0 iteration protocol would silently yield nothing on a preserved domain. - The I/O layer and chunk_projections normalize to zero-origin via translate_domain_to at their boundaries, so chunk resolution and buffer placement are unchanged; ChunkProjection.array_selection stays positional. Negative steps (supported by TensorStore) remain rejected for now — the new resolver's arithmetic already generalizes, enablement is a follow-up. Tests: oracle-parity module (47 cases incl. negative-origin domains where -1 is just another index); all lazy/view/composition/error tests rewritten to domain coordinates; property parity re-zeroes views via translate_to (exercising it on every example); superseded zero-origin/clamp-era tests removed. Assisted-by: ClaudeCode:claude-opus-4.8 --- src/zarr/core/array.py | 76 ++++- src/zarr/core/chunk_partition.py | 8 +- src/zarr/core/transforms/transform.py | 310 ++++++++---------- tests/test_lazy_indexing.py | 118 ++++--- tests/test_properties.py | 60 +++- .../test_transforms/test_chunk_resolution.py | 5 +- .../test_tensorstore_parity.py | 263 +++++++++++++++ tests/test_transforms/test_transform.py | 111 ++----- 8 files changed, 625 insertions(+), 326 deletions(-) create mode 100644 tests/test_transforms/test_tensorstore_parity.py diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5653d9b26a..777ef19d9e 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -104,7 +104,6 @@ VIndex, _iter_grid, _iter_regions, - boundscheck_indices, check_fields, check_no_multi_fields, ensure_tuple, @@ -115,7 +114,6 @@ is_scalar, pop_fields, replace_lists, - wraparound_indices, ) from zarr.core.metadata import ( ArrayMetadata, @@ -142,7 +140,6 @@ from zarr.core.transforms.output_map import ArrayMap from zarr.core.transforms.transform import ( IndexTransform, - _normalize_negative_indices, selection_to_transform, ) from zarr.errors import ( @@ -807,6 +804,24 @@ def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetada object.__setattr__(new, "_transform", transform) return new + def translate_by(self, shift: tuple[int, ...]) -> AsyncArray[T_ArrayMetadata]: + """Shift this array's domain by `shift`, preserving which cells it addresses. + + TensorStore's `translate_by`: the view's coordinate labels move; the data + does not. `a.translate_by((-10,))` gives a view whose domain starts at + -10, where coordinate -10 addresses the cell that 0 addressed before. + """ + return self._with_transform(self._transform.translate_domain_by(tuple(shift))) + + def translate_to(self, origins: tuple[int, ...]) -> AsyncArray[T_ArrayMetadata]: + """Move this array's domain so its per-dimension origins equal `origins`. + + TensorStore's `translate_to`; `view.translate_to((0,) * view.ndim)` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + return self._with_transform(self._transform.translate_domain_to(tuple(origins))) + @property def store(self) -> Store: return self.store_path.store @@ -1996,6 +2011,43 @@ def _with_transform(self, transform: IndexTransform) -> Array[T_ArrayMetadata]: new_async = self._async_array._with_transform(transform) return type(self)(new_async) + def translate_by(self, shift: tuple[int, ...]) -> Array[T_ArrayMetadata]: + """Shift this array's domain by `shift`, preserving which cells it addresses. + + TensorStore's `translate_by`: the view's coordinate labels move; the data + does not. `a.translate_by((-10,))` gives a view whose domain starts at + -10, where coordinate -10 addresses the cell that 0 addressed before. + """ + return self._with_transform(self._async_array._transform.translate_domain_by(tuple(shift))) + + def translate_to(self, origins: tuple[int, ...]) -> Array[T_ArrayMetadata]: + """Move this array's domain so its per-dimension origins equal `origins`. + + TensorStore's `translate_to`; `view.translate_to((0,) * view.ndim)` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + return self._with_transform( + self._async_array._transform.translate_domain_to(tuple(origins)) + ) + + def __iter__(self) -> Any: + """Iterate over the first axis (identity arrays only). + + Lazy views are not iterable, matching TensorStore: the Python iteration + protocol counts positions from 0, which are not valid coordinates in a + preserved (possibly non-zero-origin) domain — silently yielding nothing + or the wrong cells. Read the view first: `iter(view.result())`. + """ + if not self._async_array._is_identity: + raise TypeError( + "lazy views are not iterable; materialize first, e.g. iterate `view.result()`" + ) + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + for i in range(self.shape[0]): + yield self[i] + @classmethod def _create( cls, @@ -3024,7 +3076,6 @@ def get_basic_selection( prototype=prototype, ) ) - selection = _normalize_negative_indices(selection, self.shape) transform = selection_to_transform(selection, self._async_array._transform, "basic") return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) @@ -3136,7 +3187,6 @@ def set_basic_selection( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) return - selection = _normalize_negative_indices(selection, self.shape) transform = selection_to_transform(selection, self._async_array._transform, "basic") sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) @@ -3277,7 +3327,6 @@ def get_orthogonal_selection( # Lazy view (non-identity transform): route through the transform so the # view's offset/stride are honored. Use basic mode for plain int/slice # selections (so integer axes drop), orthogonal mode for fancy selections. - selection = _normalize_negative_indices(selection, self.shape) mode: Literal["basic", "orthogonal"] = ( "basic" if is_basic_selection(selection) else "orthogonal" ) @@ -3406,7 +3455,6 @@ def set_orthogonal_selection( # Lazy view (non-identity transform): route through the transform so the # view's offset/stride are honored. Use basic mode for plain int/slice # selections (so integer axes drop), orthogonal mode for fancy selections. - selection = _normalize_negative_indices(selection, self.shape) mode: Literal["basic", "orthogonal"] = ( "basic" if is_basic_selection(selection) else "orthogonal" ) @@ -3729,10 +3777,6 @@ def get_coordinate_selection( "(coordinate) array per dimension of the target array, " f"got {selection!r}" ) - # Handle wraparound and bounds checking - for dim_sel, dim_len in zip(sel_normalized, self.shape, strict=True): - wraparound_indices(dim_sel, dim_len) - boundscheck_indices(dim_sel, dim_len) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -3864,9 +3908,6 @@ def set_coordinate_selection( "(coordinate) array per dimension of the target array, " f"got {selection!r}" ) - for dim_sel, dim_len in zip(sel_normalized, self.shape, strict=True): - wraparound_indices(dim_sel, dim_len) - boundscheck_indices(dim_sel, dim_len) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -5856,6 +5897,10 @@ async def _get_selection_via_transform( chunk_grid: ChunkGrid | None = None, ) -> NDArrayLikeOrScalar: """Read data using an IndexTransform.""" + # The user-facing transform may carry a preserved (possibly non-zero-origin) + # domain; the I/O layer addresses output buffers positionally, so normalize + # to a zero-origin domain here. Translation preserves the cell mapping. + transform = transform.translate_domain_to((0,) * transform.input_rank) if chunk_grid is None: chunk_grid = ChunkGrid.from_metadata(metadata) @@ -5959,6 +6004,9 @@ async def _set_selection_via_transform( chunk_grid: ChunkGrid | None = None, ) -> None: """Write data using an IndexTransform.""" + # See _get_selection_via_transform: normalize to a zero-origin domain + # so value/buffer placement is positional. + transform = transform.translate_domain_to((0,) * transform.input_rank) if chunk_grid is None: chunk_grid = ChunkGrid.from_metadata(metadata) diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py index 7edc586183..712921928d 100644 --- a/src/zarr/core/chunk_partition.py +++ b/src/zarr/core/chunk_partition.py @@ -79,7 +79,13 @@ def iter_chunk_projections( chunk_grid: ChunkGrid, encode_key: Callable[[tuple[int, ...]], str], ) -> Iterator[ChunkProjection]: - """Yield a `ChunkProjection` for each stored chunk ``transform`` projects onto.""" + """Yield a `ChunkProjection` for each stored chunk ``transform`` projects onto. + + `array_selection` is positional (0-based into the view's extent), so the + transform is normalized to a zero-origin domain first — translation + preserves which cells are addressed. + """ + transform = transform.translate_domain_to((0,) * transform.input_rank) chunk_sizes = chunk_grid.chunk_sizes # per-dimension, extent-clipped for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): if chunk_grid[chunk_coords] is None: diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 9bac9babe9..473d15d0ef 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -186,6 +186,46 @@ def translate(self, shift: tuple[int, ...]) -> IndexTransform: def __getitem__(self, selection: Any) -> IndexTransform: return _apply_basic_indexing(self, selection) + def translate_domain_by(self, shift: tuple[int, ...]) -> IndexTransform: + """Shift the *input* domain by ``shift``, preserving which cells are addressed. + + TensorStore's ``translate_by``: the domain moves, and every output map is + re-offset so that new coordinate ``c`` addresses the cell that ``c - shift`` + addressed before. ArrayMaps are indexed positionally over the domain, so + their index arrays are unchanged. + """ + if len(shift) != self.input_rank: + raise ValueError(f"shift must have length {self.input_rank}, got {len(shift)}") + new_domain = self.domain.translate(shift) + new_output: list[OutputIndexMap] = [] + for m in self.output: + if isinstance(m, DimensionMap): + s = shift[m.input_dimension] + new_output.append( + DimensionMap( + input_dimension=m.input_dimension, + offset=m.offset - m.stride * s, + stride=m.stride, + ) + ) + else: + # ConstantMap: no input dependence. ArrayMap: positional over + # the domain, invariant under domain translation. + new_output.append(m) + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + def translate_domain_to(self, origins: tuple[int, ...]) -> IndexTransform: + """Move the input domain so its per-dimension origins equal ``origins``. + + TensorStore's ``translate_to``; ``translate_domain_to((0,) * rank)`` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + if len(origins) != self.input_rank: + raise ValueError(f"origins must have length {self.input_rank}, got {len(origins)}") + shift = tuple(o - m for o, m in zip(origins, self.domain.inclusive_min, strict=True)) + return self.translate_domain_by(shift) + @property def oindex(self) -> _OIndexHelper: return _OIndexHelper(self) @@ -449,10 +489,13 @@ def _reindex_array( old_dim += 1 elif isinstance(sel, slice): if old_dim < arr.ndim: - dim_size = domain.shape[old_dim] - # sel.indices gives 0-based start/stop/step for the array axis - start, stop, step = sel.indices(dim_size) - idx.append(slice(start, stop, step)) + lo = domain.inclusive_min[old_dim] + hi = domain.exclusive_max[old_dim] + # Bounds are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) old_dim += 1 result_axis += 1 @@ -478,12 +521,16 @@ def _reindex_array_oindex( for old_dim, sel in enumerate(normalized): if old_dim >= arr.ndim: break + lo = domain.inclusive_min[old_dim] if isinstance(sel, np.ndarray): - idx.append(sel) + # Values are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + idx.append(sel - lo) elif isinstance(sel, slice): - dim_size = domain.shape[old_dim] - start, stop, step = sel.indices(dim_size) - idx.append(slice(start, stop, step)) + hi = domain.exclusive_max[old_dim] + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) else: idx.append(slice(None)) @@ -531,24 +578,14 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra elif isinstance(sel, slice): lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] - dim_size = hi - lo - # Literal coordinates: bounds below the origin are out of the - # domain, not from-the-end (the eager dialect wraps them first). - _check_slice_bounds(sel, old_dim, lo, hi) - start, stop, step = _resolve_slice_literal(sel, lo, dim_size) - # start, stop, step are now relative to a 0-based range of size dim_size - - if step <= 0: - raise IndexError("slice step must be positive") - - new_size = max(0, math.ceil((stop - start) / step)) - new_inclusive_min.append(0) - new_exclusive_max.append(new_size) - - # Absolute start in the original domain coordinates - abs_start = lo + start - dim_slice_params[old_dim] = (abs_start, stop, step) + # TensorStore semantics: bounds are literal coordinates; a step-1 + # slice keeps them as the new domain, a strided slice's domain is + # [trunc(start/step), trunc(start/step) + size). + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) old_to_new_dim[old_dim] = new_dim_idx new_dim_idx += 1 old_dim += 1 @@ -570,9 +607,10 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra new_offset = m.offset + m.stride * dim_int_val[d] new_output.append(ConstantMap(offset=new_offset)) elif d in old_to_new_dim: - # Slice: update offset and stride - abs_start, _, step = dim_slice_params[d] - new_offset = m.offset + m.stride * abs_start + # Slice: new coordinate `origin + k` maps to old coordinate + # `start + k*step`, i.e. old = start - step*origin + step*new. + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) new_stride = m.stride * step new_input_dim = old_to_new_dim[d] new_output.append( @@ -669,7 +707,9 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: if isinstance(sel, np.ndarray): lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] - _check_array_in_bounds(sel, hi - lo) + # Index-array values are literal domain coordinates; the fancy dim + # they create gets a fresh zero-origin [0, n) domain (TensorStore). + _check_array_in_bounds(sel, lo, hi) dim_array[old_dim] = sel new_inclusive_min.append(0) new_exclusive_max.append(len(sel)) @@ -678,16 +718,10 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: elif isinstance(sel, slice): lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] - dim_size = hi - lo - _check_slice_bounds(sel, old_dim, lo, hi) - start, stop, step = _resolve_slice_literal(sel, lo, dim_size) - if step <= 0: - raise IndexError("slice step must be positive") - new_size = max(0, math.ceil((stop - start) / step)) - new_inclusive_min.append(0) - new_exclusive_max.append(new_size) - abs_start = lo + start - dim_slice_params[old_dim] = (abs_start, stop, step) + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) old_to_new_dim[old_dim] = new_dim_idx new_dim_idx += 1 @@ -715,8 +749,8 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: ) ) elif d in dim_slice_params: - abs_start, _, step = dim_slice_params[d] - new_offset = m.offset + m.stride * abs_start + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) new_stride = m.stride * step new_input_dim = old_to_new_dim[d] new_output.append( @@ -816,7 +850,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: if isinstance(sel, np.ndarray): lo = transform.domain.inclusive_min[i] hi = transform.domain.exclusive_max[i] - _check_array_in_bounds(sel, hi - lo) + _check_array_in_bounds(sel, lo, hi) array_dims.append(i) arrays.append(sel) else: @@ -840,23 +874,17 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: new_inclusive_min.append(0) new_exclusive_max.append(s) - # Slice dimensions + # Slice dimensions (preserved-domain literal semantics, like basic indexing) slice_dim_params: dict[int, tuple[int, int, int]] = {} for old_dim in slice_dims: sel = processed[old_dim] assert isinstance(sel, slice) lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] - dim_size = hi - lo - _check_slice_bounds(sel, old_dim, lo, hi) - start, stop, step = _resolve_slice_literal(sel, lo, dim_size) - if step <= 0: - raise IndexError("slice step must be positive") - new_size = max(0, math.ceil((stop - start) / step)) - new_inclusive_min.append(0) - new_exclusive_max.append(new_size) - abs_start = lo + start - slice_dim_params[old_dim] = (abs_start, stop, step) + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + slice_dim_params[old_dim] = (start, step, origin) new_domain = IndexDomain( inclusive_min=tuple(new_inclusive_min), @@ -886,9 +914,9 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: ) ) else: - # Slice dim - abs_start, _, step = slice_dim_params[d] - new_offset = m.offset + m.stride * abs_start + # Slice dim: new coord `origin + k` maps to old `start + k*step` + start, step, origin = slice_dim_params[d] + new_offset = m.offset + m.stride * (start - step * origin) new_stride = m.stride * step new_input_dim = n_broadcast_dims + slice_dims.index(d) new_output.append( @@ -910,140 +938,80 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: return IndexTransform(domain=new_domain, output=tuple(new_output)) -def _normalize_negative_indices(selection: Any, shape: tuple[int, ...]) -> Any: - """Convert negative indices to positive ones using the array shape. - - Only normalizes integer and array-like index components; leaves - slices, Ellipsis, None, etc. untouched. - """ - if not isinstance(selection, tuple): - selection_tuple: tuple[Any, ...] = (selection,) - else: - selection_tuple = selection - - # Count real dimensions (non-None, non-Ellipsis) to map each entry to a shape - # dim. The ellipsis covers the dims not consumed by those entries; since - # n_non_newaxis already excludes the ellipsis itself, that count is exactly - # len(shape) - n_non_newaxis (no +1). - n_non_newaxis = sum(1 for s in selection_tuple if s is not None and s is not Ellipsis) - n_ellipsis_dims = len(shape) - n_non_newaxis - - result: list[Any] = [] - dim = 0 - - for sel in selection_tuple: - if sel is Ellipsis: - result.append(sel) - dim += max(0, n_ellipsis_dims) - elif sel is None: - result.append(sel) - elif isinstance(sel, (int, np.integer)) and not isinstance(sel, bool): - idx = int(sel) - if idx < 0 and dim < len(shape): - idx = idx + shape[dim] - result.append(idx) - dim += 1 - elif isinstance(sel, np.ndarray) and sel.dtype != np.bool_: - arr = sel.copy() - if dim < len(shape): - arr = np.where(arr < 0, arr + shape[dim], arr) - result.append(arr) - dim += 1 - elif isinstance(sel, list): - # Convert lists to arrays with negative index normalization - arr = np.asarray(sel, dtype=np.intp) - if dim < len(shape): - arr = np.where(arr < 0, arr + shape[dim], arr) - result.append(arr) - dim += 1 - elif isinstance(sel, slice): - # Wrap negative slice bounds NumPy-style here (the eager dialect); - # the transform layer treats every coordinate literally and rejects - # negatives, so the wrap must happen before lowering. Negative steps - # are left untouched (rejected later with their own error); a - # wrapped bound still below 0 clamps to 0, matching NumPy. - if (sel.step is None or sel.step > 0) and dim < len(shape): - n = shape[dim] - start = sel.start if sel.start is None or sel.start >= 0 else max(0, sel.start + n) - stop = sel.stop if sel.stop is None or sel.stop >= 0 else max(0, sel.stop + n) - result.append(slice(start, stop, sel.step)) - else: - result.append(sel) - dim += 1 - else: - # bool array, or anything else: pass through - result.append(sel) - if sel is not None and sel is not Ellipsis: - dim += 1 - - if not isinstance(selection, tuple) and len(result) == 1: - return result[0] - return tuple(result) - - _LITERAL_HINT = ( "; negative indices are literal coordinates in lazy indexing, not from-the-end " - "(use `shape[dim] - k`, or the eager methods, for NumPy semantics)" + "(use `shape[dim] - k`, or materialize with `result()`, for NumPy semantics)" ) -def _check_slice_bounds(sel: slice, dim: int, lo: int, hi: int) -> None: - """Reject slice bounds below the domain origin: coordinates are literal. +def _trunc_div(a: int, b: int) -> int: + """Integer division rounded toward zero (C semantics), as TensorStore uses + for strided-slice domain origins — distinct from Python's floor division + for negative operands (``trunc(-9/2) == -4`` where ``-9 // 2 == -5``).""" + q = a // b + if q < 0 and q * b != a: + q += 1 + return q - Consistent with integer and index-array selections — a lazy selection is a - declaration in literal coordinates, so a bound below ``inclusive_min`` is - never in the domain regardless of syntactic form. For the zero-origin - domains all public views have, this is exactly "negative bounds raise". - Bounds past ``exclusive_max`` are NOT rejected: a slice denotes a range, - and ranges intersect the domain (Python/NumPy clamping), which can only - shorten the result, never silently select different data the way - wraparound would. - """ - for name, bound in (("start", sel.start), ("stop", sel.stop)): - if bound is not None and bound < lo: - raise BoundsCheckError( - f"slice {name} {bound} is out of bounds for dimension {dim} " - f"(valid indices [{lo}, {hi})){_LITERAL_HINT}" - ) +def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, int, int]: + """Resolve a slice against domain ``[lo, hi)`` with TensorStore semantics. + + Slice bounds are **literal domain coordinates** — never from-the-end, never + clamped. Rules (each verified against tensorstore 0.1.84): -def _resolve_slice_literal(sel: slice, lo: int, dim_size: int) -> tuple[int, int, int]: - """Resolve slice bounds as literal domain coordinates. + - defaults: ``start = lo``, ``stop = hi``; + - a non-empty interval must be contained in the domain (no clamping — a + NumPy-style out-of-range or negative bound is an error, not a shorter or + wrapped result); + - an **empty** interval (``start == stop``) is valid anywhere; + - reversed bounds (``start > stop`` with positive step) are an error, not + an empty result; + - the result's domain origin is ``trunc(start/step)`` (rounded toward + zero) and coordinate ``origin + k`` maps to input ``start + k*step``. - Bounds name coordinates, not positions: they are shifted into the domain's - 0-based range (an identity when ``lo == 0``, i.e. for every public view) - and then clamped by ``slice.indices`` — safe from wraparound because - ``_check_slice_bounds`` has already rejected anything below ``lo``. - Returns 0-based ``(start, stop, step)`` relative to the domain origin. + Returns ``(start, step, origin, size)`` in domain coordinates. """ - rel = slice( - None if sel.start is None else sel.start - lo, - None if sel.stop is None else sel.stop - lo, - sel.step, - ) - return rel.indices(dim_size) + step = 1 if sel.step is None else sel.step + if step <= 0: + # Negative steps are valid in TensorStore but not yet supported here; + # step 0 is invalid everywhere. + raise IndexError("slice step must be positive") + start = lo if sel.start is None else sel.start + stop = hi if sel.stop is None else sel.stop + if stop < start: + raise IndexError( + f"slice interval [{start}, {stop}) with step {step} does not specify " + f"a valid interval for dimension {dim} (start > stop)" + ) + size = -(-(stop - start) // step) # ceil((stop - start) / step) + if size > 0 and (start < lo or stop > hi): + hint = _LITERAL_HINT if (start < 0 or stop < 0) and lo >= 0 else "" + raise BoundsCheckError( + f"slice interval [{start}, {stop}) is not contained within domain " + f"[{lo}, {hi}) for dimension {dim}{hint}" + ) + origin = _trunc_div(start, step) + return start, step, origin, size -def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], dim_size: int) -> None: - """Reject index-array values outside ``[0, dim_size)``. +def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], lo: int, hi: int) -> None: + """Reject index-array values outside the domain ``[lo, hi)``. - Index-array values are literal coordinates (origin 0), so a negative value is - out of bounds rather than counting from the end — the Array layer normalizes - NumPy-style negatives before building the transform. Matches the eager - bounds-check error so out-of-range values raise instead of silently wrapping. + Index-array values are literal domain coordinates (TensorStore semantics): + a value below ``inclusive_min`` is out of bounds rather than counting from + the end. Out-of-range values raise instead of silently wrapping. """ if arr.size == 0: return lo_val, hi_val = int(arr.min()), int(arr.max()) - if lo_val < 0: - raise BoundsCheckError( - f"index {lo_val} is out of bounds for dimension with length {dim_size}{_LITERAL_HINT}" - ) - if hi_val >= dim_size: + if lo_val < lo: + hint = _LITERAL_HINT if lo_val < 0 and lo >= 0 else "" raise BoundsCheckError( - f"index {hi_val} is out of bounds for dimension with length {dim_size} " - f"(valid indices [0, {dim_size}))" + f"index {lo_val} is out of bounds (valid indices [{lo}, {hi})){hint}" ) + if hi_val >= hi: + raise BoundsCheckError(f"index {hi_val} is out of bounds (valid indices [{lo}, {hi}))") def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 5e50a80b73..825268ce4b 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -13,9 +13,8 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from unittest import mock import numpy as np @@ -27,6 +26,9 @@ from zarr.errors import BoundsCheckError, LazyViewError from zarr.storage import MemoryStore +if TYPE_CHECKING: + from collections.abc import Callable + @dataclass(frozen=True) class Config: @@ -290,10 +292,13 @@ def test_chained_views_compose(self, cfg: Config) -> None: """ a, ref = _make(cfg) n = cfg.shape[0] - view = a.lazy[1 : n - 1].lazy[1 : n - 3] - expected = ref[1 : n - 1][1 : n - 3] + view = a.lazy[1 : n - 1].lazy[2 : n - 2] + expected = ref[2 : n - 2] # literal coordinates: the inner slice re-selects assert tuple(view.shape) == expected.shape np.testing.assert_array_equal(view[...], expected) + # the composed view's domain is the literal interval it covers + t = view._async_array._transform + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (2, n - 2) class TestLazyViewMethods: @@ -312,19 +317,20 @@ def test_view_oindex_respects_transform(self, cfg: Config) -> None: n0 = cfg.shape[0] cut = n0 // 2 vslice = (slice(cut, n0), *([slice(None)] * (len(cfg.shape) - 1))) - vref = ref[vslice] v = a.lazy[vslice] - idx = np.array([0, vref.shape[0] - 1], dtype=np.intp) + idx = np.array([cut, cfg.shape[0] - 1], dtype=np.intp) # domain coordinates osel: Any = (idx, *([slice(None)] * (len(cfg.shape) - 1))) - np.testing.assert_array_equal(v.oindex[osel], vref[osel]) + np.testing.assert_array_equal(v.oindex[osel], ref[osel]) @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) - def test_view_negative_index_after_ellipsis(self, cfg: Config) -> None: - """``v[..., -1]`` on a view selects the last element of the last axis.""" + def test_view_trailing_index_after_ellipsis(self, cfg: Config) -> None: + """``v[..., k]`` uses literal coordinates; ``-1`` is out of the domain.""" a, ref = _make(cfg) v = a.lazy[1 : cfg.shape[0]] - vref = ref[1 : cfg.shape[0]] - np.testing.assert_array_equal(v[..., -1], vref[..., -1]) + last = cfg.shape[-1] - 1 + np.testing.assert_array_equal(v[..., last], ref[1:, ..., last]) + with pytest.raises(BoundsCheckError): + v[..., -1] @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) def test_view_basic_tuple_and_int(self, cfg: Config) -> None: @@ -332,11 +338,10 @@ def test_view_basic_tuple_and_int(self, cfg: Config) -> None: a, ref = _make(cfg) cut = cfg.shape[0] // 2 v = a.lazy[cut:] - vref = ref[cut:] - full: Any = (slice(0, 2), *([slice(None)] * (len(cfg.shape) - 1))) - np.testing.assert_array_equal(v[full], vref[full]) - intsel: Any = (1, *([slice(None)] * (len(cfg.shape) - 1))) - np.testing.assert_array_equal(v[intsel], vref[intsel]) # int axis drops + full: Any = (slice(cut, cut + 2), *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[full], ref[full]) + intsel: Any = (cut + 1, *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[intsel], ref[intsel]) # int axis drops @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) def test_view_vindex(self, cfg: Config) -> None: @@ -344,9 +349,12 @@ def test_view_vindex(self, cfg: Config) -> None: a, ref = _make(cfg) cut = cfg.shape[0] // 2 v = a.lazy[cut:] - vref = ref[cut:] - idx = tuple(np.array([0, 1, 2], dtype=np.intp) for _ in cfg.shape) - np.testing.assert_array_equal(v.vindex[idx], vref[idx]) + # coordinates: axis 0 lives in [cut, n); the other axes are untouched + idx = ( + np.array([cut, cut + 1, cut + 2], dtype=np.intp), + *(np.array([0, 1, 2], dtype=np.intp) for _ in cfg.shape[1:]), + ) + np.testing.assert_array_equal(v.vindex[idx], ref[idx]) def test_view_vindex_with_flat_out_buffer(self) -> None: """vindex with a multi-dim result and out= on a view uses a flat out buffer. @@ -356,10 +364,9 @@ def test_view_vindex_with_flat_out_buffer(self) -> None: """ a, ref = _make(CONFIGS[3]) # 2d-unsharded v = a.lazy[2:18] - vref = ref[2:18] - i0 = np.array([[0, 1], [2, 3]], dtype=np.intp) + i0 = np.array([[2, 3], [4, 5]], dtype=np.intp) # coordinates within [2, 18) i1 = np.array([[0, 5], [10, 15]], dtype=np.intp) - expected = vref[i0, i1] + expected = ref[i0, i1] buf = default_buffer_prototype().nd_buffer.empty( shape=(expected.size,), dtype=np.dtype("i4") ) @@ -375,9 +382,9 @@ def test_view_write_through_tuple(self, cfg: Config) -> None: cut = cfg.shape[0] // 2 v = a.lazy[cut:] expected = ref.copy() - wsel: Any = (slice(0, 2), *([slice(None)] * (len(cfg.shape) - 1))) - val = _value_like(ref[cut:][wsel]) - expected[cut:][wsel] = val + wsel: Any = (slice(cut, cut + 2), *([slice(None)] * (len(cfg.shape) - 1))) + val = _value_like(ref[wsel]) + expected[wsel] = val v[wsel] = val np.testing.assert_array_equal(a[...], expected) @@ -416,32 +423,59 @@ def test_negative_slice_bounds_are_literal(self) -> None: array). """ a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) - for sel in (slice(-3, None), slice(None, -1), slice(1, -1), slice(-24, None)): - with pytest.raises(BoundsCheckError, match="out of bounds"): + for sel in (slice(-3, None), slice(-24, None)): + with pytest.raises(BoundsCheckError, match="not contained"): a.lazy[sel] - with pytest.raises(BoundsCheckError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="not contained"): a.lazy.oindex[(sel,)] + # bounds that RESOLVE reversed (stop < start), e.g. [1, -1) or [0, -1), + # are invalid intervals — TensorStore's a[5:2] case — not empties + for sel in (slice(None, -1), slice(1, -1)): + with pytest.raises(IndexError, match="interval"): + a.lazy[sel] # ... on views too v = a.lazy[2:10] - with pytest.raises(BoundsCheckError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="not contained"): v.lazy[-2:] # ... and for writes - with pytest.raises(BoundsCheckError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="not contained"): a.lazy[-3:] = 0 - def test_positive_slice_overflow_still_clamps(self) -> None: - """Positive out-of-range slice bounds intersect the domain (Python range - semantics): a long stop shortens the range, it cannot select wrong data.""" - a, ref = _make(CONFIGS[1]) # shape (24,) - np.testing.assert_array_equal(a.lazy[5:100].result(), ref[5:100]) - assert a.lazy[100:200].shape == (0,) - - def test_eager_view_methods_keep_numpy_negative_slices(self) -> None: - """The eager dialect on a view still wraps negative slice bounds like NumPy.""" + def test_slice_bounds_strict_containment(self) -> None: + """Non-empty slice intervals must be contained in the domain — no + clamping (TensorStore semantics); empty intervals are valid anywhere; + reversed bounds are an error, not an empty result.""" + a, _ = _make(CONFIGS[1]) # shape (24,) + for sel in (slice(5, 100), slice(100, 200), slice(0, 25)): + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy[sel] + assert a.lazy[5:5].shape == (0,) + assert a.lazy[30:30].shape == (0,) # empty is valid even outside the domain + with pytest.raises(IndexError, match="interval"): + a.lazy[5:2] + + def test_one_dialect_view_methods_use_domain_coordinates(self) -> None: + """A view has ONE coordinate system: its preserved domain. Eager methods + (`v[...]`, `v[k]`) use the same literal coordinates as `.lazy` — + matching TensorStore, where indexing a view of [2, 10) with 0 or -1 is + out of bounds. Zero-based NumPy-style access is spelled explicitly: + materialize (`result()` / `np.asarray`) or re-zero with `translate_to`. + """ a, ref = _make(CONFIGS[1]) v = a.lazy[2:10] - np.testing.assert_array_equal(v[-3:], ref[2:10][-3:]) - np.testing.assert_array_equal(v[1:-1], ref[2:10][1:-1]) + np.testing.assert_array_equal(v[2:5], ref[2:5]) + np.testing.assert_array_equal(v[3], ref[3]) + bad_reads: list[Callable[[], Any]] = [ + lambda: v[-1], + lambda: v[0:3], + lambda: v[-3:], + ] + for bad in bad_reads: + with pytest.raises(BoundsCheckError): + bad() + np.testing.assert_array_equal(np.asarray(v), ref[2:10]) + z = v.translate_to((0,)) + np.testing.assert_array_equal(z[0:3], ref[2:5]) def test_lazy_bounds_errors_share_one_type(self) -> None: """All three literal-coordinate rejections raise BoundsCheckError (an @@ -453,7 +487,7 @@ def test_lazy_bounds_errors_share_one_type(self) -> None: lambda: a.lazy.oindex[(np.array([-1], dtype=np.intp),)], ] for trigger in triggers: - with pytest.raises(BoundsCheckError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="out of bounds|not contained"): trigger() def test_guard_message_names_real_apis(self) -> None: diff --git a/tests/test_properties.py b/tests/test_properties.py index d292f7ce55..c71a713ed9 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -725,6 +725,41 @@ class _ThreeDStateMachine(_IndexingStateMachine): TestIndexingStateMachine3D = pytest.mark.slow_hypothesis(_ThreeDStateMachine.TestCase) +def _canonical_nonneg(sel: Any, shape: tuple[int, ...]) -> Any: + """NumPy-normalize a selection to its non-negative canonical form. + + Wraps negative integers and integer-array values, resolves slice bounds via + Python slice semantics, and expands ellipsis — producing the literal + spelling of the same NumPy selection, as the lazy path requires. + """ + items = sel if isinstance(sel, tuple) else (sel,) + n_explicit = sum(1 for s in items if s is not Ellipsis) + out: list[Any] = [] + dim = 0 + for s in items: + if s is Ellipsis: + for _ in range(len(shape) - n_explicit): + out.append(slice(0, shape[dim], 1)) + dim += 1 + elif isinstance(s, (int, np.integer)) and not isinstance(s, (bool, np.bool_)): + v = int(s) + out.append(v + shape[dim] if v < 0 else v) + dim += 1 + elif isinstance(s, slice): + start, stop, step = s.indices(shape[dim]) + if stop < start: # NumPy empty-by-reversal -> literal empty interval + stop = start + out.append(slice(start, stop, step)) + dim += 1 + elif isinstance(s, np.ndarray) and s.dtype != np.bool_: + out.append(np.where(s < 0, s + shape[dim], s)) + dim += 1 + else: # boolean mask or anything already canonical + out.append(s) + dim += 1 + return tuple(out) if isinstance(sel, tuple) else out[0] + + @pytest.mark.skipif( not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" ) @@ -732,21 +767,26 @@ class _ThreeDStateMachine(_IndexingStateMachine): @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: - """The eager read oracle, applied to a lazy view (the lazy consumer). - - A view (built from a non-negative window) is just another `zarr.Array`, so - it flows through the same `assert_read_matches_numpy` harness. The lazy - indexing bugs found in review were all "the view path diverges from NumPy for - some (method, parameter) combination" — enumerating that surface here catches - the class. Skipped until `Array.lazy` exists, so the eager oracle can merge - ahead of the lazy feature. + """The eager read oracle, applied to a re-zeroed lazy view (the lazy consumer). + + A view is just another `zarr.Array`, so it flows through the same + `assert_read_matches_numpy` harness. Views preserve their domain + (TensorStore semantics: coordinates are literal), so to compare against a + zero-based NumPy reference the view is explicitly re-zeroed with + `translate_to` — exercising domain translation on every example. Skipped + until `Array.lazy` exists, so the eager oracle can merge ahead of the lazy + feature. """ zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) nparray = zarray[:] window = data.draw(windows(shape=nparray.shape)) - view = zarray.lazy[window] + view = zarray.lazy[window].translate_to((0,) * nparray.ndim) vref = nparray[window] mode = data.draw(st.sampled_from(_INDEX_MODES)) assume(_eligible(mode, vref.shape)) zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=vref.shape)) - assert_read_matches_numpy(view, vref, mode, zsel, npsel) + # The strategies draw NumPy-dialect selections (negatives from-the-end); + # lazy-path coordinates are literal, so canonicalize to the equivalent + # non-negative form for the view. NumPy semantics are invariant under this, + # so the reference side (npsel) is unchanged. + assert_read_matches_numpy(view, vref, mode, _canonical_nonneg(zsel, vref.shape), npsel) diff --git a/tests/test_transforms/test_chunk_resolution.py b/tests/test_transforms/test_chunk_resolution.py index 63246a9caf..893866740d 100644 --- a/tests/test_transforms/test_chunk_resolution.py +++ b/tests/test_transforms/test_chunk_resolution.py @@ -50,7 +50,10 @@ def test_multiple_chunks_2d(self) -> None: class TestChunkResolutionSliced: def test_slice_within_chunk(self) -> None: """Slice that falls within a single chunk.""" - t = IndexTransform.from_shape((100,))[5:8] + # Chunk resolution consumes zero-origin transforms: the I/O layer + # normalizes preserved (user-facing) domains via translate_domain_to + # before resolving, so mirror that contract here. + t = IndexTransform.from_shape((100,))[5:8].translate_domain_to((0,)) grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) results = list(iter_chunk_transforms(t, grid)) assert len(results) == 1 diff --git a/tests/test_transforms/test_tensorstore_parity.py b/tests/test_transforms/test_tensorstore_parity.py new file mode 100644 index 0000000000..4670cf5928 --- /dev/null +++ b/tests/test_transforms/test_tensorstore_parity.py @@ -0,0 +1,263 @@ +"""TensorStore-parity oracle tests for IndexTransform semantics. + +Every case in this module was executed against tensorstore 0.1.84 (see the +lazy-indexing design notes): the expected domains, values, and error conditions +are TensorStore's observed behavior, which zarr's lazy indexing matches by +design. Core rules pinned here: + +- **Domain preservation**: a step-1 slice keeps the literal coordinates of the + selected interval (`a[2:10]` has domain `[2, 10)`); nothing re-zeros + implicitly. Re-zeroing is explicit via `translate_to`. +- **Strided-domain rule**: for step ``k``, ``origin = trunc(start/k)`` (rounded + toward zero), ``shape = ceil((stop - start)/k)``, and coordinate + ``origin + i`` maps to base cell ``start + i*k``. +- **Strict containment**: non-empty slice intervals must lie within the domain + — no clamping, no negative-wrapping; empty intervals are valid anywhere; + reversed non-empty bounds are an error, not an empty result. +- **Fancy-dim rule**: index-array dims get fresh explicit ``[0, n)`` domains; + index-array values are absolute domain coordinates. +- **Translate rules**: ``translate_by``/``translate_to`` shift the input domain + while preserving which cells are addressed. +""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform +from zarr.errors import BoundsCheckError + + +def _identity(lo: int, hi: int) -> IndexTransform: + """Identity transform over the 1-D domain [lo, hi).""" + return IndexTransform.identity(IndexDomain(inclusive_min=(lo,), exclusive_max=(hi,))) + + +def _a() -> IndexTransform: + """The oracle's base fixture: identity over [0, 12).""" + return _identity(0, 12) + + +def _w() -> IndexTransform: + """The oracle's translated fixture: identity over [-10, 2), cell c -> base c + 10.""" + return _a().translate_domain_by((-10,)) + + +def _dim(t: IndexTransform) -> DimensionMap: + m = t.output[0] + assert isinstance(m, DimensionMap) + return m + + +def _base_cells(t: IndexTransform) -> list[int]: + """The base cells a 1-D single-DimensionMap transform addresses, in order.""" + m = _dim(t) + lo, hi = t.domain.inclusive_min[0], t.domain.exclusive_max[0] + return [m.offset + m.stride * c for c in range(lo, hi)] + + +class TestDomainPreservation: + """Oracle section 1-2: step-1 slices keep literal coordinates.""" + + def test_slice_preserves_domain(self) -> None: + t = _a()[2:10] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_integer_on_preserved_domain_is_a_coordinate(self) -> None: + v = _a()[2:10] + assert isinstance(v[3].output[0], ConstantMap) + assert v[3].output[0].offset == 3 # coordinate 3 = base cell 3 + assert v[2].output[0].offset == 2 + assert v[9].output[0].offset == 9 + + @pytest.mark.parametrize("bad", [0, -1, 10]) + def test_out_of_domain_integer_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[2, 10\)"): + _a()[2:10][bad] + + def test_slice_of_slice_is_literal(self) -> None: + v = _a()[2:10] + t = v[3:7] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((3,), (7,)) + assert _base_cells(t) == [3, 4, 5, 6] + + def test_ellipsis_preserves_domain(self) -> None: + v = _a()[2:10] + t = v[...] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + + +class TestNegativeOriginDomain: + """Oracle section 3: on domain [-10, 2), -1 is just another index.""" + + def test_translated_domain(self) -> None: + w = _w() + assert (w.domain.inclusive_min, w.domain.exclusive_max) == ((-10,), (2,)) + assert _base_cells(w) == list(range(12)) + + @pytest.mark.parametrize(("coord", "base"), [(-5, 5), (-10, 0), (-1, 9), (1, 11)]) + def test_negative_coordinates_address_cells(self, coord: int, base: int) -> None: + t = _w()[coord] + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == base + + @pytest.mark.parametrize("bad", [-11, 2]) + def test_out_of_domain_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[-10, 2\)"): + _w()[bad] + + def test_negative_slice_bounds_are_coordinates(self) -> None: + t = _w()[-5:] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((-5,), (2,)) + assert _base_cells(t) == [5, 6, 7, 8, 9, 10, 11] + t2 = _w()[-5:-2] + assert (t2.domain.inclusive_min, t2.domain.exclusive_max) == ((-5,), (-2,)) + assert _base_cells(t2) == [5, 6, 7] + + +class TestStridedDomains: + """Oracle section 5: origin = trunc(start/step), coord origin+i -> start + i*step.""" + + # (slice, expected (lo, hi), expected base cells) — verbatim oracle rows. + CASES: ClassVar[list[tuple[slice, tuple[int, int], list[int]]]] = [ + (slice(1, 10, 3), (0, 3), [1, 4, 7]), + (slice(None, None, 2), (0, 6), [0, 2, 4, 6, 8, 10]), + (slice(2, 11, 3), (0, 3), [2, 5, 8]), + (slice(0, 12, 4), (0, 3), [0, 4, 8]), + (slice(5, 12, 2), (2, 6), [5, 7, 9, 11]), + (slice(6, 12, 2), (3, 6), [6, 8, 10]), + (slice(7, 12, 3), (2, 4), [7, 10]), + ] + + @pytest.mark.parametrize(("sel", "dom", "cells"), CASES) + def test_strided_domain_and_cells( + self, sel: slice, dom: tuple[int, int], cells: list[int] + ) -> None: + t = _a()[sel] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == dom + assert _base_cells(t) == cells + + def test_strided_on_negative_origin(self) -> None: + # w[-9:2:2] -> domain [-4, 2), base cells 1,3,5,7,9,11 + t = _w()[-9:2:2] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (-4, 2) + assert _base_cells(t) == [1, 3, 5, 7, 9, 11] + # w[::2] -> domain [-5, 1), base cells 0,2,4,6,8,10 + t2 = _w()[::2] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (-5, 1) + assert _base_cells(t2) == [0, 2, 4, 6, 8, 10] + + def test_strided_composition(self) -> None: + s = _a()[1:10:3] # domain [0, 3), cells 1,4,7 + assert [s[k].output[0].offset for k in range(3)] == [1, 4, 7] + t = s[1:3] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (1, 3) + assert _base_cells(t) == [4, 7] + t2 = _a()[::2][1:4] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (1, 4) + assert _base_cells(t2) == [2, 4, 6] + t3 = _a()[::2][::2] + assert (t3.domain.inclusive_min[0], t3.domain.exclusive_max[0]) == (0, 3) + assert _base_cells(t3) == [0, 4, 8] + + @pytest.mark.parametrize("bad", [-2, -1, 3, 4]) + def test_strided_bounds(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[0, 3\)"): + _a()[1:10:3][bad] + + +class TestStrictContainment: + """Oracle section 11: no clamping, no wrapping; empty intervals valid anywhere.""" + + @pytest.mark.parametrize( + "sel", + [ + slice(5, 100), + slice(-3, None), + slice(-3, -1), + slice(0, 13), + slice(12, 14), + slice(100, 200), + ], + ) + def test_uncontained_interval_raises(self, sel: slice) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _a()[sel] + + def test_uncontained_on_negative_origin(self) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _w()[-20:] + + @pytest.mark.parametrize( + ("sel", "pos"), [(slice(5, 5), 5), (slice(0, 0), 0), (slice(13, 13), 13)] + ) + def test_empty_interval_valid_anywhere(self, sel: slice, pos: int) -> None: + t = _a()[sel] + assert t.domain.shape == (0,) + assert t.domain.inclusive_min[0] == pos + + @pytest.mark.parametrize("sel", [slice(5, 2), slice(100, 50)]) + def test_reversed_bounds_raise(self, sel: slice) -> None: + with pytest.raises(IndexError, match="valid.*interval|interval"): + _a()[sel] + + +class TestTranslate: + """Oracle sections 4 and 12: translate_by / translate_to preserve the cell mapping.""" + + def test_translate_to_zero(self) -> None: + t = _a()[2:10].translate_domain_to((0,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (8,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_translate_to_offset(self) -> None: + t = _a().translate_domain_to((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (17,)) + assert _base_cells(t) == list(range(12)) + + def test_translate_by_composes_with_stride(self) -> None: + # a[::2].translate_by[5] -> domain [5, 11), base = 2*(coord-5) + t = _a()[::2].translate_domain_by((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (11,)) + assert _base_cells(t) == [0, 2, 4, 6, 8, 10] + assert t[5].output[0].offset == 0 + assert t[10].output[0].offset == 10 + with pytest.raises(BoundsCheckError, match=r"valid indices \[5, 11\)"): + t[0] + + def test_translate_strided_to(self) -> None: + t = _a()[1:10:3].translate_domain_to((100,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((100,), (103,)) + assert _base_cells(t) == [1, 4, 7] + + +class TestFancyDims: + """Oracle section 7: fancy dims get fresh [0, n); values are absolute coordinates.""" + + def test_index_array_values_are_coordinates(self) -> None: + v = _a()[2:10] + t = v.oindex[(np.array([3, 5], dtype=np.intp),)] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (2,)) + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [3, 5]) + + def test_index_array_on_negative_origin(self) -> None: + t = _w().oindex[(np.array([-10, -1], dtype=np.intp),)] + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [0, 9]) + + def test_index_array_out_of_domain_raises(self) -> None: + v = _a()[2:10] + for bad in ([0, 3], [-1, 3], [3, 10]): + with pytest.raises(BoundsCheckError): + v.oindex[(np.array(bad, dtype=np.intp),)] diff --git a/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py index 92bb770c45..8e3810465c 100644 --- a/tests/test_transforms/test_transform.py +++ b/tests/test_transforms/test_transform.py @@ -62,14 +62,16 @@ def test_slice_identity(self) -> None: def test_slice_narrows(self) -> None: t = IndexTransform.from_shape((10, 20)) result = t[2:8, 5:15] + # Domains are preserved (TensorStore): the slice keeps its literal + # coordinates, so the map stays the identity (out = in). assert result.domain.shape == (6, 10) - assert result.domain.origin == (0, 0) + assert result.domain.origin == (2, 5) assert isinstance(result.output[0], DimensionMap) - assert result.output[0].offset == 2 + assert result.output[0].offset == 0 assert result.output[0].stride == 1 assert result.output[0].input_dimension == 0 assert isinstance(result.output[1], DimensionMap) - assert result.output[1].offset == 5 + assert result.output[1].offset == 0 assert result.output[1].input_dimension == 1 def test_strided_slice(self) -> None: @@ -149,12 +151,13 @@ def test_negative_int_valid_with_negative_origin(self) -> None: assert result.output[0].offset == -3 def test_composition_of_slices(self) -> None: - """Slicing a sliced transform should compose offsets.""" + """Slicing a sliced transform re-selects in literal domain coordinates.""" t = IndexTransform.from_shape((100,)) - result = t[10:50][5:20] + result = t[10:50][15:30] assert result.domain.shape == (15,) + assert result.domain.origin == (15,) assert isinstance(result.output[0], DimensionMap) - assert result.output[0].offset == 15 + assert result.output[0].offset == 0 assert result.output[0].stride == 1 def test_composition_of_strides(self) -> None: @@ -281,9 +284,11 @@ def test_oindex_mixed(self) -> None: result = t.oindex[idx, 5:15] assert result.input_rank == 2 assert result.domain.shape == (2, 10) + # fancy dim: fresh zero-origin; slice dim: preserved literal coords + assert result.domain.origin == (0, 5) assert isinstance(result.output[0], ArrayMap) assert isinstance(result.output[1], DimensionMap) - assert result.output[1].offset == 5 + assert result.output[1].offset == 0 def test_oindex_multiple_arrays(self) -> None: t = IndexTransform.from_shape((10, 20, 30)) @@ -348,8 +353,9 @@ def test_basic_slice(self) -> None: t = IndexTransform.from_shape((10, 20)) result = selection_to_transform((slice(2, 8), slice(5, 15)), t, "basic") assert result.domain.shape == (6, 10) + assert result.domain.origin == (2, 5) # preserved literal coordinates assert isinstance(result.output[0], DimensionMap) - assert result.output[0].offset == 2 + assert result.output[0].offset == 0 def test_basic_int(self) -> None: t = IndexTransform.from_shape((10, 20)) @@ -380,12 +386,18 @@ def test_vectorized(self) -> None: assert isinstance(result.output[1], ArrayMap) def test_composition_with_non_identity(self) -> None: - """Indexing a sliced transform composes offsets.""" + """Indexing a sliced transform uses literal domain coordinates. + + The slice [10:50] preserves its domain, so a follow-up [15:30] + re-selects coordinates 15..29 of the base (TensorStore semantics), and + the composed map stays the identity (out = in). + """ t = IndexTransform.from_shape((100,))[10:50] - result = selection_to_transform(slice(5, 20), t, "basic") - assert result.domain.shape == (15,) + result = selection_to_transform(slice(15, 30), t, "basic") + assert (result.domain.inclusive_min, result.domain.exclusive_max) == ((15,), (30,)) assert isinstance(result.output[0], DimensionMap) - assert result.output[0].offset == 15 + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 class TestIndexTransformIntersect: @@ -514,78 +526,3 @@ def test_translate_2d(self) -> None: assert result.output[0].offset == -5 assert isinstance(result.output[1], DimensionMap) assert result.output[1].offset == -10 - - -class TestNegativeOriginDomains: - """Literal-coordinate semantics must hold at ANY domain origin, not just 0. - - Public views always re-zero their domains, but the transform library admits - non-zero origins (`IndexDomain.translate`), and the literal rule — an index - names a coordinate, never a position or an offset from the end — must apply - uniformly to integers and slice bounds alike. On domain [-10, 2): coordinate - -5 is *in bounds*; slice bounds are coordinates (t[0:2] means storage - [0, 2), the last two cells, NOT the first two positions); bounds below - inclusive_min raise (the any-origin generalization of "negative raises at - origin 0"); bounds past exclusive_max clamp (range intersection). - """ - - def _t(self) -> IndexTransform: - return IndexTransform.identity(IndexDomain.from_shape((12,)).translate((-10,))) - - def test_integer_literal_in_bounds(self) -> None: - """An in-domain negative coordinate is a valid literal integer index.""" - t = self._t()[-5] - assert isinstance(t.output[0], ConstantMap) - assert t.output[0].offset == -5 - - def test_integer_below_origin_raises(self) -> None: - """A coordinate below inclusive_min is out of bounds.""" - with pytest.raises(IndexError, match=r"valid indices \[-10, 2\)"): - self._t()[-11] - - def test_slice_bounds_are_literal_coordinates(self) -> None: - """Slice bounds are coordinates, not positions: t[0:2] is storage [0, 2).""" - t = self._t()[0:2] - assert t.domain.shape == (2,) - assert isinstance(t.output[0], DimensionMap) - # first element of the result must map to storage coordinate 0, not -10 - assert t.output[0].offset == 0 - - def test_slice_with_negative_literal_bound(self) -> None: - """t[-5:] on domain [-10, 2) is the literal interval [-5, 2).""" - t = self._t()[-5:] - assert t.domain.shape == (7,) - assert isinstance(t.output[0], DimensionMap) - assert t.output[0].offset == -5 - - def test_slice_bound_below_origin_raises(self) -> None: - """A slice bound below inclusive_min raises, mirroring the origin-0 rule.""" - with pytest.raises(IndexError, match=r"valid indices \[-10, 2\)"): - self._t()[-12:] - - def test_slice_overflow_clamps(self) -> None: - """Bounds past exclusive_max clamp: ranges intersect the domain.""" - t = self._t()[-5:100] - assert t.domain.shape == (7,) - - -def test_public_view_domains_are_zero_origin() -> None: - """Every view the public `.lazy` API can produce has a zero-origin domain. - - The user-guide rule "a negative value is never in bounds" is the origin-0 - corollary of literal coordinates; this pins the premise so a future - translate-style API cannot silently invalidate the documented behavior. - """ - import zarr - - a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") - a[...] = np.arange(12, dtype="i4") - views = ( - a.lazy[2:10], - a.lazy[2:10].lazy[1:5], - a.lazy[::2], - a.lazy.oindex[np.array([1, 5, 9], dtype=np.intp)], - a.lazy.vindex[np.array([1, 5], dtype=np.intp)], - ) - for v in views: - assert all(m == 0 for m in v._async_array._transform.domain.inclusive_min) From 816810db9458248ade197ce8a9db13eefca5089a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 11:25:45 +0200 Subject: [PATCH 38/43] docs(lazy): rewrite the lazy-indexing guide around TensorStore domain semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The theory section now teaches the real model: domains are preserved (an index is a stable name for a cell, not a position), -1 is just another coordinate — demonstrably valid after translate_by((-10,)) — strict containment replaces clamping, strided views renumber by trunc-division, and every path on a view shares one coordinate system (base arrays keep NumPy semantics unchanged). Adds translate_to/translate_by patterns, eager iter() rejection on views (made eager — a generator __iter__ deferred the raise to first next()), and the positional contract of ChunkProjection.array_selection. All examples remain live (markdown-exec) and were executed end-to-end. Assisted-by: ClaudeCode:claude-opus-4.8 --- docs/user-guide/lazy_indexing.md | 184 +++++++++++++++++++++---------- src/zarr/core/array.py | 5 +- tests/test_lazy_indexing.py | 8 ++ 3 files changed, 134 insertions(+), 63 deletions(-) diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md index 8a43626c3b..7e46b6b175 100644 --- a/docs/user-guide/lazy_indexing.md +++ b/docs/user-guide/lazy_indexing.md @@ -6,6 +6,11 @@ lightweight **view** — itself a `zarr.Array` — without touching storage. Vie compose, support orthogonal and coordinate selection, write through to the backing array, and materialize on demand. +Zarr's lazy indexing follows [TensorStore's indexing +model](https://google.github.io/tensorstore/python/indexing.html): a view has a +**domain** — a box of coordinates — and every index is a **literal coordinate** +in that domain. + ```python exec="true" session="lazy" import numpy as np import zarr @@ -33,36 +38,58 @@ action, zarr can: - **plan** with it (`chunk_projections` enumerates exactly the stored chunks the declaration touches). -A view is a *window*: its shape is the selection's shape, and indexing a view is -always relative to the view itself, starting at zero — exactly like indexing a -NumPy view: +### Domains are preserved: an index is a name, not a position + +A view keeps the coordinates of the cells it selects. Slicing `[2:10]` does not +renumber anything — the view's domain *is* `[2, 10)`, and coordinate 3 still +means what it meant on the parent: ```python exec="true" session="lazy" source="above" result="ansi" -v = a.lazy[2:10] # window onto cells 2..9 -print(v.shape) # the window's shape -print(v.lazy[0].result()) # the window's first element -> base cell 2 -print(v.lazy[3:5].result()) # window cells 3,4 -> base cells 5,6 +v = a.lazy[2:10] +print(v.shape) # (8,) — eight cells... +print(v.lazy[3].result()) # ...and coordinate 3 is still base cell 3 +print(v.lazy[3:7].result()) # coordinates [3, 7) — literal, stable ``` -The `domain={ [2, 10) }` shown in a view's repr is **provenance** — the region -of the backing store the view maps onto. It is *not* the coordinate system you -index with; that is always `[0, shape)`. +This is what makes composition safe: a coordinate means the same cell no matter +how many views deep you are. `a.lazy[2:10].lazy[3:7]` and `a.lazy[3:7]` are the +same selection. -### Literal coordinates: `-1` is just another index +The price of stable names: **positions are not valid indices.** The first +element of `v` is coordinate 2, not 0 — and `v[0]` is an error, not the first +element: -A lazy selection is a declaration in **literal coordinates**. NumPy's negative -indexing ("count from the end") is sugar applied at *evaluation* time; in a -deferred, composable setting it would silently re-bind meaning as views compose. -Zarr therefore treats every coordinate in a lazy selection literally, and since -view domains start at 0, **a negative value is never in bounds — no matter the -syntactic form** (integer, slice bound, or index array): +```python exec="true" session="lazy" source="above" result="ansi" +try: + v[0] +except IndexError as e: + print(e) +``` + +To renumber a view explicitly, move its domain with `translate_to` (or shift it +with `translate_by`) — the data does not move, only the labels: + +```python exec="true" session="lazy" source="above" result="ansi" +z = v.translate_to((0,)) # same cells, coordinates now [0, 8) +print(z.lazy[0].result()) # coordinate 0 -> base cell 2 +w = a.translate_by((-10,)) # a view of `a` labeled [-10, 2) +print(w.lazy[-1].result()) # -1 is just another index: base cell 9 +``` + +### `-1` is just another index + +Because indices are literal coordinates, a negative index is *not* "from the +end" — it names the coordinate `-1`, which your domain may or may not contain. +On the translated view above, `-1` was perfectly valid. On a fresh array (domain +`[0, n)`), it is out of bounds — in **every** syntactic form: integers, slice +bounds, and index arrays are treated identically. ```python exec="true" session="lazy" source="above" result="ansi" for make in (lambda: a.lazy[-1], lambda: a.lazy[-3:], lambda: a.lazy.oindex[[-1]]): try: make() except IndexError as e: - print(type(e).__name__, "-", e) + print(type(e).__name__, "-", str(e).split(";")[0]) ``` To select from the end, say what you mean literally: @@ -72,29 +99,46 @@ n = a.shape[0] print(a.lazy[n - 3 :].result()) # the last three elements ``` -Positive out-of-range *slice* bounds are still fine — a slice denotes a range, -and ranges intersect the domain (as in Python and NumPy). Clamping can only -shorten a result; it can never silently select different data the way -wraparound can: +### No clamping: intervals must fit the domain + +A slice interval must be contained in the domain — an out-of-range bound is an +error, not a silently shorter result. Empty intervals are the one exception: +they are valid anywhere. Reversed bounds are an error, not an empty result. ```python exec="true" session="lazy" source="above" result="ansi" -print(a.lazy[5:100].result()) # clamps to [5, 12) +for sel in (slice(5, 100), slice(5, 2)): + try: + a.lazy[sel] + except IndexError as e: + print(type(e).__name__, "-", str(e).split(";")[0]) +print(a.lazy[5:5].shape) # empty is fine, anywhere ``` -### Two dialects on one object +### Strided views renumber by division + +A strided slice produces a domain in *strided units*: for step `k`, the new +origin is `start / k` rounded toward zero, and coordinate `origin + i` maps to +base cell `start + i*k` (TensorStore's rule): -A view is an ordinary `zarr.Array`, so it also has the ordinary *eager* methods -— and those keep full NumPy semantics, negatives included. The two dialects -split cleanly by intent: +```python exec="true" session="lazy" source="above" result="ansi" +s = a.lazy[1:10:3] # base cells 1, 4, 7 +print(s) # domain [0, 3) +print(s.lazy[1].result()) # coordinate 1 -> base cell 4 +``` + +### One coordinate system per view -| You are... | Spelling | Negative indices | -| ----------- | --------------------------- | ----------------------- | -| declaring | `v.lazy[...]` (a new view) | literal — raise | -| accessing | `v[...]`, `v.oindex[...]` | NumPy — from the end | +Every way of indexing a view — `v[...]`, `v.lazy[...]`, `v.oindex`, `v.vindex` +— uses the same domain coordinates. (The base array's ordinary `a[...]` API is +unchanged: it keeps full NumPy semantics, negatives and all. The literal rules +apply to *views* and to the `.lazy` accessor.) NumPy-style zero-based access to +a view's data is spelled explicitly: materialize with `result()` / +`np.asarray`, or renumber with `translate_to`. ```python exec="true" session="lazy" source="above" result="ansi" -print(v[-1]) # eager access: NumPy dialect -> base cell 9 -print(v.lazy[v.shape[0] - 1].result()) # lazy declaration: same cell, said literally +print(v[3], v.lazy[3].result()) # same coordinate, same cell +print(np.asarray(v)[0]) # materialized: NumPy rules apply +print(a[-1]) # base arrays keep NumPy semantics ``` ## Common patterns @@ -105,8 +149,8 @@ print(v.lazy[v.shape[0] - 1].result()) # lazy declaration: same cell, said lite img = zarr.create_array(store="memory://lazy-img", shape=(100, 100), chunks=(10, 10), dtype="float64") img[...] = np.arange(100 * 100).reshape(100, 100) -crop = img.lazy[25:75, 25:75] # 50x50 window, no I/O -inner = crop.lazy[10:40, 10:40] # 30x30 window of the window +crop = img.lazy[25:75, 25:75] # no I/O; domain [25,75) x [25,75) +inner = crop.lazy[35:65, 35:65] # coordinates are literal: this is img[35:65, 35:65] print(crop.shape, inner.shape) print(float(np.mean(inner))) # I/O happens here, for the inner crop only ``` @@ -117,28 +161,31 @@ Assignment through the accessor, or through a view, routes values back to storage — including strided and composed selections: ```python exec="true" session="lazy" source="above" result="ansi" -img.lazy[30:50, 40:60] = 0.0 # region write, no read-back needed here +img.lazy[30:50, 40:60] = 0.0 # region write tile = img.lazy[30:50, 40:60] -tile[0:5, 0:5] = 7.0 # eager write through the view -img.lazy[::2, ::2] = -1.0 # strided write, NumPy-equivalent +tile[30:35, 40:45] = 7.0 # write through the view, same coordinates +img.lazy[::2, ::2] = -1.0 # strided write, NumPy-equivalent cells print(img[29:33, 39:43]) ``` ### Orthogonal and coordinate selection `lazy.oindex` selects an outer product per axis; `lazy.vindex` selects points. -Both return views that compose: +Index-array *values* are domain coordinates; the dimension a fancy selection +*creates* gets a fresh `[0, n)` domain (there is no meaningful coordinate to +preserve for "the i-th pick"): ```python exec="true" session="lazy" source="above" result="ansi" -rows = img.lazy.oindex[[3, 17, 42], :] # three full rows -> shape (3, 100) -sub = rows.lazy[:, 10:20] # then a column window of those rows +rows = img.lazy.oindex[[3, 17, 42], :] # picked dim: domain [0, 3); row dim preserved +sub = rows.lazy[:, 10:20] # column window, literal coords print(rows.shape, sub.shape) -pts = img.lazy.vindex[[3, 17], [5, 9]] # two points -> shape (2,) +pts = img.lazy.vindex[[3, 17], [5, 9]] # two points -> fresh domain [0, 2) print(pts.result()) ``` -Boolean masks are array selections, so they go through `oindex`/`vindex`: +Boolean masks are array selections, so they go through `oindex`/`vindex`; the +positions of `True` values become coordinates: ```python exec="true" session="lazy" source="above" result="ansi" mask = np.zeros(12, dtype=bool) @@ -148,12 +195,17 @@ print(a.lazy.oindex[mask].result()) ### Materializing -`view.result()`, `view[...]`, and `np.asarray(view)` are equivalent reads of the -whole view; views also work directly with NumPy reductions: +`view.result()`, `view[...]`, and `np.asarray(view)` are equivalent whole-view +reads; views also work directly with NumPy reductions. Views are **not** +iterable (iterate the materialized result instead): ```python exec="true" session="lazy" source="above" result="ansi" -w = a.lazy[3:9] -print(w.result(), np.asarray(w), float(np.mean(w))) +w2 = a.lazy[3:9] +print(w2.result(), float(np.mean(w2))) +try: + iter(w2) +except TypeError as e: + print(e) ``` ### Chunk-aware processing @@ -161,8 +213,9 @@ print(w.result(), np.asarray(w), float(np.mean(w))) `chunk_projections` enumerates the stored chunks a view touches: which store object (`coord`, `key`), its stored `shape`, the region *within the chunk* the view covers (`chunk_selection`), the region *of the view* it maps to -(`array_selection`), and whether the chunk is only partially covered -(`is_partial` — a partial write requires a read-modify-write): +(`array_selection`, positional — 0-based into the view's extent), and whether +the chunk is only partially covered (`is_partial` — a partial write requires a +read-modify-write): ```python exec="true" session="lazy" source="above" result="ansi" for p in a.lazy[2:10].chunk_projections(): @@ -172,14 +225,16 @@ print(a.lazy[3:9].is_chunk_aligned()) # starts and ends on chunk boundaries ``` This is the supported way to partition any selection for parallel or -chunk-at-a-time work — compose the selection through `.lazy`, then project: +chunk-at-a-time work — compose the selection through `.lazy`, then project. +Since `array_selection` is positional, re-zero the view (or materialize) to use +it: ```python exec="true" session="lazy" source="above" result="ansi" +crop0 = img.lazy[25:75, 25:75].translate_to((0, 0)) total = 0.0 -crop = img.lazy[25:75, 25:75] -for p in crop.chunk_projections(): - total += float(np.sum(crop[p.array_selection])) -print(total == float(np.sum(crop))) +for p in crop0.chunk_projections(): + total += float(np.sum(crop0[tuple(slice(s.start, s.stop) for s in p.array_selection)])) +print(total == float(np.sum(crop0))) ``` For sharded arrays, pass `unit="write"` to enumerate at shard (write-unit) @@ -204,21 +259,28 @@ array. ## Coming from NumPy -- **Negative indices raise on the lazy path** — in every form (integer, slice - bound, index array). Use `shape[dim] - k`, or the eager methods, which keep - NumPy semantics. -- **No negative steps**: `a.lazy[::-1]` raises; reversal is not supported by - zarr indexing generally. +- **A view's indices are coordinates, not positions.** `a.lazy[2:10]` is + indexed with 2..9, not 0..7. Renumber explicitly with + `view.translate_to((0, ...))` if you want positions. +- **Negative indices are not from-the-end** — in any form (integer, slice + bound, index array). They name literal coordinates, which fresh arrays' + domains (starting at 0) do not contain. Use `shape[dim] - k`, or translate + the domain so negative coordinates exist. +- **No clamping**: out-of-range slice bounds raise; reversed bounds raise; + only empty intervals are allowed anywhere. +- **No negative steps**: `a.lazy[::-1]` raises; reversal is not yet supported. - **No `newaxis`**: `a.lazy[None]` raises; insert axes on the materialized result instead. - **The basic accessor takes basic selections only** (integers, slices, ellipsis). Lists, arrays, and boolean masks go through `lazy.oindex` / `lazy.vindex`. -- **Scalar reads return NumPy scalars**: `a.lazy[3].result()` is `np.int32(3)`, - matching `a[3]`'s value with scalar type. +- **Views are not iterable**; iterate `view.result()`. +- **Base arrays are unchanged**: `a[-1]`, `a[5:100]`, and friends keep full + NumPy semantics on non-view arrays. ## Current limitations +- Negative slice steps (reversal) are not yet supported. - Integer indexing a dimension *created by* an `oindex`/`vindex` selection (e.g. `rows.lazy[0]` after `rows = a.lazy.oindex[[3, 17, 42], :]`) is not yet supported reliably; slice the view instead (`rows.lazy[0:1]`). diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 777ef19d9e..fcd0b5f873 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2045,8 +2045,9 @@ def __iter__(self) -> Any: ) if self.ndim == 0: raise TypeError("iteration over a 0-d array") - for i in range(self.shape[0]): - yield self[i] + # A plain generator function would defer these raises to the first + # next(); returning an inner generator keeps them eager at iter(). + return (self[i] for i in range(self.shape[0])) @classmethod def _create( diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 825268ce4b..537960a960 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -499,6 +499,14 @@ def test_guard_message_names_real_apis(self) -> None: assert "chunk_projections" in str(ei.value) assert "chunk_layout" not in str(ei.value) + def test_views_are_not_iterable(self) -> None: + """iter() on a view raises eagerly (TensorStore parity): the getitem + protocol counts positions from 0, which are not domain coordinates.""" + a, ref = _make(CONFIGS[1]) + with pytest.raises(TypeError, match="not iterable"): + iter(a.lazy[2:10]) + assert [int(x) for x in a][:3] == [int(v) for v in ref[:3]] # base unchanged + def test_fancy_view_repr_does_not_crash(self) -> None: """repr of an integer-indexed fancy view must not raise (0-d index array in selection_repr).""" From ba9760316a45688e60b99d9427e56b9139bc0d98 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 2 Jul 2026 20:57:37 +0200 Subject: [PATCH 39/43] test(lazy): pin mask True-positions as absolute coordinates (roborev #358) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit roborev flagged mask-on-non-zero-origin-view as reading "wrong" cells, assuming NumPy positional semantics. Verified against the executed tensorstore 0.1.84 oracle: TensorStore treats mask True-positions as absolute coordinates counted from 0, origin-blind (a mask is sugar for the coordinate array of its True positions), so a True at position 3 on a domain-[2,10) view addresses cell 3 — exactly our behavior; a True below the origin is out of domain (we reject eagerly where TensorStore defers to read — the documented timing deviation). The real gap was coverage: nothing pinned this on a translated view, so the reviewer could not infer intent. Add the oracle-pinning regression test (read + write + below-origin rejection) and sharpen the docs callout. Assisted-by: ClaudeCode:claude-opus-4.8 --- docs/user-guide/lazy_indexing.md | 6 +++++- tests/test_lazy_indexing.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md index 7e46b6b175..19951e0922 100644 --- a/docs/user-guide/lazy_indexing.md +++ b/docs/user-guide/lazy_indexing.md @@ -185,7 +185,11 @@ print(pts.result()) ``` Boolean masks are array selections, so they go through `oindex`/`vindex`; the -positions of `True` values become coordinates: +positions of `True` values become **coordinates, counted from 0** — not offsets +from the view's origin. On a view whose domain starts at 2, a mask `True` at +position 3 addresses coordinate 3, and a `True` at position 0 or 1 is out of +the domain (matching TensorStore, where a mask is sugar for the coordinate +array of its `True` positions): ```python exec="true" session="lazy" source="above" result="ansi" mask = np.zeros(12, dtype=bool) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 537960a960..2bb4592912 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -499,6 +499,25 @@ def test_guard_message_names_real_apis(self) -> None: assert "chunk_projections" in str(ei.value) assert "chunk_layout" not in str(ei.value) + def test_mask_positions_are_absolute_coordinates(self) -> None: + """Boolean-mask True-positions are absolute coordinates counted from 0, + NOT offsets from the view's origin (TensorStore semantics, verified + against tensorstore 0.1.84: on a domain-[2,10) view, a mask with True + at {3,5,7} addresses cells 3,5,7 — and a True below the origin is out + of the domain, rejected eagerly here where TensorStore defers to read). + """ + a, ref = _make(CONFIGS[1]) # shape (24,) + v = a.lazy[2:10] # domain [2, 10) + mask = np.zeros(8, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(v.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + v.lazy.oindex[mask] = 99 # write path: same coordinates + assert list(np.flatnonzero(a[...] == 99)) == [3, 5, 7] + below_origin = np.zeros(8, dtype=bool) + below_origin[0] = True # coordinate 0 is not in [2, 10) + with pytest.raises(BoundsCheckError, match="out of bounds"): + v.lazy.oindex[below_origin] + def test_views_are_not_iterable(self) -> None: """iter() on a view raises eagerly (TensorStore parity): the getitem protocol counts positions from 0, which are not domain coordinates.""" From 8d273d2e596823fb8d4a639cec8879da27036b3e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 15 Jul 2026 22:40:20 +0200 Subject: [PATCH 40/43] docs: specify index transform corrections Assisted-by: Codex:GPT-5 --- ...tore-index-transform-corrections-design.md | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md diff --git a/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md new file mode 100644 index 0000000000..4755cf8311 --- /dev/null +++ b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md @@ -0,0 +1,103 @@ +# TensorStore-Aligned Index Transform Corrections + +## Goal + +Correct lazy and transform-based indexing so that basic, orthogonal, and vectorized selections preserve NumPy/Zarr semantics through composition, chunk resolution, reads, writes, and serialization. + +## Scope + +This change fixes five reviewed regressions: + +1. Legacy advanced-indexing fallbacks bypass a lazy view's transform. +2. Multiple independent orthogonal index arrays are mistaken for correlated vectorized arrays. +3. Lazy selections do not consistently normalize negative indices or validate bounds. +4. Zero-dimensional identity transforms are incorrectly flattened. +5. Multidimensional vectorized reads cannot safely write directly into a caller-provided `out` buffer. + +The public indexing API remains unchanged. Structured-dtype field selection remains on the legacy codec path, but it must either operate on the correct composed storage selection or reject unsupported transformed access explicitly; it must never silently access the wrong storage region. + +## Transform Representation + +Follow TensorStore's normalized index-transform model. Every `ArrayMap.index_array` is interpreted over the complete input domain of its containing `IndexTransform`. + +- The index array rank equals the transform input rank. +- An axis on which an index map does not vary is represented as a singleton/broadcast axis. +- Independent orthogonal maps vary along different input dimensions. +- Correlated vectorized maps vary along the same broadcast input dimensions. +- Zero-rank transforms use zero-rank index arrays where applicable. + +For an orthogonal selection with row indices of length 2 and column indices of length 3, the normalized maps have shapes `(2, 1)` and `(1, 3)`. For a pairwise vectorized selection of length 2, both maps have shape `(2,)`. + +Correlation must be derived from the maps' input-dimension dependencies. Code must not infer vectorized semantics merely from the presence of two or more `ArrayMap` instances. + +## Construction and Composition + +Basic, orthogonal, and vectorized selection constructors normalize their inputs before producing transforms: + +- Python and NumPy integer scalars use normal Zarr negative-index wraparound. +- Integer arrays normalize negative elements and reject values outside the selected view's bounds. +- Boolean masks must match the dimensions they consume. +- Slice steps remain positive-only, consistent with current Zarr behavior. +- Orthogonal integer arrays remain one-dimensional at the public API boundary, then become full-rank broadcast arrays in the normalized transform. +- Vectorized integer arrays are broadcast together before being expanded to full input rank. + +Composition preserves full-rank dependency information. When an outer transform indexes an inner `ArrayMap`, the resulting array is evaluated and broadcast over the outer input domain without squeezing dependency axes. Composition may simplify a map to `ConstantMap` or `DimensionMap` only when the resulting mapping is provably equivalent. + +## Intersection and Chunk Resolution + +Intersection classifies maps by their dependency axes: + +- Correlated maps are filtered with one joint mask so a coordinate survives only when all mapped storage coordinates are inside the chunk. +- Independent maps are filtered per dependency axis and retain outer-product semantics. +- Mixed `ArrayMap`, `DimensionMap`, and `ConstantMap` transforms retain their input-domain axis ordering. + +Chunk resolution produces output selectors consistent with the normalized input domain. Correlated maps share scatter coordinates; independent maps produce orthogonal selectors. The resolver must not rely on array-map count as a proxy for correlation. + +## Array Read and Write Paths + +All non-field lazy reads and writes use the composed `IndexTransform` path. Eager advanced operations may continue using legacy indexers when operating on an identity transform, but a transformed view must not fall back to storage-relative indexing that ignores its transform. + +Structured-field operations continue using the existing field-aware codec path. If a transformed structured-field selection cannot be expressed correctly through that path, raise a clear `NotImplementedError` instead of reading or writing the wrong region. + +Buffer handling follows transform semantics: + +- Flatten only when the transform contains correlated array maps whose chunk resolution emits flat scatter indices. +- Never treat an empty output-map tuple as vectorized; zero-dimensional identity reads and writes retain shape `()`. +- When a multidimensional vectorized read supplies `out`, allocate a flat temporary buffer, perform scatter reads into it, reshape it to the selection shape, and copy it into `out`. +- Writes validate broadcasting against the visible selection domain before any vectorized flattening. + +## Serialization + +JSON round-tripping preserves full-rank index-array shapes, including singleton dependency axes. No transform-wide `orthogonal` or `vectorized` mode flag is added. Existing JSON that already satisfies the rank invariant continues to load unchanged. + +If compatibility with branch-local JSON fixtures containing lower-rank arrays is required, the loader may normalize them only when their dependency alignment is unambiguous. Ambiguous lower-rank arrays must be rejected rather than guessed. + +## Error Handling + +- Scalar and array indices outside explicit view bounds raise `IndexError` before chunk access. +- Invalid boolean mask shapes raise `IndexError` during transform construction. +- Incompatible assignment shapes raise `ValueError` before storage mutation. +- Unsupported transformed structured-field access raises `NotImplementedError` with an explanation. +- All validation is relative to the current lazy view domain, not the root array shape. + +## Test Strategy + +Add focused regression tests before implementation and verify that each fails for the reviewed reason: + +- A sliced lazy view followed by advanced orthogonal read and write accesses the composed storage region. +- Multiple orthogonal index arrays with unequal lengths produce an outer-product result across chunks. +- Multiple vectorized arrays retain pairwise/broadcast semantics. +- Negative scalar and array indices work relative to a lazy view; overly negative and positive out-of-range values raise. +- Invalid Boolean mask shapes raise. +- Zero-dimensional reads return a scalar and zero-dimensional writes persist it. +- Multidimensional coordinate selection with a caller-provided `out` buffer matches NumPy. +- Transform and JSON unit tests assert full-rank singleton-axis normalization and round-trip preservation. + +After focused tests pass, run the complete transform, lazy-indexing, indexing, and array test modules, followed by the repository's configured static checks for the modified files. + +## Non-Goals + +- Negative slice steps. +- New public indexing modes or APIs. +- General TensorStore dimension labels, implicit bounds, or arbitrary-origin domains. +- Unrelated restructuring of the legacy indexer or codec pipeline. From e1ba6f1c4d7af33ef6c7b4325c4070541c4978ad Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 15 Jul 2026 22:44:11 +0200 Subject: [PATCH 41/43] docs: add stateful indexing test design Assisted-by: Codex:GPT-5 --- ...tore-index-transform-corrections-design.md | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md index 4755cf8311..b2edb86e6f 100644 --- a/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md +++ b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md @@ -95,6 +95,28 @@ Add focused regression tests before implementation and verify that each fails fo After focused tests pass, run the complete transform, lazy-indexing, indexing, and array test modules, followed by the repository's configured static checks for the modified files. +### Stateful Indexing Properties + +Extend the property-based test architecture from single eager selections to indexing programs. A generated example consists of: + +1. A NumPy array and equivalent chunked Zarr array. +2. A valid sequence of one or more indexing operations chosen from basic, orthogonal, and vectorized modes. +3. An execution mode chosen from lazy materialization, eager access on a lazy view, caller-provided `out`, scalar assignment, and array assignment. + +Apply the same program to a NumPy reference and to Zarr. Each operation is relative to the visible domain produced by the preceding operations. The generator must preserve whether index arrays are independent (`oindex`) or correlated (`vindex`) and must deliberately cover unequal orthogonal index lengths, broadcast-compatible vectorized shapes, negative indices, zero-rank results, empty results, and selections spanning multiple chunks. + +Read properties compare all of the following where applicable: + +- result shape; +- scalar versus array result kind; +- dtype; +- element values; +- contents of a caller-provided `out` buffer. + +Write properties compare the entire NumPy and Zarr backing arrays after the operation, not merely the selected view. This detects writes that produce the expected selected values while mutating the wrong storage region. + +Keep focused example-based regressions for each reviewed bug. The stateful properties complement those tests by exploring compositions and shape interactions; they do not replace precise regression cases. + ## Non-Goals - Negative slice steps. From 49fd195d654e939a5f7b7407689207f17566a4f4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 15 Jul 2026 22:46:43 +0200 Subject: [PATCH 42/43] docs: plan index transform corrections Assisted-by: Codex:GPT-5 --- ...tensorstore-index-transform-corrections.md | 319 ++++++++++++++++++ 1 file changed, 319 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md diff --git a/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md b/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md new file mode 100644 index 0000000000..438fbea082 --- /dev/null +++ b/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md @@ -0,0 +1,319 @@ +# TensorStore-Aligned Index Transform Corrections Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make transform-based Zarr indexing preserve eager NumPy/Zarr semantics across lazy composition, orthogonal and vectorized indexing, zero-rank arrays, caller-provided output buffers, and writes. + +**Architecture:** Normalize every `ArrayMap.index_array` over the complete transform input domain, following TensorStore. Derive independent versus correlated indexing from array dependency axes, preserve those axes through composition and chunk intersection, and keep public Zarr negative-index normalization at the array-selection boundary. Verify behavior with focused regressions plus generated indexing programs compared against NumPy. + +**Tech Stack:** Python 3.12+, NumPy, pytest, Hypothesis, Zarr transform and codec pipeline internals. + +## Global Constraints + +- Keep the public indexing API unchanged. +- Use full-input-rank index arrays; do not add a transform-wide indexing-mode flag. +- Preserve TensorStore-style literal coordinates in the general `IndexTransform` API while applying NumPy-style negative wraparound at the public Zarr array boundary. +- Never silently bypass a lazy view transform. +- Keep negative slice steps unsupported. +- Follow test-first red/green/refactor cycles. + +--- + +### Task 1: Normalize ArrayMap Dependency Axes + +**Files:** +- Modify: `src/zarr/core/transforms/transform.py` +- Modify: `src/zarr/core/transforms/composition.py` +- Modify: `src/zarr/core/transforms/json.py` +- Test: `tests/test_transforms/test_transform.py` +- Test: `tests/test_transforms/test_composition.py` +- Test: `tests/test_transforms/test_json.py` + +**Interfaces:** +- Consumes: `IndexTransform(domain, output)`, `_apply_oindex`, `_apply_vindex`, `compose`. +- Produces: full-input-rank `ArrayMap.index_array` values and an internal dependency-axis classifier used by later tasks. + +- [ ] **Step 1: Write failing transform-construction tests** + +Add assertions that orthogonal maps have shapes `(2, 1)` and `(1, 3)`, while two vectorized maps share shape `(2,)`. Assert JSON round-tripping preserves singleton axes. + +```python +def test_oindex_multiple_arrays_preserves_independent_axes() -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.oindex[np.array([1, 3]), np.array([2, 4, 6])] + assert result.domain.shape == (2, 3) + assert result.output[0].index_array.shape == (2, 1) + assert result.output[1].index_array.shape == (1, 3) + + +def test_vindex_multiple_arrays_preserves_shared_axes() -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.vindex[np.array([1, 3]), np.array([2, 4])] + assert result.domain.shape == (2,) + assert result.output[0].index_array.shape == (2,) + assert result.output[1].index_array.shape == (2,) +``` + +- [ ] **Step 2: Run the focused tests and confirm the orthogonal shape assertion fails** + +Run: + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_transform.py tests/test_transforms/test_composition.py tests/test_transforms/test_json.py -q +``` + +Expected: the new orthogonal singleton-axis assertions fail against squeezed arrays. + +- [ ] **Step 3: Implement full-rank normalization** + +Update `_apply_oindex` to reshape each public one-dimensional integer array with one varying axis and singleton axes elsewhere. Update `_apply_vindex` to align every broadcast array with the broadcast dimensions plus singleton slice dimensions. Add a private helper with this contract: + +```python +def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, ...]: + """Return input axes on which a normalized index array varies.""" + return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) +``` + +Preserve full rank in basic indexing, composition, and JSON conversion. Reject an `ArrayMap` whose rank differs from its containing transform input rank, except during a clearly defined compatibility normalization at construction. + +- [ ] **Step 4: Run focused transform tests and refactor without changing behavior** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 1** + +```bash +git add src/zarr/core/transforms tests/test_transforms +git commit -m "fix(indexing): preserve array map dependency axes" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 2: Resolve Independent and Correlated Array Maps Correctly + +**Files:** +- Modify: `src/zarr/core/transforms/transform.py` +- Modify: `src/zarr/core/transforms/chunk_resolution.py` +- Test: `tests/test_transforms/test_transform.py` +- Test: `tests/test_transforms/test_chunk_resolution.py` +- Test: `tests/test_lazy_indexing.py` + +**Interfaces:** +- Consumes: normalized full-rank `ArrayMap` values and `_array_map_dependency_axes` from Task 1. +- Produces: chunk intersections and output selectors that preserve outer-product or pairwise semantics. + +- [ ] **Step 1: Write failing orthogonal resolution tests** + +Create a chunked 2-D array and compare a lazy orthogonal selection with unequal index lengths to `np.ix_`. Cover both reads and writes across more than one chunk. + +```python +def test_lazy_oindex_multiple_arrays_outer_product(arr, data) -> None: + rows = np.array([1, 11]) + cols = np.array([2, 12, 22]) + actual = arr.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) +``` + +- [ ] **Step 2: Run focused tests and confirm failure in intersection or placement** + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_chunk_resolution.py tests/test_lazy_indexing.py -q +``` + +Expected: the new multi-array orthogonal regression fails. + +- [ ] **Step 3: Implement dependency-aware intersection and selectors** + +Replace the `len(array_dims) >= 2` correlation test with dependency analysis. Jointly mask maps only when they share the same non-singleton dependency axes. Filter independent maps along their own axes and retain singleton axes. In `sub_transform_to_selections`, emit shared scatter indices only for correlated maps; build orthogonal selectors for independent maps. + +- [ ] **Step 4: Verify orthogonal and vectorized behavior** + +Run: + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_transform.py tests/test_transforms/test_chunk_resolution.py tests/test_lazy_indexing.py -q +``` + +Expected: PASS, including existing vectorized cases. + +- [ ] **Step 5: Commit Task 2** + +```bash +git add src/zarr/core/transforms tests/test_transforms tests/test_lazy_indexing.py +git commit -m "fix(indexing): distinguish orthogonal array maps" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 3: Preserve Lazy Views and Normalize Public Indices + +**Files:** +- Modify: `src/zarr/core/array.py` +- Modify: `src/zarr/core/transforms/transform.py` +- Test: `tests/test_lazy_indexing.py` +- Test: `tests/test_indexing.py` + +**Interfaces:** +- Consumes: `selection_to_transform(selection, transform, mode)`. +- Produces: array-boundary selection normalization relative to the current visible view and transform-aware fallback behavior. + +- [ ] **Step 1: Write failing composed-view and validation tests** + +Cover `arr.lazy[10:20].oindex[[0]]`, the corresponding write, `arr.lazy[-1]`, negative integer arrays, overly negative indices, positive out-of-range indices, and invalid Boolean masks. Assert writes by comparing the complete root array with NumPy. + +- [ ] **Step 2: Confirm the regressions fail for the reviewed reasons** + +```bash +.venv/bin/python -m pytest tests/test_lazy_indexing.py tests/test_indexing.py -q +``` + +Expected: composed advanced operations access the wrong offset or validation differs from eager Zarr. + +- [ ] **Step 3: Add a public-selection normalization boundary** + +Implement a helper in `array.py` that normalizes integer scalars and arrays against `self.shape`, validates bounds and mask shapes, and returns a selection suitable for `selection_to_transform`. Use it in `_LazyIndexAccessor`, `_LazyOIndex`, and `_LazyVIndex` before composing transforms. + +For advanced access on a non-identity view, compose with `self._transform` instead of creating a legacy indexer against storage. Allow the legacy path only where it is guaranteed to operate on the identity transform. Raise `NotImplementedError` for transformed structured-field operations that cannot preserve field-aware semantics. + +- [ ] **Step 4: Run focused lazy and eager indexing tests** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 3** + +```bash +git add src/zarr/core/array.py src/zarr/core/transforms/transform.py tests/test_lazy_indexing.py tests/test_indexing.py +git commit -m "fix(indexing): compose advanced lazy selections" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 4: Correct Zero-Rank and Caller-Provided Buffer Handling + +**Files:** +- Modify: `src/zarr/core/array.py` +- Test: `tests/test_indexing.py` +- Test: `tests/test_lazy_indexing.py` + +**Interfaces:** +- Consumes: normalized transforms and chunk output selectors. +- Produces: scalar-shaped zero-rank I/O and safe multidimensional vectorized `out` behavior. + +- [ ] **Step 1: Write failing result-kind and `out` tests** + +Assert that `arr[...]` and `arr.lazy[...].result()` for a zero-dimensional array have shape `()` rather than `(1,)`. Add multidimensional coordinate arrays and an `NDBuffer` whose shape equals the broadcast selection shape; assert returned values and buffer contents match NumPy. + +- [ ] **Step 2: Run tests and observe the exact shape/placement failures** + +```bash +.venv/bin/python -m pytest tests/test_indexing.py::test_get_basic_selection_0d tests/test_indexing.py::test_get_selection_out tests/test_lazy_indexing.py -q +``` + +Expected: strict zero-rank shape assertions and multidimensional `out` regression fail. + +- [ ] **Step 3: Implement correlation-aware flattening** + +Define flattening as `bool(correlated_array_maps)` rather than `all(...)`. Keep `buffer_shape == ()` for a zero-rank identity transform. When `out` is supplied for correlated multidimensional indexing, read into a flat temporary, reshape to `out_shape`, and copy into `out`; return the caller-visible result without converting backend buffers unnecessarily. + +- [ ] **Step 4: Verify focused buffer tests** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 4** + +```bash +git add src/zarr/core/array.py tests/test_indexing.py tests/test_lazy_indexing.py +git commit -m "fix(indexing): preserve transform result shapes" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 5: Add Stateful Property-Based Indexing Programs + +**Files:** +- Modify: `src/zarr/testing/strategies.py` +- Modify: `tests/test_properties.py` + +**Interfaces:** +- Consumes: array shapes and existing basic/orthogonal/vectorized strategies. +- Produces: `indexing_programs(shape)` yielding valid operation sequences and execution modes for NumPy/Zarr differential testing. + +- [ ] **Step 1: Add the program model and failing composed examples** + +Define small immutable operation records: + +```python +@dataclass(frozen=True) +class IndexOperation: + mode: Literal["basic", "orthogonal", "vectorized"] + selection: Any + + +@dataclass(frozen=True) +class IndexingProgram: + operations: tuple[IndexOperation, ...] + execution: Literal["materialize", "eager_on_lazy", "out", "set_scalar", "set_array"] +``` + +Generate one to three operations while carrying the visible shape forward. Keep sizes small enough for 300 Hypothesis examples. Add deterministic `@example` cases for slice-then-oindex, unequal orthogonal lengths, multidimensional vindex with `out`, and a zero-rank result. + +- [ ] **Step 2: Run the new property with deterministic regression examples enabled** + +```bash +.venv/bin/python -m pytest tests/test_properties.py -k indexing_program -q +``` + +Expected on the fixed tree: PASS. Temporarily reverse the Task 2 correlation predicate in the working tree, rerun the slice-then-oindex deterministic example, and confirm it fails before restoring the Task 2 code. This mutation check proves the generated-program oracle detects the original semantic error. + +- [ ] **Step 3: Implement the NumPy/Zarr program runner** + +For reads, compare exact shape, scalar-versus-array kind, dtype, values, and `out`. For writes, retain the NumPy root array, apply the equivalent indexed assignment, and compare the entire Zarr root array after the write. Use `np.ix_` to implement NumPy orthogonal semantics and normal NumPy indexing for vectorized semantics. + +- [ ] **Step 4: Run property tests under the CI profile** + +```bash +HYPOTHESIS_PROFILE=ci .venv/bin/python -m pytest --run-slow-hypothesis tests/test_properties.py -k "indexing_program or basic_indexing or oindex or vindex" -q +``` + +Expected: PASS with 300 deterministic examples per property. + +- [ ] **Step 5: Commit Task 5** + +```bash +git add src/zarr/testing/strategies.py tests/test_properties.py +git commit -m "test(indexing): generate composed indexing programs" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 6: Full Verification and Review + +**Files:** +- Verify all modified production and test files. + +**Interfaces:** +- Consumes: Tasks 1-5. +- Produces: evidence that the branch is ready for review. + +- [ ] **Step 1: Run the complete affected test modules** + +```bash +.venv/bin/python -m pytest tests/test_transforms tests/test_lazy_indexing.py tests/test_indexing.py tests/test_array.py -q +``` + +Expected: PASS with only existing documented xfails. + +- [ ] **Step 2: Run the focused Hypothesis suite** + +```bash +HYPOTHESIS_PROFILE=ci .venv/bin/python -m pytest --run-slow-hypothesis tests/test_properties.py -q +``` + +Expected: PASS. + +- [ ] **Step 3: Run formatting and static checks** + +Use the repository's configured commands for Ruff, mypy, and pre-commit on modified files. Expected: no new errors. + +- [ ] **Step 4: Inspect the final branch diff** + +```bash +git diff --check origin/main...HEAD +git status --short +``` + +Expected: no whitespace errors and no uncommitted implementation changes. + +- [ ] **Step 5: Request a branch code review** + +Run the branch-review workflow and address any correctness findings before declaring completion. From d0abd4bec98e3a83f8502fd4916cab8d999eab63 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 09:03:14 +0200 Subject: [PATCH 43/43] fix(indexing): handle empty and negative view selections Assisted-by: Codex:GPT-5 --- src/zarr/core/array.py | 15 +++++++++++++++ tests/test_lazy_indexing.py | 22 ++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index fcd0b5f873..f27ef4195c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3778,6 +3778,12 @@ def get_coordinate_selection( "(coordinate) array per dimension of the target array, " f"got {selection!r}" ) + sel_normalized = tuple( + np.where(np.asarray(s) < 0, np.asarray(s) + hi, s) + for s, hi in zip( + sel_normalized, self._async_array._transform.domain.exclusive_max, strict=True + ) + ) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -3909,6 +3915,12 @@ def set_coordinate_selection( "(coordinate) array per dimension of the target array, " f"got {selection!r}" ) + sel_normalized = tuple( + np.where(np.asarray(s) < 0, np.asarray(s) + hi, s) + for s, hi in zip( + sel_normalized, self._async_array._transform.domain.exclusive_max, strict=True + ) + ) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -6057,6 +6069,9 @@ async def _set_selection_via_transform( f"The selection has shape {sel_shape}, but the value has shape {value.shape}." ) from None + if product(sel_shape) == 0: + return + # When the transform has ArrayMap outputs, chunk resolution produces # flat scatter indices (out_indices). The value buffer must be 1D # during the write, matching the flat index layout. diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 2bb4592912..67403d1ea6 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -259,6 +259,15 @@ def test_multi_axis_write_sharded_unsupported(self, cfg: Config) -> None: class TestLazyVIndex: + def test_empty_writes_are_noops(self) -> None: + a, ref = _make(CONFIGS[1]) + view = a.lazy[2:10] + + view.lazy.vindex[(np.array([], dtype=np.intp),)] = np.array([], dtype="i4") + view.set_mask_selection(np.zeros(view.shape, dtype=bool), np.array([], dtype="i4")) + + np.testing.assert_array_equal(a[...], ref) + @pytest.mark.parametrize("cfg", ND_CASES) def test_read(self, cfg: Config) -> None: """Lazy vectorized indexing matches NumPy's point (coordinate) selection.""" @@ -310,6 +319,19 @@ class TestLazyViewMethods: through ``Array``'s own dispatch, which had correctness gaps. """ + def test_coordinate_methods_wrap_negative_indices(self) -> None: + a, ref = _make(CONFIGS[1]) + view = a.lazy[2:10] + + np.testing.assert_array_equal( + view.get_coordinate_selection((np.array([-1], dtype=np.intp),)), ref[[9]] + ) + view.set_coordinate_selection((np.array([-1], dtype=np.intp),), 999) + + expected = ref.copy() + expected[9] = 999 + np.testing.assert_array_equal(a[...], expected) + @pytest.mark.parametrize("cfg", ND_CASES) def test_view_oindex_respects_transform(self, cfg: Config) -> None: """``v.oindex[idx]`` on a sub-view reads from the view, not the base array."""