Repository navigation
Unsigned integer index arrays are mishandled (IndexError on unsorted, and on any uint64) #4285
Description
hi here is an AI report on the bug I found while working on zarrs/zarrs-python#189
Zarr version
3.3.1.dev34+gd44f9f92
Numcodecs version
0.16.5
Python Version
3.14.6
Operating System
macOS 26.5.2, arm64 (Apple silicon)
Installation
uv, from a git checkout of main (reproducer run via `uv run` with PEP 723 inline metadata)
Description
An integer array selection whose dtype is unsigned fails, where the identical values in a
signed dtype work. is_integer_array() explicitly admits unsigned dtypes
(x.dtype.kind in "ui"), so these selections pass validation and are then handled
incorrectly downstream.
z[np.array([3, 0], dtype="int64"), :] # works
z[np.array([3, 0], dtype="uint8"), :] # IndexErrorSame array, same values — only the index dtype differs.
There are two distinct failures, both stemming from the index dtype never being normalized
to intp:
- unsorted index + any unsigned dtype — misclassified as increasing, so the indices are
never regrouped by chunk. uint64, sorted or not, through[]and.vindexalike —uint64combined with a
signed offset promotes tofloat64.
Neither returns wrong data; both raise.
Expected behaviour
All unsigned index dtypes should behave exactly as the signed ones do, matching NumPy, which
accepts uint8…uint64 fancy indices without complaint:
values[np.array([3, 0], dtype="uint64"), :] # numpy: fineWhy the current behaviour is a problem
Unsigned index arrays are not exotic — they arrive routinely from np.argsort-adjacent code,
from libraries that store row ids compactly (uint32 cell indices, uint8 category codes),
and from any downstream consumer that narrows dtypes to save memory. Such an array indexes a
NumPy array fine and indexes a Zarr array fine as long as it happens to be sorted, so the
failure surfaces only on some inputs and looks like a data problem rather than a dtype
problem. The error message (index 3 is out of bounds for axis 0 with size 2) points at the
chunk shape and gives no hint that the index dtype is what matters, which makes it
expensive to diagnose.
The three conditions in §Conditions below also mean a test suite can easily pass while the
bug is live: 1-D arrays, single-chunk arrays, and sorted indices all work.
Why the suggested fix is better
Normalizing the index dtype to intp restores an invariant the code already declares —
IntArrayDimIndexer annotates dim_sel: npt.NDArray[np.intp], it just never enforced it.
Doing the cast once at the indexer boundary fixes both failure modes for every caller,
rather than guarding each arithmetic site downstream. Classifying order by comparison
instead of subtraction removes a wraparound hazard that cannot be reintroduced by a future
dtype.
Root cause
(1) Order.check differences in the index's own dtype — src/zarr/core/indexing.py:
diff = np.diff(a)
diff_positive = diff >= 0On an unsigned array a subtraction cannot go negative, so it wraps:
np.diff(np.array([3, 0], dtype="uint8")) # -> [253] -> Order.INCREASING
np.diff(np.array([255, 0], dtype="uint8")) # -> [1] -> Order.INCREASINGA descending selection is therefore classified INCREASING, IntArrayDimIndexer leaves
dim_out_sel = None, and the argsort that groups indices by chunk is skipped. Index 3 is
then looked up in the chunk that only holds rows 0..1.
(2) uint64 promotes to float — in IntArrayDimIndexer.__iter__,
self.dim_sel[start:stop] - dim_offset promotes uint64 against a signed offset to
float64, which is no longer a valid index. CoordinateIndexer has the same gap: it casts
scalar ints to intp but leaves supplied arrays at their original dtype.
Conditions
Failure (1) needs all four of:
- an unsigned index dtype —
int8/int32are fine, so this is signedness, not width - an unsorted index —
np.sort(rows)always works - a second axis — the same index on a 1-D array is fine
- an index spanning more than one chunk — a single-chunk array is fine
Failure (2) needs only uint64.
Suggested fix
- classify with comparisons in
Order.check(a[1:] >= a[:-1]), which cannot wrap - normalize the index dtype once, where the
intpinvariant is claimed:
dim_sel.astype(np.intp, copy=False)inIntArrayDimIndexer.__init__, and the same cast
inCoordinateIndexerafteris_coordinate_selectionvalidation, so a float index
still errors
Happy to open a PR if this looks right.
Steps to reproduce
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "zarr@git+https://github.com/zarr-developers/zarr-python.git@main",
# ]
# ///
#
# This script automatically imports the development branch of zarr to check for issues
import tempfile
from pathlib import Path
import numpy as np
import zarr
from zarr.core.indexing import Order
values = np.arange(8, dtype="float32").reshape(4, 2)
store = Path(tempfile.mkdtemp()) / "a.zarr"
z = zarr.create_array(store, shape=values.shape, dtype="float32", chunks=(2, 1))
z[:] = values
print(f"expected: {values[[3, 0], :].ravel()}\n")
print("unsorted index [3, 0], by index dtype:")
for dtype in ("int64", "int8", "uint8", "uint16", "uint32", "uint64"):
rows = np.array([3, 0], dtype=dtype)
try:
print(f" {dtype:>7}: {z[rows, :].ravel()}")
except Exception as exc: # noqa: BLE001 - the point is to show what escapes
print(f" {dtype:>7}: {type(exc).__name__}: {exc}")
print("\nuint64 fails even sorted, and through .vindex:")
for label, call in (
("z[sorted uint64, :]", lambda: z[np.array([0, 3], dtype="uint64"), :].ravel()),
(
"z.vindex[uint64, uint64]",
lambda: z.vindex[np.array([0, 3], dtype="uint64"), np.array([1, 0], dtype="uint64")],
),
):
try:
print(f" {label:>24}: {call()}")
except Exception as exc: # noqa: BLE001
print(f" {label:>24}: {type(exc).__name__}: {exc}")
print("\nthe misclassification, directly:")
for a in (
np.array([3, 0], dtype="int64"),
np.array([3, 0], dtype="uint8"),
np.array([255, 0], dtype="uint8"),
):
print(f" np.diff({a.dtype}{a.tolist()}) = {np.diff(a)} -> {Order.check(a)}")
zarr.print_debug_info()Output on main (d44f9f92):
expected: [6. 7. 0. 1.]
unsorted index [3, 0], by index dtype:
int64: [6. 7. 0. 1.]
int8: [6. 7. 0. 1.]
uint8: IndexError: index 3 is out of bounds for axis 0 with size 2
uint16: IndexError: index 3 is out of bounds for axis 0 with size 2
uint32: IndexError: index 3 is out of bounds for axis 0 with size 2
uint64: IndexError: arrays used as indices must be of integer (or boolean) type
uint64 fails even sorted, and through .vindex:
z[sorted uint64, :]: IndexError: arrays used as indices must be of integer (or boolean) type
z.vindex[uint64, uint64]: IndexError: arrays used as indices must be of integer (or boolean) type
the misclassification, directly:
np.diff(int64[3, 0]) = [-3] -> Order.DECREASING
np.diff(uint8[3, 0]) = [253] -> Order.INCREASING
np.diff(uint8[255, 0]) = [1] -> Order.INCREASING
Additional output
platform: macOS-26.5.2-arm64-arm-64bit-Mach-O
python: 3.14.6
zarr: 3.3.1.dev34+gd44f9f92
**Required dependencies:**
packaging: 26.3
numpy: 2.5.2
numcodecs: 0.16.5
typing_extensions: 4.16.0
donfig: 0.8.1.post1
**Optional dependencies:**
numcodecs: 0.16.5
**Not Installed:**
botocore, cupy-cuda12x, fsspec, s3fs, gcsfs, universal-pathlib, obstore
Activity