Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions packages/zarr-indexing/changes/4333.bugfix.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Direct `IndexTransform.oindex` and `IndexTransform.vindex` selections now reject
integer array coordinates outside the `np.intp` range before conversion.
Previously, oversized `uint64` values could wrap to negative coordinates and
silently select a different location in a domain containing negative coordinates.
12 changes: 8 additions & 4 deletions packages/zarr-indexing/src/zarr_indexing/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -1541,7 +1541,7 @@ def _normalize_oindex_selection(
(indices,) = np.nonzero(sel)
result.append(indices.astype(np.intp))
elif isinstance(sel, np.ndarray):
result.append(sel.astype(np.intp))
result.append(checked_affine(0, 1, sel))
elif isinstance(sel, slice):
result.append(sel)
elif (scalar := as_scalar_index(sel)) is not None:
Expand All @@ -1553,7 +1553,9 @@ def _normalize_oindex_selection(
(indices,) = np.nonzero(array)
result.append(indices.astype(np.intp))
else:
result.append(np.asarray(sel, dtype=np.intp))
# Advanced selection validation has already checked the element types.
integer_array = cast("npt.NDArray[np.integer[Any]]", array)
result.append(checked_affine(0, 1, integer_array))
else:
result.append(sel)

Expand Down Expand Up @@ -1743,9 +1745,11 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform:
indices_tuple = np.nonzero(boolean_array)
processed.extend(indices.astype(np.intp) for indices in indices_tuple)
elif isinstance(sel, np.ndarray):
processed.append(sel.astype(np.intp))
processed.append(checked_affine(0, 1, sel))
elif isinstance(sel, (list, tuple)):
processed.append(np.asarray(sel, dtype=np.intp))
# Advanced selection validation has already checked the element types.
integer_array = cast("npt.NDArray[np.integer[Any]]", np.asarray(sel))
processed.append(checked_affine(0, 1, integer_array))
elif (scalar := as_scalar_index(sel)) is not None:
processed.append(np.array([scalar], dtype=np.intp))
else:
Expand Down
26 changes: 26 additions & 0 deletions packages/zarr-indexing/tests/test_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,6 +763,32 @@ def test_vindex_multiple_arrays_preserves_shared_axes(self) -> None:
assert result.output[1].index_array.shape == (2,)


@pytest.mark.parametrize("mode", ["oindex", "vindex"])
@pytest.mark.parametrize("as_list", [False, True])
@pytest.mark.parametrize("value", [2**63, 2**64 - 1])
def test_direct_advanced_index_rejects_unsigned_overflow(
mode: str, as_list: bool, value: int
) -> None:
"""Narrowing must not turn huge positive coordinates into valid negative ones."""
transform = IndexTransform.identity(
IndexDomain(inclusive_min=(np.iinfo(np.intp).min,), exclusive_max=(0,))
)
selector = [value] if as_list else np.array([value], dtype=np.uint64)
with pytest.raises(OverflowError, match="outside np.intp range"):
getattr(transform, mode)[selector]


@pytest.mark.parametrize("mode", ["oindex", "vindex"])
@pytest.mark.parametrize("dtype", ["intp", "uint8", "uint64"])
def test_direct_advanced_index_preserves_valid_coordinates(mode: str, dtype: str) -> None:
transform = IndexTransform.from_shape((5,))
selector = np.array([3, 0, 3], dtype=dtype)
selector.flags.writeable = False
result = getattr(transform, mode)[selector]
np.testing.assert_array_equal(result.apply_many(np.arange(3).reshape(3, 1)), [[3], [0], [3]])
np.testing.assert_array_equal(selector, [3, 0, 3])


@pytest.mark.parametrize("mode", ["oindex", "vindex"])
def test_direct_advanced_index_rejects_float_arrays(mode: str) -> None:
helper = getattr(IndexTransform.from_shape((5,)), mode)
Expand Down
Loading