diff --git a/changes/4316.bugfix.md b/changes/4316.bugfix.md new file mode 100644 index 0000000000..e4b9bd8d94 --- /dev/null +++ b/changes/4316.bugfix.md @@ -0,0 +1 @@ +Fixed a `ValueError` when setting an orthogonal selection on a sharded array that mixes an integer index with two or more array indices, such as `a.oindex[[3, 1, 2], 1, [0, 2]] = value`. The fix for the array-only case in #4284 reshaped the value only when its shape matched the coordinate selection exactly; the sharding codec now also ravels a value that is the selection's shape minus the integer-indexed axes. Values of any other rank are left alone, so a write that is invalid on an unsharded array fails the same way on a sharded one. Both partial-encode paths share one helper for this. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 7a4a339f13..bd0760f7e3 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -57,6 +57,7 @@ from zarr.core.indexing import ( BasicIndexer, ChunkProjection, + CoordinateIndexer, SelectorTuple, SliceDimIndexer, _lexicographic_order, @@ -399,6 +400,25 @@ def to_dict_vectorized(self) -> dict[tuple[int, ...], Buffer | None]: return result +def _drops_only_unit_axes(shape: tuple[int, ...], full: tuple[int, ...]) -> bool: + """Return whether ``shape`` is ``full`` with zero or more length-1 axes removed. + + ``(2, 2)`` is ``(2, 1, 2)`` minus its unit axis, and ``(1, 2)`` is + ``(1, 2, 1)`` minus its last; ``(2, 2)`` is not ``(4,)``, and + ``(3, 2, 1)`` is not ``(3, 2)`` because it adds an axis. + """ + remaining = iter(full) + for size in shape: + for full_size in remaining: + if full_size == size: + break + if full_size != 1: + return False + else: + return False + return all(full_size == 1 for full_size in remaining) + + @dataclass(frozen=True) class ShardingCodec( ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin @@ -791,23 +811,11 @@ def _encode_partial_sync( Loads the existing shard, merges the written region into the affected inner chunks, and rewrites the whole shard. """ - shard_shape = shard_spec.shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) - shard_indexer = get_indexer( - selection, - shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), - ) - # A coordinate indexer flattens the selection, so its projections address - # `value` as 1-D while the caller shaped it like `sel_shape`. Mirrors the - # reshape `_encode_partial_single` applies on the async path. - sel_shape = getattr(shard_indexer, "sel_shape", None) - if sel_shape is not None and value.shape == sel_shape: - value = value.reshape(shard_indexer.shape) - indexer = list(shard_indexer) + indexer, value = self._get_shard_indexer_and_value(selection, shard_spec, value) is_complete = self._is_complete_shard_write(indexer, chunks_per_shard) @@ -1359,23 +1367,10 @@ async def _encode_partial_single( selection: SelectorTuple, shard_spec: ArraySpec, ) -> None: - shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) - shard_indexer = get_indexer( - selection, - shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), - ) - # A coordinate indexer flattens the selection, so its projections address - # `shard_array` as 1-D while the caller shaped it like `sel_shape`. This - # mirrors the reshape `_decode_partial_single` applies on the way out. - sel_shape = getattr(shard_indexer, "sel_shape", None) - if sel_shape is not None and shard_array.shape == sel_shape: - shard_array = shard_array.reshape(shard_indexer.shape) - indexer = list(shard_indexer) + indexer, shard_array = self._get_shard_indexer_and_value(selection, shard_spec, shard_array) if self._is_complete_shard_write(indexer, chunks_per_shard): shard_dict = dict.fromkeys(lexicographic_order_coords(chunks_per_shard)) @@ -1433,6 +1428,45 @@ async def _encode_shard_dict( index_bytes, buffers, buffer_prototype, chunks_per_shard=chunks_per_shard ) + def _get_shard_indexer_and_value( + self, selection: SelectorTuple, shard_spec: ArraySpec, value: NDBuffer + ) -> tuple[list[ChunkProjection], NDBuffer]: + """Index ``selection`` over the inner chunk grid, flattening ``value`` to match. + + ``get_indexer`` classifies a tuple of integer arrays as a coordinate + selection, and a ``CoordinateIndexer`` addresses the value buffer as + 1-D. An ``OrthogonalIndexer`` with two or more array-indexed axes hands + down an ``np.ix_`` tuple, so ``sel_shape`` is N-D, while the caller + shaped ``value`` like the orthogonal result: the broadcast shape minus + the integer-indexed axes, which ``np.ix_`` keeps as length-1 axes. That + value has the element count and C order of the flattened projections, + so ravel it. + + Only that value shape is ravelled. A mask or coordinate selection + arrives with a 1-D ``sel_shape`` and a value that must already be flat, + and an orthogonal value with an axis the selection does not have is + invalid. Both are left alone so the write fails the same way it does + on an unsharded array; the shard-level selection alone cannot tell + orthogonal from mask indexing, the value shape can. Scalars pass + through and are broadcast downstream. + + The partial-decode paths apply the inverse reshape, to + ``indexer.sel_shape``, on the way out. + """ + shard_shape = shard_spec.shape + indexer = get_indexer( + selection, + shape=shard_shape, + chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), + ) + if ( + isinstance(indexer, CoordinateIndexer) + and len(value.shape) > 1 + and _drops_only_unit_axes(value.shape, indexer.sel_shape) + ): + value = value.reshape(indexer.shape) + return list(indexer), value + def _is_total_shard( self, all_chunk_coords: set[tuple[int, ...]], chunks_per_shard: tuple[int, ...] ) -> bool: diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 6679dbcee4..f7ddca969e 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -328,11 +328,14 @@ def arrays( else: chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") - if all(s > c > 1 for s, c in zip(nparray.shape, chunks_param, strict=True)): + # Any chunk that fits the array can be sharded: shard_shapes draws a + # whole number of chunks per axis, one inner chunk included. + if all(s >= c >= 1 for s, c in zip(nparray.shape, chunks_param, strict=True)): shard_shape = draw( st.none() | shard_shapes(shape=nparray.shape, chunk_shape=chunks_param), label="shard shape", ) + event("sharded" if shard_shape is not None else "unsharded") if shard_shape is not None: subchunk_write_order = draw(subchunk_write_orders) inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") @@ -572,20 +575,30 @@ def basic_indices( @st.composite def orthogonal_indices( draw: st.DrawFn, *, shape: tuple[int, ...] -) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: +) -> tuple[tuple[int | slice | np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: """ Strategy that returns - (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays. - (2) a tuple of integer arrays that can be used for equivalent indexing of numpy arrays + (1) a tuple of per-axis selectors (integer array, slice, or bare integer) for + orthogonal indexing of Zarr arrays. + (2) a tuple of broadcast integer arrays that index a numpy array to the same + result. A bare integer drops its axis, as ``oindex`` does, so it is + given as a 0-d array and does not contribute a result dimension. """ - zindexer = [] - npindexer = [] - ndim = len(shape) + zindexer: list[int | slice | np.ndarray[Any, Any]] = [] + kept: list[tuple[int, np.ndarray[Any, Any]]] = [] + npindexer: dict[int, np.ndarray[Any, Any]] = {} for axis, size in enumerate(shape): if size != 0: - strategy = npst.integer_array_indices( - shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1) - ) | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + strategy = ( + npst.integer_array_indices( + shape=(size,), + result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1), + ) + | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + # basic_indices(min_dims=1) never yields a bare integer, so draw + # one explicitly: it is the only selector that drops an axis. + | st.integers(min_value=-size, max_value=size - 1) + ) else: strategy = basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) @@ -597,19 +610,25 @@ def orthogonal_indices( .filter(bool) ) (idxr,) = val - if isinstance(idxr, int): - idxr = np.array([idxr]) zindexer.append(idxr) + if isinstance(idxr, int): + npindexer[axis] = np.array(idxr) + continue if isinstance(idxr, slice): idxr = np.arange(*idxr.indices(size)) - elif isinstance(idxr, (tuple, int)): + elif isinstance(idxr, tuple): idxr = np.array(idxr) - newshape = [1] * ndim - newshape[axis] = idxr.size - npindexer.append(idxr.reshape(newshape)) + kept.append((axis, idxr)) + + for pos, (axis, idxr) in enumerate(kept): + newshape = [1] * len(kept) + newshape[pos] = idxr.size + npindexer[axis] = idxr.reshape(newshape) # casting the output of broadcast_arrays is needed for numpy < 2 - return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) + return tuple(zindexer), tuple( + np.broadcast_arrays(*(npindexer[axis] for axis in range(len(shape)))) + ) @st.composite diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 80fea63760..42a88ddc7a 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1,4 +1,5 @@ import enum +import math import pickle import warnings from typing import Any, cast, get_args @@ -1272,31 +1273,58 @@ def test_shard_reader_to_dict_vectorized(chunks_per_shard: tuple[int, ...]) -> N ], ) @pytest.mark.parametrize("nested", [False, True], ids=["single", "nested"]) -def test_sharding_orthogonal_set_multiple_array_dims(nested: bool, pipeline_path: str) -> None: +@pytest.mark.parametrize( + "selection", + [ + pytest.param((np.array([3, 1, 2]), np.array([0, 2])), id="2d-arr-arr"), + pytest.param((np.array([3, 0]), np.array([2, 0])), id="2d-arr-arr-unsorted-two-shards"), + pytest.param((np.array([3, 1, 2]), 1, np.array([0, 2])), id="3d-arr-int-arr"), + pytest.param((1, np.array([0, 2]), np.array([1, 3])), id="3d-int-arr-arr"), + pytest.param((np.array([3, 1]), np.array([0, 2]), 2), id="3d-arr-arr-int"), + pytest.param( + (np.array([3, 1, 2]), np.array([0, 2]), np.array([1, 3])), id="3d-arr-arr-arr" + ), + pytest.param((np.array([3]), np.array([0, 2]), 1), id="3d-arr1-arr-int"), + ], +) +def test_sharding_orthogonal_set_multiple_array_dims( + selection: tuple[int | npt.NDArray[np.intp], ...], nested: bool, pipeline_path: str +) -> None: """Orthogonal set with more than one array-indexed dimension. - ``OrthogonalIndexer`` converts such a chunk selection to an ``np.ix_`` pair + ``OrthogonalIndexer`` converts such a chunk selection to an ``np.ix_`` tuple of broadcastable arrays before handing it to the codec pipeline. The sharding codec re-derives an indexer from that selection and gets a ``CoordinateIndexer``, whose projections address the value buffer flat. Regression test for the resulting shape mismatch on write. + An integer index alongside the arrays is the case that a shape-equality + guard misses: ``OrthogonalIndexer`` drops that axis from the value but + ``np.ix_`` keeps it as a length-1 axis in the chunk selection, so the value + and the re-derived indexer's ``sel_shape`` differ in rank while agreeing in + element count. + Parametrized over both pipelines because the partial-encode path is written twice -- ``_encode_partial_single`` for ``BatchedCodecPipeline`` - and ``_encode_partial_sync`` for ``FusedCodecPipeline`` -- and each - derives its own indexer. + and ``_encode_partial_sync`` for ``FusedCodecPipeline``. """ - inner = ShardingCodec(chunk_shape=(1, 1), codecs=(BytesCodec(),)) - serializer = ShardingCodec(chunk_shape=(2, 2), codecs=((inner,) if nested else (BytesCodec(),))) - base = np.arange(16, dtype="int32").reshape(4, 4) - selection = (np.array([3, 1, 2]), np.array([0, 2])) - value = np.arange(6, dtype="int32").reshape(3, 2) + 100 + ndim = len(selection) + shape = (4,) * ndim + inner = ShardingCodec(chunk_shape=(1,) * ndim, codecs=(BytesCodec(),)) + serializer = ShardingCodec( + chunk_shape=(2,) * ndim, codecs=((inner,) if nested else (BytesCodec(),)) + ) + base = np.arange(4**ndim, dtype="int32").reshape(shape) + ix = np.ix_(*(np.atleast_1d(s) for s in selection)) + # The value is shaped like the orthogonal result: integer axes dropped. + value_shape = tuple(len(s) for s in selection if not isinstance(s, int)) + value = np.arange(math.prod(value_shape), dtype="int32").reshape(value_shape) + 100 with zarr.config.set({"codec_pipeline.path": pipeline_path}): a = zarr.create_array( MemoryStore(), - shape=base.shape, - chunks=(2, 4), + shape=shape, + chunks=(2,) + (4,) * (ndim - 1), dtype=base.dtype, serializer=serializer, compressors=None, @@ -1306,6 +1334,6 @@ def test_sharding_orthogonal_set_multiple_array_dims(nested: bool, pipeline_path a.oindex[selection] = value expected = base.copy() - expected[np.ix_(*selection)] = value + expected[ix] = value.reshape(expected[ix].shape) assert np.array_equal(a[:], expected) assert np.array_equal(a.oindex[selection], value) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 04fbdad8c6..4a1a4172d2 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -51,12 +51,14 @@ def zarr_array_from_numpy_array( store: StorePath, a: npt.NDArray[Any], chunk_shape: tuple[int, ...] | None = None, + shards: tuple[int, ...] | None = None, ) -> zarr.Array: z = zarr.create_array( store=store / str(uuid4()), shape=a.shape, dtype=a.dtype, chunks=chunk_shape or a.shape, + shards=shards, chunk_key_encoding={"name": "v2", "separator": "."}, ) z[()] = a @@ -2103,23 +2105,51 @@ def test_zero_sized_chunks(store: StorePath, shape: list[int]) -> None: assert_array_equal(z[...], np.zeros(shape, dtype="f8")) -@pytest.mark.parametrize("store", ["memory"], indirect=["store"]) -def test_vectorized_indexing_incompatible_shape(store) -> None: - """Regression for GH2469: vectorized set-indexing raises ValueError when the value shape is incompatible with the indexer shape.""" - # GH2469 - shape = (4, 4) - chunks = (2, 2) - fill_value = 32767 - arr = zarr.create( - shape, - store=store, - chunks=chunks, - dtype=np.int16, - fill_value=fill_value, - codecs=[zarr.codecs.BytesCodec(), zarr.codecs.BloscCodec()], - ) - with pytest.raises(ValueError, match="Attempting to set"): - arr[np.array([1, 2]), np.array([1, 2])] = np.array([[-1, -2], [-3, -4]]) +@pytest.mark.parametrize( + "pipeline_path", + [ + "zarr.core.codec_pipeline.BatchedCodecPipeline", + "zarr.core.codec_pipeline.FusedCodecPipeline", + ], + ids=["batched", "fused"], +) +@pytest.mark.parametrize("shards", [None, (4, 4)], ids=["chunked", "sharded"]) +@pytest.mark.parametrize( + ("kind", "selection", "value_shape"), + [ + # GH2469: array-level check, the value has twice the selected elements. + pytest.param("vindex", (np.array([1, 2]), np.array([1, 2])), (2, 2), id="coord-2d-value"), + # Right element count, wrong rank: a mask takes a flat value. + pytest.param("vindex", np.eye(4, dtype=bool), (2, 2), id="mask-2d-value"), + # Right element count, an axis the selection does not have. + pytest.param( + "oindex", (np.array([3, 1, 2]), np.array([0, 2])), (3, 2, 1), id="oindex-extra-axis" + ), + ], +) +def test_set_selection_rejects_value_with_wrong_rank( + store: StorePath, + kind: str, + selection: Any, + value_shape: tuple[int, ...], + shards: tuple[int, ...] | None, + pipeline_path: str, +) -> None: + """A value whose rank does not fit the selection raises regardless of storage layout. + + The sharding codec re-derives an indexer from the selection it is handed + and ravels the value when it is the selection's broadcast shape minus + integer-indexed axes. Any other rank must fail on a sharded array exactly + as it does on a chunked one; an element count that happens to match is + not grounds to accept it. Only the rejection is asserted: a write that + fails inside the chunk merge may already have touched other chunks. + """ + a = np.zeros((4, 4), dtype=np.int32) + value = np.arange(np.prod(value_shape), dtype=np.int32).reshape(value_shape) + with zarr.config.set({"codec_pipeline.path": pipeline_path}): + z = zarr_array_from_numpy_array(store, a, chunk_shape=(2, 2), shards=shards) + with pytest.raises(ValueError, match="Attempting to set|shape mismatch"): + getattr(z, kind)[selection] = value def test_iter_chunk_regions(): diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..d57475df9e 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -176,9 +176,8 @@ async def test_oindex(data: st.DataObject) -> None: 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: + # sync set + for idxr in zindexer: 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) @@ -217,13 +216,14 @@ async def test_vindex(data: st.DataObject) -> None: 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[:]) + points = np.stack([idxr.ravel() for idxr in np.broadcast_arrays(*indexer)], axis=-1) + if len(np.unique(points, axis=0)) != len(points): + # behaviour of setitem with repeated coordinates is not guaranteed in practice + assume(False) + new_data = data.draw(numpy_arrays(shapes=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 @@ -243,7 +243,6 @@ def test_mask_indexing(data: st.DataObject) -> None: 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) @@ -270,8 +269,7 @@ def test_block_indexing(data: st.DataObject) -> None: assert_array_equal(expected, zarray.blocks[block_indexer]) assert_array_equal(expected, zarray.get_block_selection(block_indexer)) - # sync set, via both interfaces; sharded set is broken upstream (GH2834) - assume(zarray.shards is None) + # sync set, via both interfaces new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) nparray[array_indexer] = new_data zarray.blocks[block_indexer] = new_data