Skip to content

Commit

Permalink
Only use oindex for advanced indexing (int arrays) (#440)
Browse files Browse the repository at this point in the history
This change means that more of the array API is supported for
storage backends that don't have oindex (e.g. the new Zarr v3
implementation).
  • Loading branch information
tomwhite authored Apr 3, 2024
1 parent 7b4f7e2 commit c2283f3
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions cubed/core/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -489,6 +489,7 @@ def merged_chunk_len_for_indexer(s):
extra_projected_mem=extra_projected_mem,
target_chunks=target_chunks,
selection=selection,
advanced_indexing=len(where_list) > 0,
)

# merge chunks for any dims with step > 1 so they are
Expand All @@ -509,10 +510,19 @@ def merged_chunk_len_for_indexer(s):
return out


def _read_index_chunk(x, *arrays, target_chunks=None, selection=None, block_id=None):
array = arrays[0]
def _read_index_chunk(
x,
*arrays,
target_chunks=None,
selection=None,
advanced_indexing=None,
block_id=None,
):
array = arrays[0].zarray
if advanced_indexing:
array = array.oindex
idx = block_id
out = array.zarray.oindex[_target_chunk_selection(target_chunks, idx, selection)]
out = array[_target_chunk_selection(target_chunks, idx, selection)]
out = numpy_array_to_backend_array(out)
return out

Expand Down

0 comments on commit c2283f3

Please sign in to comment.