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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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/60] =?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/60] 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/60] 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/60] 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/60] =?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/60] 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/60] 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/60] 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/60] 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/60] 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/60] 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.""" From 0804d4eae2d4b05b13b1ff8f5029b0db9947771d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 15 Jul 2026 20:14:29 +0200 Subject: [PATCH 44/60] feat(lazy): add .lazy.shape and dask duck-typing compatibility - _LazyIndexAccessor.shape mirrors the wrapped array/view shape - LazyViewError now also subclasses AttributeError so hasattr/getattr probes (e.g. dask.array.from_array) treat guarded grid members as absent instead of crashing - add a dask.array.from_array interop test and an interop-tests dependency group wired into the optional hatch matrix Assisted-by: ClaudeCode:claude-fable-5 --- pyproject.toml | 6 ++++ src/zarr/core/array.py | 5 +++ src/zarr/errors.py | 7 ++-- tests/test_lazy_indexing.py | 24 +++++++++++++ uv.lock | 67 +++++++++++++++++++++++++++++++++++++ 5 files changed, 107 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6a7238ff8f..f2c7a619e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,11 @@ remote-tests = [ "moto[s3,server]==5.2.2", "requests==2.34.2", ] +interop-tests = [ + # Third-party array consumers whose duck-typing contract we test against; + # left floating so the `optional` matrix exercises their latest releases. + "dask[array]", +] release = [ "towncrier==25.8.0", ] @@ -190,6 +195,7 @@ matrix.deps.features = [ ] matrix.deps.dependency-groups = [ {value = "remote-tests", if = ["optional"]}, + {value = "interop-tests", if = ["optional"]}, ] [tool.hatch.envs.test.scripts] diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f27ef4195c..7d5f00853a 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4453,6 +4453,11 @@ 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 shape(self) -> tuple[int, ...]: + """Shape of the array (or view) this accessor indexes into.""" + return self._array.shape + @property def oindex(self) -> _LazyOIndex: return _LazyOIndex(self._array) diff --git a/src/zarr/errors.py b/src/zarr/errors.py index d3bbb5ff74..b4480ef035 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -152,7 +152,7 @@ class BoundsCheckError(IndexError): ... class ArrayIndexError(IndexError): ... -class LazyViewError(NotImplementedError): +class LazyViewError(NotImplementedError, AttributeError): """Raised when an operation that assumes an array fills its chunk grid is used on a non-identity lazy view (created via ``Array.lazy[...]``). @@ -161,7 +161,10 @@ class LazyViewError(NotImplementedError): 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. + ``NotImplementedError`` so existing consumers that catch it keep working, and + ``AttributeError`` so duck-typing probes (`hasattr(view, "chunks")`, + `getattr(x, "chunks", None)` — e.g. `dask.array.from_array`) treat guarded + members as absent instead of crashing. """ diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 67403d1ea6..8ba7e9416a 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -411,6 +411,30 @@ def test_view_write_through_tuple(self, cfg: Config) -> None: np.testing.assert_array_equal(a[...], expected) +class TestLazyAccessorSurface: + def test_shape(self) -> None: + """`.lazy.shape` mirrors the wrapped array's shape, including on views.""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded (20, 30) + assert a.lazy.shape == (20, 30) + v = a.lazy[2:8, 5:15] + assert v.lazy.shape == (6, 10) + + +class TestDaskInterop: + def test_from_array_lazy_view(self) -> None: + """A zero-origin lazy view works as a `dask.array.from_array` source. + + This is the dask-free-wrapper use case (napari): the view exposes + `shape`/`dtype`/`ndim` and eager `__getitem__`, which is all dask needs. + """ + da = pytest.importorskip("dask.array") + a, ref = _make(CONFIGS[3]) # 2d-unsharded (20, 30) + v = a.lazy[2:18, 5:25].translate_to((0, 0)) + d = da.from_array(v, chunks=(8, 10)) + assert d.shape == (16, 20) + np.testing.assert_array_equal(d.compute(scheduler="synchronous"), ref[2:18, 5:25]) + + class TestLazyErrors: def test_negative_step_raises(self) -> None: """A negative slice step is unsupported and raises on the lazy path.""" diff --git a/uv.lock b/uv.lock index 799ea6e45a..3802eaad67 100644 --- a/uv.lock +++ b/uv.lock @@ -655,6 +655,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl", hash = "sha256:a2bf429bb3033c89fa4936ffb35d5cb471e3719e1f3c8a7c3fff0b8314305613", size = 110502, upload-time = "2026-04-22T15:11:25.044Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -853,6 +862,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/80/5e05de89ba61df072aab6f8a6ee3ffeec57db68a0a456825b3b4ce608426/cupy_cuda12x-14.1.1-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:238080487174268d0f09770fe518de7c5b206527bef5c6792aef7ba0626a1c48", size = 132635338, upload-time = "2026-06-01T04:53:35.229Z" }, ] +[[package]] +name = "dask" +version = "2026.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "cloudpickle" }, + { name = "fsspec" }, + { name = "packaging" }, + { name = "partd" }, + { name = "pyyaml" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/39/cbd21c9133d02b4e60899ed466ad5e876553ea68ebffee6209e7dcafc8d4/dask-2026.7.1.tar.gz", hash = "sha256:5727484427665f051e86bf87d021a64d6411141cdc8a20bfe3c1ad2968cc06b7", size = 11548794, upload-time = "2026-07-14T01:06:22.46Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/5f/7c22733da92b3a6cc4dddcaa8731089d213c2790bbc997e51c429a4e8f8b/dask-2026.7.1-py3-none-any.whl", hash = "sha256:985ffd6c5e9d7979ede515e84ae8d39b647d6aa64f77600f15714ff65f578fe6", size = 1496882, upload-time = "2026-07-14T01:06:20.341Z" }, +] + +[package.optional-dependencies] +array = [ + { name = "numpy" }, +] + [[package]] name = "debugpy" version = "1.8.20" @@ -1515,6 +1547,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/62/b40b382fa0c66fee1478073eb8db352a4a6beda4a1adccf1df911d8c289c/librt-0.11.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dee008f20b542e3cd162ba338a7f9ec0f6d23d395f66fe8aeeec3c9d067ea253", size = 102572, upload-time = "2026-05-10T18:17:06.809Z" }, ] +[[package]] +name = "locket" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" }, +] + [[package]] name = "markdown" version = "3.10.2" @@ -2385,6 +2426,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" }, ] +[[package]] +name = "partd" +version = "1.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "locket" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, +] + [[package]] name = "pathable" version = "0.5.0" @@ -3587,6 +3641,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/43/8bd850ee71a191bf072e31302c73a66be413fecdd98fdcd111ecbcce13ca/tomlkit-0.15.0-py3-none-any.whl", hash = "sha256:4dbc8f0fc024412b57ced8757ac7461305126a648ff8c2c807fcb8e133a78738", size = 41328, upload-time = "2026-05-10T07:38:23.517Z" }, ] +[[package]] +name = "toolz" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, +] + [[package]] name = "tornado" version = "6.5.7" @@ -4033,6 +4096,9 @@ docs = [ { name = "s3fs" }, { name = "towncrier" }, ] +interop-tests = [ + { name = "dask", extra = ["array"] }, +] release = [ { name = "towncrier" }, ] @@ -4141,6 +4207,7 @@ docs = [ { name = "s3fs", specifier = ">=2023.10.0" }, { name = "towncrier", specifier = "==25.8.0" }, ] +interop-tests = [{ name = "dask", extras = ["array"] }] release = [{ name = "towncrier", specifier = "==25.8.0" }] remote-tests = [ { name = "botocore" }, From c370f7923ed64ba3aed3d720709c18bc2ce3835e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 09:34:15 +0200 Subject: [PATCH 45/60] fix(indexing): preserve array map dependency axes Normalize ArrayMap index arrays to the full input rank of their enclosing transform: an axis the array varies over is full-sized, every other axis is a singleton. Dependency axes (orthogonal vs vectorized) become derivable from the shape via _array_map_dependency_axes; input_dimension is retained as a compatibility shim for not-yet-migrated consumers. Full rank is preserved through basic indexing, composition, and JSON. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/transforms/composition.py | 4 ++ src/zarr/core/transforms/json.py | 10 ++++- src/zarr/core/transforms/output_map.py | 30 +++++++------- src/zarr/core/transforms/transform.py | 52 ++++++++++++++++++++++--- tests/test_transforms/test_json.py | 20 +++++++++- tests/test_transforms/test_transform.py | 52 ++++++++++++++++++++++++- 6 files changed, 145 insertions(+), 23 deletions(-) diff --git a/src/zarr/core/transforms/composition.py b/src/zarr/core/transforms/composition.py index 9d07bd3324..e33a94c470 100644 --- a/src/zarr/core/transforms/composition.py +++ b/src/zarr/core/transforms/composition.py @@ -62,10 +62,14 @@ def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> Output ) if isinstance(outer_map, ArrayMap): + # Affine post-composition leaves the index array (and hence its full + # input rank and dependency axes) untouched; carry the orthogonal + # binding through unchanged. return ArrayMap( index_array=outer_map.index_array, offset=offset_i + stride_i * outer_map.offset, stride=stride_i * outer_map.stride, + input_dimension=outer_map.input_dimension, ) raise TypeError(f"Unknown output map type: {type(outer_map)}") # pragma: no cover diff --git a/src/zarr/core/transforms/json.py b/src/zarr/core/transforms/json.py index 79fe6e58b8..11e824a61c 100644 --- a/src/zarr/core/transforms/json.py +++ b/src/zarr/core/transforms/json.py @@ -19,7 +19,7 @@ from __future__ import annotations -from typing import Required, TypedDict +from typing import Any, Required, TypedDict import numpy as np @@ -31,6 +31,12 @@ # TypedDict definitions (JSON shapes) # --------------------------------------------------------------------------- +# An ``index_array`` serializes via ``ndarray.tolist()``, so it is a nested list +# of ints whose nesting depth equals the array rank. Normalized arrays carry the +# full input rank (with singleton axes), so the nesting can be arbitrarily deep — +# ``list[int] | list[list[int]]`` would under-describe a 3-D+ orthogonal array. +NestedIntList = list[Any] + class IndexDomainJSON(TypedDict, total=False): """JSON representation of an IndexDomain.""" @@ -52,7 +58,7 @@ class OutputIndexMapJSON(TypedDict, total=False): offset: int stride: int input_dimension: int - index_array: list[int] | list[list[int]] + index_array: NestedIntList class IndexTransformJSON(TypedDict, total=False): diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py index c53402be4d..5aa0f32892 100644 --- a/src/zarr/core/transforms/output_map.py +++ b/src/zarr/core/transforms/output_map.py @@ -74,19 +74,23 @@ 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. + Freshly constructed maps are normalized to the **full input rank** of their + enclosing transform: `index_array` has the enclosing domain's rank, sized + fully on the axes it varies over and singleton (size 1) elsewhere. The + dependency axes are therefore derivable from the shape (see + `transform._array_map_dependency_axes`), which distinguishes the two flavours + of multi-array fancy indexing: + + - **orthogonal** (`oindex`): each array varies along a single, *distinct* + axis (all others singleton); the result is their outer product. + - **vectorized** (`vindex`): the arrays are correlated and share the same + non-singleton (broadcast) axes; the result is a pointwise scatter. + + `input_dimension` records the single axis an orthogonal array varies over + (`None` for vectorized), binding it the way `DimensionMap` is bound. It is + now redundant with the shape-derived classifier and is retained as a + compatibility shim for consumers not yet migrated to the shape-derived + dependency axes. """ index_array: npt.NDArray[np.intp] diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 473d15d0ef..d5ab50f9f7 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -637,6 +637,33 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra return IndexTransform(domain=new_domain, output=tuple(new_output)) +def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, ...]: + """Return the input axes on which a normalized index array varies. + + Normalized `ArrayMap` index arrays carry the full input rank of their + enclosing transform: an axis the array varies over has its full size, while + an axis the array is independent of is a singleton (size 1). The dependency + axes are therefore exactly the non-singleton axes. An orthogonal (`oindex`) + array depends on a single axis; a vectorized (`vindex`) array depends on all + of the (shared) broadcast axes. + """ + return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) + + +def _reshape_to_axis( + values: np.ndarray[Any, np.dtype[np.intp]], axis: int, ndim: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Reshape a 1-D selection to full rank ``ndim`` varying only along ``axis``. + + The result has ``values`` laid out along ``axis`` and singleton (size-1) axes + everywhere else, so its dependency axis is derivable from its shape. + """ + flat = np.asarray(values, dtype=np.intp).ravel() + shape = [1] * ndim + shape[axis] = flat.shape[0] + return flat.reshape(shape) + + class _OIndexHelper: """Helper that provides orthogonal (outer) indexing via ``transform.oindex[...]``.""" @@ -737,15 +764,21 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: elif isinstance(m, DimensionMap): d = m.input_dimension if d in dim_array: + new_axis = old_to_new_dim[d] + # Normalize to full input rank: the selection varies along its own + # new axis and is singleton on every other axis. The dependency + # axis is then derivable from the shape (a single non-singleton + # axis marks the selection orthogonal / outer-product rather than + # vectorized). `input_dimension` is kept populated as a + # compatibility shim for consumers not yet migrated to the + # shape-derived classifier. + full_arr = _reshape_to_axis(dim_array[d], new_axis, new_dim_idx) new_output.append( ArrayMap( - index_array=dim_array[d], + index_array=full_arr, 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], + input_dimension=new_axis, ) ) elif d in dim_slice_params: @@ -906,9 +939,16 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: elif isinstance(m, DimensionMap): d = m.input_dimension if d in array_dim_to_broadcast: + # Normalize to full input rank: the broadcast (correlated) axes + # come first, followed by a singleton axis per slice dimension. + # Every vectorized array shares the same broadcast axes, so the + # dependency axes derived from the shape coincide — the signature + # of a pointwise scatter rather than an outer product. + broadcast_arr = array_dim_to_broadcast[d] + full_arr = broadcast_arr.reshape(broadcast_shape + (1,) * len(slice_dims)) new_output.append( ArrayMap( - index_array=array_dim_to_broadcast[d], + index_array=full_arr, offset=m.offset, stride=m.stride, ) diff --git a/tests/test_transforms/test_json.py b/tests/test_transforms/test_json.py index db582a9561..325ccd4502 100644 --- a/tests/test_transforms/test_json.py +++ b/tests/test_transforms/test_json.py @@ -158,9 +158,27 @@ def test_with_array(self) -> None: 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) + # Orthogonal arrays are normalized to full input rank with a singleton + # axis on the dimension they do not vary over. + assert restored.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(restored.output[0].index_array, idx.reshape(3, 1)) assert isinstance(restored.output[1], DimensionMap) + def test_roundtrip_preserves_singleton_axes(self) -> None: + """Full-rank orthogonal arrays keep their singleton axes across JSON.""" + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + restored = index_transform_from_json(index_transform_to_json(t)) + orig0, orig1 = t.output[0], t.output[1] + rest0, rest1 = restored.output[0], restored.output[1] + assert isinstance(orig0, ArrayMap) + assert isinstance(orig1, ArrayMap) + assert isinstance(rest0, ArrayMap) + assert isinstance(rest1, ArrayMap) + assert rest0.index_array.shape == (2, 1) + assert rest1.index_array.shape == (1, 3) + np.testing.assert_array_equal(rest0.index_array, orig0.index_array) + np.testing.assert_array_equal(rest1.index_array, orig1.index_array) + def test_with_labels(self) -> None: domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) t = IndexTransform.identity(domain) diff --git a/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py index 8e3810465c..dec5192ce1 100644 --- a/tests/test_transforms/test_transform.py +++ b/tests/test_transforms/test_transform.py @@ -262,7 +262,9 @@ def test_oindex_int_array(self) -> None: 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) + # Full input rank: the array varies along its own axis (0), singleton on 1. + assert result.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(result.output[0].index_array, idx.reshape(3, 1)) assert result.output[0].offset == 0 assert result.output[0].stride == 1 assert isinstance(result.output[1], DimensionMap) @@ -301,6 +303,15 @@ def test_oindex_multiple_arrays(self) -> None: assert isinstance(result.output[1], DimensionMap) assert isinstance(result.output[2], ArrayMap) + def test_oindex_multiple_arrays_preserves_independent_axes(self) -> 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 isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2, 1) + assert result.output[1].index_array.shape == (1, 3) + class TestIndexTransformVindex: def test_vindex_single_array(self) -> None: @@ -347,6 +358,15 @@ def test_vindex_broadcast_different_shapes(self) -> None: assert result.input_rank == 2 assert result.domain.shape == (2, 3) + def test_vindex_multiple_arrays_preserves_shared_axes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.vindex[np.array([1, 3]), np.array([2, 4])] + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2,) + assert result.output[1].index_array.shape == (2,) + class TestSelectionToTransform: def test_basic_slice(self) -> None: @@ -526,3 +546,33 @@ 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 TestArrayMapDependencyAxes: + """`_array_map_dependency_axes` derives the input axes an array varies on + from its (full-rank) shape: non-singleton axes vary, singleton axes do not.""" + + def test_orthogonal_single_axis(self) -> None: + from zarr.core.transforms.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (1,) + + def test_vectorized_shares_axes(self) -> None: + from zarr.core.transforms.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3]), np.array([2, 4])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (0,) + + def test_scalar_array_has_no_dependency(self) -> None: + from zarr.core.transforms.transform import _array_map_dependency_axes + + assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () From 3b182a59f8b208da85fbd012d034d819f35eb4d8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 10:22:01 +0200 Subject: [PATCH 46/60] fix(indexing): distinguish orthogonal array maps Resolve independent (oindex) and correlated (vindex) ArrayMaps by their dependency axes during chunk intersection instead of assuming >=2 arrays are correlated or raveling full-rank arrays to 1-D. - _intersect routes on input_dimension is None (correlated marker): orthogonal maps are filtered along their own dependency axis at full input rank; correlated maps are jointly masked over the shared broadcast block while residual DimensionMap dims are narrowed against the chunk bounds and preserved (previously a rank-1 domain was paired with untouched DimensionMap outputs, raising ValueError at read time for any vindex with >=2 arrays plus a residual slice dim, including partial-rank boolean masks). - The correlated intersection now emits a flat row-major scatter index covering (surviving points, residual slice block), and sub_transform_to_selections consumes it as the single out-selection. - _reindex_array applies basic selections only along an ArrayMap's dependency axes; selections on singleton (broadcast) axes narrow the domain without touching the array's values, fixing basic slicing of the non-fancy axis of an oindex view. - Degenerate length-1 selections (all-singleton shapes) use input_dimension as the tie-breaker so they are not misclassified. - iter_chunk_transforms returns early for empty fancy selections instead of crashing on chunk_ids.min(). Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/transforms/chunk_resolution.py | 61 ++- src/zarr/core/transforms/output_map.py | 9 +- src/zarr/core/transforms/transform.py | 358 +++++++++++++----- tests/test_lazy_indexing.py | 107 ++++++ .../test_transforms/test_chunk_resolution.py | 60 +++ tests/test_transforms/test_transform.py | 50 +++ 6 files changed, 528 insertions(+), 117 deletions(-) diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py index e57d71d408..0b045d39cd 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -96,6 +96,9 @@ def iter_chunk_transforms( elif isinstance(m, ArrayMap): storage = m.offset + m.stride * m.index_array flat = storage.ravel().astype(np.intp) + if flat.size == 0: + # Empty fancy selection: no coordinates, so no chunks are touched. + return chunk_ids = dg.indices_to_chunks(flat) first = int(chunk_ids.min()) last = int(chunk_ids.max()) @@ -187,13 +190,49 @@ def sub_transform_to_selections( 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]]] = [] + # Correlated (vindex) sub-transforms carry ArrayMaps with ``input_dimension`` + # None. They scatter through a single flat index (``out_indices``) into the + # row-major-flattened output buffer; the chunk selection reads a + # (points, residual-slice) block via the raveled coordinate arrays and any + # residual DimensionMap slices. + correlated = any( + isinstance(m, ArrayMap) and m.input_dimension is None for m in sub_transform.output + ) + if correlated: + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + d = m.input_dimension + start = m.offset + m.stride * inclusive_min[d] + stop = m.offset + m.stride * exclusive_max[d] + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + else: # ArrayMap + idx = m.index_array.reshape(-1) + chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) + # Chunk resolution always supplies the flat scatter index for a + # correlated transform. Absent one (a bare sub-transform), fall back to an + # identity scatter over the whole flattened output buffer. + # ``out_indices`` is narrowed to a flat scatter array or None here (the + # per-dimension dict is an orthogonal outer product, handled above). + out_scatter: slice | np.ndarray[Any, np.dtype[np.intp]] + if out_indices is None: + n = 1 + for s in sub_transform.domain.shape: + n *= s + out_scatter = slice(0, n) + else: + out_scatter = out_indices + return tuple(chunk_sel), (out_scatter,), () + + chunk_sel = [] # annotated in the correlated branch above (same function scope) out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] - # Single-pass build for the basic / single-array / vectorized cases. + # Single-pass build for the basic / single-orthogonal-array cases. # ConstantMap dims are dropped (no out_sel entry). - n_array_maps = 0 - for m in sub_transform.output: if isinstance(m, ConstantMap): chunk_sel.append(m.offset) @@ -207,17 +246,13 @@ def sub_transform_to_selections( start, stop = stop + 1, start + 1 chunk_sel.append(slice(start, stop, m.stride)) out_sel.append(slice(dim_lo, dim_hi)) - else: # ArrayMap - n_array_maps += 1 + else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) + idx = m.index_array.reshape(-1) if m.offset == 0 and m.stride == 1: - chunk_sel.append(m.index_array) + chunk_sel.append(idx) else: - chunk_sel.append((m.offset + m.stride * m.index_array).astype(np.intp)) + chunk_sel.append((m.offset + m.stride * idx).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))) - - # 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] + out_sel.append(out_indices if out_indices is not None else slice(0, idx.size)) return tuple(chunk_sel), tuple(out_sel), () diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py index 5aa0f32892..80696a6c65 100644 --- a/src/zarr/core/transforms/output_map.py +++ b/src/zarr/core/transforms/output_map.py @@ -88,9 +88,12 @@ class ArrayMap: `input_dimension` records the single axis an orthogonal array varies over (`None` for vectorized), binding it the way `DimensionMap` is bound. It is - now redundant with the shape-derived classifier and is retained as a - compatibility shim for consumers not yet migrated to the shape-derived - dependency axes. + usually redundant with the shape-derived classifier, but stays authoritative + for the shapes the classifier cannot distinguish: a length-1 orthogonal + selection normalizes to an all-singleton array (no non-singleton axis), and + length-1 vectorized arrays are equally degenerate. `None` therefore marks a + map as correlated, and an integer pins the dependency axis of a degenerate + orthogonal map (see `transform._array_map_dependent_axis`). """ index_array: npt.NDArray[np.intp] diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index d5ab50f9f7..8d98767f0c 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -67,6 +67,14 @@ def __post_init__(self) -> None: f"is out of range for input rank {self.domain.ndim}" ) elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim: + # ArrayMap index arrays produced by indexing and chunk resolution + # are normalized to the full input rank (an axis the array varies + # over is full-sized, every other axis a singleton). A rank + # *exceeding* the domain is always a bug. A rank *below* it is + # tolerated: TensorStore-format JSON (external input) may supply a + # lower-rank index array that broadcasts against the input domain, + # and `_array_map_dependency_axes` treats any missing leading axes + # as singleton dependencies. raise ValueError( f"output[{i}].index_array has {m.index_array.ndim} dims " f"but input domain has {self.domain.ndim} dims" @@ -247,16 +255,23 @@ def _intersect( """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]). + ``[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. + Two flavours of fancy indexing require different treatment, distinguished by + the ArrayMaps' dependency axes (see `_array_map_dependency_axes`): - 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. + - **orthogonal** (`oindex`): each ArrayMap varies over a single, distinct + input axis, forming an outer product. Every output dimension is intersected + independently and the input domain narrowed per axis. + - **correlated** (`vindex`): the ArrayMaps share their (broadcast) dependency + axes and scatter through a single flat index. A point survives only if ALL + its storage coordinates fall within the output domain; residual slice + dimensions are intersected independently, as in the orthogonal case. - Returns None if the intersection is empty. + A `None` `input_dimension` marks a correlated map, so any such map routes the + whole transform through the correlated intersection. + + Returns ``None`` if the intersection is empty. """ if output_domain.ndim != transform.output_rank: raise ValueError( @@ -264,19 +279,58 @@ def _intersect( f"transform output rank ({transform.output_rank})" ) - # 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)] - vectorized_dims = [ - i for i in array_dims if cast("ArrayMap", transform.output[i]).input_dimension is None + correlated_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is None ] - if len(vectorized_dims) >= 2: - return _intersect_vectorized(transform, output_domain, vectorized_dims) + if correlated_dims: + return _intersect_correlated(transform, output_domain, correlated_dims) + return _intersect_orthogonal(transform, output_domain) + + +def _intersect_dimension_map( + m: DimensionMap, input_lo: int, input_hi: int, lo: int, hi: int +) -> tuple[int, int] | None: + """Narrow a DimensionMap's input range to storage coordinates in ``[lo, hi)``. + + ``input_lo``/``input_hi`` are the current (possibly already narrowed) input + range for the map's axis. Returns the new ``(input_lo, input_hi)`` or ``None`` + if no input produces an in-bounds storage coordinate. + """ + if input_lo >= input_hi: + return None + 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 + return new_input_lo, new_input_hi - # 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. + +def _intersect_orthogonal( + transform: IndexTransform, output_domain: IndexDomain +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): + """Intersect a transform with no correlated ArrayMaps. + + Every output dimension is intersected independently. Multiple ArrayMaps bound + to distinct input dimensions form an outer product, so each array'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] = [] @@ -294,50 +348,35 @@ def _intersect( elif isinstance(m, DimensionMap): d = m.input_dimension - input_lo = new_min[d] - input_hi = new_max[d] - if input_lo >= input_hi: + narrowed = _intersect_dimension_map(m, new_min[d], new_max[d], lo, hi) + if narrowed is None: 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_min[d], new_max[d] = narrowed new_output.append(m) elif isinstance(m, ArrayMap): + # Orthogonal: the array varies over a single axis (its dependency + # axis, or `input_dimension` for a degenerate length-1 array). Filter + # along that axis and keep the array at full input rank so the + # singleton axes it broadcasts over are preserved. + d = _array_map_dependent_axis(m) storage = m.offset + m.stride * m.index_array mask = (storage >= lo) & (storage < hi) - survivors = np.nonzero(mask.ravel())[0].astype(np.intp) + # The array is singleton on every axis but ``d``, so its mask reduces + # to a 1-D vector along ``d``. + survivors = np.nonzero(mask.reshape(-1))[0].astype(np.intp) if survivors.size == 0: return None - filtered = m.index_array.ravel()[survivors] + filtered = np.take(m.index_array, survivors, axis=d) new_output.append( ArrayMap( - index_array=filtered, + index_array=np.asarray(filtered, dtype=np.intp), 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) + new_max[d] = new_min[d] + int(survivors.size) out_positions[out_dim] = survivors new_domain = IndexDomain( @@ -361,63 +400,143 @@ def _intersect( return (result, out_indices) -def _intersect_vectorized( +def _intersect_correlated( 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. + correlated_dims: list[int], +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]]] | None: + """Intersect a correlated (vindex) transform with an output domain. + + The correlated ArrayMaps share their broadcast (dependency) axes; a broadcast + point survives only if ALL its storage coordinates fall within the output + domain. Residual DimensionMap dimensions are intersected independently (as in + the orthogonal case) and preserved, so a partial vindex — e.g. two coordinate + arrays over a 3-D array, leaving one slice dimension — resolves correctly. + + The surviving broadcast axes collapse to a single axis; the returned + ``out_indices`` is the flat scatter index into the (row-major flattened) + output buffer, of shape ``(surviving_points,) + (residual slice sizes)``. """ - # Compute storage coords per array dim and check bounds simultaneously - masks: list[np.ndarray[Any, np.dtype[np.bool_]]] = [] + corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] + + # Mixing correlated and orthogonal ArrayMaps in one transform is not produced + # by any single selection and is not supported here. + orthogonal_array_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is not None + ] + if orthogonal_array_dims: + raise NotImplementedError( + "intersecting a transform with both correlated and orthogonal " + "ArrayMaps is not supported" + ) + + # The broadcast (dependency) axes are shared by every correlated map; they are + # the leading axes of the domain, followed by the residual slice axes. + broadcast_axes = _array_map_dependency_axes(corr_maps[0].index_array) + broadcast_shape = tuple(corr_maps[0].index_array.shape[a] for a in broadcast_axes) - for out_dim in array_dims: - m = transform.output[out_dim] - assert isinstance(m, ArrayMap) - storage = m.offset + m.stride * m.index_array + # Joint bounds mask over the broadcast block. + combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None + for out_dim in correlated_dims: + cm = cast("ArrayMap", transform.output[out_dim]) + storage = cm.offset + cm.stride * cm.index_array lo = output_domain.inclusive_min[out_dim] hi = output_domain.exclusive_max[out_dim] - masks.append((storage >= lo) & (storage < hi)) - - # 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): + mask = (storage >= lo) & (storage < hi) + combined = mask if combined is None else (combined & mask) + assert combined is not None + # The correlated maps are singleton on every non-broadcast axis, so the mask + # collapses (C-order) to the broadcast block. + combined_bcast = combined.reshape(broadcast_shape) + surviving = np.nonzero(combined_bcast.reshape(-1))[0].astype(np.intp) + if surviving.size == 0: return None - surviving = np.nonzero(combined_mask.ravel())[0].astype(np.intp) - - # Build new output maps + # Intersect residual (slice / constant) dimensions independently. Slice dims + # are ordered by input dimension so their flat-buffer strides are row-major. + slice_dims: list[tuple[int, int, int, int, DimensionMap]] = [] # (in_dim, lo, hi, full, m) + for out_dim, m in enumerate(transform.output): + if out_dim in correlated_dims: + continue + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + if isinstance(m, ConstantMap): + if not (lo <= m.offset < hi): + return None + elif isinstance(m, DimensionMap): + d = m.input_dimension + input_lo = transform.domain.inclusive_min[d] + input_hi = transform.domain.exclusive_max[d] + narrowed = _intersect_dimension_map(m, input_lo, input_hi, lo, hi) + if narrowed is None: + return None + slice_dims.append((d, narrowed[0], narrowed[1], input_hi - input_lo, m)) + slice_dims.sort(key=lambda item: item[0]) + + n_points = int(surviving.size) + n_slice = len(slice_dims) + corr_values = { + out_dim: cast("ArrayMap", transform.output[out_dim]) + .index_array.reshape(broadcast_shape) + .reshape(-1)[surviving] + for out_dim in correlated_dims + } + + # New domain: the collapsed broadcast axis, then one axis per residual slice. + new_min = [0] + new_max = [n_points] + new_input_dim_of = {} + for new_axis, (d, nlo, nhi, _full, _m) in enumerate(slice_dims, start=1): + new_min.append(nlo) + new_max.append(nhi) + new_input_dim_of[d] = new_axis + new_domain = IndexDomain(inclusive_min=tuple(new_min), exclusive_max=tuple(new_max)) + + corr_shape = (n_points,) + (1,) * n_slice new_output: list[OutputIndexMap] = [] for out_dim, m in enumerate(transform.output): - if isinstance(m, ArrayMap): - filtered = m.index_array.ravel()[surviving] + if out_dim in correlated_dims: + corr = cast("ArrayMap", m) new_output.append( ArrayMap( - index_array=filtered, - offset=m.offset, - stride=m.stride, - input_dimension=m.input_dimension, + index_array=corr_values[out_dim].reshape(corr_shape).astype(np.intp), + offset=corr.offset, + stride=corr.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),)) + else: + assert isinstance(m, DimensionMap) + new_output.append( + DimensionMap( + input_dimension=new_input_dim_of[m.input_dimension], + offset=m.offset, + stride=m.stride, + ) + ) result = IndexTransform(domain=new_domain, output=tuple(new_output)) - return (result, surviving) + + # Flat scatter index into the row-major output buffer of shape + # (broadcast points, residual slice sizes...): flat = point * prod(slice) + + # (row-major offset within the slice block). + prod_slice = 1 + for _d, _lo, _hi, full, _m in slice_dims: + prod_slice *= full + out_indices: np.ndarray[Any, np.dtype[np.intp]] = (surviving * prod_slice).reshape( + (n_points,) + (1,) * n_slice + ) + running = 1 + for j in range(n_slice - 1, -1, -1): + _d, nlo, nhi, full, _m = slice_dims[j] + coords = np.arange(nlo, nhi, dtype=np.intp) * running + shape = [1] * (1 + n_slice) + shape[1 + j] = coords.size + out_indices = out_indices + coords.reshape(shape) + running *= full + return (result, out_indices.astype(np.intp)) def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | None, ...]: @@ -461,16 +580,26 @@ def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | def _reindex_array( - arr: np.ndarray[Any, np.dtype[np.intp]], + m: ArrayMap, 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. + over the domain shape). Each axis is either a **dependency axis** — the array + genuinely varies with that input dimension — or a **singleton** axis it + broadcasts over. Integer indexing, slicing, or newaxis is applied to the + array only along its dependency axes; a selection on a singleton axis does not + touch the array's values (it just narrows or drops that broadcast axis). """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + # Degenerate length-1 orthogonal selection: the recorded axis is a + # dependency even though its size (1) makes it look singleton. + dependent.add(m.input_dimension) + arr = m.index_array + # Build a numpy indexing tuple: one entry per old input dimension idx: list[Any] = [] old_dim = 0 @@ -483,19 +612,27 @@ def _reindex_array( 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) + if old_dim in dependent: + # Convert absolute domain coordinate to 0-based array index + idx.append(sel - domain.inclusive_min[old_dim]) + else: + # Broadcast axis: keep the single element and drop the axis. + idx.append(0) old_dim += 1 elif isinstance(sel, slice): if old_dim < arr.ndim: - 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)) + if old_dim in dependent: + 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)) + else: + # Broadcast axis: preserve the singleton (it still broadcasts + # over the narrowed domain), regardless of the slice bounds. + idx.append(slice(None)) old_dim += 1 result_axis += 1 @@ -621,7 +758,7 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra else: raise RuntimeError(f"unexpected: dimension {d} not handled") elif isinstance(m, ArrayMap): - new_arr = _reindex_array(m.index_array, normalized, transform.domain) + new_arr = _reindex_array(m, normalized, transform.domain) 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) @@ -650,6 +787,25 @@ def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) +def _array_map_dependent_axis(m: ArrayMap) -> int: + """Return the single input axis an orthogonal `ArrayMap` varies over. + + Normally this is the array's one non-singleton axis. A degenerate length-1 + orthogonal selection normalizes to an all-singleton shape (its dependency + axes are empty and indistinguishable by shape from a scalar), so + `input_dimension` breaks the tie — it records the axis the map binds. + """ + dep = _array_map_dependency_axes(m.index_array) + if len(dep) == 1: + return dep[0] + if m.input_dimension is not None: + return m.input_dimension + raise ValueError( + f"orthogonal ArrayMap must vary over exactly one axis; got dependency " + f"axes {dep} with input_dimension={m.input_dimension}" + ) + + def _reshape_to_axis( values: np.ndarray[Any, np.dtype[np.intp]], axis: int, ndim: int ) -> np.ndarray[Any, np.dtype[np.intp]]: diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 8ba7e9416a..e5eea43bda 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -807,3 +807,110 @@ def test_invalid_unit_rejected(self) -> None: a = self._array((6,), (2,)) with pytest.raises(ValueError, match="unit"): a.chunk_projections(unit="bogus") # type: ignore[arg-type] + + +class TestLazyArrayMapResolution: + """Chunk resolution of orthogonal (outer-product) and correlated (vectorized) + ArrayMaps: intersections and output selectors must preserve the distinct + semantics of the two flavours, across chunk boundaries, and for the + length-1 shapes that are degenerate under the shape-derived classifier. + """ + + def _make(self, shape: tuple[int, ...], chunks: tuple[int, ...]) -> tuple[Any, Any]: + a = zarr.create_array(store={}, shape=shape, chunks=chunks, dtype="i4") + data = np.arange(int(np.prod(shape)), dtype="i4").reshape(shape) + a[...] = data + return a, data + + def test_oindex_multiple_arrays_outer_product(self) -> None: + """Orthogonal selection with unequal index lengths equals ``np.ix_`` (read).""" + a, data = self._make((20, 30), (5, 10)) + rows = np.array([1, 11], dtype=np.intp) + cols = np.array([2, 12, 22], dtype=np.intp) + actual = a.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) + + def test_oindex_multiple_arrays_outer_product_write(self) -> None: + """Orthogonal outer-product write spanning several chunks (unsharded).""" + a, data = self._make((20, 30), (5, 10)) + rows = np.array([1, 11], dtype=np.intp) + cols = np.array([2, 12, 22], dtype=np.intp) + expected = data.copy() + val = (np.arange(6, dtype="i4").reshape(2, 3) + 1) * 100 + expected[np.ix_(rows, cols)] = val + a.lazy.oindex[rows, cols] = val + np.testing.assert_array_equal(a[...], expected) + + def test_oindex_length1_and_length3_axes(self) -> None: + """A length-1 orthogonal axis (all-singleton shape) must not be + misclassified as scalar/correlated: the outer product still holds.""" + a, data = self._make((6, 6), (2, 2)) + rows = np.array([2], dtype=np.intp) + cols = np.array([1, 3, 5], dtype=np.intp) + actual = a.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) + assert actual.shape == (1, 3) + + def test_vindex_length1_pair(self) -> None: + """Two length-1 correlated arrays (degenerate (1,) shape) stay a single + pointwise scatter, not an outer product.""" + a, data = self._make((6, 6), (2, 2)) + actual = a.lazy.vindex[np.array([2]), np.array([3])].result() + np.testing.assert_array_equal(actual, data[np.array([2]), np.array([3])]) + + def test_vindex_two_arrays_with_residual_slice(self) -> None: + """Vectorized selection with two correlated arrays plus a residual slice + dim (partial indexing of a 3-D array) matches NumPy fancy indexing.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + rows = np.array([1, 3], dtype=np.intp) + cols = np.array([2, 0], dtype=np.intp) + actual = a.lazy.vindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[rows, cols]) + assert actual.shape == (2, 5) + + def test_vindex_two_arrays_with_residual_slice_write(self) -> None: + """Write-through of the correlated + residual-slice case.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + rows = np.array([1, 3], dtype=np.intp) + cols = np.array([2, 0], dtype=np.intp) + expected = data.copy() + val = (np.arange(10, dtype="i4").reshape(2, 5) + 1) * 100 + expected[rows, cols] = val + a.lazy.vindex[rows, cols] = val + np.testing.assert_array_equal(a[...], expected) + + def test_vindex_partial_rank_bool_mask(self) -> None: + """A partial-rank boolean mask (rank < array ndim) leaves a residual slice + dim; the vectorized read matches NumPy.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + mask = np.zeros((4, 3), dtype=bool) + mask[1, 2] = mask[3, 0] = True + actual = a.lazy.vindex[mask].result() + np.testing.assert_array_equal(actual, data[mask]) + assert actual.shape == (2, 5) + + def test_basic_slice_after_oindex_independent_axis(self) -> None: + """Basic-slicing an axis the ArrayMap is independent of broadcasts the + map (does not slice its values); the slice narrows only that axis.""" + a, data = self._make((6, 6), (2, 2)) + cols = a.lazy.oindex[:, np.array([1, 3, 5], dtype=np.intp)] + ref = data[:, [1, 3, 5]] + np.testing.assert_array_equal(cols.lazy[2:4, :].result(), ref[2:4, :]) + + def test_basic_slice_after_oindex_fancy_axis(self) -> None: + """Basic-slicing the fancy axis of an ArrayMap slices the map's values.""" + a, data = self._make((6, 6), (2, 2)) + cols = a.lazy.oindex[:, np.array([1, 3, 5], dtype=np.intp)] + ref = data[:, [1, 3, 5]] + np.testing.assert_array_equal(cols.lazy[:, 0:2].result(), ref[:, 0:2]) + + def test_empty_oindex_chunk_projections(self) -> None: + """An empty fancy selection yields no chunk projections (no crash).""" + a = zarr.create_array(store={}, shape=(10,), chunks=(3,), dtype="i4") + assert list(a.lazy.oindex[np.array([], dtype=np.intp)].chunk_projections()) == [] + + def test_empty_oindex_read(self) -> None: + """An empty fancy selection reads back an empty array.""" + a, _ = self._make((10,), (3,)) + actual = a.lazy.oindex[np.array([], dtype=np.intp)].result() + assert actual.shape == (0,) diff --git a/tests/test_transforms/test_chunk_resolution.py b/tests/test_transforms/test_chunk_resolution.py index 893866740d..4e24ad6065 100644 --- a/tests/test_transforms/test_chunk_resolution.py +++ b/tests/test_transforms/test_chunk_resolution.py @@ -179,3 +179,63 @@ def test_mixed_maps_2d(self) -> None: 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 == () + + +class TestChunkResolutionArrayMapFlavours: + """Chunk resolution must yield outer-product (np.ix_) selectors for + orthogonal ArrayMaps and shared flat-scatter selectors for correlated ones, + and must return early for empty fancy selections.""" + + def test_empty_array_selection_yields_nothing(self) -> None: + """An empty ArrayMap selection produces no chunk transforms (no crash).""" + t = IndexTransform( + domain=IndexDomain.from_shape((0,)), + output=(ArrayMap(index_array=np.array([], dtype=np.intp)),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=3, extent=10),)) + assert list(iter_chunk_transforms(t, grid)) == [] + + def test_orthogonal_outer_product_selectors(self) -> None: + """Two independent arrays produce np.ix_-style (mesh) chunk/out selectors.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + grid = ChunkGrid( + dimensions=(FixedDimension(size=10, extent=10), FixedDimension(size=10, extent=10)) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) + # np.ix_ produces one 2-D open-mesh selector per axis, for both sides. + assert len(chunk_sel) == 2 + assert len(out_sel) == 2 + assert isinstance(chunk_sel[0], np.ndarray) + assert isinstance(chunk_sel[1], np.ndarray) + assert chunk_sel[0].shape == (2, 1) + assert chunk_sel[1].shape == (1, 3) + assert drop_axes == () + + def test_correlated_scatter_with_residual_slice(self) -> None: + """Correlated arrays + a residual slice dim scatter through a single flat + index whose shape matches the (points, slice) block read from the chunk.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4), + FixedDimension(size=3, extent=3), + FixedDimension(size=5, extent=5), + ) + ) + # One chunk holds everything: both points survive, slice dim spans [0,5). + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, _drop = sub_transform_to_selections(sub_t, out_indices) + # Chunk side: flat coordinate arrays for the two correlated dims plus a + # slice for the residual dim. + assert len(chunk_sel) == 3 + np.testing.assert_array_equal(np.asarray(chunk_sel[0]), [1, 3]) + np.testing.assert_array_equal(np.asarray(chunk_sel[1]), [2, 0]) + assert chunk_sel[2] == slice(0, 5, 1) + # Output side: a single flat scatter index of shape (points, slice) = (2, 5). + assert len(out_sel) == 1 + assert np.asarray(out_sel[0]).shape == (2, 5) diff --git a/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py index dec5192ce1..786eb81575 100644 --- a/tests/test_transforms/test_transform.py +++ b/tests/test_transforms/test_transform.py @@ -576,3 +576,53 @@ def test_scalar_array_has_no_dependency(self) -> None: from zarr.core.transforms.transform import _array_map_dependency_axes assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () + + +class TestIntersectArrayMapClassification: + """`_intersect` must distinguish orthogonal (outer-product) ArrayMaps from + correlated (vectorized) ones by their dependency axes, keep surviving arrays + at full input rank, and preserve residual (slice) dimensions.""" + + def test_orthogonal_outer_product_keeps_full_rank(self) -> None: + """Two arrays on distinct axes narrow independently and stay full rank; + out_indices is a per-output-dim dict of surviving positions.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3, 8]), np.array([2, 6, 9])] + # Chunk covering storage [0,5) x [0,5): rows 1,3 survive (out pos 0,1), + # cols 2 survives (out pos 0). + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(5, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + assert isinstance(restricted.output[0], ArrayMap) + assert isinstance(restricted.output[1], ArrayMap) + # Full input rank preserved (not raveled to 1-D). + assert restricted.output[0].index_array.ndim == 2 + assert restricted.output[1].index_array.ndim == 2 + assert restricted.domain.ndim == 2 + assert isinstance(out_indices, dict) + np.testing.assert_array_equal(out_indices[0], np.array([0, 1])) + np.testing.assert_array_equal(out_indices[1], np.array([0])) + + def test_correlated_with_residual_slice_preserves_slice_dim(self) -> None: + """A vindex transform with two correlated arrays plus a residual slice + dim intersects without a rank error and keeps the DimensionMap.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + # Chunk covering storage [0,2) x [2,3) x [0,5): only point (1,2,*) is in + # bounds on both array dims -> one surviving broadcast point. + chunk = IndexDomain(inclusive_min=(0, 2, 0), exclusive_max=(2, 3, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + # A DimensionMap for the residual slice dim survives (no post-init error). + assert any(isinstance(m, DimensionMap) for m in restricted.output) + assert out_indices is not None + + def test_length1_orthogonal_not_treated_as_correlated(self) -> None: + """A length-1 orthogonal array (all-singleton shape) is still an outer + product with the length-3 axis: out_indices is a dict, not a flat array.""" + t = IndexTransform.from_shape((6, 6)).oindex[np.array([2]), np.array([1, 3, 5])] + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(6, 6)) + result = t.intersect(chunk) + assert result is not None + _restricted, out_indices = result + assert isinstance(out_indices, dict) From a0f48f9469a0783fcbffbd7dbbd77119401f0362 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 11:02:49 +0200 Subject: [PATCH 47/60] fix(indexing): compose advanced lazy selections Normalize public selections at the Array boundary: wrap NumPy-style negatives against the current view's domain (positive indices stay literal domain coordinates), reject boolean scalar indices, and route basic-slice-then-fancy composition through the view's transform. Generalizes the coordinate-selection wrapping from d0abd4bec into one helper used by the lazy accessors and the non-identity view methods; the transform layer stays literal. Also reject fields= on a non-identity view (NotImplementedError) in all eight get/set_*_selection methods before any storage access, fixing silent corruption where the legacy indexer ignored the transform. Assisted-by: ClaudeCode:claude-fable-5 --- docs/user-guide/lazy_indexing.md | 67 ++++--- src/zarr/core/array.py | 128 ++++++++++-- src/zarr/core/transforms/transform.py | 4 +- tests/test_lazy_indexing.py | 275 ++++++++++++++++++++------ tests/test_properties.py | 9 +- 5 files changed, 370 insertions(+), 113 deletions(-) diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md index 19951e0922..e44c15a919 100644 --- a/docs/user-guide/lazy_indexing.md +++ b/docs/user-guide/lazy_indexing.md @@ -8,8 +8,9 @@ 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. +**domain** — a box of coordinates — and a positive index is a **literal +coordinate** in that domain (a negative index counts from the domain's end, as +in NumPy). ```python exec="true" session="lazy" import numpy as np @@ -72,31 +73,36 @@ 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 +### Negative indices count from the end -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. +*Positive* indices are literal coordinates (so positions still are not indices — +see above), but a **negative** index counts from the end of the current view's +domain, exactly like NumPy: `k` maps to `exclusive_max + k`. This holds in every +syntactic form — integers, slice bounds, and index arrays — at the public +boundary (`a.lazy[...]`, `v.lazy[...]`, `v[...]`, `.oindex`, `.vindex`, and the +`get_/set_*_selection` methods). For a fresh, zero-origin array both rules +coincide with NumPy in full: ```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__, "-", str(e).split(";")[0]) +print(a.lazy[-1].result()) # last element, like NumPy +print(a.lazy[-3:].result()) # the last three +print(a.lazy.oindex[[-1, -2]].result()) # index arrays wrap too ``` -To select from the end, say what you mean literally: +On a view the wrap is relative to the *view's* end: on the domain-`[2, 10)` +view `v`, `-1` names coordinate `9`, and `-8` names the first cell (coordinate +`2`). Only an out-of-range index raises — a positive past the end, or a +negative below `-size` (one that would wrap past the domain origin): ```python exec="true" session="lazy" source="above" result="ansi" n = a.shape[0] -print(a.lazy[n - 3 :].result()) # the last three elements +for make in (lambda: a.lazy[n], lambda: a.lazy[-n - 1], lambda: a.lazy.oindex[[-n - 1]]): + try: + make() + except IndexError as e: + print(type(e).__name__, "-", str(e).split(";")[0]) ``` ### No clamping: intervals must fit the domain @@ -129,11 +135,11 @@ print(s.lazy[1].result()) # coordinate 1 -> base cell 4 ### One coordinate system per view 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`. +— uses the same rule: positive indices are literal domain coordinates, and +negative indices wrap from the domain's end. (The base array's ordinary `a[...]` +API is unchanged: it keeps full NumPy semantics.) NumPy-style zero-based access +to a view's *positive* range is spelled explicitly: materialize with `result()` +/ `np.asarray`, or renumber with `translate_to`. ```python exec="true" session="lazy" source="above" result="ansi" print(v[3], v.lazy[3].result()) # same coordinate, same cell @@ -263,15 +269,16 @@ array. ## Coming from NumPy -- **A view's indices are coordinates, not positions.** `a.lazy[2:10]` is - indexed with 2..9, not 0..7. Renumber explicitly with +- **A view's positive 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. +- **Negative indices count from the end**, exactly like NumPy — in every form + (integer, slice bound, index array). `k` maps to `exclusive_max + k` in the + current view's domain; only a negative below `-size` (or a positive past the + end) is out of bounds. +- **No clamping**: out-of-range slice bounds raise (including a negative that + would wrap past the domain origin); 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. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 7d5f00853a..fc011c3c7a 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2944,6 +2944,68 @@ def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None: else: self.set_basic_selection(cast("BasicSelection", pure_selection), value, fields=fields) + def _reject_bool_scalar(self, sel: Any) -> None: + # `isinstance(True, int)` is True, so a bare Python/NumPy bool scalar would + # otherwise pass the integer validators and silently read element 0/1 where + # eager NumPy indexing raises. Boolean *arrays* (masks) remain valid. + if isinstance(sel, (bool, np.bool_)): + raise IndexError( + "boolean scalar indices are not supported; use a boolean array " + "(mask) for boolean indexing" + ) + + def _normalize_public_selection(self, selection: Any) -> Any: + """Normalize a user selection at the public Array boundary. + + Rejects boolean *scalar* indices, and wraps negative integers, negative + index-array elements, and negative slice bounds to literal coordinates in + the current view's domain: an index `k` with `-size <= k < 0` on axis `d` + becomes `exclusive_max[d] + k` (from the end of the view's domain — exactly + NumPy for a zero-origin domain). Positive indices are left as literal domain + coordinates. The transform layer stays literal; it performs the final bounds + check, so an out-of-range value (a positive past the end, or `k < -size`) is + rejected there before any chunk access. + """ + exclusive_max = self._async_array._transform.domain.exclusive_max + ndim = len(exclusive_max) + + is_tuple = isinstance(selection, tuple) + items = list(selection) if is_tuple else [selection] + + for item in items: + self._reject_bool_scalar(item) + + n_real = sum(1 for s in items if s is not Ellipsis and s is not None) + # Too many indices: leave the selection untouched and let the transform + # layer raise its canonical error with the correct message. + if n_real > ndim: + return selection + + normalized: list[Any] = [] + axis = 0 + for item in items: + if item is Ellipsis: + axis += ndim - n_real # ellipsis fills the unaddressed axes + normalized.append(item) + elif item is None: + normalized.append(item) # newaxis consumes no domain axis + else: + normalized.append(_wrap_negative_index(item, exclusive_max[axis])) + axis += 1 + return tuple(normalized) if is_tuple else normalized[0] + + def _reject_fields_on_view(self, fields: Fields | None) -> None: + # Routing `fields=` on a non-identity view through the legacy indexer + # ignores the transform and reads/writes the wrong storage region. The + # transform layer cannot preserve field-aware semantics, so reject before + # any storage access rather than corrupt data. + if fields is not None and not self._async_array._is_identity: + raise NotImplementedError( + "field selection (`fields=`) is not supported on a lazy view " + "(non-identity transform); materialize the view first " + "(e.g. zarr.array(view[...]))" + ) + def get_basic_selection( self, selection: BasicSelection = Ellipsis, @@ -3066,6 +3128,7 @@ def get_basic_selection( if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) 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. @@ -3077,6 +3140,7 @@ def get_basic_selection( prototype=prototype, ) ) + selection = self._normalize_public_selection(selection) transform = selection_to_transform(selection, self._async_array._transform, "basic") return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) @@ -3180,6 +3244,7 @@ def set_basic_selection( """ if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) 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. @@ -3188,6 +3253,7 @@ def set_basic_selection( self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) return + selection = self._normalize_public_selection(selection) transform = selection_to_transform(selection, self._async_array._transform, "basic") sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) @@ -3316,6 +3382,7 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) 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. @@ -3328,6 +3395,7 @@ 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 = self._normalize_public_selection(selection) mode: Literal["basic", "orthogonal"] = ( "basic" if is_basic_selection(selection) else "orthogonal" ) @@ -3445,6 +3513,7 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) 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. @@ -3456,6 +3525,7 @@ 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 = self._normalize_public_selection(selection) mode: Literal["basic", "orthogonal"] = ( "basic" if is_basic_selection(selection) else "orthogonal" ) @@ -3545,6 +3615,7 @@ def get_mask_selection( if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) if fields is not None or self._async_array._is_identity: indexer = MaskIndexer(mask, self.shape, self._chunk_grid) return sync( @@ -3650,6 +3721,7 @@ def set_mask_selection( """ if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) if fields is not None or self._async_array._is_identity: indexer = MaskIndexer(mask, self.shape, self._chunk_grid) sync( @@ -3755,6 +3827,7 @@ def get_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() + self._reject_fields_on_view(fields) if fields is not None or self._async_array._is_identity: indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) out_array = sync( @@ -3778,12 +3851,9 @@ 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 - ) - ) + # Wrap negative coordinates against the view's domain (see + # _normalize_public_selection); positive coordinates stay literal. + sel_normalized = self._normalize_public_selection(sel_normalized) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -3881,6 +3951,7 @@ def set_coordinate_selection( # no fields); keep the exact-emptiness check rather than a falsy one. if fields == [] or fields == (): fields = None + self._reject_fields_on_view(fields) 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): @@ -3915,12 +3986,9 @@ 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 - ) - ) + # Wrap negative coordinates against the view's domain (see + # _normalize_public_selection); positive coordinates stay literal. + sel_normalized = self._normalize_public_selection(sel_normalized) transform = selection_to_transform( sel_normalized, self._async_array._transform, "vectorized" ) @@ -4403,6 +4471,36 @@ async def _shards_initialized( type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"] +def _wrap_negative_index(item: Any, exclusive_max: int) -> Any: + """Wrap negative values in a single per-axis selection element to literal + coordinates against `exclusive_max` (from-the-end within the axis's domain). + + A negative integer `k` becomes `exclusive_max + k`; negative slice bounds and + negative integer-array elements are wrapped the same way. Positive values, + `None` slice bounds, boolean masks, and non-integer items pass through + unchanged. This performs no bounds checking — the transform layer rejects an + out-of-range value (a positive past the end, or one that wraps below the + domain origin) before any chunk access. + """ + if isinstance(item, slice): + start = item.start + stop = item.stop + if start is not None and start < 0: + start += exclusive_max + if stop is not None and stop < 0: + stop += exclusive_max + return slice(start, stop, item.step) + if isinstance(item, (int, np.integer)): + k = int(item) + return k + exclusive_max if k < 0 else k + if isinstance(item, (list, np.ndarray)): + arr = np.asarray(item) + if arr.dtype == np.bool_ or not np.issubdtype(arr.dtype, np.integer): + return item # boolean mask or non-integer array: no wrapping + return np.where(arr < 0, arr + exclusive_max, arr) + return item + + class _LazyOIndex: """Lazy orthogonal indexing via ``array.lazy.oindex[...]``.""" @@ -4412,10 +4510,12 @@ def __init__(self, array: Array[Any]) -> None: self._array = array def __getitem__(self, selection: Any) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) 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: + selection = self._array._normalize_public_selection(selection) new_t = selection_to_transform(selection, self._array._async_array._transform, "orthogonal") self._array._with_transform(new_t)[...] = value @@ -4429,10 +4529,12 @@ def __init__(self, array: Array[Any]) -> None: self._array = array def __getitem__(self, selection: Any) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) 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: + selection = self._array._normalize_public_selection(selection) new_t = selection_to_transform(selection, self._array._async_array._transform, "vectorized") self._array._with_transform(new_t)[...] = value @@ -4446,10 +4548,12 @@ def __init__(self, array: Array[Any]) -> None: self._array = array def __getitem__(self, selection: Selection) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) 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: + selection = self._array._normalize_public_selection(selection) new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") self._array._with_transform(new_t)[...] = value diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 8d98767f0c..98014196af 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -1135,8 +1135,8 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: _LITERAL_HINT = ( - "; negative indices are literal coordinates in lazy indexing, not from-the-end " - "(use `shape[dim] - k`, or materialize with `result()`, for NumPy semantics)" + "; within this transform layer, indices are literal domain coordinates (the " + "public Array boundary wraps NumPy-style negatives before they reach here)" ) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index e5eea43bda..25b5f547da 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -295,9 +295,8 @@ class TestLazyComposition: def test_chained_views_compose(self, cfg: Config) -> None: """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). + Positive bounds are literal domain coordinates, so the inner slice + re-selects in the outer view's preserved coordinate system. """ a, ref = _make(cfg) n = cfg.shape[0] @@ -346,13 +345,13 @@ def test_view_oindex_respects_transform(self, cfg: Config) -> None: @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) def test_view_trailing_index_after_ellipsis(self, cfg: Config) -> None: - """``v[..., k]`` uses literal coordinates; ``-1`` is out of the domain.""" + """``v[..., k]``: positive ``k`` is a literal domain coordinate; a negative + ``k`` wraps from the end of the view's domain (NumPy parity at the boundary).""" a, ref = _make(cfg) v = a.lazy[1 : cfg.shape[0]] last = cfg.shape[-1] - 1 np.testing.assert_array_equal(v[..., last], ref[1:, ..., last]) - with pytest.raises(BoundsCheckError): - v[..., -1] + np.testing.assert_array_equal(v[..., -1], ref[1:, ..., -1]) @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) def test_view_basic_tuple_and_int(self, cfg: Config) -> None: @@ -442,50 +441,50 @@ 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. + def test_accessor_negative_index_wraps(self) -> None: + """The public boundary wraps negative indices NumPy-style against the + current view's ``exclusive_max`` (positive indices stay literal domain + coordinates). Only out-of-range negatives (``k < -size``) or positive + indices past the end raise — before any chunk access. """ - 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)): + a, ref = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + np.testing.assert_array_equal(a.lazy[-1].result(), ref[-1]) + np.testing.assert_array_equal( + a.lazy.oindex[(np.array([-1], dtype=np.intp),)].result(), ref[[-1]] + ) + np.testing.assert_array_equal( + a.lazy.vindex[(np.array([-1], dtype=np.intp),)].result(), ref[[-1]] + ) + # too-negative (< -size) and positive-past-the-end still raise + for sel in (np.array([-25], dtype=np.intp), np.array([24], dtype=np.intp)): with pytest.raises(IndexError, match="out of bounds"): - _ = a.lazy.oindex[(sel,)][...] + _ = a.lazy.oindex[(sel,)].result() 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. + _ = a.lazy.vindex[(sel,)].result() - 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). + def test_negative_slice_bounds_wrap(self) -> None: + """Negative slice bounds wrap from the end of the current view's domain + (NumPy parity at the boundary). Bounds that wrap out of range still raise, + and bounds that resolve reversed are invalid intervals. """ - a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) - for sel in (slice(-3, None), slice(-24, None)): - with pytest.raises(BoundsCheckError, match="not contained"): - a.lazy[sel] - 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="not contained"): - v.lazy[-2:] - # ... and for writes + a, ref = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + for sel in (slice(-3, None), slice(-24, None), slice(None, -1), slice(1, -1)): + np.testing.assert_array_equal(a.lazy[sel].result(), ref[sel]) + np.testing.assert_array_equal(a.lazy.oindex[(sel,)].result(), ref[sel]) + # on a view: wraps against the view's exclusive_max, not the base's + v = a.lazy[2:10] # domain [2, 10) + np.testing.assert_array_equal(v.lazy[-2:].result(), ref[8:10]) + # writes wrap too + a.lazy[-3:] = 0 + expected = ref.copy() + expected[-3:] = 0 + np.testing.assert_array_equal(a[...], expected) + # out-of-range wrap (start < -size) still raises with pytest.raises(BoundsCheckError, match="not contained"): - a.lazy[-3:] = 0 + a.lazy[-25:] + # bounds that RESOLVE reversed (stop < start after wrapping) are invalid + with pytest.raises(IndexError, match="interval"): + a.lazy[-1:1] def test_slice_bounds_strict_containment(self) -> None: """Non-empty slice intervals must be contained in the domain — no @@ -500,23 +499,21 @@ def test_slice_bounds_strict_containment(self) -> None: 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`. + def test_view_positive_literal_negative_wraps(self) -> None: + """A view's positive indices are literal domain coordinates; negative + indices wrap from the view's ``exclusive_max`` (NumPy parity at the + boundary). Positive coordinates below the domain origin remain out of + bounds — the view keeps its preserved coordinate system for positives. """ a, ref = _make(CONFIGS[1]) - v = a.lazy[2:10] - 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: + v = a.lazy[2:10] # domain [2, 10) + np.testing.assert_array_equal(v[2:5], ref[2:5]) # positive literal + np.testing.assert_array_equal(v[3], ref[3]) # positive literal + np.testing.assert_array_equal(v[-1], ref[9]) # wrap: exclusive_max 10 - 1 + np.testing.assert_array_equal(v[-3:], ref[7:10]) # wrap + # positive coordinates below the origin are still out of bounds + below_origin: list[Callable[[], Any]] = [lambda: v[1], lambda: v[0:3]] + for bad in below_origin: with pytest.raises(BoundsCheckError): bad() np.testing.assert_array_equal(np.asarray(v), ref[2:10]) @@ -524,13 +521,15 @@ def test_one_dialect_view_methods_use_domain_coordinates(self) -> None: 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 - IndexError), with one message shape naming the valid range.""" - a, _ = _make(CONFIGS[1]) + """Out-of-range indices — positive past the end, or negative past + ``-size`` — all raise BoundsCheckError (an IndexError), with one message + shape naming the valid range.""" + a, _ = _make(CONFIGS[1]) # shape (24,) triggers: list[Callable[[], Any]] = [ - lambda: a.lazy[-1], - lambda: a.lazy[-3:], - lambda: a.lazy.oindex[(np.array([-1], dtype=np.intp),)], + lambda: a.lazy[24], + lambda: a.lazy[-25], + lambda: a.lazy[20:30], + lambda: a.lazy.oindex[(np.array([-25], dtype=np.intp),)], ] for trigger in triggers: with pytest.raises(BoundsCheckError, match="out of bounds|not contained"): @@ -581,6 +580,152 @@ def test_fancy_view_repr_does_not_crash(self) -> None: assert isinstance(repr(ov.lazy[0]), str) +class TestLazyComposedAdvanced: + """Basic-slice-then-fancy composition routes through the view's transform + (the view's domain origin is preserved), with negatives wrapped at the + boundary against the view's ``exclusive_max``.""" + + def test_slice_then_oindex_read(self) -> None: + """``v.oindex[idx]`` on a sliced view reads in the view's domain + coordinates; negatives wrap from the view's end.""" + a, ref = _make(CONFIGS[1]) # shape (24,) + v = a.lazy[10:20] # domain [10, 20) + np.testing.assert_array_equal(v.oindex[np.array([10, 15, 19])], ref[[10, 15, 19]]) + np.testing.assert_array_equal(v.oindex[np.array([-1, -10])], ref[[19, 10]]) + + def test_slice_then_lazy_oindex_read(self) -> None: + """The lazy accessor chain ``v.lazy.oindex[idx]`` composes the same way.""" + a, ref = _make(CONFIGS[1]) + v = a.lazy[10:20] + np.testing.assert_array_equal(v.lazy.oindex[np.array([-1, -10])].result(), ref[[19, 10]]) + + def test_slice_then_oindex_write(self) -> None: + """Writes through the composed view land at the wrapped domain coordinates.""" + a, ref = _make(CONFIGS[1]) + v = a.lazy[10:20] + v.oindex[np.array([-1, -10])] = np.array([777, 888], dtype="i4") + expected = ref.copy() + expected[19] = 777 + expected[10] = 888 + np.testing.assert_array_equal(a[...], expected) + + def test_slice_then_oindex_out_of_bounds(self) -> None: + """Coordinates outside the view's ``[10, 20)`` domain raise before I/O: + a positive below the origin, a negative wrapping below the origin, and a + positive past the end.""" + a, _ = _make(CONFIGS[1]) + v = a.lazy[10:20] + for sel in (np.array([0]), np.array([-11]), np.array([20])): + with pytest.raises(BoundsCheckError): + _ = v.oindex[sel] + + +class TestLazyBoolScalar: + """`isinstance(True, int)` is True, so a bare boolean scalar would otherwise + pass the integer validators and silently read element 0/1. Reject it at the + boundary; boolean *arrays* (masks) stay valid.""" + + @pytest.mark.parametrize("accessor", ["basic", "oindex", "vindex"]) + @pytest.mark.parametrize("val", [True, False, np.True_, np.False_]) + def test_bool_scalar_rejected(self, accessor: str, val: Any) -> None: + a, _ = _make(CONFIGS[1]) + acc: Any = {"basic": a.lazy, "oindex": a.lazy.oindex, "vindex": a.lazy.vindex}[accessor] + with pytest.raises(IndexError): + _ = acc[val] + + def test_bool_array_still_valid(self) -> None: + """A boolean *array* is a mask, not a scalar, and remains a valid index.""" + a, ref = _make(CONFIGS[1]) + mask = np.zeros(24, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(a.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + + +_GET_FIELD_METHODS = [ + "get_basic_selection", + "get_orthogonal_selection", + "get_mask_selection", + "get_coordinate_selection", +] +_SET_FIELD_METHODS = [ + "set_basic_selection", + "set_orthogonal_selection", + "set_mask_selection", + "set_coordinate_selection", +] + + +class TestLazyFieldsOnView: + """`fields=` on a non-identity view routed through the legacy indexer, which + ignores the transform and reads/writes the wrong storage region. It must now + raise NotImplementedError before any storage access in all eight + get/set_*_selection methods (and the ``v['f0', :]`` sugar), leaving the base + array bit-identical.""" + + @staticmethod + def _struct_array() -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + dt = np.dtype([("f0", "i4"), ("f1", "i4")]) + a = zarr.create_array(MemoryStore(), shape=(10,), chunks=(10,), dtype=dt, zarr_format=2) + base = np.zeros(10, dtype=dt) + base["f0"] = np.arange(10, dtype="i4") + base["f1"] = np.arange(10, dtype="i4") + 100 + a[...] = base + return a, base + + @staticmethod + def _selection_for(method: str) -> Any: + if "mask" in method: + return np.zeros(5, dtype=bool) # view shape is (5,) + if "coordinate" in method: + return (np.array([5, 6, 7], dtype=np.intp),) + return Ellipsis + + @pytest.mark.parametrize("method", _GET_FIELD_METHODS) + def test_get_fields_on_view_raises(self, method: str) -> None: + a, base = self._struct_array() + v = a.lazy[5:10] + with pytest.raises(NotImplementedError, match="field"): + getattr(v, method)(self._selection_for(method), fields="f0") + np.testing.assert_array_equal(a[...], base) + + @pytest.mark.parametrize("method", _SET_FIELD_METHODS) + def test_set_fields_on_view_raises(self, method: str) -> None: + a, _base = self._struct_array() + before = np.asarray(a[...]).copy() + v = a.lazy[5:10] + with pytest.raises(NotImplementedError, match="field"): + getattr(v, method)(self._selection_for(method), np.arange(3, dtype="i4"), fields="f0") + np.testing.assert_array_equal(a[...], before) + + def test_getitem_field_sugar_on_view_raises(self) -> None: + a, base = self._struct_array() + v = a.lazy[5:10] + selections: list[Any] = [("f0", slice(None)), "f0"] + for sel in selections: + with pytest.raises(NotImplementedError, match="field"): + _ = v[sel] + np.testing.assert_array_equal(a[...], base) + + def test_setitem_field_sugar_on_view_raises(self) -> None: + a, _base = self._struct_array() + before = np.asarray(a[...]).copy() + v = a.lazy[5:10] + field_sel: Any = ("f0", slice(None)) + with pytest.raises(NotImplementedError, match="field"): + v[field_sel] = np.arange(5, dtype="i4") + np.testing.assert_array_equal(a[...], before) + + def test_fields_on_identity_array_still_work(self) -> None: + """Field selection on the base (identity) array still routes to the legacy + path — the view guard must not block it. (Only the write path is checked + here; structured-dtype field *reads* are broken independently of this + change.)""" + a, _ = self._struct_array() + a.set_basic_selection(Ellipsis, np.arange(10, dtype="i4") + 1, fields="f0") + result: Any = np.asarray(a[...]) + np.testing.assert_array_equal(result["f0"], np.arange(10, dtype="i4") + 1) + + 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 diff --git a/tests/test_properties.py b/tests/test_properties.py index c71a713ed9..c5634e4cf9 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -785,8 +785,9 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: 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)) - # 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. + # The strategies draw NumPy-dialect selections. The lazy boundary wraps + # negatives against the (here zero-origin) view, matching NumPy, but positive + # indices are literal domain coordinates; canonicalizing to the equivalent + # non-negative form keeps the comparison valid for both rules. 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) From 43bcf7174a6a18d0836e99c91c88e4e433a5ed00 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 11:16:01 +0200 Subject: [PATCH 48/60] fix(indexing): validate boolean mask shapes at the lazy boundary A boolean mask must exactly match the view domain's shape on the dimensions it consumes (NumPy parity: 'boolean index did not match indexed array'); under/over-length masks previously truncated silently and a too-high-rank mask raised ValueError from the transform layer. Validate in _normalize_public_selection against the current view domain, raising IndexError before transform construction. Also pin the pre-existing eager fields= sibling-field corruption with a strict xfail so the identity fields test documents rather than hides it. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 76 +++++++++++++++++++++++++++++-------- tests/test_lazy_indexing.py | 68 +++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 15 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index fc011c3c7a..0ad6d8afdf 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2957,16 +2957,20 @@ def _reject_bool_scalar(self, sel: Any) -> None: def _normalize_public_selection(self, selection: Any) -> Any: """Normalize a user selection at the public Array boundary. - Rejects boolean *scalar* indices, and wraps negative integers, negative - index-array elements, and negative slice bounds to literal coordinates in - the current view's domain: an index `k` with `-size <= k < 0` on axis `d` - becomes `exclusive_max[d] + k` (from the end of the view's domain — exactly - NumPy for a zero-origin domain). Positive indices are left as literal domain - coordinates. The transform layer stays literal; it performs the final bounds - check, so an out-of-range value (a positive past the end, or `k < -size`) is - rejected there before any chunk access. - """ - exclusive_max = self._async_array._transform.domain.exclusive_max + Rejects boolean *scalar* indices, validates boolean-mask shapes, and wraps + negative integers, negative index-array elements, and negative slice bounds + to literal coordinates in the current view's domain: an index `k` with + `-size <= k < 0` on axis `d` becomes `exclusive_max[d] + k` (from the end of + the view's domain — exactly NumPy for a zero-origin domain). Positive indices + are left as literal domain coordinates. A boolean mask consumes `mask.ndim` + dimensions and must exactly match the view domain's shape on them (NumPy + parity), else IndexError — never a silent truncation. The transform layer + stays literal; it performs the final bounds check, so an out-of-range value + (a positive past the end, or `k < -size`) is rejected there before any chunk + access. + """ + domain = self._async_array._transform.domain + exclusive_max = domain.exclusive_max ndim = len(exclusive_max) is_tuple = isinstance(selection, tuple) @@ -2975,10 +2979,29 @@ def _normalize_public_selection(self, selection: Any) -> Any: for item in items: self._reject_bool_scalar(item) - n_real = sum(1 for s in items if s is not Ellipsis and s is not None) - # Too many indices: leave the selection untouched and let the transform - # layer raise its canonical error with the correct message. + # A boolean mask consumes one dimension per mask axis; everything else + # consumes exactly one (newaxis/Ellipsis consume none). + n_real = 0 + has_mask = False + for item in items: + if item is Ellipsis or item is None: + continue + mask = _as_bool_mask(item) + if mask is not None: + has_mask = True + n_real += mask.ndim + else: + n_real += 1 if n_real > ndim: + if has_mask: + # The transform layer would raise an unhelpful ValueError for a + # too-high-rank mask; raise the canonical IndexError here. + raise IndexError( + f"too many indices for array: array has {ndim} dimension(s), " + f"but {n_real} were indexed" + ) + # Otherwise leave the selection untouched and let the transform layer + # raise its canonical error with the correct message. return selection normalized: list[Any] = [] @@ -2990,8 +3013,21 @@ def _normalize_public_selection(self, selection: Any) -> Any: elif item is None: normalized.append(item) # newaxis consumes no domain axis else: - normalized.append(_wrap_negative_index(item, exclusive_max[axis])) - axis += 1 + mask = _as_bool_mask(item) + if mask is not None: + expected = domain.shape[axis : axis + mask.ndim] + for offset, (want, got) in enumerate(zip(expected, mask.shape, strict=True)): + if want != got: + raise IndexError( + "boolean index did not match indexed array along " + f"dimension {axis + offset}; dimension is {want} but " + f"corresponding boolean dimension is {got}" + ) + normalized.append(item) + axis += mask.ndim + else: + normalized.append(_wrap_negative_index(item, exclusive_max[axis])) + axis += 1 return tuple(normalized) if is_tuple else normalized[0] def _reject_fields_on_view(self, fields: Fields | None) -> None: @@ -4471,6 +4507,16 @@ async def _shards_initialized( type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"] +def _as_bool_mask(item: Any) -> np.ndarray[Any, np.dtype[np.bool_]] | None: + """Return `item` as a boolean ndarray if it is a boolean mask (ndarray or + list), else None. Boolean *scalars* are not masks (rejected separately).""" + if isinstance(item, (list, np.ndarray)): + arr = np.asarray(item) + if arr.dtype == np.bool_: + return arr + return None + + def _wrap_negative_index(item: Any, exclusive_max: int) -> Any: """Wrap negative values in a single per-axis selection element to literal coordinates against `exclusive_max` (from-the-end within the axis's domain). diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 25b5f547da..13b76769a7 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -641,6 +641,60 @@ def test_bool_array_still_valid(self) -> None: np.testing.assert_array_equal(a.lazy.oindex[mask].result(), ref[[3, 5, 7]]) +class TestLazyMaskShape: + """A boolean mask must exactly match the view domain's shape on the + dimension(s) it consumes (NumPy parity: "boolean index did not match indexed + array"). Anything else raises IndexError at the boundary, before transform + construction — never a silent truncation.""" + + def test_correct_shape_masks_work(self) -> None: + """Exact-shape masks select as before, through oindex, vindex, and on a + non-zero-origin view (whose domain shape, not the base's, is the ruler).""" + a, ref = _make(CONFIGS[1]) # shape (24,) + mask = np.zeros(24, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(a.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + np.testing.assert_array_equal(a.lazy.vindex[mask].result(), ref[[3, 5, 7]]) + v = a.lazy[2:10] # domain [2, 10) -> shape (8,) + vmask = np.zeros(8, dtype=bool) + vmask[[3, 5]] = True # True positions are absolute coordinates + np.testing.assert_array_equal(v.lazy.oindex[vmask].result(), ref[[3, 5]]) + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + def test_under_length_mask_raises(self, accessor: str) -> None: + a, _ = _make(CONFIGS[1]) # shape (24,) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = acc[np.zeros(5, dtype=bool)] + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + @pytest.mark.parametrize("fill", [False, True]) + def test_over_length_mask_raises(self, accessor: str, fill: bool) -> None: + """Over-length masks raise the shape-mismatch error regardless of where + their True values fall (not a bounds error on the True positions).""" + a, _ = _make(CONFIGS[1]) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = acc[np.full(30, fill, dtype=bool)] + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + def test_wrong_ndim_mask_raises(self, accessor: str) -> None: + """A 2-D mask on a 1-D array is an IndexError (it consumes two + dimensions), not a ValueError from deeper in the transform layer.""" + a, _ = _make(CONFIGS[1]) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="too many indices"): + _ = acc[np.zeros((4, 6), dtype=bool)] + + def test_view_mask_checked_against_view_domain(self) -> None: + """On a view, the mask ruler is the view's domain shape: the base + array's shape is the wrong length there.""" + a, _ = _make(CONFIGS[1]) + v = a.lazy[2:10] # domain shape (8,) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = v.lazy.oindex[np.zeros(24, dtype=bool)] + + _GET_FIELD_METHODS = [ "get_basic_selection", "get_orthogonal_selection", @@ -725,6 +779,20 @@ def test_fields_on_identity_array_still_work(self) -> None: result: Any = np.asarray(a[...]) np.testing.assert_array_equal(result["f0"], np.arange(10, dtype="i4") + 1) + @pytest.mark.xfail( + reason="the eager fields= write path corrupts sibling fields (pre-existing " + "upstream bug in the legacy identity path, not introduced or fixed by the " + "lazy-view guard); strict xfail flips when the eager path is fixed", + strict=True, + ) + def test_fields_write_preserves_sibling_fields(self) -> None: + """A `fields='f0'` write must leave `f1` untouched. Documents the known + eager-path corruption (mirrors the TestKnownFancyIntBugs pinning style).""" + a, base = self._struct_array() + a.set_basic_selection(Ellipsis, np.arange(10, dtype="i4") + 1, fields="f0") + result: Any = np.asarray(a[...]) + np.testing.assert_array_equal(result["f1"], base["f1"]) + class TestKnownFancyIntBugs: """Strict-xfail pins for the int-on-fancy-picked-dim defect (see review notes): From 8c6d152bc6dd84e5f992bd7e9b538ac9807cb893 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 11:32:42 +0200 Subject: [PATCH 49/60] fix(indexing): preserve transform result shapes Zero-rank correlated reads (e.g. an integer index into a vindex-created view) collapsed a flat (1,) working buffer with as_scalar, which returns shape-(1,) data; collapse to 0-d first so the caller gets a true scalar. The caller-visible `out` contract for vectorized reads is now the broadcast selection shape rather than a flat buffer: a flat temporary is used internally and the reshaped result is copied into the supplied buffer. Efficiency: reshape the freshly-allocated contiguous working buffer as a view instead of np.array(...)-copying, and drop the now-redundant second reshape in get_coordinate_selection's transform branch. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 65 +++++++++++++++++++++++-------------- tests/test_lazy_indexing.py | 35 +++++++++++--------- 2 files changed, 60 insertions(+), 40 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 0ad6d8afdf..4f3d0794d1 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3893,16 +3893,10 @@ def get_coordinate_selection( 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 = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(sel_shape)) - return out_array + # `_get_selection_via_transform` already returns the broadcast selection + # shape (transform.domain.shape == the broadcast shape of the coordinate + # arrays), so no further reshape/copy is needed here. + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_coordinate_selection( self, @@ -6083,26 +6077,34 @@ async def _get_selection_via_transform( out_shape = transform.domain.shape - # Vectorized indexing (any coordinate ArrayMap, input_dimension None — even a + # Vectorized indexing (any correlated 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. + # the working 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. A zero-rank domain + # still needs a flat (1,) working buffer for the scatter, but its single cell + # maps to a scalar on return (see the out_shape == () branch below). needs_flat_buffer = _is_vectorized_transform(transform) buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape - # 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. + # Setup output buffer. The caller-visible `out` contract is the broadcast + # selection shape (out_shape). Vectorized reads still scatter through a flat + # index internally, so when `out` is supplied for a vectorized selection we + # read into a flat temporary and copy the reshaped result into `out` below. + copy_into_out: NDBuffer | None = None 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 != buffer_shape: + if out.shape != out_shape: raise ValueError( - f"shape of out argument doesn't match. Expected {buffer_shape}, got {out.shape}" + f"shape of out argument doesn't match. Expected {out_shape}, got {out.shape}" ) - out_buffer = out + if needs_flat_buffer: + copy_into_out = out + out_buffer = prototype.nd_buffer.empty(shape=buffer_shape, dtype=dtype, order=order) + else: + out_buffer = out else: out_buffer = prototype.nd_buffer.empty(shape=buffer_shape, dtype=dtype, order=order) @@ -6150,13 +6152,26 @@ async def _get_selection_via_transform( f"Missing chunks:\n{chunks_str}" ) - # Return scalar for 0-d results + # Return scalar for 0-d results. A correlated read used a flat (1,) working + # buffer for the scatter, so collapse it to a 0-d value before extracting. if out_shape == (): - return out_buffer.as_scalar() + if copy_into_out is not None: + copy_into_out[...] = out_buffer.as_ndarray_like().reshape(()) + return copy_into_out.as_scalar() + # Mirror NDBuffer.as_scalar (GPU-safe conversion), but collapse the flat + # (1,) correlated working buffer to 0-d first. + return cast("NDArrayLikeOrScalar", out_buffer.as_numpy_array().reshape(())[()]) out_result = out_buffer.as_ndarray_like() - # Reshape if we flattened for array indexing + # Reshape the flat working buffer back to the broadcast selection shape. The + # buffer is freshly allocated and contiguous, so reshape returns a view — no + # copy — and preserves the backend array type. if needs_flat_buffer and hasattr(out_result, "reshape"): - out_result = np.array(out_result).reshape(out_shape) + out_result = out_result.reshape(out_shape) + # A vectorized read filled a flat temporary; copy the reshaped result into + # the caller-provided out buffer, whose shape is the broadcast selection shape. + if copy_into_out is not None: + copy_into_out[...] = out_result + return copy_into_out.as_ndarray_like() return out_result diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 13b76769a7..efb77afc67 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -318,6 +318,15 @@ class TestLazyViewMethods: through ``Array``'s own dispatch, which had correctness gaps. """ + def test_zero_rank_result_is_scalar(self) -> None: + """A zero-dimensional array reads back as a scalar (shape ``()``) eagerly and lazily.""" + a = zarr.create_array({}, shape=(), chunks=(), dtype="i4") + a[...] = 5 + assert np.shape(a[...]) == () + assert np.shape(a[()]) == () + assert np.shape(a.lazy[...].result()) == () + assert a.lazy[...].result() == 5 + def test_coordinate_methods_wrap_negative_indices(self) -> None: a, ref = _make(CONFIGS[1]) view = a.lazy[2:10] @@ -377,24 +386,23 @@ def test_view_vindex(self, cfg: Config) -> None: ) 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. + def test_view_vindex_with_broadcast_out_buffer(self) -> None: + """vindex with a multi-dim result and out= on a view takes a broadcast-shaped 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). + Vectorized indexing scatters through a single flat index internally, but + the caller-visible `out` contract is the broadcast selection shape: the + flat temporary is an implementation detail. The out buffer must therefore + match `expected.shape`, and the results must land in it. """ a, ref = _make(CONFIGS[3]) # 2d-unsharded v = a.lazy[2:18] 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 = ref[i0, i1] - buf = default_buffer_prototype().nd_buffer.empty( - shape=(expected.size,), dtype=np.dtype("i4") - ) + assert expected.shape == (2, 2) + buf = default_buffer_prototype().nd_buffer.empty(shape=expected.shape, 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 - ) + np.testing.assert_array_equal(np.asarray(buf.as_ndarray_like()), expected) @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) def test_view_write_through_tuple(self, cfg: Config) -> None: @@ -813,17 +821,14 @@ def test_int_read_on_oindex_view(self) -> None: 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]) == () + assert np.shape(pts.lazy[0].result()) == () + assert pts[0] == a[1] def test_block_selection_on_view_rejected(self) -> None: """Block selection is ill-defined on a lazy view and must raise, not corrupt.""" From 6b20925b2c09ee2468ae7b5d2ed663260aa03fec Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 23:03:56 +0200 Subject: [PATCH 50/60] test(indexing): generate composed indexing programs Property-based differential oracle for the index-transform layer: an IndexingProgram composes one to three indexing operations (basic prefix, optional trailing orthogonal/vectorized step) on a small array and checks five execution recipes (materialize, eager-on-lazy, out=, scalar and array writes) against a NumPy model applied in lockstep. Deterministic @example cases pin slice-then-oindex, unequal orthogonal lengths, multidimensional vindex with out=, a zero-rank result, and an empty result. Also restores zero-extent coverage lost with the deleted pre-branch tests: _indexing_array and the lazy-view parity test draw min_side=0 again (mask mode verified to work on zero-size shapes and now admitted by _eligible), and assert_read_matches_numpy learns the transform-path out= convention (broadcast result shape) so lazy-view vindex reads with multidimensional results are checked correctly. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/testing/strategies.py | 134 +++++++++++++++ tests/test_properties.py | 300 +++++++++++++++++++++++++++++++-- 2 files changed, 421 insertions(+), 13 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 60c99e92a1..a37013bbf5 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -2,6 +2,7 @@ import math import sys from collections.abc import Callable, Mapping +from dataclasses import dataclass from typing import Any, Literal import hypothesis.extra.numpy as npst @@ -669,6 +670,139 @@ def numpy_array_indexers( raise ValueError(f"unknown indexing mode: {mode!r}") +# --- Composed indexing programs ------------------------------------------------- +# +# A *program* is a short sequence of indexing operations composed on a single +# array, plus an execution recipe describing how to realize the composed +# selection (read whole, read a sub-selection eagerly, read into an ``out=`` +# buffer, or write). It exists to differentially test the index-transform layer: +# every operation composes a real transform, and the runner (in +# ``tests/test_properties.py``) checks the composed result against a NumPy oracle +# that applies the equivalent selections in lockstep. +# +# The mode vocabulary here is the *public dialect* ("orthogonal"/"vectorized"), +# not the internal accessor names ("oindex"/"vindex"); the runner maps between +# them. +ProgramMode = Literal["basic", "orthogonal", "vectorized"] +ProgramExecution = Literal["materialize", "eager_on_lazy", "out", "set_scalar", "set_array"] + + +@dataclass(frozen=True) +class IndexOperation: + """One indexing step in an :class:`IndexingProgram`. + + ``selection`` carries exactly what the runner needs to reproduce the step on + *both* a zarr array and its NumPy reference: + + - ``"basic"`` — a single object used verbatim for both (slice / int / + ellipsis / tuple thereof). + - ``"orthogonal"`` — a ``(zarr_selection, numpy_selection)`` pair, where the + NumPy side is the ``np.ix_`` outer-product spelling of the per-axis + selection. + - ``"vectorized"`` — a single coordinate selection used verbatim for both + (plain NumPy fancy-indexing semantics). + """ + + mode: ProgramMode + selection: Any + + +@dataclass(frozen=True) +class IndexingProgram: + """A composed sequence of :class:`IndexOperation` plus an execution recipe.""" + + operations: tuple[IndexOperation, ...] + execution: ProgramExecution + + +# Small arrays keep 300 Hypothesis examples cheap while still crossing chunk +# boundaries (the runner chunks each axis into ~2 pieces). ``min_side=0`` lets a +# program start from — or a basic step produce — a zero-extent axis, which is the +# required "empty result" coverage class. +program_shapes = npst.array_shapes(min_dims=1, max_dims=3, min_side=0, max_side=4) + + +def _basic_result_shape(shape: tuple[int, ...], selection: Any) -> tuple[int, ...]: + """The NumPy result shape of applying a *basic* ``selection`` to ``shape``. + + Used to carry the visible shape forward while generating a program, so each + step's selection is drawn against the shape the previous steps leave behind. + Uses a tiny placeholder array (program shapes are capped small). + """ + result_shape: tuple[int, ...] = np.empty(shape, dtype=np.int8)[selection].shape + return result_shape + + +@st.composite +def indexing_programs(draw: st.DrawFn, *, shape: tuple[int, ...]) -> IndexingProgram: + """Generate a composable indexing program over an array of ``shape``. + + Structure (one to three operations): a prefix of zero-to-two **basic** steps, + optionally followed by a single **fancy** (orthogonal/vectorized) step as the + *last* operation. The visible shape is carried forward so every step is valid + against what its predecessors produce. + + **Why fancy is confined to the last position** — the exclusions here are + deliberate and documented so they can shrink as the underlying bugs are + fixed: + + - *Fancy-after-fancy* composition is broken in ``_reindex_array_oindex`` + (applying oindex/vindex to a view that already carries an ArrayMap axis), so + at most one fancy step is generated. + - *Basic-after-fancy* is excluded wholesale. Integer basic indexing on an + oindex-picked axis is a strict xfail + (``TestKnownFancyIntBugs::test_int_read_on_oindex_view``); rather than track + which composed axes became "fancy-picked" and admit only the working + basic-on-non-fancy-axis subset, we keep fancy last and let the basic prefix + cover multi-step basic composition. (When those bugs are fixed this + restriction can be relaxed to interleave basic and fancy freely.) + + A fancy step is only emitted when the current shape has rank ``>= 1`` and no + zero-length axis (integer/orthogonal selections cannot address an empty axis); + otherwise the program stays basic-only. If nothing else was generated, a + single basic step is appended so a program always has at least one operation. + """ + n_basic = draw(st.integers(min_value=0, max_value=2)) + want_fancy = draw(st.booleans()) + + operations: list[IndexOperation] = [] + cur = shape + for _ in range(n_basic): + if len(cur) == 0: + break # a rank-0 (scalar) view has no further indexing to compose + sel = draw(basic_indices(shape=cur)) + operations.append(IndexOperation("basic", sel)) + cur = _basic_result_shape(cur, sel) + + fancy_ok = len(cur) > 0 and all(s > 0 for s in cur) + if want_fancy and fancy_ok: + mode = draw(st.sampled_from(("orthogonal", "vectorized"))) + if mode == "orthogonal": + zsel, npsel = draw(orthogonal_indices(shape=cur)) + operations.append(IndexOperation("orthogonal", (zsel, npsel))) + else: + idx = draw( + npst.integer_array_indices( + shape=cur, + result_shape=npst.array_shapes(min_side=1, max_side=4, max_dims=2), + ) + ) + operations.append(IndexOperation("vectorized", idx)) + + if not operations: + operations.append(IndexOperation("basic", draw(basic_indices(shape=cur)))) + + executions: tuple[ProgramExecution, ...] = ( + "materialize", + "eager_on_lazy", + "out", + "set_scalar", + "set_array", + ) + execution = draw(st.sampled_from(executions)) + return IndexingProgram(tuple(operations), execution) + + @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 c5634e4cf9..6f111e1a6f 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -2,7 +2,7 @@ import json import numbers from collections.abc import Generator -from typing import Any +from typing import Any, Literal import numpy as np import pytest @@ -15,7 +15,7 @@ import hypothesis.extra.numpy as npst import hypothesis.strategies as st -from hypothesis import assume, given, settings +from hypothesis import assume, example, given, settings from hypothesis.stateful import RuleBasedStateMachine, invariant, rule from tests.conftest import Expect @@ -24,15 +24,19 @@ from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.sync import sync from zarr.testing.strategies import ( + IndexingProgram, IndexMode, + IndexOperation, array_metadata, arrays, basic_indices, block_indices, block_test_arrays, complex_rectilinear_arrays, + indexing_programs, numpy_array_indexers, numpy_arrays, + program_shapes, rectilinear_arrays, simple_arrays, stores, @@ -431,23 +435,39 @@ def _n_array_axes(zsel: Any) -> int: 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 - (oindex/vindex/mask via integer/boolean arrays) can't handle zero-size axes. + Rank-0 arrays have no interesting selections. `basic` and `mask` handle + zero-length axes (a boolean mask of the array's shape simply selects nothing, + which zarr and NumPy agree on — verified against `set_mask_selection` too). + The integer-array fancy modes (`oindex`/`vindex`) cannot address an empty + axis, so they require every axis to be non-empty. """ if len(shape) == 0: return False - return mode == "basic" or all(s > 0 for s in shape) + if mode in ("basic", "mask"): + return True + return all(s > 0 for s in shape) def assert_read_matches_numpy( - target: zarr.Array, ref: np.ndarray[Any, Any], mode: IndexMode, zsel: Any, npsel: Any + target: zarr.Array, + ref: np.ndarray[Any, Any], + mode: IndexMode, + zsel: Any, + npsel: Any, + *, + out_layout: Literal["eager", "transform"] = "eager", ) -> None: """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 - is flat for the vectorized vindex/mask modes (which scatter through a flat - index) and full-shape otherwise. + view) and `ref` is the NumPy array it must behave like. The two consumers + disagree on the `out=` buffer shape for the vectorized vindex/mask modes, so + `out_layout` names which convention `target` implements: + + - `"eager"` (legacy path): vindex/mask scatter through a flat index, so the + buffer is flat — `(number of selected points,)`. + - `"transform"` (lazy views, Task 4): `out=` takes the broadcast selection + shape for every mode, so the buffer always has the result shape. """ expected = np.asarray(ref[npsel]) actual = np.asarray(_get(target, mode, zsel)) @@ -458,7 +478,10 @@ def assert_read_matches_numpy( 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 + if out_layout == "eager" and mode in _VECTORIZED_MODES: + buf_shape: tuple[int, ...] = (expected.size,) + else: + buf_shape = 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) @@ -468,7 +491,9 @@ 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)), + # min_side=0 restores zero-extent-array coverage: the `_eligible` helper + # routes zero-size shapes to the modes that tolerate them (basic/mask). + simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=0)), rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), ) ) @@ -777,7 +802,9 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: 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))) + # min_side=0 exercises the `windows()` zero-extent branch and, via `_eligible`, + # zero-size lazy-view reads in the basic/mask modes that tolerate them. + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=0))) nparray = zarray[:] window = data.draw(windows(shape=nparray.shape)) view = zarray.lazy[window].translate_to((0,) * nparray.ndim) @@ -790,4 +817,251 @@ async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: # indices are literal domain coordinates; canonicalizing to the equivalent # non-negative form keeps the comparison valid for both rules. 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) + # Lazy views take out= in the broadcast result shape for every mode (Task 4) + # — *except* when the composed transform is an identity, where the read + # methods fast-path to the legacy implementation and its flat-vindex out= + # convention applies. Mirror that dispatch here. + out_layout: Literal["eager", "transform"] = ( + "eager" if view._async_array._is_identity else "transform" + ) + assert_read_matches_numpy( + view, vref, mode, _canonical_nonneg(zsel, vref.shape), npsel, out_layout=out_layout + ) + + +# --------------------------------------------------------------------------- +# Composed indexing programs (Task 5): a differential oracle over sequences of +# composed indexing operations. Where the single-shot parity tests above check +# one selection at a time, these compose one to three operations through the real +# index-transform layer and check the composed result — and each of five +# execution recipes — against a NumPy model applied in lockstep. +# --------------------------------------------------------------------------- + +# Public program dialect -> the internal accessor mode used by _get/_setitem. +_PROGRAM_INDEX_MODE: dict[str, IndexMode] = { + "basic": "basic", + "orthogonal": "oindex", + "vectorized": "vindex", +} + + +def _program_selections(op: IndexOperation) -> tuple[Any, Any]: + """The ``(zarr_selection, numpy_selection)`` pair for one program step. + + Only ``orthogonal`` needs distinct spellings (per-axis vs the ``np.ix_`` + outer-product form); ``basic`` and ``vectorized`` index both sides alike. + """ + if op.mode == "orthogonal": + zsel, npsel = op.selection + return zsel, npsel + return op.selection, op.selection + + +def _apply_lazy_step(view: zarr.Array, mode: str, zsel: Any) -> zarr.Array: + """Compose one lazy indexing step onto ``view`` and re-zero the result. + + Re-zeroing (``translate_to``) keeps the next step's NumPy-dialect selection + valid against a zero-origin domain (positive coordinates stay literal, so they + match NumPy) while still exercising domain translation on every step. + """ + if mode == "basic": + out = view.lazy[zsel] + elif mode == "orthogonal": + out = view.lazy.oindex[zsel] + elif mode == "vectorized": + out = view.lazy.vindex[zsel] + else: + raise AssertionError(mode) + if out.ndim: + out = out.translate_to((0,) * out.ndim) + return out + + +def _compare_read(actual_raw: Any, expected: np.ndarray[Any, Any]) -> None: + """Assert a program read matches NumPy in scalar/array kind, shape, dtype, values.""" + actual = np.asarray(actual_raw) + # Scalar-versus-array kind: a zero-rank read must come back as a scalar, a + # ranked read as an array. np.ndim reads 0 for python/numpy scalars. + assert np.ndim(actual_raw) == expected.ndim, ("kind", np.ndim(actual_raw), expected.ndim) + if expected.ndim >= 1: # skip 0-d to avoid python-vs-numpy scalar dtype noise + assert actual.dtype == expected.dtype, (actual.dtype, expected.dtype) + assert actual.shape == expected.shape, (actual.shape, expected.shape) + assert_array_equal(actual, expected) + + +def _compose_base( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], operations: tuple[IndexOperation, ...] +) -> tuple[zarr.Array, np.ndarray[Any, Any]]: + """Compose every operation *except the last* as lazy views, in lockstep with NumPy. + + Returns the composed (zero-origin) lazy base view and the matching NumPy + array; their shapes are asserted equal at each step. The prefix operations are + always basic (fancy steps only ever occur last), so the NumPy side stays a + view — which matters for the write executions, where assignment through it must + reach the root array. + """ + base = zarray + ref = nparray + for op in operations[:-1]: + zsel, npsel = _program_selections(op) + base = _apply_lazy_step(base, op.mode, _canonical_nonneg(zsel, base.shape)) + ref = ref[npsel] + assert base.shape == ref.shape, (base.shape, ref.shape) + return base, ref + + +def _run_program_read( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], program: IndexingProgram +) -> None: + base, ref = _compose_base(zarray, nparray, program.operations) + last = program.operations[-1] + imode = _PROGRAM_INDEX_MODE[last.mode] + zsel, npsel = _program_selections(last) + czsel = _canonical_nonneg(zsel, base.shape) + expected = np.asarray(ref[npsel]) + + if program.execution == "eager_on_lazy": + # Apply the final step as an *eager* indexing read on the composed view. + _compare_read(_get(base, imode, czsel), expected) + return + + # materialize / out: compose the final step lazily too, then read the whole view. + view = _apply_lazy_step(base, last.mode, czsel) + _compare_read(view[...], expected) + if program.execution == "out" and expected.ndim >= 1 and expected.size > 0: + # A whole-view basic read: the transform path takes the full result shape + # (Task 4), for every mode, so the buffer is never the flat vectorized form. + buf = default_buffer_prototype().nd_buffer.empty(shape=expected.shape, dtype=expected.dtype) + view.get_basic_selection(Ellipsis, out=buf) + assert_array_equal(np.asarray(buf.as_ndarray_like()), expected) + + +def _run_program_write( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], program: IndexingProgram +) -> None: + root = nparray.copy() # the NumPy model we mutate and compare against + base, sub = _compose_base(zarray, root, program.operations) + last = program.operations[-1] + imode = _PROGRAM_INDEX_MODE[last.mode] + zsel, npsel = _program_selections(last) + czsel = _canonical_nonneg(zsel, base.shape) + if last.mode in ("orthogonal", "vectorized"): + # A fancy write hitting a cell twice is order-dependent and has no oracle + # (same limitation NumPy has); program arrays are unsharded so GH2834 (the + # only other write restriction) cannot arise. + assume(_write_is_unambiguous(imode, czsel, base.shape)) + + result_shape = sub[npsel].shape + value: Any + if program.execution == "set_scalar": + value = np.array(-7, dtype=root.dtype) + else: # set_array: distinct (negative) values expose scatter/order bugs + value = (-(np.arange(int(np.prod(result_shape)), dtype="i8") + 1)).reshape(result_shape) + value = value.astype(root.dtype) + + sub[npsel] = value # writes through the basic-view chain into `root` + _setitem(base, imode, czsel, value) + assert_array_equal(zarray[:], root) + + +def _build_program_array(shape: tuple[int, ...]) -> tuple[zarr.Array, np.ndarray[Any, Any]]: + """A deterministic i4 array filled with ``arange``, chunked ~2-per-axis. + + Fixed dtype and a regular (unsharded) grid keep the program oracle focused on + composition and execution correctness — dtype breadth and rectilinear/sharded + grids are covered by the single-shot parity tests and the state machines. The + ~2-chunks-per-axis geometry makes composed selections routinely cross chunk + boundaries. + """ + chunks = tuple(max(1, (s + 1) // 2) for s in shape) + nparray = np.arange(int(np.prod(shape, dtype=int)), dtype="i4").reshape(shape) + zarray = zarr.create_array({}, shape=shape, chunks=chunks, dtype="i4") + zarray[...] = nparray + return zarray, nparray + + +@st.composite +def _indexing_scenarios(draw: st.DrawFn) -> tuple[tuple[int, ...], IndexingProgram]: + """Draw a ``(shape, program)`` scenario. Kept together so `@example` can pin both.""" + shape = draw(program_shapes) + program = draw(indexing_programs(shape=shape)) + return shape, program + + +def _orth(*axes: Any) -> tuple[tuple[np.ndarray[Any, Any], ...], Any]: + """Build an orthogonal selection ``(zarr_per_axis, numpy_ix_form)`` for `@example`.""" + arrs = tuple(np.asarray(a) for a in axes) + return arrs, np.ix_(*arrs) + + +@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") +# slice-then-oindex: the canonical composed-fancy regression. Reversing the Task 2 +# correlation predicate makes this example fail (mutation check). +@example( + scenario=( + (4, 4), + IndexingProgram( + ( + IndexOperation("basic", (slice(1, 4), slice(None))), + IndexOperation("orthogonal", _orth([0, 2], [1, 3])), + ), + "materialize", + ), + ) +) +# unequal orthogonal lengths: the outer product must not assume square selections. +@example( + scenario=( + (4, 4), + IndexingProgram((IndexOperation("orthogonal", _orth([0, 1, 2], [0, 1])),), "materialize"), + ) +) +# multidimensional vindex read into an out= buffer. +@example( + scenario=( + (4, 4), + IndexingProgram( + ( + IndexOperation( + "vectorized", (np.array([[0, 1], [2, 3]]), np.array([[1, 0], [3, 2]])) + ), + ), + "out", + ), + ) +) +# zero-rank result: a full integer selection collapses to a scalar. +@example( + scenario=( + (3, 3), + IndexingProgram((IndexOperation("basic", (1, 2)),), "materialize"), + ) +) +# empty result: a 0-length slice yields a zero-size read. +@example( + scenario=( + (4, 4), + IndexingProgram((IndexOperation("basic", (slice(0, 0), slice(None))),), "materialize"), + ) +) +@given(scenario=_indexing_scenarios()) +def test_indexing_program_parity(scenario: tuple[tuple[int, ...], IndexingProgram]) -> None: + """A composed indexing program matches NumPy under every execution recipe. + + Reads (`materialize`, `eager_on_lazy`, `out`) check scalar/array kind, shape, + dtype and values; writes (`set_scalar`, `set_array`) apply the equivalent + assignment to a NumPy root and compare the whole zarr array afterwards. The + generator composes at most one fancy step and always keeps it last — see + `indexing_programs` for the documented exclusions (fancy-after-fancy and + basic-after-fancy are known-broken and deliberately not generated). + """ + shape, program = scenario + zarray, nparray = _build_program_array(shape) + if program.execution in ("set_scalar", "set_array"): + _run_program_write(zarray, nparray, program) + else: + _run_program_read(zarray, nparray, program) From 15d1ffbac7796fd6468e14660d7a1f69cec31fa0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 16 Jul 2026 23:21:04 +0200 Subject: [PATCH 51/60] fix(lazy): guard the async surface against silent transform bypass The async surface built legacy indexers from metadata.shape without consulting a lazy view's index transform, silently reading/writing the wrong region (getitem/setitem/selection methods, oindex/vindex accessors, and from_array's copy loop). Guard every such entry point with a new AsyncArray._require_identity_for_async helper that raises LazyViewError on non-identity views, and precheck from_array up front so no half-created target array is left behind. Steer users to the sync surface / .result(). Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 41 +++++++++++++- src/zarr/core/indexing.py | 2 + tests/test_lazy_indexing.py | 103 ++++++++++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 1 deletion(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4f3d0794d1..44d6cc4985 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -882,6 +882,24 @@ def _require_identity(self, name: str) -> None: f"`metadata` / `chunk_grid`." ) + def _require_identity_for_async(self, name: str) -> None: + """Raise `LazyViewError` if this async entry point is used on a lazy view. + + The async surface does not yet route selections through a view's index + transform: it builds a legacy indexer from `metadata.shape`, which for a + non-identity view would silently read or write the *wrong* region. Until + async transform routing lands, guard these entry points and steer users + to the synchronous surface. `name` is the member being guarded. + """ + if not self._is_identity: + raise LazyViewError( + f"`{name}` on the async surface does not yet support lazy views " + f"(sliced or indexed arrays with a non-identity transform); it would " + f"silently read or write the wrong region. Use the synchronous `Array` " + f"surface, or materialize the view first with `.result()` (for example, " + f"index `view.result()` instead of `await view.async_array...`)." + ) + @property def _is_sharded(self) -> bool: """Whether the array stores inner chunks inside shards (a sharding codec).""" @@ -1606,7 +1624,7 @@ async def getitem( >>> asyncio.run(example()) np.int32(0) """ - + self._require_identity_for_async("getitem") return await _getitem( self.store_path, self.metadata, @@ -1625,6 +1643,7 @@ async def get_orthogonal_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_orthogonal_selection") return await _get_orthogonal_selection( self.store_path, self.metadata, @@ -1645,6 +1664,7 @@ async def get_mask_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_mask_selection") return await _get_mask_selection( self.store_path, self.metadata, @@ -1665,6 +1685,7 @@ async def get_coordinate_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_coordinate_selection") return await _get_coordinate_selection( self.store_path, self.metadata, @@ -1778,6 +1799,7 @@ async def setitem( - This method is asynchronous and should be awaited. - Supports basic indexing, where the selection is contiguous and does not involve advanced indexing. """ + self._require_identity_for_async("setitem") return await _setitem( self.store_path, self.metadata, @@ -4816,6 +4838,23 @@ async def from_array( [0, 0]]) """ mode: Literal["a"] = "a" + # Reject lazy-view sources up front, before touching the store or creating the + # target array. The copy loop reads via `data.async_array.getitem`, which does + # not yet route through a view's index transform (it would copy the wrong + # region); guard here so the user gets one clear error and no half-created + # target is left behind. + _source_async: AsyncArray[Any] | None = None + if isinstance(data, Array): + _source_async = data._async_array + elif isinstance(data, AsyncArray): + _source_async = data + if _source_async is not None and not _source_async._is_identity: + raise LazyViewError( + "Creating an array from a lazy view (a sliced or indexed array with a " + "non-identity transform) via `from_array` is not yet supported. " + "Materialize the view first, e.g. `zarr.from_array(store, data=view.result())`, " + "or use the synchronous copy idiom." + ) config_parsed = parse_array_config(config) store_path = await make_store_path(store, path=name, mode=mode, storage_options=storage_options) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 26eb06c996..006dec36fc 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1029,6 +1029,7 @@ class AsyncOIndex[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: async def getitem(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: from zarr.core.array import Array + self.array._require_identity_for_async("oindex.getitem") # if input is a Zarr array, we materialize it now. if isinstance(selection, Array): selection = _zarr_array_to_int_or_bool_array(selection) @@ -1377,6 +1378,7 @@ async def getitem( # TODO requires solving this circular sync issue: https://github.com/zarr-developers/zarr-python/pull/3083#discussion_r2230737448 from zarr.core.array import Array + self.array._require_identity_for_async("vindex.getitem") # if input is a Zarr array, we materialize it now. if isinstance(selection, Array): selection = _zarr_array_to_int_or_bool_array(selection) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index efb77afc67..0719306815 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.core.sync import sync from zarr.errors import BoundsCheckError, LazyViewError from zarr.storage import MemoryStore @@ -1132,3 +1133,105 @@ def test_empty_oindex_read(self) -> None: a, _ = self._make((10,), (3,)) actual = a.lazy.oindex[np.array([], dtype=np.intp)].result() assert actual.shape == (0,) + + +# --------------------------------------------------------------------------- +# Async-surface guards +# --------------------------------------------------------------------------- +# +# The async surface (``AsyncArray.getitem`` / ``setitem`` / selection methods and +# the ``AsyncOIndex`` / ``AsyncVIndex`` accessors) does not yet route selections +# through a view's index transform: it builds a legacy indexer from +# ``metadata.shape``, which for a lazy view silently reads or writes the *wrong* +# region. Until async transform routing lands, every such entry point must raise +# ``LazyViewError`` for a non-identity view, steering users to the synchronous +# ``Array`` surface (or ``.result()``). Identity arrays must be unaffected. + + +def _read_ops() -> list[Any]: + """Async read entry points that build a legacy indexer, keyed by name.""" + coord = (np.array([0, 1], dtype=np.intp),) + mask = np.zeros((10,), dtype=bool) + mask[0] = mask[2] = True + return [ + pytest.param(lambda av: av.getitem(slice(None)), id="getitem"), + pytest.param( + lambda av: av.get_orthogonal_selection((slice(0, 3),)), id="get_orthogonal_selection" + ), + pytest.param(lambda av: av.get_coordinate_selection(coord), id="get_coordinate_selection"), + pytest.param(lambda av: av.get_mask_selection(mask), id="get_mask_selection"), + pytest.param(lambda av: av.oindex.getitem((slice(0, 3),)), id="oindex.getitem"), + pytest.param(lambda av: av.vindex.getitem(coord), id="vindex.getitem"), + ] + + +class TestAsyncSurfaceGuards: + def _view(self) -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + """A base array and a non-identity async view ``arr.lazy[2:12]``'s async_array.""" + a = zarr.create_array(MemoryStore(), shape=(24,), chunks=(4,), dtype="i4") + ref = np.arange(24, dtype="i4") + a[...] = ref + return a, ref + + @pytest.mark.parametrize("op", _read_ops()) + def test_read_entry_points_raise(self, op: Callable[[Any], Any]) -> None: + """Every async read entry point raises ``LazyViewError`` on a lazy view.""" + a, _ = self._view() + av = a.lazy[2:12]._async_array + assert not av._is_identity + with pytest.raises(LazyViewError): + sync(op(av)) + + def test_setitem_raises_and_leaves_base_unchanged(self) -> None: + """``AsyncArray.setitem`` on a view raises and writes nothing to the base.""" + a, ref = self._view() + av = a.lazy[5:10]._async_array + with pytest.raises(LazyViewError): + sync(av.setitem(slice(0, 5), 99)) + np.testing.assert_array_equal(a[...], ref) + + def test_from_array_view_explicit_kwargs_raises_and_writes_nothing(self) -> None: + """``from_array`` with a lazy-view source raises and leaves the target empty.""" + a, _ = self._view() + view = a.lazy[5:10] + target = MemoryStore() + with pytest.raises(LazyViewError): + zarr.from_array(target, data=view, chunks=(5,), shards=None, overwrite=False) + assert target._store_dict == {} + + def test_from_array_view_default_kwargs_raises_by_design(self) -> None: + """The default ('keep') kwarg path raises the same clear error, not by accident.""" + a, _ = self._view() + view = a.lazy[5:10] + target = MemoryStore() + with pytest.raises(LazyViewError): + zarr.from_array(target, data=view) + assert target._store_dict == {} + + def test_translate_by_async_view_getitem_raises(self) -> None: + """A view produced by ``AsyncArray.translate_by`` is guarded too.""" + a, _ = self._view() + av = a._async_array.translate_by((3,)) + assert not av._is_identity + with pytest.raises(LazyViewError): + sync(av.getitem(slice(None))) + + # --- identity regression guards: none of the above may affect eager arrays --- + + def test_identity_getitem_still_reads(self) -> None: + a, ref = self._view() + got = sync(a._async_array.getitem(slice(2, 12))) + np.testing.assert_array_equal(got, ref[2:12]) + + def test_identity_setitem_still_writes(self) -> None: + a, ref = self._view() + sync(a._async_array.setitem(slice(0, 5), 99)) + expected = ref.copy() + expected[0:5] = 99 + np.testing.assert_array_equal(a[...], expected) + + def test_identity_from_array_still_copies(self) -> None: + a, ref = self._view() + target = MemoryStore() + out = zarr.from_array(target, data=a) + np.testing.assert_array_equal(out[...], ref) From cd33fb0f9417f7b161258714f6d6ee60a0dd83c8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 09:50:29 +0200 Subject: [PATCH 52/60] fix(indexing): correct sharded detection, chunk-coverage parity, and out-of-grid guardrail `_is_sharded` keyed off `len(codecs) == 1`, misclassifying a valid sharded array with a trailing bytes-bytes codec (e.g. an outer compressor) as unsharded, so `chunk_projections(unit="read")` silently returned shard-granularity projections instead of raising. `_covers_full_chunk` returned partial for any integer chunk_selection entry, diverging from the write path's `_is_complete_chunk` for a constant over a chunk dim of extent 1. The transform read/write loops and `iter_chunk_projections` silently `continue`d on an out-of-grid chunk coordinate instead of raising, risking uninitialized-buffer reads and silently dropped writes if a transform composition bug ever produced one. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 16 ++- src/zarr/core/chunk_partition.py | 26 +++- tests/test_lazy_indexing.py | 213 ++++++++++++++++++++++++++++++- 3 files changed, 243 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 44d6cc4985..b3fa7d627d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -902,11 +902,19 @@ def _require_identity_for_async(self, name: str) -> None: @property def _is_sharded(self) -> bool: - """Whether the array stores inner chunks inside shards (a sharding codec).""" + """Whether the array stores inner chunks inside shards (a sharding codec). + + A sharding codec is the array-bytes codec in the chain: array-array + filters may precede it and bytes-bytes codecs (e.g. an outer compressor) + may follow it, so its presence — not the chain's length — is what makes + an array sharded. `validate_codecs` permits such trailing codecs, and + they round-trip data correctly; keying off `len(codecs) == 1` would + misclassify them as unsharded. + """ from zarr.codecs.sharding import ShardingCodec codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) - return len(codecs) == 1 and isinstance(codecs[0], ShardingCodec) + return any(isinstance(codec, ShardingCodec) for codec in codecs) def chunk_projections( self, *, unit: Literal["read", "write"] = "read" @@ -6160,7 +6168,7 @@ async def _get_selection_via_transform( ): chunk_spec = chunk_grid[chunk_coords] if chunk_spec is None: - continue + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) drop_axes = da # same for all chunks batch_info.append( @@ -6311,7 +6319,7 @@ async def _set_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 + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) drop_axes = da # same for all chunks batch_info.append( diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py index 712921928d..61f98db9f8 100644 --- a/src/zarr/core/chunk_partition.py +++ b/src/zarr/core/chunk_partition.py @@ -60,16 +60,28 @@ class ChunkProjection: 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. + A per-dimension full `0:size:1` slice is a full cover. An integer entry is + also a full cover *iff* the chunk's extent along that dimension is 1 — fixing + the only element of a length-1 dimension covers it exactly, the same as an + equivalent length-1 slice would. This mirrors the write path's + `_is_complete_chunk` (`ConstantMap` over a dimension of size 1 is complete); + keeping the two in sync matters because a divergence would make + `chunk_projections` plan needless read-modify-writes for coverage the write + path already treats as complete. An array selection always touches a subset + (or requires per-element comparison this cheap check does not attempt), so it + is never a full cover. """ 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): + if isinstance(sel, slice): + start, stop, step = sel.indices(size) + if not (start == 0 and stop == size and step == 1): + return False + elif isinstance(sel, int): + if size > 1: + return False + else: return False return True @@ -89,7 +101,7 @@ def iter_chunk_projections( 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 + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") chunk_selection, array_selection, _drop_axes = sub_transform_to_selections( sub_transform, out_indices ) diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 0719306815..5ee2b505c7 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -22,9 +22,15 @@ import pytest import zarr +from zarr.codecs.bytes import BytesCodec +from zarr.codecs.gzip import GzipCodec +from zarr.codecs.sharding import ShardingCodec from zarr.core.buffer import default_buffer_prototype from zarr.core.sync import sync -from zarr.errors import BoundsCheckError, LazyViewError +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap +from zarr.core.transforms.transform import IndexTransform +from zarr.errors import BoundsCheckError, LazyViewError, ZarrUserWarning from zarr.storage import MemoryStore if TYPE_CHECKING: @@ -1021,6 +1027,29 @@ def test_sharded_write_unit(self) -> None: with pytest.raises(NotImplementedError): a.chunk_projections(unit="read") + def test_multi_codec_sharded_read_unit_raises(self) -> None: + """A sharded array whose codec chain has more than one entry (a `ShardingCodec` + followed by a trailing bytes-bytes codec) is still sharded: `unit='read'` must + raise the same `NotImplementedError` as a plain single-codec sharded array.""" + with pytest.warns(ZarrUserWarning, match="sharding_indexed"): + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ + ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()]), + GzipCodec(), + ], + ) + 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,)) @@ -1028,6 +1057,188 @@ def test_invalid_unit_rejected(self) -> None: a.chunk_projections(unit="bogus") # type: ignore[arg-type] +class TestIsSharded: + """`_is_sharded` must key off the presence of a `ShardingCodec` in the codec + chain, not the chain's length: a valid sharded array may carry additional + codecs (e.g. a trailing bytes-bytes compressor after the sharding codec).""" + + def test_true_for_multi_codec_sharded_chain(self) -> None: + with pytest.warns(ZarrUserWarning, match="sharding_indexed"): + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ + ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()]), + GzipCodec(), + ], + ) + assert a._async_array._is_sharded is True + + def test_true_for_single_codec_sharded_chain(self) -> None: + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()])], + ) + assert a._async_array._is_sharded is True + + def test_false_for_plain_bytes_chain(self) -> None: + a = zarr.create( + store={}, shape=(12,), chunks=(6,), dtype="i4", zarr_format=3, codecs=[BytesCodec()] + ) + assert a._async_array._is_sharded is False + + def test_false_for_plain_compressed_chain(self) -> None: + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[BytesCodec(), GzipCodec()], + ) + assert a._async_array._is_sharded is False + + +class TestChunkCoverageParity: + """`chunk_projections`' `is_partial` (backed by `_covers_full_chunk`) must agree + with the write path's `_is_complete_chunk`: both derive coverage from the same + `iter_chunk_transforms` enumeration and must reach the same verdict — including + the case a naive selection-kind check misses: an integer index on a + chunk-extent-1 dimension covers that dimension exactly as its equivalent + length-1 slice does.""" + + @pytest.mark.parametrize( + "sel", + [ + (slice(None), slice(None)), # full slice, both dims -> aligned + (slice(1, 3), slice(None)), # partial slice on dim 0 (chunk extent 1) + (2, slice(None)), # int on a chunk-extent-1 dim -> aligned + (slice(2, 3), slice(None)), # equivalent length-1 slice -> aligned + (slice(None), 2), # int on a chunk-extent-4 dim -> partial + (slice(None), slice(2, 3)), # length-1 slice on the same dim -> partial + ], + ) + def test_is_partial_matches_is_complete_chunk(self, sel: tuple[Any, ...]) -> None: + """For every stored chunk a selection touches, `is_partial` (computed via + `chunk_projections`) must equal `not _is_complete_chunk(...)` computed + directly from the same `iter_chunk_transforms` enumeration.""" + from zarr.core.array import _is_complete_chunk + from zarr.core.transforms.chunk_resolution import iter_chunk_transforms + + a = zarr.create_array({}, shape=(4, 4), chunks=(1, 4), dtype="i4") + a[...] = np.arange(16, dtype="i4").reshape(4, 4) + view = a.lazy[sel] + async_view = view._async_array + transform = async_view._transform.translate_domain_to( + (0,) * async_view._transform.input_rank + ) + chunk_grid = async_view._chunk_grid + + projections = {p.coord: p for p in view.chunk_projections()} + assert len(projections) > 0 + + seen: set[tuple[int, ...]] = set() + for chunk_coords, sub_transform, _out_indices in iter_chunk_transforms( + transform, chunk_grid + ): + spec = chunk_grid[chunk_coords] + if spec is None: + continue + seen.add(chunk_coords) + expected_partial = not _is_complete_chunk(sub_transform, spec) + assert projections[chunk_coords].is_partial == expected_partial + + assert seen == set(projections) + + def test_int_and_length1_slice_report_identically(self) -> None: + """An integer selection and its equivalent length-1 slice, on a + chunk-extent-1 dimension, must report the same is_partial/aligned verdict — + both cover the same data, so a read-modify-write plan should not depend on + which selection spelling was used.""" + a = zarr.create_array({}, shape=(4, 4), chunks=(1, 4), dtype="i4") + a[...] = np.arange(16, dtype="i4").reshape(4, 4) + + int_view = a.lazy[2, :] + slice_view = a.lazy[2:3, :] + + int_partial = [p.is_partial for p in int_view.chunk_projections()] + slice_partial = [p.is_partial for p in slice_view.chunk_projections()] + assert int_partial == slice_partial + assert int_view.is_chunk_aligned() == slice_view.is_chunk_aligned() + assert int_view.is_chunk_aligned() is True # chunk extent along dim 0 is 1 + + +class TestOutOfGridGuardrail: + """The transform I/O loops must raise loudly on an out-of-grid chunk + coordinate rather than silently `continue`. A read buffer is allocated with + `nd_buffer.empty()` (uninitialized), so a bug anywhere in transform + composition that yields an out-of-grid storage coordinate would otherwise + return garbage memory on reads and silently drop data on writes. + + The public indexing APIs already bounds-check selections before a transform + reaches chunk resolution, so an out-of-grid coordinate can only arise from a + transform built directly (simulating a composition bug upstream). Each test + hand-builds an `IndexTransform` whose `ArrayMap` storage coordinates + exceed the backing array's chunk grid. + """ + + @staticmethod + def _out_of_grid_transform() -> IndexTransform: + """For shape=(20,), chunks=(5,): storage coords 1 and 2 are in-grid (chunk + 0), but 23 and 24 exceed the array's extent and fall in chunk id 4, which + does not exist (the grid only has chunks 0-3).""" + idx = np.array([1, 2, 23, 24], dtype=np.intp) + return IndexTransform( + domain=IndexDomain.from_shape((4,)), output=(ArrayMap(index_array=idx),) + ) + + def test_get_selection_via_transform_raises(self) -> None: + """The read pipeline (`_get_selection_via_transform`) must raise instead + of silently returning whatever was in the uninitialized output buffer.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + a[...] = np.arange(20, dtype="i4") + transform = self._out_of_grid_transform() + with pytest.raises(IndexError, match="out of bounds"): + sync(a._async_array._get_selection_t(transform, prototype=default_buffer_prototype())) + + def test_set_selection_via_transform_raises(self) -> None: + """The write pipeline (`_set_selection_via_transform`) must raise instead + of silently dropping the out-of-grid entries -- and must not apply a + partial write for the in-grid entries either.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + original = np.arange(20, dtype="i4") + a[...] = original + transform = self._out_of_grid_transform() + value = np.array([100, 200, 300, 400], dtype="i4") + with pytest.raises(IndexError, match="out of bounds"): + sync( + a._async_array._set_selection_t( + transform, value, prototype=default_buffer_prototype() + ) + ) + # The batch is built (and validated) before any write is issued, so a + # rejected batch must leave the array untouched -- not even the in-grid + # entries (indices 1, 2) should have been written. + np.testing.assert_array_equal(a[...], original) + + def test_iter_chunk_projections_raises(self) -> None: + """`chunk_projections` (backed by `iter_chunk_projections`) must raise + instead of silently omitting the out-of-grid chunk from the enumeration.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + a[...] = np.arange(20, dtype="i4") + transform = self._out_of_grid_transform() + view = a._with_transform(transform) + with pytest.raises(IndexError, match="out of bounds"): + list(view.chunk_projections()) + + class TestLazyArrayMapResolution: """Chunk resolution of orthogonal (outer-product) and correlated (vectorized) ArrayMaps: intersections and output selectors must preserve the distinct From 06ff95ab23139965296fea10eeed225481efdc4a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 10:15:02 +0200 Subject: [PATCH 53/60] =?UTF-8?q?chore:=20review=20housekeeping=20?= =?UTF-8?q?=E2=80=94=20truthiness,=20mkdocs=20docstrings,=20changelog,=200?= =?UTF-8?q?-d=20pin,=20ChunkProjection=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address five small code-review items on the lazy-indexing branch: - Convert branch-introduced bare collection truthiness checks to explicit len() comparisons (array.py, transforms/transform.py, testing/strategies.py, test_properties.py). - Convert branch-introduced RST-style ``double-backtick`` markup and :role:`target` references to plain markdown backticks, since docs render via mkdocs/mkdocstrings (transforms/*.py, chunk_partition.py, errors.py, array.py). - Add changes/3906.feature.md documenting the new .lazy accessor and chunk_projections. - Pin the intentional 0-d array iteration behavior change (raises TypeError eagerly at iter(), matching numpy, instead of the old silent empty sequence-protocol iteration). - Document zarr.ChunkProjection under docs/api/zarr/array.md, which was exported in zarr.__all__ but missing from the API docs. Assisted-by: ClaudeCode:claude-fable-5 --- changes/3906.feature.md | 1 + docs/api/zarr/array.md | 1 + src/zarr/core/array.py | 38 ++++---- src/zarr/core/chunk_partition.py | 12 +-- src/zarr/core/transforms/__init__.py | 10 +-- src/zarr/core/transforms/chunk_resolution.py | 32 +++---- src/zarr/core/transforms/composition.py | 6 +- src/zarr/core/transforms/domain.py | 10 +-- src/zarr/core/transforms/json.py | 10 +-- src/zarr/core/transforms/output_map.py | 36 ++++---- src/zarr/core/transforms/transform.py | 92 ++++++++++---------- src/zarr/errors.py | 10 +-- src/zarr/testing/strategies.py | 2 +- tests/test_array.py | 13 +++ tests/test_properties.py | 2 +- 15 files changed, 145 insertions(+), 130 deletions(-) create mode 100644 changes/3906.feature.md diff --git a/changes/3906.feature.md b/changes/3906.feature.md new file mode 100644 index 0000000000..bc8518bd0f --- /dev/null +++ b/changes/3906.feature.md @@ -0,0 +1 @@ +Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged. Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. diff --git a/docs/api/zarr/array.md b/docs/api/zarr/array.md index ff61cb1fe2..2445ee58d3 100644 --- a/docs/api/zarr/array.md +++ b/docs/api/zarr/array.md @@ -1,2 +1,3 @@ ::: zarr.Array ::: zarr.AsyncArray +::: zarr.ChunkProjection diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b3fa7d627d..0e6907cdd5 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -858,10 +858,10 @@ 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 + `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`) + transform resolution. Lazy views (created via `_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. """ @@ -4306,8 +4306,8 @@ 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, and forwards - ``prototype``. Named after `tensorstore.Future.result`. + Equivalent to `self[...]` but more explicit for lazy views, and forwards + `prototype`. Named after `tensorstore.Future.result`. """ return self.get_basic_selection(Ellipsis, prototype=prototype) @@ -4572,7 +4572,7 @@ def _wrap_negative_index(item: Any, exclusive_max: int) -> Any: class _LazyOIndex: - """Lazy orthogonal indexing via ``array.lazy.oindex[...]``.""" + """Lazy orthogonal indexing via `array.lazy.oindex[...]`.""" __slots__ = ("_array",) @@ -4591,7 +4591,7 @@ def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: class _LazyVIndex: - """Lazy vectorized indexing via ``array.lazy.vindex[...]``.""" + """Lazy vectorized indexing via `array.lazy.vindex[...]`.""" __slots__ = ("_array",) @@ -4610,7 +4610,7 @@ def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: class _LazyIndexAccessor: - """Provides lazy indexing via ``array.lazy[...]``.""" + """Provides lazy indexing via `array.lazy[...]`.""" __slots__ = ("_array",) @@ -5990,9 +5990,9 @@ def _array_spec_from_chunk_spec( ) -> ArraySpec: """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`. + Split out from `_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 `_is_complete_chunk`. """ return ArraySpec( shape=spec.codec_shape, @@ -6020,23 +6020,23 @@ def _get_chunk_spec( def _is_vectorized_transform(transform: IndexTransform) -> bool: """Return True for a vectorized (vindex/coordinate) transform. - Coordinate ArrayMaps (``input_dimension is None``) are jointly indexed over + 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), + 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]``). + (e.g. `oindex[i0, i1]`). """ 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: - """Return True if ``transform`` is the identity over the full storage domain. + """Return True if `transform` is the identity over the full storage domain. - An identity transform maps input coordinate ``i`` to storage coordinate ``i`` + 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 @@ -6063,8 +6063,8 @@ def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, 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`). + `spec` is the chunk's already-resolved `ChunkSpec` (the caller looks + it up once and shares it with `_array_spec_from_chunk_spec`). """ from zarr.core.transforms.output_map import ConstantMap, DimensionMap @@ -6190,7 +6190,7 @@ async def _get_selection_via_transform( if result["status"] == "missing": coords_path = batch_info[i][0] missing_info.append(f" chunk at '{coords_path}'") - if missing_info: + if len(missing_info) > 0: chunks_str = "\n".join(missing_info) raise ChunkNotFoundError( f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" @@ -6448,7 +6448,7 @@ async def _get_selection( coords = indexed_chunks[i][0] key = metadata.encode_chunk_key(coords) missing_info.append(f" chunk '{key}' (grid position {coords})") - if missing_info: + if len(missing_info) > 0: chunks_str = "\n".join(missing_info) raise ChunkNotFoundError( f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py index 61f98db9f8..d165cdba88 100644 --- a/src/zarr/core/chunk_partition.py +++ b/src/zarr/core/chunk_partition.py @@ -1,10 +1,10 @@ """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 +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 +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). """ @@ -58,7 +58,7 @@ class ChunkProjection: def _covers_full_chunk(chunk_selection: tuple[Any, ...], shape: tuple[int, ...]) -> bool: - """Whether ``chunk_selection`` selects every element of a chunk of ``shape``. + """Whether `chunk_selection` selects every element of a chunk of `shape`. A per-dimension full `0:size:1` slice is a full cover. An integer entry is also a full cover *iff* the chunk's extent along that dimension is 1 — fixing @@ -91,7 +91,7 @@ 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 diff --git a/src/zarr/core/transforms/__init__.py b/src/zarr/core/transforms/__init__.py index ec98e02d4d..283ff834e6 100644 --- a/src/zarr/core/transforms/__init__.py +++ b/src/zarr/core/transforms/__init__.py @@ -7,11 +7,11 @@ 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 +- `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 diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py index 0b045d39cd..21d9e509b2 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -1,7 +1,7 @@ """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: +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? @@ -12,17 +12,17 @@ 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 + `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)``. + via `transform.translate(-chunk_origin)`. -4. **Yield** — produce ``(chunk_coords, local_transform, surviving_indices)`` +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 +`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. """ @@ -59,12 +59,12 @@ def iter_chunk_transforms( ) -> Iterator[ChunkTransformResult]: """Resolve a composed IndexTransform against a ChunkGrid. - Yields ``(chunk_coords, sub_transform, out_indices)`` triples: + 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. + - `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 @@ -161,7 +161,7 @@ def sub_transform_to_selections( Returns ------- tuple - ``(chunk_selection, out_selection, drop_axes)`` + `(chunk_selection, out_selection, drop_axes)` """ inclusive_min = sub_transform.domain.inclusive_min exclusive_max = sub_transform.domain.exclusive_max @@ -190,8 +190,8 @@ def sub_transform_to_selections( out_arrays.append(out_indices[out_dim]) return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) - # Correlated (vindex) sub-transforms carry ArrayMaps with ``input_dimension`` - # None. They scatter through a single flat index (``out_indices``) into the + # Correlated (vindex) sub-transforms carry ArrayMaps with `input_dimension` + # None. They scatter through a single flat index (`out_indices`) into the # row-major-flattened output buffer; the chunk selection reads a # (points, residual-slice) block via the raveled coordinate arrays and any # residual DimensionMap slices. @@ -216,7 +216,7 @@ def sub_transform_to_selections( # Chunk resolution always supplies the flat scatter index for a # correlated transform. Absent one (a bare sub-transform), fall back to an # identity scatter over the whole flattened output buffer. - # ``out_indices`` is narrowed to a flat scatter array or None here (the + # `out_indices` is narrowed to a flat scatter array or None here (the # per-dimension dict is an orthogonal outer product, handled above). out_scatter: slice | np.ndarray[Any, np.dtype[np.intp]] if out_indices is None: diff --git a/src/zarr/core/transforms/composition.py b/src/zarr/core/transforms/composition.py index e33a94c470..865478abf0 100644 --- a/src/zarr/core/transforms/composition.py +++ b/src/zarr/core/transforms/composition.py @@ -9,11 +9,11 @@ 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). + `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``. + Precondition: `outer.output_rank == inner.domain.ndim`. """ if outer.output_rank != inner.domain.ndim: raise ValueError( diff --git a/src/zarr/core/transforms/domain.py b/src/zarr/core/transforms/domain.py index bca6488ad1..f20d5bf7bd 100644 --- a/src/zarr/core/transforms/domain.py +++ b/src/zarr/core/transforms/domain.py @@ -1,13 +1,13 @@ """Index domains — rectangular regions in N-dimensional integer space. -An ``IndexDomain`` represents the set of valid coordinates for an array or +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 +`arr[5:10]`, the result has origin 5 and shape 5 — coordinates 5 through 9 are valid. This follows the TensorStore convention. """ @@ -22,7 +22,7 @@ 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[d], exclusive_max[d])` for each dimension `d`. """ inclusive_min: tuple[int, ...] @@ -30,8 +30,8 @@ class IndexDomain: 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 ``()``. + # 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: diff --git a/src/zarr/core/transforms/json.py b/src/zarr/core/transforms/json.py index 11e824a61c..78c09664b5 100644 --- a/src/zarr/core/transforms/json.py +++ b/src/zarr/core/transforms/json.py @@ -31,10 +31,10 @@ # TypedDict definitions (JSON shapes) # --------------------------------------------------------------------------- -# An ``index_array`` serializes via ``ndarray.tolist()``, so it is a nested list +# An `index_array` serializes via `ndarray.tolist()`, so it is a nested list # of ints whose nesting depth equals the array rank. Normalized arrays carry the # full input rank (with singleton axes), so the nesting can be arbitrarily deep — -# ``list[int] | list[list[int]]`` would under-describe a 3-D+ orthogonal array. +# `list[int] | list[list[int]]` would under-describe a 3-D+ orthogonal array. NestedIntList = list[Any] @@ -50,9 +50,9 @@ 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": 5}` — constant + - `{"offset": 0, "stride": 1, "input_dimension": 0}` — dimension + - `{"offset": 0, "stride": 1, "index_array": [...]}` — array """ offset: int diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py index 80696a6c65..581229bd22 100644 --- a/src/zarr/core/transforms/output_map.py +++ b/src/zarr/core/transforms/output_map.py @@ -4,18 +4,18 @@ 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}`` +- `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): +`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}`` + `{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}`` + `{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 @@ -23,13 +23,13 @@ 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) +- `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 +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. +them by chunk, when `DimensionMap` does it with three integers. """ from __future__ import annotations @@ -46,7 +46,7 @@ class ConstantMap: """A singleton set: one storage coordinate. - Represents ``{offset}``. Arises from integer indexing (e.g., ``arr[5]`` + Represents `{offset}`. Arises from integer indexing (e.g., `arr[5]` fixes one dimension to coordinate 5). """ @@ -57,9 +57,9 @@ class ConstantMap: 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). + 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 @@ -71,8 +71,8 @@ class DimensionMap: 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). + Represents `{offset + stride * index_array[i] : i in input_range}`. + Arises from fancy indexing (e.g., `arr[[1, 5, 9]]` or boolean masks). Freshly constructed maps are normalized to the **full input rank** of their enclosing transform: `index_array` has the enclosing domain's rank, sized diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py index 98014196af..5d7add487c 100644 --- a/src/zarr/core/transforms/transform.py +++ b/src/zarr/core/transforms/transform.py @@ -1,13 +1,13 @@ """Index transforms — composable, lazy coordinate mappings. -An ``IndexTransform`` pairs an **input domain** (the coordinates a user sees) +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 +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]``) — +- **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. @@ -18,11 +18,11 @@ - **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``. +- **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.lazy[...]`` composes a new transform lazily. Reading resolves the +chunk-level I/O. Every `Array` holds a transform (identity by default). +`Array.lazy[...]` composes a new transform lazily. Reading resolves the transform against the chunk grid via intersect + translate. """ @@ -43,15 +43,15 @@ class IndexTransform: """A composable mapping from input coordinates to storage coordinates. - An ``IndexTransform`` has: + An `IndexTransform` has: - - ``domain``: an ``IndexDomain`` describing the valid input coordinates + - `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 + - `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 + coordinate `i` maps to storage coordinate `i`. Indexing operations compose new transforms without I/O. """ @@ -99,10 +99,10 @@ def from_shape(cls, shape: tuple[int, ...]) -> IndexTransform: @property def selection_repr(self) -> str: - """Compact domain string, e.g. ``'{ [2, 8), [0, 10) }'``. + """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. + 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. """ @@ -155,9 +155,9 @@ def intersect( ): """Restrict this transform to storage coordinates within output_domain. - Returns ``(restricted_transform, out_indices)`` or None if empty. + Returns `(restricted_transform, out_indices)` or None if empty. - ``out_indices`` carries the surviving output positions: ``None`` when all + `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). @@ -165,7 +165,7 @@ def intersect( return _intersect(self, output_domain) def translate(self, shift: tuple[int, ...]) -> IndexTransform: - """Shift all output coordinates by ``shift``.""" + """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] = [] @@ -195,10 +195,10 @@ 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. + """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`` + 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. """ @@ -223,9 +223,9 @@ def translate_domain_by(self, shift: tuple[int, ...]) -> IndexTransform: 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``. + """Move the input domain so its per-dimension origins equal `origins`. - TensorStore's ``translate_to``; ``translate_domain_to((0,) * rank)`` + TensorStore's `translate_to`; `translate_domain_to((0,) * rank)` re-zeros a view's coordinate system without changing which cells it addresses. """ @@ -255,7 +255,7 @@ def _intersect( """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])``. + `[output_domain.inclusive_min[d], output_domain.exclusive_max[d])`. Two flavours of fancy indexing require different treatment, distinguished by the ArrayMaps' dependency axes (see `_array_map_dependency_axes`): @@ -271,7 +271,7 @@ def _intersect( A `None` `input_dimension` marks a correlated map, so any such map routes the whole transform through the correlated intersection. - Returns ``None`` if the intersection is empty. + Returns `None` if the intersection is empty. """ if output_domain.ndim != transform.output_rank: raise ValueError( @@ -284,7 +284,7 @@ def _intersect( for i, m in enumerate(transform.output) if isinstance(m, ArrayMap) and m.input_dimension is None ] - if correlated_dims: + if len(correlated_dims) > 0: return _intersect_correlated(transform, output_domain, correlated_dims) return _intersect_orthogonal(transform, output_domain) @@ -292,10 +292,10 @@ def _intersect( def _intersect_dimension_map( m: DimensionMap, input_lo: int, input_hi: int, lo: int, hi: int ) -> tuple[int, int] | None: - """Narrow a DimensionMap's input range to storage coordinates in ``[lo, hi)``. + """Narrow a DimensionMap's input range to storage coordinates in `[lo, hi)`. - ``input_lo``/``input_hi`` are the current (possibly already narrowed) input - range for the map's axis. Returns the new ``(input_lo, input_hi)`` or ``None`` + `input_lo`/`input_hi` are the current (possibly already narrowed) input + range for the map's axis. Returns the new `(input_lo, input_hi)` or `None` if no input produces an in-bounds storage coordinate. """ if input_lo >= input_hi: @@ -362,8 +362,8 @@ def _intersect_orthogonal( d = _array_map_dependent_axis(m) storage = m.offset + m.stride * m.index_array mask = (storage >= lo) & (storage < hi) - # The array is singleton on every axis but ``d``, so its mask reduces - # to a 1-D vector along ``d``. + # The array is singleton on every axis but `d`, so its mask reduces + # to a 1-D vector along `d`. survivors = np.nonzero(mask.reshape(-1))[0].astype(np.intp) if survivors.size == 0: return None @@ -414,8 +414,8 @@ def _intersect_correlated( arrays over a 3-D array, leaving one slice dimension — resolves correctly. The surviving broadcast axes collapse to a single axis; the returned - ``out_indices`` is the flat scatter index into the (row-major flattened) - output buffer, of shape ``(surviving_points,) + (residual slice sizes)``. + `out_indices` is the flat scatter index into the (row-major flattened) + output buffer, of shape `(surviving_points,) + (residual slice sizes)`. """ corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] @@ -426,7 +426,7 @@ def _intersect_correlated( for i, m in enumerate(transform.output) if isinstance(m, ArrayMap) and m.input_dimension is not None ] - if orthogonal_array_dims: + if len(orthogonal_array_dims) > 0: raise NotImplementedError( "intersecting a transform with both correlated and orthogonal " "ArrayMaps is not supported" @@ -809,9 +809,9 @@ def _array_map_dependent_axis(m: ArrayMap) -> int: def _reshape_to_axis( values: np.ndarray[Any, np.dtype[np.intp]], axis: int, ndim: int ) -> np.ndarray[Any, np.dtype[np.intp]]: - """Reshape a 1-D selection to full rank ``ndim`` varying only along ``axis``. + """Reshape a 1-D selection to full rank `ndim` varying only along `axis`. - The result has ``values`` laid out along ``axis`` and singleton (size-1) axes + The result has `values` laid out along `axis` and singleton (size-1) axes everywhere else, so its dependency axis is derivable from its shape. """ flat = np.asarray(values, dtype=np.intp).ravel() @@ -821,7 +821,7 @@ def _reshape_to_axis( class _OIndexHelper: - """Helper that provides orthogonal (outer) indexing via ``transform.oindex[...]``.""" + """Helper that provides orthogonal (outer) indexing via `transform.oindex[...]`.""" def __init__(self, transform: IndexTransform) -> None: self._transform = transform @@ -967,7 +967,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: class _VIndexHelper: - """Helper that provides vectorized (fancy) indexing via ``transform.vindex[...]``.""" + """Helper that provides vectorized (fancy) indexing via `transform.vindex[...]`.""" def __init__(self, transform: IndexTransform) -> None: self._transform = transform @@ -1047,7 +1047,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: # Broadcast all arrays together broadcast_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] - if arrays: + if len(arrays) > 0: broadcast_arrays = list(np.broadcast_arrays(*arrays)) broadcast_shape = broadcast_arrays[0].shape else: @@ -1143,7 +1143,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: 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``).""" + for negative operands (`trunc(-9/2) == -4` where `-9 // 2 == -5`).""" q = a // b if q < 0 and q * b != a: q += 1 @@ -1151,22 +1151,22 @@ def _trunc_div(a: int, b: int) -> int: 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. + """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): - - defaults: ``start = lo``, ``stop = hi``; + - 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** 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``. + - the result's domain origin is `trunc(start/step)` (rounded toward + zero) and coordinate `origin + k` maps to input `start + k*step`. - Returns ``(start, step, origin, size)`` in domain coordinates. + Returns `(start, step, origin, size)` in domain coordinates. """ step = 1 if sel.step is None else sel.step if step <= 0: @@ -1192,10 +1192,10 @@ def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, 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)``. + """Reject index-array values outside the domain `[lo, hi)`. Index-array values are literal domain coordinates (TensorStore semantics): - a value below ``inclusive_min`` is out of bounds rather than counting from + 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: diff --git a/src/zarr/errors.py b/src/zarr/errors.py index b4480ef035..53e8e8535c 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -154,15 +154,15 @@ class ArrayIndexError(IndexError): ... class LazyViewError(NotImplementedError, AttributeError): """Raised when an operation that assumes an array fills its chunk grid is used - on a non-identity lazy view (created via ``Array.lazy[...]``). + 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 + 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_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, and - ``AttributeError`` so duck-typing probes (`hasattr(view, "chunks")`, + `NotImplementedError` so existing consumers that catch it keep working, and + `AttributeError` so duck-typing probes (`hasattr(view, "chunks")`, `getattr(x, "chunks", None)` — e.g. `dask.array.from_array`) treat guarded members as absent instead of crashing. """ diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index a37013bbf5..289b3f166e 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -789,7 +789,7 @@ def indexing_programs(draw: st.DrawFn, *, shape: tuple[int, ...]) -> IndexingPro ) operations.append(IndexOperation("vectorized", idx)) - if not operations: + if len(operations) == 0: operations.append(IndexOperation("basic", draw(basic_indices(shape=cur)))) executions: tuple[ProgramExecution, ...] = ( diff --git a/tests/test_array.py b/tests/test_array.py index 991120e9a0..1991c5ead0 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1704,6 +1704,19 @@ def test_scalar_array(value: Any, zarr_format: ZarrFormat) -> None: assert isinstance(arr[()], NDArrayLikeOrScalar) +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_iter_0d_array_raises(zarr_format: ZarrFormat) -> None: + """Iterating a 0-d array raises TypeError, matching NumPy. + + This pins an intentional behavior change: previously `Array.__iter__` fell + through to the generic (length-based) sequence protocol and silently produced + an empty iterator for a 0-d array; it now raises eagerly, like `numpy.ndarray`. + """ + arr = zarr.array(np.array(1), zarr_format=zarr_format) + with pytest.raises(TypeError, match="iteration over a 0-d array"): + iter(arr) + + @pytest.mark.parametrize("store", ["local"], indirect=True) @pytest.mark.parametrize("store2", ["local"], indirect=["store2"]) @pytest.mark.parametrize("src_format", [2, 3]) diff --git a/tests/test_properties.py b/tests/test_properties.py index 6f111e1a6f..da97a0cc96 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -415,7 +415,7 @@ def _write_is_unambiguous(mode: IndexMode, zsel: Any, shape: tuple[int, ...]) -> for axis, s in enumerate(sel) if isinstance(s, np.ndarray) ] - if not norm: + if len(norm) == 0: 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)) From 7ea8899efefd60be13ed2d4de948bb9ed06eacaa Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 11:20:37 +0200 Subject: [PATCH 54/60] refactor(transforms): extract index-transform algebra into zarr-transforms package Move the TensorStore-style index-transform algebra out of zarr core into a new numpy-only uv-workspace subpackage packages/zarr-transforms (import name zarr_transforms), mirroring packages/zarr-metadata. This is a pure move: behavior is byte-for-byte unchanged and the full suite stays green. The package depends on numpy + stdlib only and must not import zarr. Two couplings are broken: - errors: the canonical BoundsCheckError / VindexInvalidSelectionError class definitions now live in zarr_transforms.errors (both still subclass IndexError). zarr.errors re-exports the same objects, so zarr.errors.BoundsCheckError is zarr_transforms.errors.BoundsCheckError and every existing catch site is unaffected. - chunk grid: chunk_resolution is typed against new structural Protocols (ChunkGridLike / DimensionGridLike) in zarr_transforms.grid instead of importing zarr's concrete ChunkGrid, which satisfies them structurally. zarr now declares a runtime dependency on zarr-transforms, resolved to the in-tree package via a uv workspace source. The package tests are collected by the root test suite (they exercise chunk resolution against zarr's ChunkGrid, so they need zarr importable) rather than run in isolation. The package __init__ promotes the names the zarr integration layer consumes (iter_chunk_transforms, sub_transform_to_selections, selection_to_transform) to the public surface alongside the existing exports. Assisted-by: ClaudeCode:claude-opus-4.8 --- packages/zarr-transforms/CHANGELOG.md | 3 + packages/zarr-transforms/LICENSE.txt | 21 ++++ packages/zarr-transforms/README.md | 30 +++++ packages/zarr-transforms/changes/README.md | 25 +++++ packages/zarr-transforms/pyproject.toml | 103 ++++++++++++++++++ .../src/zarr_transforms}/__init__.py | 26 ++++- .../src/zarr_transforms}/chunk_resolution.py | 10 +- .../src/zarr_transforms}/composition.py | 4 +- .../src/zarr_transforms}/domain.py | 0 .../src/zarr_transforms/errors.py | 21 ++++ .../src/zarr_transforms/grid.py | 39 +++++++ .../src/zarr_transforms}/json.py | 6 +- .../src/zarr_transforms}/output_map.py | 0 .../src/zarr_transforms/py.typed | 0 .../src/zarr_transforms}/transform.py | 6 +- .../tests}/test_chunk_resolution.py | 10 +- .../tests}/test_composition.py | 8 +- .../zarr-transforms/tests}/test_domain.py | 2 +- .../zarr-transforms/tests}/test_json.py | 8 +- .../zarr-transforms/tests}/test_output_map.py | 2 +- .../tests}/test_tensorstore_parity.py | 8 +- .../zarr-transforms/tests}/test_transform.py | 12 +- pyproject.toml | 18 ++- src/zarr/core/array.py | 16 +-- src/zarr/core/chunk_partition.py | 5 +- src/zarr/errors.py | 11 +- tests/test_lazy_indexing.py | 9 +- uv.lock | 26 +++++ 28 files changed, 365 insertions(+), 64 deletions(-) create mode 100644 packages/zarr-transforms/CHANGELOG.md create mode 100644 packages/zarr-transforms/LICENSE.txt create mode 100644 packages/zarr-transforms/README.md create mode 100644 packages/zarr-transforms/changes/README.md create mode 100644 packages/zarr-transforms/pyproject.toml rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/__init__.py (56%) rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/chunk_resolution.py (97%) rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/composition.py (96%) rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/domain.py (100%) create mode 100644 packages/zarr-transforms/src/zarr_transforms/errors.py create mode 100644 packages/zarr-transforms/src/zarr_transforms/grid.py rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/json.py (96%) rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/output_map.py (100%) rename tests/test_transforms/__init__.py => packages/zarr-transforms/src/zarr_transforms/py.typed (100%) rename {src/zarr/core/transforms => packages/zarr-transforms/src/zarr_transforms}/transform.py (99%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_chunk_resolution.py (97%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_composition.py (96%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_domain.py (99%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_json.py (97%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_output_map.py (94%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_tensorstore_parity.py (97%) rename {tests/test_transforms => packages/zarr-transforms/tests}/test_transform.py (98%) diff --git a/packages/zarr-transforms/CHANGELOG.md b/packages/zarr-transforms/CHANGELOG.md new file mode 100644 index 0000000000..7c4bc92cad --- /dev/null +++ b/packages/zarr-transforms/CHANGELOG.md @@ -0,0 +1,3 @@ +# Release notes + + diff --git a/packages/zarr-transforms/LICENSE.txt b/packages/zarr-transforms/LICENSE.txt new file mode 100644 index 0000000000..1e8da4d242 --- /dev/null +++ b/packages/zarr-transforms/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/zarr-transforms/README.md b/packages/zarr-transforms/README.md new file mode 100644 index 0000000000..cf3620831b --- /dev/null +++ b/packages/zarr-transforms/README.md @@ -0,0 +1,30 @@ +# zarr-transforms + +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 +- `compose` — chain two transforms into one + +The package depends only on NumPy and the standard library; it does not import +`zarr`. It is developed in the [zarr-python](https://github.com/zarr-developers/zarr-python) +repository and consumed by `zarr` to resolve array indexing operations. + +## Installation + +```bash +pip install zarr-transforms +``` + +## License + +MIT diff --git a/packages/zarr-transforms/changes/README.md b/packages/zarr-transforms/changes/README.md new file mode 100644 index 0000000000..a543987fbf --- /dev/null +++ b/packages/zarr-transforms/changes/README.md @@ -0,0 +1,25 @@ +Writing a changelog entry for `zarr-transforms` +----------------------------------------------- + +Fragments in **this** directory are release notes for the `zarr-transforms` +package only — kept separate from the parent zarr-python `changes/` +directory so a PR touching only `packages/zarr-transforms/` produces a +release note for this package only. + +Please put a new file in this directory named `xxxx..md`, where + +- `xxxx` is the pull request number associated with this entry +- `` is one of: + - feature + - bugfix + - doc + - removal + - misc + +Inside the file, please write a short description of what you have +changed, and how it impacts users of `zarr-transforms`. + +A `zarr-transforms` release runs `towncrier build` in `packages/zarr-transforms/`, +which consumes the fragments here and updates `CHANGELOG.md`. Fragments +that describe parent zarr-python changes (not the transforms package) +belong in the top-level `changes/` directory, not here. diff --git a/packages/zarr-transforms/pyproject.toml b/packages/zarr-transforms/pyproject.toml new file mode 100644 index 0000000000..36d3829b1b --- /dev/null +++ b/packages/zarr-transforms/pyproject.toml @@ -0,0 +1,103 @@ +[build-system] +requires = ["hatchling>=1.29.0", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +name = "zarr-transforms" +dynamic = ["version"] +description = "Composable, lazy coordinate transforms for Zarr array indexing." +readme = "README.md" +requires-python = ">=3.11" +license = "MIT" +license-files = ["LICENSE.txt"] +authors = [ + { name = "Davis Bennett", email = "davis.v.bennett@gmail.com" }, +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Information Technology", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", + "Topic :: Software Development :: Libraries :: Python Modules", + "Typing :: Typed", +] +keywords = ["zarr"] +dependencies = [ + "numpy>=2", +] + +[project.urls] +Homepage = "https://github.com/zarr-developers/zarr-python" +Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-transforms" +Issues = "https://github.com/zarr-developers/zarr-python/issues" +Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-transforms/CHANGELOG.md" +Documentation = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-transforms/README.md" + +[dependency-groups] +# The transform tests exercise chunk resolution against zarr's ChunkGrid +# (tests/test_chunk_resolution.py) and are collected by the parent zarr-python +# test suite, which already has zarr installed. `zarr` is intentionally NOT +# listed here to avoid a workspace dependency cycle; run these tests from the +# repo root (`uv run pytest packages/zarr-transforms/tests`), not in isolation. +test = ["pytest"] + +[tool.hatch.version] +source = "vcs" +tag-pattern = '^zarr_transforms-v(?P.+)$' +# `git_describe_command` ensures we get the zarr_transforms tags instead of latest. +# `local_scheme` strips the git commit info so the appending info is just a counter from latest tag. +# test-pypi doesn't accept git commit info in tags, and the count should be enough to distinguish unique runs. +raw-options = { root = "../..", git_describe_command = "git describe --dirty --tags --long --match zarr_transforms-v*", local_scheme = "no-local-version" } + +[tool.hatch.build.targets.wheel] +packages = ["src/zarr_transforms"] + +[tool.ruff] +extend = "../../pyproject.toml" +target-version = "py311" + +[tool.pytest.ini_options] +minversion = "7" +testpaths = ["tests"] +xfail_strict = true +addopts = ["-ra", "--strict-config", "--strict-markers"] +filterwarnings = [ + "error", +] + +[tool.pyright] +include = ["src"] +enableExperimentalFeatures = true +typeCheckingMode = "strict" +pythonVersion = "3.11" + +[tool.numpydoc_validation] +checks = [ + "GL10", + "SS04", + "PR02", + "PR03", + "PR05", + "PR06", +] + +[tool.towncrier] +# Fragments for this package live alongside the package source, separate +# from the parent zarr-python `changes/` directory, so a PR touching only +# `packages/zarr-transforms/` produces a release note for this package only. +directory = "changes" +filename = "CHANGELOG.md" +package = "zarr_transforms" +underlines = ["", "", ""] +title_format = "## {version} ({project_date})" +issue_format = "[#{issue}](https://github.com/zarr-developers/zarr-python/issues/{issue})" +start_string = "\n" diff --git a/src/zarr/core/transforms/__init__.py b/packages/zarr-transforms/src/zarr_transforms/__init__.py similarity index 56% rename from src/zarr/core/transforms/__init__.py rename to packages/zarr-transforms/src/zarr_transforms/__init__.py index 283ff834e6..458fb60e97 100644 --- a/src/zarr/core/transforms/__init__.py +++ b/packages/zarr-transforms/src/zarr_transforms/__init__.py @@ -12,11 +12,22 @@ - `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 + +The chunk-resolution helpers (`iter_chunk_transforms`, +`sub_transform_to_selections`) and `selection_to_transform` are also exported +here: they form the surface the zarr integration layer (array indexing) depends +on. The `*Like` grid Protocols describe the chunk-grid surface chunk resolution +consumes without importing zarr. """ -from zarr.core.transforms.composition import compose -from zarr.core.transforms.domain import IndexDomain -from zarr.core.transforms.json import ( +from zarr_transforms.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) +from zarr_transforms.composition import compose +from zarr_transforms.domain import IndexDomain +from zarr_transforms.grid import ChunkGridLike, DimensionGridLike +from zarr_transforms.json import ( IndexDomainJSON, IndexTransformJSON, OutputIndexMapJSON, @@ -25,12 +36,14 @@ 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 +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_transforms.transform import IndexTransform, selection_to_transform __all__ = [ "ArrayMap", + "ChunkGridLike", "ConstantMap", + "DimensionGridLike", "DimensionMap", "IndexDomain", "IndexDomainJSON", @@ -43,4 +56,7 @@ "index_domain_to_json", "index_transform_from_json", "index_transform_to_json", + "iter_chunk_transforms", + "selection_to_transform", + "sub_transform_to_selections", ] diff --git a/src/zarr/core/transforms/chunk_resolution.py b/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py similarity index 97% rename from src/zarr/core/transforms/chunk_resolution.py rename to packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py index 21d9e509b2..9848a19bdf 100644 --- a/src/zarr/core/transforms/chunk_resolution.py +++ b/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py @@ -33,14 +33,14 @@ 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 +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform if TYPE_CHECKING: from collections.abc import Iterator - from zarr.core.chunk_grids import ChunkGrid + from zarr_transforms.grid import ChunkGridLike OutIndices = ( dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None @@ -55,7 +55,7 @@ def iter_chunk_transforms( transform: IndexTransform, - chunk_grid: ChunkGrid, + chunk_grid: ChunkGridLike, ) -> Iterator[ChunkTransformResult]: """Resolve a composed IndexTransform against a ChunkGrid. diff --git a/src/zarr/core/transforms/composition.py b/packages/zarr-transforms/src/zarr_transforms/composition.py similarity index 96% rename from src/zarr/core/transforms/composition.py rename to packages/zarr-transforms/src/zarr_transforms/composition.py index 865478abf0..900d96ce1a 100644 --- a/src/zarr/core/transforms/composition.py +++ b/packages/zarr-transforms/src/zarr_transforms/composition.py @@ -2,8 +2,8 @@ import numpy as np -from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap -from zarr.core.transforms.transform import IndexTransform +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_transforms.transform import IndexTransform def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: diff --git a/src/zarr/core/transforms/domain.py b/packages/zarr-transforms/src/zarr_transforms/domain.py similarity index 100% rename from src/zarr/core/transforms/domain.py rename to packages/zarr-transforms/src/zarr_transforms/domain.py diff --git a/packages/zarr-transforms/src/zarr_transforms/errors.py b/packages/zarr-transforms/src/zarr_transforms/errors.py new file mode 100644 index 0000000000..f0cf7a6d18 --- /dev/null +++ b/packages/zarr-transforms/src/zarr_transforms/errors.py @@ -0,0 +1,21 @@ +"""Canonical index-error types raised by the transform algebra. + +These are the authoritative class definitions. `zarr.errors` re-exports the +same objects (`from zarr_transforms.errors import ...`) so that, e.g., +`zarr.errors.BoundsCheckError is zarr_transforms.errors.BoundsCheckError`. +Both subclass the built-in `IndexError`, so existing `except IndexError` (or +`except zarr.errors.BoundsCheckError`) catch sites keep working unchanged. +""" + +from __future__ import annotations + +__all__ = [ + "BoundsCheckError", + "VindexInvalidSelectionError", +] + + +class VindexInvalidSelectionError(IndexError): ... + + +class BoundsCheckError(IndexError): ... diff --git a/packages/zarr-transforms/src/zarr_transforms/grid.py b/packages/zarr-transforms/src/zarr_transforms/grid.py new file mode 100644 index 0000000000..6db1025720 --- /dev/null +++ b/packages/zarr-transforms/src/zarr_transforms/grid.py @@ -0,0 +1,39 @@ +"""Structural typing for the chunk-grid surface used by chunk resolution. + +`chunk_resolution` needs only a narrow slice of a chunk grid: the per-dimension +mapping between storage indices and chunk coordinates. Rather than import +zarr's concrete `ChunkGrid`, we type against these Protocols. zarr's +`ChunkGrid` / `DimensionGrid` satisfy them structurally, so no zarr import is +needed here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from collections.abc import Sequence + + import numpy as np + import numpy.typing as npt + + +class DimensionGridLike(Protocol): + """The per-dimension chunk-mapping surface consumed by chunk resolution.""" + + def index_to_chunk(self, idx: int) -> int: ... + def chunk_offset(self, chunk_ix: int) -> int: ... + def chunk_size(self, chunk_ix: int) -> int: ... + def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... + + +class ChunkGridLike(Protocol): + """A chunk grid exposing its per-dimension grids via `_dimensions`. + + Typed as a read-only, covariant `Sequence` so a concrete grid whose + `_dimensions` is a `tuple` of a more specific dimension type structurally + satisfies this Protocol. + """ + + @property + def _dimensions(self) -> Sequence[DimensionGridLike]: ... diff --git a/src/zarr/core/transforms/json.py b/packages/zarr-transforms/src/zarr_transforms/json.py similarity index 96% rename from src/zarr/core/transforms/json.py rename to packages/zarr-transforms/src/zarr_transforms/json.py index 78c09664b5..c8d1f07660 100644 --- a/src/zarr/core/transforms/json.py +++ b/packages/zarr-transforms/src/zarr_transforms/json.py @@ -23,9 +23,9 @@ 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 +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_transforms.transform import IndexTransform # --------------------------------------------------------------------------- # TypedDict definitions (JSON shapes) diff --git a/src/zarr/core/transforms/output_map.py b/packages/zarr-transforms/src/zarr_transforms/output_map.py similarity index 100% rename from src/zarr/core/transforms/output_map.py rename to packages/zarr-transforms/src/zarr_transforms/output_map.py diff --git a/tests/test_transforms/__init__.py b/packages/zarr-transforms/src/zarr_transforms/py.typed similarity index 100% rename from tests/test_transforms/__init__.py rename to packages/zarr-transforms/src/zarr_transforms/py.typed diff --git a/src/zarr/core/transforms/transform.py b/packages/zarr-transforms/src/zarr_transforms/transform.py similarity index 99% rename from src/zarr/core/transforms/transform.py rename to packages/zarr-transforms/src/zarr_transforms/transform.py index 5d7add487c..f1ee5737db 100644 --- a/src/zarr/core/transforms/transform.py +++ b/packages/zarr-transforms/src/zarr_transforms/transform.py @@ -34,9 +34,9 @@ import numpy as np -from zarr.core.transforms.domain import IndexDomain -from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap -from zarr.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_transforms.domain import IndexDomain +from zarr_transforms.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap @dataclass(frozen=True, slots=True) diff --git a/tests/test_transforms/test_chunk_resolution.py b/packages/zarr-transforms/tests/test_chunk_resolution.py similarity index 97% rename from tests/test_transforms/test_chunk_resolution.py rename to packages/zarr-transforms/tests/test_chunk_resolution.py index 4e24ad6065..557014f905 100644 --- a/tests/test_transforms/test_chunk_resolution.py +++ b/packages/zarr-transforms/tests/test_chunk_resolution.py @@ -1,12 +1,12 @@ 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 + +from zarr_transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform class TestChunkResolutionIdentity: diff --git a/tests/test_transforms/test_composition.py b/packages/zarr-transforms/tests/test_composition.py similarity index 96% rename from tests/test_transforms/test_composition.py rename to packages/zarr-transforms/tests/test_composition.py index e96616933d..50a9b01302 100644 --- a/tests/test_transforms/test_composition.py +++ b/packages/zarr-transforms/tests/test_composition.py @@ -3,10 +3,10 @@ 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 +from zarr_transforms.composition import compose +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform class TestComposeConstantInner: diff --git a/tests/test_transforms/test_domain.py b/packages/zarr-transforms/tests/test_domain.py similarity index 99% rename from tests/test_transforms/test_domain.py rename to packages/zarr-transforms/tests/test_domain.py index 5a222a548c..533186bd20 100644 --- a/tests/test_transforms/test_domain.py +++ b/packages/zarr-transforms/tests/test_domain.py @@ -2,7 +2,7 @@ import pytest -from zarr.core.transforms.domain import IndexDomain +from zarr_transforms.domain import IndexDomain class TestIndexDomainConstruction: diff --git a/tests/test_transforms/test_json.py b/packages/zarr-transforms/tests/test_json.py similarity index 97% rename from tests/test_transforms/test_json.py rename to packages/zarr-transforms/tests/test_json.py index 325ccd4502..e689917980 100644 --- a/tests/test_transforms/test_json.py +++ b/packages/zarr-transforms/tests/test_json.py @@ -2,8 +2,8 @@ import numpy as np -from zarr.core.transforms.domain import IndexDomain -from zarr.core.transforms.json import ( +from zarr_transforms.domain import IndexDomain +from zarr_transforms.json import ( IndexTransformJSON, index_domain_from_json, index_domain_to_json, @@ -12,8 +12,8 @@ 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 +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform class TestIndexDomainJSON: diff --git a/tests/test_transforms/test_output_map.py b/packages/zarr-transforms/tests/test_output_map.py similarity index 94% rename from tests/test_transforms/test_output_map.py rename to packages/zarr-transforms/tests/test_output_map.py index 57dbebaf44..979ca3eb1b 100644 --- a/tests/test_transforms/test_output_map.py +++ b/packages/zarr-transforms/tests/test_output_map.py @@ -2,7 +2,7 @@ import numpy as np -from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap class TestConstantMap: diff --git a/tests/test_transforms/test_tensorstore_parity.py b/packages/zarr-transforms/tests/test_tensorstore_parity.py similarity index 97% rename from tests/test_transforms/test_tensorstore_parity.py rename to packages/zarr-transforms/tests/test_tensorstore_parity.py index 4670cf5928..3804392934 100644 --- a/tests/test_transforms/test_tensorstore_parity.py +++ b/packages/zarr-transforms/tests/test_tensorstore_parity.py @@ -27,10 +27,10 @@ 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 +from zarr_transforms.domain import IndexDomain +from zarr_transforms.errors import BoundsCheckError +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform def _identity(lo: int, hi: int) -> IndexTransform: diff --git a/tests/test_transforms/test_transform.py b/packages/zarr-transforms/tests/test_transform.py similarity index 98% rename from tests/test_transforms/test_transform.py rename to packages/zarr-transforms/tests/test_transform.py index 786eb81575..efc1eb8613 100644 --- a/tests/test_transforms/test_transform.py +++ b/packages/zarr-transforms/tests/test_transform.py @@ -3,9 +3,9 @@ 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 +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_transforms.transform import IndexTransform, selection_to_transform class TestIndexTransformConstruction: @@ -553,7 +553,7 @@ class TestArrayMapDependencyAxes: from its (full-rank) shape: non-singleton axes vary, singleton axes do not.""" def test_orthogonal_single_axis(self) -> None: - from zarr.core.transforms.transform import _array_map_dependency_axes + from zarr_transforms.transform import _array_map_dependency_axes t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] m0, m1 = t.output[0], t.output[1] @@ -563,7 +563,7 @@ def test_orthogonal_single_axis(self) -> None: assert _array_map_dependency_axes(m1.index_array) == (1,) def test_vectorized_shares_axes(self) -> None: - from zarr.core.transforms.transform import _array_map_dependency_axes + from zarr_transforms.transform import _array_map_dependency_axes t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3]), np.array([2, 4])] m0, m1 = t.output[0], t.output[1] @@ -573,7 +573,7 @@ def test_vectorized_shares_axes(self) -> None: assert _array_map_dependency_axes(m1.index_array) == (0,) def test_scalar_array_has_no_dependency(self) -> None: - from zarr.core.transforms.transform import _array_map_dependency_axes + from zarr_transforms.transform import _array_map_dependency_axes assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () diff --git a/pyproject.toml b/pyproject.toml index f2c7a619e6..b393cea1c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ 'google-crc32c>=1.5', 'typing_extensions>=4.14', 'donfig>=0.8', + 'zarr-transforms>=0.1.0', ] dynamic = [ @@ -152,6 +153,16 @@ dev = [ "mypy==2.1.0", ] +# The repo is a uv workspace so that `zarr` can depend on the in-tree +# `packages/zarr-transforms` during development and CI. Released wheels depend on +# the published zarr-transforms distribution; the workspace source below only +# affects resolution from a checkout (uv.lock, `uv run`, hatch envs using uv). +[tool.uv.workspace] +members = ["packages/zarr-transforms"] + +[tool.uv.sources] +zarr-transforms = { workspace = true } + [tool.coverage.report] exclude_also = [ 'if TYPE_CHECKING:', @@ -434,7 +445,12 @@ ignore_errors = true [tool.pytest.ini_options] minversion = "7" -testpaths = ["src", "tests", "docs/user-guide"] +# `packages/zarr-transforms/tests` is collected here rather than run in an +# isolated package env: those tests exercise chunk resolution against zarr's +# concrete `ChunkGrid`, so they need `zarr` importable (which the root env +# provides). Keeping them in testpaths means the main test suite / `run-coverage` +# continues to exercise the transform algebra after the move to the subpackage. +testpaths = ["src", "tests", "docs/user-guide", "packages/zarr-transforms/tests"] log_cli_level = "INFO" log_level = "INFO" xfail_strict = true diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 0e6907cdd5..5764ca7cf8 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -19,6 +19,12 @@ import numpy as np from typing_extensions import Sentinel, deprecated +from zarr_transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr_transforms.output_map import ArrayMap +from zarr_transforms.transform import ( + IndexTransform, + selection_to_transform, +) import zarr from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec @@ -136,12 +142,6 @@ parse_node_type_array, ) 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, - selection_to_transform, -) from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, @@ -6044,7 +6044,7 @@ def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, 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 + from zarr_transforms.output_map import DimensionMap domain = transform.domain ndim = len(storage_shape) @@ -6066,7 +6066,7 @@ def _is_complete_chunk(sub_transform: IndexTransform, spec: ChunkSpec) -> bool: `spec` is the chunk's already-resolved `ChunkSpec` (the caller looks it up once and shares it with `_array_spec_from_chunk_spec`). """ - from zarr.core.transforms.output_map import ConstantMap, DimensionMap + from zarr_transforms.output_map import ConstantMap, DimensionMap shape = spec.shape for out_dim, m in enumerate(sub_transform.output): diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py index d165cdba88..1b77a4b85e 100644 --- a/src/zarr/core/chunk_partition.py +++ b/src/zarr/core/chunk_partition.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any -from zarr.core.transforms.chunk_resolution import ( +from zarr_transforms.chunk_resolution import ( iter_chunk_transforms, sub_transform_to_selections, ) @@ -22,8 +22,9 @@ if TYPE_CHECKING: from collections.abc import Callable, Iterator + from zarr_transforms.transform import IndexTransform + from zarr.core.chunk_grids import ChunkGrid - from zarr.core.transforms.transform import IndexTransform __all__ = ["ChunkProjection", "iter_chunk_projections"] diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 53e8e8535c..418553b119 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -1,3 +1,8 @@ +# `BoundsCheckError` and `VindexInvalidSelectionError` are defined canonically in +# `zarr_transforms.errors` (the transform algebra raises them) and re-exported here +# as aliases so `zarr.errors.BoundsCheckError` stays the identical class object. +from zarr_transforms.errors import BoundsCheckError, VindexInvalidSelectionError + __all__ = [ "ArrayIndexError", "ArrayNotFoundError", @@ -140,15 +145,9 @@ class ZarrRuntimeWarning(RuntimeWarning): """ -class VindexInvalidSelectionError(IndexError): ... - - class NegativeStepError(IndexError): ... -class BoundsCheckError(IndexError): ... - - class ArrayIndexError(IndexError): ... diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 5ee2b505c7..197f5e2c98 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -20,6 +20,9 @@ import numpy as np import numpy.typing as npt import pytest +from zarr_transforms.domain import IndexDomain +from zarr_transforms.output_map import ArrayMap +from zarr_transforms.transform import IndexTransform import zarr from zarr.codecs.bytes import BytesCodec @@ -27,9 +30,6 @@ from zarr.codecs.sharding import ShardingCodec from zarr.core.buffer import default_buffer_prototype from zarr.core.sync import sync -from zarr.core.transforms.domain import IndexDomain -from zarr.core.transforms.output_map import ArrayMap -from zarr.core.transforms.transform import IndexTransform from zarr.errors import BoundsCheckError, LazyViewError, ZarrUserWarning from zarr.storage import MemoryStore @@ -1129,8 +1129,9 @@ def test_is_partial_matches_is_complete_chunk(self, sel: tuple[Any, ...]) -> Non """For every stored chunk a selection touches, `is_partial` (computed via `chunk_projections`) must equal `not _is_complete_chunk(...)` computed directly from the same `iter_chunk_transforms` enumeration.""" + from zarr_transforms.chunk_resolution import iter_chunk_transforms + from zarr.core.array import _is_complete_chunk - from zarr.core.transforms.chunk_resolution import iter_chunk_transforms a = zarr.create_array({}, shape=(4, 4), chunks=(1, 4), dtype="i4") a[...] = np.arange(16, dtype="i4").reshape(4, 4) diff --git a/uv.lock b/uv.lock index 3802eaad67..e950c21c2b 100644 --- a/uv.lock +++ b/uv.lock @@ -6,6 +6,12 @@ resolution-markers = [ "python_full_version < '3.15'", ] +[manifest] +members = [ + "zarr", + "zarr-transforms", +] + [[package]] name = "aiobotocore" version = "3.7.0" @@ -4023,6 +4029,7 @@ dependencies = [ { name = "numpy" }, { name = "packaging" }, { name = "typing-extensions" }, + { name = "zarr-transforms" }, ] [package.optional-dependencies] @@ -4151,6 +4158,7 @@ requires-dist = [ { name = "typer", marker = "extra == 'cli'" }, { name = "typing-extensions", specifier = ">=4.14" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, + { name = "zarr-transforms", editable = "packages/zarr-transforms" }, ] provides-extras = ["cast-value-rs", "cli", "gpu", "optional", "remote"] @@ -4243,3 +4251,21 @@ test = [ { name = "tomlkit", specifier = "==0.15.0" }, { name = "uv", specifier = "==0.11.20" }, ] + +[[package]] +name = "zarr-transforms" +source = { editable = "packages/zarr-transforms" } +dependencies = [ + { name = "numpy" }, +] + +[package.dev-dependencies] +test = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [{ name = "numpy", specifier = ">=2" }] + +[package.metadata.requires-dev] +test = [{ name = "pytest" }] From f72df94a4b7d4d1c88c5472e1cadc28204e60fb5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 11:20:43 +0200 Subject: [PATCH 55/60] docs(transforms): note zarr-transforms package in 3906 changelog Assisted-by: ClaudeCode:claude-opus-4.8 --- changes/3906.feature.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changes/3906.feature.md b/changes/3906.feature.md index bc8518bd0f..c7526617ea 100644 --- a/changes/3906.feature.md +++ b/changes/3906.feature.md @@ -1 +1 @@ -Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged. Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. +Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged. Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. The index-transform machinery that powers this lives in a new NumPy-only `zarr-transforms` workspace package (import name `zarr_transforms`), which `zarr` now depends on at runtime. From 420956abed54badfe98ae24c9f056e4de67154fe Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 11:33:37 +0200 Subject: [PATCH 56/60] fix(coverage): measure zarr-transforms package source in coverage runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hatch run-coverage / run-coverage-html / run-hypothesis (and gputest run-coverage) scripts passed --source=src to coverage run, which excluded the index-transform algebra after its move to packages/zarr-transforms/src — it was measured before the move. Declare the measured source trees in [tool.coverage.run] source instead, so every coverage run invocation measures both src and the workspace package consistently, and drop the CLI flag from all four scripts. Assisted-by: ClaudeCode:claude-opus-4.8 --- pyproject.toml | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b393cea1c8..1f38dbf38b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -169,6 +169,15 @@ exclude_also = [ ] [tool.coverage.run] +# Measured source lives in two trees: zarr's own `src` and the in-tree +# `zarr-transforms` workspace package. Configured here (not via `--source=src` +# on the CLI) so every `coverage run` invocation — the hatch run-coverage / +# run-coverage-html / run-hypothesis scripts and ad-hoc local runs — measures +# both consistently. +source = [ + "src", + "packages/zarr-transforms/src", +] omit = [ "bench/compress_normal.py", ] @@ -211,17 +220,17 @@ matrix.deps.dependency-groups = [ [tool.hatch.envs.test.scripts] run-coverage = [ - "coverage run --source=src -m pytest --ignore tests/benchmarks --junitxml=junit.xml -o junit_family=legacy {args:}", + "coverage run -m pytest --ignore tests/benchmarks --junitxml=junit.xml -o junit_family=legacy {args:}", "coverage xml", ] run-coverage-html = [ - "coverage run --source=src -m pytest --ignore tests/benchmarks {args:}", + "coverage run -m pytest --ignore tests/benchmarks {args:}", "coverage html", ] run = "pytest --ignore tests/benchmarks" run-verbose = "run-coverage --verbose" run-hypothesis = [ - "coverage run --source=src -m pytest -nauto --run-slow-hypothesis tests/test_properties.py tests/test_store/test_stateful* {args:}", + "coverage run -m pytest -nauto --run-slow-hypothesis tests/test_properties.py tests/test_store/test_stateful* {args:}", "coverage xml", ] run-benchmark = "pytest --benchmark-enable tests/benchmarks" @@ -244,7 +253,7 @@ python = ["3.12", "3.13"] [tool.hatch.envs.gputest.scripts] run-coverage = [ - "coverage run --source=src -m pytest -m gpu --junitxml=junit.xml -o junit_family=legacy --ignore tests/benchmarks {args:}", + "coverage run -m pytest -m gpu --junitxml=junit.xml -o junit_family=legacy --ignore tests/benchmarks {args:}", "coverage xml", ] run = "pytest -m gpu --ignore tests/benchmarks" From f6675be735276b7ea727cd9fad77a2a08a091b90 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 12:09:56 +0200 Subject: [PATCH 57/60] fix(lazy): raise NotImplementedError for fancy-after-fancy composition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Applying an oindex/vindex selection to a view that already carries an orthogonal ArrayMap could land a fancy index on a broadcast (singleton) axis of that map, leaking a raw numpy IndexError ("index N is out of bounds for axis ... with size 1") at resolve time — reading like a user error rather than the implementation gap it is. Add `_guard_fancy_after_fancy`, invoked from `_apply_oindex` and `_apply_vindex`, which raises a clear NotImplementedError at composition time when a fancy axis targets a non-dependency axis of an existing ArrayMap. It names the limitation and the workaround (materialize via .result() then index, or reorder so the fancy step is last). Compositions that keep the fancy step on the dependency axis (or on a correlated vindex view) are unaffected and still resolve correctly. Adds TestFancyAfterFancy, documents the limitation in the lazy-indexing user guide, and refreshes the now-stale generator/runner comments that described this class as a silent bug. Assisted-by: ClaudeCode:claude-fable-5 --- docs/user-guide/lazy_indexing.md | 5 ++ .../src/zarr_transforms/transform.py | 29 ++++++++++ src/zarr/testing/strategies.py | 5 +- tests/test_lazy_indexing.py | 55 +++++++++++++++++++ tests/test_properties.py | 5 +- 5 files changed, 95 insertions(+), 4 deletions(-) diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md index e44c15a919..25b06658bf 100644 --- a/docs/user-guide/lazy_indexing.md +++ b/docs/user-guide/lazy_indexing.md @@ -295,6 +295,11 @@ array. - 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]`). +- Composing a fancy (`oindex`/`vindex`) selection onto a view that already has a + fancy-indexed axis (fancy-after-fancy — e.g. `a.lazy.oindex[[0, 2, 5], :]` + then `.oindex[:, [5]]`) is not supported and raises `NotImplementedError`. + Materialize the view first with `.result()` and index the array, or reorder + the selections so the fancy step is applied last. - `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 diff --git a/packages/zarr-transforms/src/zarr_transforms/transform.py b/packages/zarr-transforms/src/zarr_transforms/transform.py index f1ee5737db..3b70045992 100644 --- a/packages/zarr-transforms/src/zarr_transforms/transform.py +++ b/packages/zarr-transforms/src/zarr_transforms/transform.py @@ -644,6 +644,33 @@ def _reindex_array( return np.asarray(result, dtype=np.intp) +_FANCY_AFTER_FANCY_MSG = ( + "applying a fancy (orthogonal/vectorized) selection to a view that already " + "has a fancy-indexed axis is not supported (fancy-after-fancy composition): " + "the new coordinates would index a broadcast axis of the existing selection. " + "Materialize the view first with `.result()` and index the array, or reorder " + "the selections so the fancy step is applied last." +) + + +def _guard_fancy_after_fancy(m: ArrayMap, fancy_dims: set[int] | list[int]) -> None: + """Reject a fancy step that lands on a broadcast axis of an existing ArrayMap. + + A new orthogonal/vectorized selection can only be absorbed into an existing + ArrayMap along the axes that map genuinely varies over (its dependency axes, + plus the recorded `input_dimension` for a degenerate length-1 orthogonal + selection). A fancy index targeting any other axis — a singleton axis the map + merely broadcasts over — cannot be reindexed and used to leak a raw NumPy + `IndexError` at resolve time. Raise a clear `NotImplementedError` instead. + """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + dependent.add(m.input_dimension) + for d in fancy_dims: + if d < m.index_array.ndim and d not in dependent: + raise NotImplementedError(_FANCY_AFTER_FANCY_MSG) + + def _reindex_array_oindex( arr: np.ndarray[Any, np.dtype[np.intp]], normalized: tuple[Any, ...] | list[Any], @@ -950,6 +977,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: else: raise RuntimeError(f"unexpected: dimension {d} not handled") elif isinstance(m, ArrayMap): + _guard_fancy_after_fancy(m, list(dim_array.keys())) new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) array_input_dim: int | None = None if m.input_dimension is not None: @@ -1121,6 +1149,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: ) ) elif isinstance(m, ArrayMap): + _guard_fancy_after_fancy(m, array_dims) new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) new_output.append( ArrayMap( diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 289b3f166e..8792601b33 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -746,8 +746,9 @@ def indexing_programs(draw: st.DrawFn, *, shape: tuple[int, ...]) -> IndexingPro deliberate and documented so they can shrink as the underlying bugs are fixed: - - *Fancy-after-fancy* composition is broken in ``_reindex_array_oindex`` - (applying oindex/vindex to a view that already carries an ArrayMap axis), so + - *Fancy-after-fancy* composition (applying oindex/vindex to a view that + already carries an orthogonal ArrayMap axis) is unsupported: it raises a + clear ``NotImplementedError`` at composition time rather than resolving, so at most one fancy step is generated. - *Basic-after-fancy* is excluded wholesale. Integer basic indexing on an oindex-picked axis is a strict xfail diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py index 197f5e2c98..bf55231349 100644 --- a/tests/test_lazy_indexing.py +++ b/tests/test_lazy_indexing.py @@ -809,6 +809,61 @@ def test_fields_write_preserves_sibling_fields(self) -> None: np.testing.assert_array_equal(result["f1"], base["f1"]) +class TestFancyAfterFancy: + """Composing a fancy (orthogonal/vectorized) selection onto a view that + already carries an orthogonal ArrayMap axis is not supported: the reindexing + machinery would try to index a broadcast (singleton) axis of the existing + ArrayMap with the new coordinates. This used to leak a raw NumPy + ``IndexError`` (``index N is out of bounds for axis ... with size 1``) at + resolve time, which read like a user error; it now raises + ``NotImplementedError`` naming the limitation and the workaround at + composition time. Compositions that keep the fancy step on the existing + dependency axis (or on a correlated vindex view) still work — see + ``test_fancy_on_same_axis_still_works``.""" + + def test_oindex_then_oindex_other_axis_raises(self) -> None: + """``v.lazy.oindex[:, [5]]`` after ``v = a.lazy.oindex[[0, 2, 5], :]`` + raises at composition time, not a raw IndexError at resolve time.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], :] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.oindex[:, [5]] + + def test_oindex_then_vindex_other_axis_raises(self) -> None: + """The vindex accessor closes the same gap: a coordinate selection over a + broadcast axis of an oindex view raises rather than crashing in numpy.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], :] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.vindex[[0, 1], [5, 6]] + + def test_oindex_both_axes_then_oindex_raises(self) -> None: + """An outer-product view (two ArrayMaps) rejects any further fancy step: + the new axis is a broadcast axis of the *other* ArrayMap.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], [1, 3, 5]] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.oindex[[0, 1], :] + + def test_fancy_on_same_axis_still_works(self) -> None: + """A fancy step on the existing dependency axis is absorbed correctly and + must keep working; likewise a fancy step on a correlated vindex view.""" + 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[[0, 2, 5], :] + np.testing.assert_array_equal( + rows.lazy.oindex[[0, 1], :].result(), ref[[0, 2, 5], :][[0, 1], :] + ) + pts = b.lazy.vindex[[0, 2, 5], [1, 2, 3]] + expected = ref[[0, 2, 5], [1, 2, 3]] + np.testing.assert_array_equal(pts.lazy.oindex[[0, 1]].result(), expected[[0, 1]]) + np.testing.assert_array_equal(pts.lazy.vindex[[0, 1]].result(), expected[[0, 1]]) + + 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 diff --git a/tests/test_properties.py b/tests/test_properties.py index da97a0cc96..fe6949dda1 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -1056,8 +1056,9 @@ def test_indexing_program_parity(scenario: tuple[tuple[int, ...], IndexingProgra dtype and values; writes (`set_scalar`, `set_array`) apply the equivalent assignment to a NumPy root and compare the whole zarr array afterwards. The generator composes at most one fancy step and always keeps it last — see - `indexing_programs` for the documented exclusions (fancy-after-fancy and - basic-after-fancy are known-broken and deliberately not generated). + `indexing_programs` for the documented exclusions (fancy-after-fancy is + unsupported and raises `NotImplementedError`; basic-after-fancy is + known-broken; both are deliberately not generated). """ shape, program = scenario zarray, nparray = _build_program_array(shape) From 0d390cb2648cde36e17b9560ea82be440a9b9031 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 12:10:15 +0200 Subject: [PATCH 58/60] ci(zarr-transforms): add test/release workflows and disclose eager changes zarr-transforms is a hard runtime dependency of zarr with no CI. Add `zarr-transforms.yml` (push/PR path-filtered test/ruff/pyright, mirroring zarr-metadata; the test job runs from the repo root because the transform tests import zarr) and `zarr-transforms-release.yml` (tag-triggered publish on `zarr_transforms-v*`). Extend `check_changelogs.yml` to check the package's changes directory. Bump the package `requires-python` floor to >=3.12 (consistent with the repo; nothing tested 3.11) and update its classifiers, ruff target, and pyright pythonVersion to match. Extend the root hatch `--match v*` comment to note the `zarr_transforms-v*` tags it also excludes. Disclose the two deliberate eager-path changes in the 3906 changelog: the Array repr `domain={...}` suffix and the 0-d iteration TypeError matching NumPy. Fix a stale comment in array.py that claimed pop_fields yields [] for no fields (it now yields None). Assisted-by: ClaudeCode:claude-fable-5 --- .github/workflows/check_changelogs.yml | 3 + .github/workflows/zarr-transforms-release.yml | 117 ++++++++++++++++++ .github/workflows/zarr-transforms.yml | 102 +++++++++++++++ changes/3906.feature.md | 2 +- packages/zarr-transforms/pyproject.toml | 7 +- pyproject.toml | 7 +- src/zarr/core/array.py | 6 +- 7 files changed, 234 insertions(+), 10 deletions(-) create mode 100644 .github/workflows/zarr-transforms-release.yml create mode 100644 .github/workflows/zarr-transforms.yml diff --git a/.github/workflows/check_changelogs.yml b/.github/workflows/check_changelogs.yml index 25034b868d..abb6a8a455 100644 --- a/.github/workflows/check_changelogs.yml +++ b/.github/workflows/check_changelogs.yml @@ -29,3 +29,6 @@ jobs: - name: Check zarr-metadata changelog entries run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-metadata/changes + + - name: Check zarr-transforms changelog entries + run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-transforms/changes diff --git a/.github/workflows/zarr-transforms-release.yml b/.github/workflows/zarr-transforms-release.yml new file mode 100644 index 0000000000..ceadcc5632 --- /dev/null +++ b/.github/workflows/zarr-transforms-release.yml @@ -0,0 +1,117 @@ +name: zarr-transforms release + +on: + workflow_dispatch: + push: + tags: + - 'zarr_transforms-v*' + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + build: + name: Build wheel and sdist + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-transforms + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + fetch-depth: 0 # hatch-vcs needs full history + tags + + - name: Install Hatch + uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc + with: + version: '1.16.5' + + - name: Build + run: hatch build + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: zarr-transforms-dist + path: packages/zarr-transforms/dist + + test_artifacts: + name: Test built artifacts + needs: [build] + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-transforms-dist + path: dist + + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: false + + - name: Set up Python + run: uv python install 3.12 + + - name: Install built wheel and run import smoke test + run: | + wheel=$(ls dist/*.whl) + uv run --with "${wheel}" --python 3.12 --no-project \ + python -c "import zarr_transforms; print('zarr_transforms', zarr_transforms.__version__)" + + upload_pypi: + name: Upload to PyPI + needs: [build, test_artifacts] + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/zarr_transforms-v') + runs-on: ubuntu-latest + environment: + name: zarr-transforms-releases + url: https://pypi.org/p/zarr-transforms + permissions: + id-token: write # required for OIDC trusted publishing + attestations: write # required for artifact attestations + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-transforms-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@59d89421af93a897026c735860bf21b6eb4f7b26 # v4.1.0 + with: + subject-path: dist/* + + - name: Publish package to PyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + + upload_testpypi: + name: Upload to TestPyPI + needs: [build, test_artifacts] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + environment: + name: zarr-transforms-releases-test + url: https://test.pypi.org/p/zarr-transforms + permissions: + id-token: write + attestations: write + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-transforms-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@59d89421af93a897026c735860bf21b6eb4f7b26 # v4.1.0 + with: + subject-path: dist/* + + - name: Publish package to TestPyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + with: + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/zarr-transforms.yml b/.github/workflows/zarr-transforms.yml new file mode 100644 index 0000000000..1304988fa9 --- /dev/null +++ b/.github/workflows/zarr-transforms.yml @@ -0,0 +1,102 @@ +name: zarr-transforms + +on: + push: + branches: [main] + paths: + - 'packages/zarr-transforms/**' + - '.github/workflows/zarr-transforms.yml' + pull_request: + paths: + - 'packages/zarr-transforms/**' + - '.github/workflows/zarr-transforms.yml' + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: pytest py=${{ matrix.python-version }} + runs-on: ubuntu-latest + defaults: + run: + shell: bash + strategy: + fail-fast: false + matrix: + python-version: ['3.12', '3.13', '3.14'] + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: true + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + # The transform tests exercise chunk resolution against zarr's ChunkGrid, + # so they run from the repo root against the full workspace (which installs + # both `zarr` and the `zarr-transforms` workspace member) rather than in + # package isolation. + - name: Sync test dependency group + run: uv sync --group test --python ${{ matrix.python-version }} + - name: Run pytest + run: uv run --group test pytest packages/zarr-transforms/tests + + ruff: + name: ruff + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-transforms + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + - name: Run ruff + run: uvx ruff check . + + pyright: + name: pyright + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-transforms + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: true + - name: Set up Python + run: uv python install 3.12 + - name: Sync test dependency group + run: uv sync --group test --python 3.12 + - name: Run pyright + run: uv run --group test --with pyright pyright src + + zarr-transforms-complete: + name: zarr-transforms complete + needs: [test, ruff, pyright] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check failure + if: | + contains(needs.*.result, 'failure') || + contains(needs.*.result, 'cancelled') + run: exit 1 + - name: Success + run: echo Success! diff --git a/changes/3906.feature.md b/changes/3906.feature.md index c7526617ea..d34c03ee4b 100644 --- a/changes/3906.feature.md +++ b/changes/3906.feature.md @@ -1 +1 @@ -Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged. Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. The index-transform machinery that powers this lives in a new NumPy-only `zarr-transforms` workspace package (import name `zarr_transforms`), which `zarr` now depends on at runtime. +Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged, with two deliberate exceptions on `Array`: its `repr` now carries a `domain={...}` suffix describing the view's coordinate box, and iterating a 0-d array now raises `TypeError("iteration over a 0-d array")` to match NumPy (previously it silently yielded nothing). Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. The index-transform machinery that powers this lives in a new NumPy-only `zarr-transforms` workspace package (import name `zarr_transforms`), which `zarr` now depends on at runtime. diff --git a/packages/zarr-transforms/pyproject.toml b/packages/zarr-transforms/pyproject.toml index 36d3829b1b..c436036931 100644 --- a/packages/zarr-transforms/pyproject.toml +++ b/packages/zarr-transforms/pyproject.toml @@ -7,7 +7,7 @@ name = "zarr-transforms" dynamic = ["version"] description = "Composable, lazy coordinate transforms for Zarr array indexing." readme = "README.md" -requires-python = ">=3.11" +requires-python = ">=3.12" license = "MIT" license-files = ["LICENSE.txt"] authors = [ @@ -22,7 +22,6 @@ classifiers = [ "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", @@ -63,7 +62,7 @@ packages = ["src/zarr_transforms"] [tool.ruff] extend = "../../pyproject.toml" -target-version = "py311" +target-version = "py312" [tool.pytest.ini_options] minversion = "7" @@ -78,7 +77,7 @@ filterwarnings = [ include = ["src"] enableExperimentalFeatures = true typeCheckingMode = "strict" -pythonVersion = "3.11" +pythonVersion = "3.12" [tool.numpydoc_validation] checks = [ diff --git a/pyproject.toml b/pyproject.toml index 1f38dbf38b..5f71063571 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -186,9 +186,10 @@ omit = [ version.source = "vcs" # Only consider zarr-python's own `v*` tags when deriving the version. Without # this filter `git describe` matches the most recent tag of any shape, -# including the `zarr_metadata-v*` tags used to release the zarr-metadata -# subpackage — which would make a from-source build report a `0.2.x` version -# instead of `3.x`. +# including the `zarr_metadata-v*` and `zarr_transforms-v*` tags used to release +# the in-tree subpackages — which would make a from-source build report a +# `0.x` subpackage version instead of `3.x`. (`--match v*` excludes both, since +# neither tag starts with a bare `v`.) version.raw-options = { git_describe_command = "git describe --dirty --tags --long --match v*" } [tool.hatch.build] diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5764ca7cf8..fcb4d2dc9c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4007,8 +4007,10 @@ def set_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - # Normalize an empty fields list/tuple to None (pop_fields yields [] for - # no fields); keep the exact-emptiness check rather than a falsy one. + # Normalize a caller-supplied empty fields list/tuple to None, so that + # `fields is not None` uniformly means "a field was requested" (matching + # pop_fields, which yields None for no fields); keep the exact-emptiness + # check rather than a falsy one. if fields == [] or fields == (): fields = None self._reject_fields_on_view(fields) From d6d4d3d7a16d9621bc294a1c81c4f446de2e3cd0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 12:14:53 +0200 Subject: [PATCH 59/60] fix(zarr-transforms): expose __version__ for the release smoke test Mirrors zarr_metadata's importlib.metadata idiom; the release workflow's isolated-wheel check imports it. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-transforms/src/zarr_transforms/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/packages/zarr-transforms/src/zarr_transforms/__init__.py b/packages/zarr-transforms/src/zarr_transforms/__init__.py index 458fb60e97..cef076bf3b 100644 --- a/packages/zarr-transforms/src/zarr_transforms/__init__.py +++ b/packages/zarr-transforms/src/zarr_transforms/__init__.py @@ -20,6 +20,8 @@ consumes without importing zarr. """ +from importlib.metadata import version + from zarr_transforms.chunk_resolution import ( iter_chunk_transforms, sub_transform_to_selections, @@ -39,6 +41,8 @@ from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_transforms.transform import IndexTransform, selection_to_transform +__version__ = version("zarr-transforms") + __all__ = [ "ArrayMap", "ChunkGridLike", @@ -51,6 +55,7 @@ "IndexTransformJSON", "OutputIndexMap", "OutputIndexMapJSON", + "__version__", "compose", "index_domain_from_json", "index_domain_to_json", From 3c276b7451a0a4e47b5dc8230c3dd2f210ba17ff Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 17 Jul 2026 12:45:25 +0200 Subject: [PATCH 60/60] perf(zarr-transforms): enumerate only touched chunks in iter_chunk_transforms The ArrayMap (fancy) branch of iter_chunk_transforms enumerated the dense range(min_chunk, max_chunk + 1) bounding box and ran transform.intersect against every candidate chunk, making sparse fancy/vindex chunk resolution scale O(n_chunks) instead of O(n_touched). The exact touched chunk ids were already computed and then discarded in favor of min/max. Enumerate each fancy dimension's distinct touched chunk ids (np.unique) instead; the cartesian product then spans only touched-per-dimension combinations. Constant/Dimension dims keep their contiguous ranges. Semantics are unchanged: the dense range only ever added empty intersections, which intersect already skipped. Sparse vindex (2 far-apart coords) goes from 15.9x slower than eager at 1k chunks / 65.5x at 4k to ~1.1x at both, flat in grid size. Dense fancy selections are unchanged. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_transforms/chunk_resolution.py | 40 +++++-- .../tests/test_chunk_resolution.py | 109 ++++++++++++++++++ 2 files changed, 138 insertions(+), 11 deletions(-) diff --git a/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py b/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py index 9848a19bdf..485341218a 100644 --- a/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py +++ b/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py @@ -38,7 +38,7 @@ from zarr_transforms.transform import IndexTransform if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterator, Sequence from zarr_transforms.grid import ChunkGridLike @@ -68,16 +68,35 @@ def iter_chunk_transforms( """ 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] = [] + # Enumerate candidate chunks via the cartesian product of per-dimension + # candidate chunk ids, then for each candidate intersect the transform with + # the chunk domain (`transform.intersect` handles orthogonal and vectorized + # cases alike, filtering out combinations it does not actually touch). + # + # Each dimension contributes exactly the chunk ids it can touch: + # + # - `ConstantMap`/`DimensionMap` dims contribute a contiguous `range` — a + # single chunk for a constant, and the span between the first and last + # chunk for a slice. These are already tight (or nearly so). + # - `ArrayMap` (fancy) dims contribute only the *distinct* chunk ids the + # index array actually lands in (`np.unique`), never the dense + # `range(min_chunk, max_chunk + 1)` between them. A sparse fancy selection + # (e.g. two far-apart coordinates) would otherwise enumerate every chunk + # in the bounding box, making resolution scale with grid size instead of + # with the number of selected coordinates. + # + # For >= 2 correlated (vindex) ArrayMaps the per-dimension distinct sets + # over-approximate the *joint* touched set (their cartesian product includes + # combinations no single point lands in), but `intersect` filters those out, + # so the yielded chunks are identical either way — and the work stays bounded + # by the per-dimension distinct chunk counts, not the grid size. + chunk_candidates: list[Sequence[int]] = [] 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)) + chunk_candidates.append((c,)) elif isinstance(m, DimensionMap): d = m.input_dimension dim_lo = transform.domain.inclusive_min[d] @@ -92,7 +111,7 @@ def iter_chunk_transforms( 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)) + chunk_candidates.append(range(first, last + 1)) elif isinstance(m, ArrayMap): storage = m.offset + m.stride * m.index_array flat = storage.ravel().astype(np.intp) @@ -100,13 +119,12 @@ def iter_chunk_transforms( # Empty fancy selection: no coordinates, so no chunks are touched. return chunk_ids = dg.indices_to_chunks(flat) - first = int(chunk_ids.min()) - last = int(chunk_ids.max()) - chunk_ranges.append(range(first, last + 1)) + # Enumerate only the distinct chunks the coordinates land in. + chunk_candidates.append([int(c) for c in np.unique(chunk_ids)]) import itertools - for chunk_coords_tuple in itertools.product(*chunk_ranges): + for chunk_coords_tuple in itertools.product(*chunk_candidates): chunk_coords = tuple(int(c) for c in chunk_coords_tuple) # Build the chunk domain in storage space diff --git a/packages/zarr-transforms/tests/test_chunk_resolution.py b/packages/zarr-transforms/tests/test_chunk_resolution.py index 557014f905..15ac8496e9 100644 --- a/packages/zarr-transforms/tests/test_chunk_resolution.py +++ b/packages/zarr-transforms/tests/test_chunk_resolution.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np from zarr.core.chunk_grids import ChunkGrid, FixedDimension @@ -8,6 +10,9 @@ from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_transforms.transform import IndexTransform +if TYPE_CHECKING: + import pytest + class TestChunkResolutionIdentity: def test_single_chunk(self) -> None: @@ -105,6 +110,110 @@ def test_array_index(self) -> None: assert (2,) in coords_list +def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]: + """Wrap `IndexTransform.intersect` with a call counter. + + Returns a mutable dict whose `"n"` entry is the number of times + `intersect` is invoked. Used to assert that candidate-chunk enumeration is + proportional to the *touched* chunks, not the dense bounding box between the + min and max touched chunk. + """ + calls = {"n": 0} + original = IndexTransform.intersect + + def counting(self: IndexTransform, output_domain: IndexDomain) -> object: + calls["n"] += 1 + return original(self, output_domain) + + monkeypatch.setattr(IndexTransform, "intersect", counting) + return calls + + +class TestChunkResolutionTouchedOnly: + """`iter_chunk_transforms` must enumerate only the chunks a fancy selection + actually touches — never the dense `range(min_chunk, max_chunk + 1)` bounding + box. These guard against a regression to bounding-box enumeration, whose cost + scales with grid size rather than with the number of selected coordinates. + """ + + def test_1d_sparse_vindex_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two far-apart coordinates on a 1000-chunk grid touch exactly 2 chunks. + + A dense bounding-box enumeration would intersect ~1000 candidate chunks; + touched-only enumeration intersects exactly 2. + """ + # 4000 elements, chunk size 4 -> 1000 chunks. coords 1 and 3997 land in + # chunk 0 and chunk 999 respectively (998 empty chunks between them). + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=4000),)) + t = IndexTransform.from_shape((4000,)).vindex[np.array([1, 3997], dtype=np.intp)] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0,), (999,)] + # Exactly the touched chunks, independent of the 1000-chunk grid size. + assert calls["n"] == 2 + + def test_2d_orthogonal_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Orthogonal outer product of two 2-coordinate arrays touches 2x2 chunks. + + Per-dimension distinct touched chunks: {0, 999} on each axis. The outer + product is 2*2 = 4 candidate chunks (all survive), versus ~1e6 for a + dense 1000x1000 bounding box. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).oindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (0, 999), (999, 0), (999, 999)] + assert calls["n"] == 4 + + def test_2d_correlated_vindex_enumerates_per_dim_distinct_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two correlated (vindex) coordinate arrays scatter to 2 diagonal chunks. + + The two points (1, 2) and (3997, 3998) touch chunks (0, 0) and + (999, 999). Per-dimension distinct touched chunks are {0, 999} on each + axis, so enumeration intersects the 2x2 = 4 combinations; the two + off-diagonal combinations are filtered out by `intersect`, leaving 2 + surviving chunks. The key guarantee is that the work is bounded by the + per-dimension distinct touched chunks (4), not the dense 1e6 grid. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).vindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (999, 999)] + # 2x2 per-dim-distinct combinations enumerated; 2 survive intersection. + assert calls["n"] == 4 + + class TestSubTransformToSelections: def test_constant_map(self) -> None: """ConstantMap produces int selection + drop axis."""