Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/3830.misc.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Optimize the performance of indexing operations when using an array-like indexer on a single dimension.
24 changes: 15 additions & 9 deletions src/zarr/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -966,15 +966,21 @@ def __iter__(self) -> Iterator[ChunkProjection]:

# handle advanced indexing arrays orthogonally
if self.is_advanced:
# N.B., numpy doesn't support orthogonal indexing directly as yet,
# so need to work around via np.ix_. Also np.ix_ does not support a
# mixture of arrays and slices or integers, so need to convert slices
# and integers into ranges.
chunk_selection = ix_(chunk_selection, self.chunk_shape)

# special case for non-monotonic indices
if not is_basic_selection(out_selection):
out_selection = ix_(out_selection, self.shape)
# NumPy can handle a single array-indexed dimension directly,
# which keeps slice dimensions as slices and avoids an
Comment thread
slevang marked this conversation as resolved.
Outdated
# unnecessary advanced-indexing copy. Integer-indexed
# dimensions still need the ix_ path for downstream squeezing.
Comment thread
slevang marked this conversation as resolved.
n_array_dims = sum(isinstance(sel, np.ndarray) for sel in chunk_selection)

if n_array_dims > 1 or self.drop_axes:
# N.B., numpy doesn't support orthogonal indexing directly
# for multiple array-indexed dimensions, so we need to
# convert the orthogonal selection into coordinate arrays.
chunk_selection = ix_(chunk_selection, self.chunk_shape)

# special case for non-monotonic indices
if not is_basic_selection(out_selection):
out_selection = ix_(out_selection, self.shape)

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