diff --git a/packages/zarr-indexing/changes/4333.bugfix.md b/packages/zarr-indexing/changes/4333.bugfix.md new file mode 100644 index 0000000000..0347c2d95d --- /dev/null +++ b/packages/zarr-indexing/changes/4333.bugfix.md @@ -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. diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index cb29776476..975494490c 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -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: @@ -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) @@ -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: diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index a13eaf6e28..9e8b24bf4a 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -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)