Skip to content

Add vectorized array_indices_to_chunk_dim to eliminate loops - #5

Merged
jhamman merged 1 commit into
jhamman:feature/rectilinear-chunk-gridfrom
maxrjones:vectorize
Feb 25, 2026
Merged

Add vectorized array_indices_to_chunk_dim to eliminate loops#5
jhamman merged 1 commit into
jhamman:feature/rectilinear-chunk-gridfrom
maxrjones:vectorize

Conversation

@maxrjones

Copy link
Copy Markdown

zarr-developers#3534 refactored IntArrayDimIndexer and CoordinateIndexer to use chunk_grid.array_index_to_chunk_coord() for mapping array indices to chunk coordinates. This replaced the old vectorized dim_sel // dim_chunk_len with a Python for-loop calling array_index_to_chunk_coord per element, which introduces significant overhead for both regular and rectilinear chunk grids.

This PR adds array_indices_to_chunk_dim() — a vectorized method that maps an entire array of indices to chunk coordinates along a single dimension:

  • RegularChunkGrid: indices // chunk_size (single numpy operation, O(n))
  • RectilinearChunkGrid: np.searchsorted(cumsum, indices, side='right') - 1 (vectorized binary search, O(n log m))

The two hot loops in indexing.py are replaced with one-line calls to this method, restoring the original performance characteristics for regular chunks while providing efficient vectorized indexing for rectilinear chunks.

New tests

  • test_chunk_grids/test_common.py — 4 tests for _is_nested_sequence() covering basic sequences, non-sequences, ChunkGrid instances, and empty iterables
  • test_chunk_grids/test_rectilinear.py — 5 end-to-end indexing tests for rectilinear arrays: slice at exact chunk boundaries, strided slicing, fancy indexing with oindex/vindex, boolean masks, and
    block indexing
  • test_indexing.py — TestRectilinearIndexing class with 16 tests covering basic selection (1D/2D, strided), orthogonal selection (boolean/integer, mixed), coordinate/vindex selection, block selection,
    and set selection on rectilinear arrays with chunks=[[5, 10, 15], [10, 20, 30, 40]]

@jhamman
jhamman merged commit 1ce7c60 into jhamman:feature/rectilinear-chunk-gridFeb 25, 2026
5 checks passed
@maxrjones
maxrjones deleted the vectorize branch February 25, 2026 18:32
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@maxrjones@jhamman