Skip to content

fix(indexing): read a Boolean array index the way NumPy does - #443

Open
d-v-b wants to merge 3 commits into
mainfrom
fix/bool-ndarray-masks
Open

d-v-b wants to merge 3 commits into
mainfrom
fix/bool-ndarray-masks

Conversation

@d-v-b

@d-v-b d-v-b commented Nov 6, 2026 •

Copy link
Copy Markdown
Owner

🤖 AI text below 🤖

Indexing an array with a Boolean NumPy array now selects what NumPy selects. That covers a mask over the leading axes (for example a (2, 3) mask on a (2, 3, 4) array), an empty mask whose length is 0 along an axis, and a full-shape mask on an array with a zero-length axis. Orthogonal selection also accepts an empty Boolean array and no longer fails on a zero-length dimension. A Boolean array with more than one dimension inside a selection tuple, which is not supported, now raises an error that says so instead of one about orthogonal selection.

What changed

expand_bool_index in src/zarr/core/indexing.py rewrites a Boolean array used as the whole index into a selection that the existing indexers already handle, following NumPy's reading of a mask m as m.nonzero():

  • A mask with the array's full shape stays a mask selection. This includes empty masks on arrays with a zero-length axis, which work on main.
  • A 1-D mask over the first of several axes becomes an orthogonal selection.
  • Any other mask, including an empty one, becomes a coordinate selection that broadcasts against the remaining axes.
  • A mask that does not match the axes it covers raises IndexError. As in NumPy, an empty mask may have length 0 along an axis; any other mismatch raises.

Array.__getitem__, Array.__setitem__ and get_indexer call it before dispatch. BoolArrayDimIndexer accepts a zero-length Boolean array and computes its item count without indexing an empty cumulative sum, which crashed on a zero-length dimension on main. The vindex contract (coordinate arrays, or one full-shape mask) is unchanged.

Verification

  • One parametrized test compares reads and writes against NumPy, and one test covers the wrong-shape error; 18 of the 27 cases fail on main.
  • A differential fuzz of 800 random cases against NumPy 2.5.4 (up to 3 dimensions, zero-length axes, random chunks, sharded and unsharded, reads plus scalar and array writes) found about 280 reads that fail on main and pass here, no regressions, and matching shapes and values wherever both return data.
  • The full test suite and mypy pass locally.

Known limits

  • A mask over two or more leading axes of a higher-dimensional array builds one coordinate per selected element, about 22 times the result's size in one measurement. The full-shape mask path costs the same.
  • Writes match NumPy for a scalar value or a value with the selection's full shape. zarr does not broadcast a smaller value on any selection path, so z[m2] = np.arange(4) on a (2, 3, 4) array still raises.
  • A mask with two or more dimensions inside a tuple, such as z[m2, 1] or z[m2, :], still raises, now with an error that says such a mask is only supported as the whole index (it used to talk about orthogonal selection). Supporting it would mean turning every slice in the tuple into coordinates, so it is left for a move onto zarr-indexing, whose vindex keeps those axes lazy.
  • A 0-d Boolean index (z[np.True_]) still raises; NumPy adds a length-1 axis.
  • Several Boolean arrays in one tuple are still read orthogonally, as on main, so z[E, E] with two empty masks gives (0, 0) where NumPy gives (0,).
  • a[m, ...] with an index array before an Ellipsis still fails; that is fixed separately.
  • Field selection on structured arrays is broken on main for every selection type and is not touched here.
Bookkeeping

This replaces zarr-developers#4519, which fixed the 1-D partial mask but sent empty full-shape masks on arrays with a zero-length axis down the orthogonal path, where they fail; those work on main. The Ellipsis case is zarr-developers#4595. If this is upstreamed, rename the changelog fragment to the upstream PR number.

🤖 Generated with Claude Code

d-v-b added 3 commits November 6, 2026 22:10
A Boolean array used as the whole index now selects what NumPy selects.
A 1-D mask over the first of several axes is read as an orthogonal
selection, an N-D mask over the leading axes as coordinates broadcast
against the trailing axes, and an empty mask selects nothing whatever
its shape. Full-shape masks keep the mask selection path, including on
arrays with a zero-length axis. Orthogonal selection also accepts an
empty Boolean array, and no longer fails on a zero-length dimension.

Assisted-by: ClaudeCode:claude-opus-5-5
Assisted-by: ClaudeCode:claude-opus-5-5
`z[m2, :]` with a 2-D mask used to fail with "Boolean arrays in an
orthogonal selection must be 1-dimensional only", which names an indexing
mode the caller never asked for. The error now says that such an array is
only supported as the whole index and suggests alternatives.

Assisted-by: ClaudeCode:claude-opus-5-5

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant