From 1c3229ef8a9b8bfe385baeda93fae6917c117ab1 Mon Sep 17 00:00:00 2001 From: XXXXXQ-0206 Date: Sat, 17 Oct 2026 03:36:58 +0800 Subject: [PATCH 1/2] fix(indexing): dispatch partial Boolean array masks to orthogonal indexing --- src/zarr/core/indexing.py | 13 +++++--- tests/test_indexing.py | 68 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 8ea59eb960..17970b6a6d 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -281,8 +281,11 @@ def is_pure_fancy_indexing(selection: Any, ndim: int) -> bool: with boolean or slices). """ if is_bool_array(selection): - # is mask selection - return True + # A full-shape Boolean array is a mask selection, which is vectorized + # indexing. A shorter one is a partial Boolean index along the leading + # axes, which NumPy resolves like orthogonal indexing, and a zero-length + # one selects nothing at all; both are handled by the orthogonal path. + return selection.ndim == ndim and selection.size > 0 if ndim == 1 and ( is_integer_list(selection) or is_integer_array(selection) or is_bool_list(selection) @@ -642,8 +645,10 @@ def __init__( if not is_bool_array(dim_sel, 1): raise IndexError("Boolean arrays in an orthogonal selection must be 1-dimensional only") - # check shape - if dim_sel.shape[0] != dim_len: + # check shape: a zero-length Boolean array selects nothing, which NumPy + # allows even when the axis is non-empty (it is treated as an empty + # integer selection); any other mismatch is an error. + if dim_sel.shape[0] != dim_len and dim_sel.shape[0] != 0: raise IndexError( f"Boolean array has the wrong length for dimension; expected {dim_len}, got {dim_sel.shape[0]}" ) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 785f089513..b2c5df51cb 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -38,6 +38,7 @@ oindex_set, replace_ellipsis, ) +from zarr.errors import VindexInvalidSelectionError from zarr.registry import get_ndbuffer_class from zarr.storage import MemoryStore, StorePath @@ -428,6 +429,73 @@ def test_get_basic_selection_2d_rejects_list_in_tuple(store: StorePath) -> None: z.get_basic_selection((slice(None), [0, 1])) +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_partial_boolean_array_mask_matches_numpy(zarr_format: ZarrFormat) -> None: + """A Boolean array shorter than the array is a partial mask over the leading axes. + + NumPy resolves such a mask with orthogonal (axis-wise) indexing, selects nothing + for a zero-length mask, and rejects any other length mismatch. A Boolean array + with the full shape stays a mask selection. + """ + expected = np.arange(16, dtype="i4").reshape(4, 4) + arr = zarr.create_array({}, data=expected, chunks=(2, 2), zarr_format=zarr_format) + mask = np.array([True, False, True, False]) + empty = np.array([], dtype=bool) + all_true = np.ones(4, dtype=bool) + all_false = np.zeros(4, dtype=bool) + + assert_array_equal(arr[mask], expected[mask]) + assert_array_equal(arr[all_true], expected[all_true]) + assert_array_equal(arr[all_false], expected[all_false]) + assert_array_equal(arr[empty], expected[empty]) + assert_array_equal(arr[mask, :], expected[mask, :]) + assert_array_equal(arr[:, mask], expected[:, mask]) + assert_array_equal(arr[empty, :], expected[empty, :]) + + # A full-shape Boolean array is still a mask selection. + full_mask = np.eye(4, dtype=bool) + assert_array_equal(arr[full_mask], expected[full_mask]) + + # Orthogonal indexing agrees. + assert_array_equal(arr.oindex[mask], expected[mask]) + assert_array_equal(arr.oindex[empty], expected[empty]) + + # 1-D and 3-D arrays. + expected_1d = np.arange(4, dtype="i4") + arr_1d = zarr.create_array({}, data=expected_1d, chunks=(2,), zarr_format=zarr_format) + assert_array_equal(arr_1d[mask], expected_1d[mask]) + assert_array_equal(arr_1d[empty], expected_1d[empty]) + + expected_3d = np.arange(24, dtype="i4").reshape(2, 3, 4) + arr_3d = zarr.create_array({}, data=expected_3d, chunks=(1, 2, 2), zarr_format=zarr_format) + leading = np.array([True, False]) + assert_array_equal(arr_3d[leading], expected_3d[leading]) + assert_array_equal(arr_3d[empty], expected_3d[empty]) + + # A Boolean list keeps working. + assert_array_equal(arr[[True, False, True, False]], expected[[True, False, True, False]]) + + # A non-empty mask of the wrong length still raises, as it does in NumPy. + wrong_length = np.array([True, False, True]) + with pytest.raises(IndexError): + arr[wrong_length] + with pytest.raises(IndexError): + expected[wrong_length] + with pytest.raises(IndexError): + arr.oindex[wrong_length] + + # Writes through a partial mask select the same rows; an empty mask writes nothing. + arr[mask] = 99 + expected[mask] = 99 + assert_array_equal(arr[:], expected) + arr[empty] = 1 + assert_array_equal(arr[:], expected) + + # vindex keeps its own contract: coordinate arrays or a single full mask. + with pytest.raises(VindexInvalidSelectionError): + arr.vindex[mask] + + def test_fancy_indexing_fallback_on_get_setitem(store: StorePath) -> None: """Paired integer-list indexing falls back to vectorized (fancy) get and set via `__getitem__`/`__setitem__`.""" z = zarr_array_from_numpy_array(store, np.zeros((20, 20))) From 7a3ec172bfc5c75c01c1f2f9614869aaeacfbc49 Mon Sep 17 00:00:00 2001 From: XXXXXQ-0206 Date: Sat, 17 Oct 2026 10:16:25 +0800 Subject: [PATCH 2/2] docs: add the changelog entry for #4519 --- changes/4519.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/4519.bugfix.md diff --git a/changes/4519.bugfix.md b/changes/4519.bugfix.md new file mode 100644 index 0000000000..dd276f3603 --- /dev/null +++ b/changes/4519.bugfix.md @@ -0,0 +1 @@ +A boolean ``ndarray`` used as a bare mask now behaves like the equivalent boolean list and like NumPy: ``a[np.array([True, False, True, False])]`` selects the matching rows instead of raising ``VindexInvalidSelectionError``, and an empty boolean array selects nothing (``(0, n)``). Masks whose length does not match the axis still raise, as in NumPy.