Skip to content

Comments

Add vectorized array_indices_to_chunk_dim to eliminate loops#5

Open
maxrjones wants to merge 1 commit intojhamman:feature/rectilinear-chunk-gridfrom
maxrjones:vectorize
Open

Add vectorized array_indices_to_chunk_dim to eliminate loops#5
maxrjones wants to merge 1 commit intojhamman:feature/rectilinear-chunk-gridfrom
maxrjones:vectorize

Conversation

@maxrjones
Copy link

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]]

Sign up for free to 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.

1 participant