diff --git a/changes/4286.bugfix.md b/changes/4286.bugfix.md index ee0a60a82c..e1ebb0dee2 100644 --- a/changes/4286.bugfix.md +++ b/changes/4286.bugfix.md @@ -4,3 +4,11 @@ the order check used `np.diff`, which wraps on unsigned dtypes and misclassified descending selection as increasing. Separately, a `uint64` index raised `IndexError` on both `array[...]` and `array.vindex[...]` — sorted or not — because `uint64` promotes to `float64` against a signed chunk offset. Index arrays are now cast to `intp`. + +Unsigned indices are bounds-checked before this conversion, so values such as +`np.uint64(2**64 - 1)` are rejected rather than wrapping to a negative index and +reading or overwriting an element at the end of the array. + +Negative-index normalization copies indices before modifying them, preserving +caller-owned arrays and supporting read-only index arrays. Reusing one index +array across axes of different lengths now normalizes each axis independently. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index b791a9b6c9..2213b69e44 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -729,10 +729,13 @@ def check(a: npt.NDArray[Any]) -> Order: return order -def wraparound_indices(x: npt.NDArray[Any], dim_len: int) -> None: +def wraparound_indices(x: npt.NDArray[Any], dim_len: int) -> npt.NDArray[Any]: + """Normalize negative indices, copying only when normalization is needed.""" loc_neg = x < 0 if np.any(loc_neg): + x = x.copy() x[loc_neg] += dim_len + return x def boundscheck_indices(x: npt.NDArray[Any], dim_len: int) -> None: @@ -769,6 +772,9 @@ def __init__( dim_sel = np.asanyarray(dim_sel) if not is_integer_array(dim_sel, 1): raise IndexError("integer arrays in an orthogonal selection must be 1-dimensional only") + # Check unsigned values before narrowing: uint64 can wrap to a valid negative index. + if boundscheck and dim_sel.dtype.kind == "u": + boundscheck_indices(dim_sel, dim_len) # uint64 promotes to float against the signed chunk offset dim_sel = dim_sel.astype(np.intp, copy=False) @@ -778,7 +784,7 @@ def __init__( # handle wraparound if wraparound: - wraparound_indices(dim_sel, dim_len) + dim_sel = wraparound_indices(dim_sel, dim_len) # handle out of bounds if boundscheck: @@ -1209,6 +1215,10 @@ def __init__( "(coordinate) array per dimension of the target array, " f"got {selection!r}" ) + # Check unsigned values before narrowing can turn an out-of-bounds value negative. + for dim_sel, dim_len in zip(selection_normalized, shape, strict=True): + if dim_sel.dtype.kind == "u": + boundscheck_indices(dim_sel, dim_len) # keep indices integral: uint64 against a signed offset promotes to float selection_normalized = cast( "CoordinateSelectionNormalized", @@ -1268,12 +1278,12 @@ def __init__( object.__setattr__(self, "drop_axes", ()) return - # handle wraparound, boundscheck - for dim_sel, dim_len in zip(selection_normalized, shape, strict=True): - # handle wraparound + # Normalize each axis independently without modifying caller-owned index arrays. + selection_normalized = tuple( wraparound_indices(dim_sel, dim_len) - - # handle out of bounds + for dim_sel, dim_len in zip(selection_normalized, shape, strict=True) + ) + for dim_sel, dim_len in zip(selection_normalized, shape, strict=True): boundscheck_indices(dim_sel, dim_len) # compute chunk index for each point in the selection diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 4cf6a0b662..e09d403081 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -40,6 +40,7 @@ from zarr.abc.store import ByteRequest from zarr.core.buffer import BufferPrototype from zarr.core.buffer.core import Buffer + from zarr.core.common import ZarrFormat @pytest.fixture @@ -949,6 +950,94 @@ def test_unsorted_index_unsigned_dtype(store: StorePath, dtype: str) -> None: assert_array_equal(a[[3, 0], [1, 0]], z.vindex[rows, np.array([1, 0], dtype=dtype)]) +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("indexer", ["oindex", "vindex"]) +@pytest.mark.parametrize("operation", ["read", "write"]) +@pytest.mark.parametrize("index", [4, 2**63, 2**64 - 4, 2**64 - 1]) +def test_unsigned_index_out_of_bounds( + zarr_format: ZarrFormat, indexer: str, operation: str, index: int +) -> None: + """Unsigned indices must be checked before narrowing can turn them negative.""" + data = np.arange(16).reshape(4, 4) + arr = zarr.create_array({}, data=data, chunks=(2, 2), zarr_format=zarr_format) + selection = (np.array([0, 1]), np.array([0, index], dtype="uint64")) + accessor = getattr(arr, indexer) + + if operation == "read": + with pytest.raises(IndexError, match="out of bounds"): + accessor[selection] + else: + with pytest.raises(IndexError, match="out of bounds"): + accessor[selection] = 99 + + assert_array_equal(arr[:], data) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("indexer", ["oindex", "vindex"]) +@pytest.mark.parametrize("dtype", ["uint8", "uint16", "uint32", "uint64"]) +@pytest.mark.parametrize("indices", [[], [3, 0], [0, 3], [0] * 16 + [3] * 16]) +def test_unsigned_index_roundtrip( + zarr_format: ZarrFormat, indexer: str, dtype: str, indices: list[int] +) -> None: + """Valid unsigned reads and writes retain the behavior fixed by #4286.""" + expected = np.arange(4) + arr = zarr.create_array({}, data=expected, chunks=(2,), zarr_format=zarr_format) + selection = np.array(indices, dtype=dtype) + accessor = getattr(arr, indexer) + + assert_array_equal(accessor[selection], expected[selection]) + accessor[selection] = 99 + expected[selection] = 99 + assert_array_equal(arr[:], expected) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("indexer", ["oindex", "vindex"]) +@pytest.mark.parametrize("readonly", [False, True]) +@pytest.mark.parametrize("indices", [[-1, 0], [0, 3]]) +def test_index_array_preserved( + zarr_format: ZarrFormat, indexer: str, readonly: bool, indices: list[int] +) -> None: + """Shared, strided index arrays remain unchanged after reads and writes.""" + expected = np.arange(20).reshape(4, 5) + arr = zarr.create_array({}, data=expected, chunks=(2, 2), zarr_format=zarr_format) + backing = np.array([indices[0], 99, indices[1], 99], dtype=np.intp) + original = backing.copy() + selection = backing[::2] + selection.flags.writeable = not readonly + accessor = getattr(arr, indexer) + # Reuse the same array on axes of different lengths: -1 must wrap independently. + numpy_selection = ( + np.ix_(selection, selection) if indexer == "oindex" else (selection, selection) + ) + assert_array_equal(accessor[selection, selection], expected[numpy_selection]) + assert_array_equal(backing, original) + + accessor[selection, selection] = 99 + expected[numpy_selection] = 99 + assert_array_equal(arr[:], expected) + assert_array_equal(backing, original) + + +@pytest.mark.parametrize("indexer", ["oindex", "vindex"]) +@pytest.mark.parametrize("operation", ["read", "write"]) +def test_invalid_negative_index_array_preserved(indexer: str, operation: str) -> None: + """Even a rejected selection must not modify caller-owned indices or array data.""" + expected = np.arange(4) + arr = zarr.create_array({}, data=expected, chunks=(2,)) + selection = np.array([-5, 0], dtype=np.intp) + accessor = getattr(arr, indexer) + if operation == "read": + with pytest.raises(IndexError, match="out of bounds"): + accessor[selection] + else: + with pytest.raises(IndexError, match="out of bounds"): + accessor[selection] = 99 + assert_array_equal(selection, [-5, 0]) + assert_array_equal(arr[:], expected) + + def _test_set_orthogonal_selection( v: npt.NDArray[np.int_], a: npt.NDArray[Any], z: Array, selection: OrthogonalSelection ) -> None: diff --git a/tests/test_properties.py b/tests/test_properties.py index 2b82d54e2e..8c3b864b86 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -177,9 +177,10 @@ async def test_oindex(data: st.DataObject) -> None: assert_array_equal(nparray[npindexer], actual) # sync set - for idxr in zindexer: - if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr).size: + for idxr, size in zip(zindexer, nparray.shape, strict=True): + if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr % size).size: # behaviour of setitem with repeated indices is not guaranteed in practice + # Negative and positive spellings of the same index are duplicates too. assume(False) new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) nparray[npindexer] = new_data @@ -216,7 +217,15 @@ async def test_vindex(data: st.DataObject) -> None: assert_array_equal(nparray[indexer], actual) # sync set - points = np.stack([idxr.ravel() for idxr in np.broadcast_arrays(*indexer)], axis=-1) + # Reads preserve the supplied indices; normalize negative indices explicitly when + # detecting repeated points rather than relying on a read to mutate the indexer. + points = np.stack( + [ + (idxr % size).ravel() + for idxr, size in zip(np.broadcast_arrays(*indexer), nparray.shape, strict=True) + ], + axis=-1, + ) if len(np.unique(points, axis=0)) != len(points): # behaviour of setitem with repeated coordinates is not guaranteed in practice assume(False)