Skip to content

Commit 03e70f5

Browse files
slevangdcherian
andauthored
perf: oindex optimization (#3830)
* oindex single dim optimization * changelog * changelog type * Apply suggestions from code review Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> * lint --------- Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
1 parent 5d92e85 commit 03e70f5

File tree

2 files changed

+17
-9
lines changed

2 files changed

+17
-9
lines changed

changes/3830.misc.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Optimize the performance of indexing operations when using an array-like indexer on a single dimension.

src/zarr/core/indexing.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -966,15 +966,22 @@ def __iter__(self) -> Iterator[ChunkProjection]:
966966

967967
# handle advanced indexing arrays orthogonally
968968
if self.is_advanced:
969-
# N.B., numpy doesn't support orthogonal indexing directly as yet,
970-
# so need to work around via np.ix_. Also np.ix_ does not support a
971-
# mixture of arrays and slices or integers, so need to convert slices
972-
# and integers into ranges.
973-
chunk_selection = ix_(chunk_selection, self.chunk_shape)
974-
975-
# special case for non-monotonic indices
976-
if not is_basic_selection(out_selection):
977-
out_selection = ix_(out_selection, self.shape)
969+
# NumPy can handle a single array-indexed dimension directly,
970+
# which preserves full slices and avoids an
971+
# unnecessary advanced-indexing copy. Integer-indexed
972+
# dimensions still need the ix_ path for downstream squeezing.
973+
# Example: we skip `ix_` for array[:, :, [1, 2, 3]]
974+
n_array_dims = sum(isinstance(sel, np.ndarray) for sel in chunk_selection)
975+
976+
if n_array_dims > 1 or self.drop_axes:
977+
# N.B., numpy doesn't support orthogonal indexing directly
978+
# for multiple array-indexed dimensions, so we need to
979+
# convert the orthogonal selection into coordinate arrays.
980+
chunk_selection = ix_(chunk_selection, self.chunk_shape)
981+
982+
# special case for non-monotonic indices
983+
if not is_basic_selection(out_selection):
984+
out_selection = ix_(out_selection, self.shape)
978985

979986
is_complete_chunk = all(p.is_complete_chunk for p in dim_projections)
980987
yield ChunkProjection(chunk_coords, chunk_selection, out_selection, is_complete_chunk)

0 commit comments

Comments
 (0)