Skip to content

Commit f879470

Browse files
authored
Merge branch 'main' into ig/spec0_py314
2 parents eb89a55 + 03e70f5 commit f879470

2 files changed

Lines changed: 17 additions & 9 deletions

File tree

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
@@ -961,15 +961,22 @@ def __iter__(self) -> Iterator[ChunkProjection]:
961961

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

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

0 commit comments

Comments
 (0)