(feat): fall back to pure python indexing in case of unhandled rust indexing for read

python/zarrs_python/pipeline.py

@@ -25,6 +25,7 @@
 
 from ._internal import CodecPipelineImpl
 from .utils import (
+    CollapsedDimensionError,
     DiscontiguousArrayError,
     get_max_threads,
     make_chunk_info_for_rust,
@@ -100,9 +101,9 @@ async def read(
         if not out.dtype.isnative:
             raise RuntimeError("Non-native byte order not supported")
         try:
-            chunks_desc = make_chunk_info_for_rust_with_indices(batch_info)
+            chunks_desc = make_chunk_info_for_rust_with_indices(batch_info, drop_axes)
             index_in_rust = True
-        except DiscontiguousArrayError:
+        except (DiscontiguousArrayError, CollapsedDimensionError):
             chunks_desc = make_chunk_info_for_rust(batch_info)
             index_in_rust = False
         if index_in_rust:
@@ -142,7 +143,7 @@ async def write(
             value = np.ascontiguousarray(value, dtype=value.dtype.newbyteorder("="))
         elif not value.flags.c_contiguous:
             value = np.ascontiguousarray(value)
-        chunks_desc = make_chunk_info_for_rust_with_indices(batch_info)
+        chunks_desc = make_chunk_info_for_rust_with_indices(batch_info, drop_axes)
         await asyncio.to_thread(
             self.impl.store_chunks_with_indices,
             chunks_desc,

python/zarrs_python/utils.py

@@ -8,6 +8,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
+    from types import EllipsisType
 
     from zarr.abc.store import ByteGetter, ByteSetter
     from zarr.core.array_spec import ArraySpec
@@ -23,7 +24,11 @@ class DiscontiguousArrayError(BaseException):
     pass
 
 
-# This is a copy of the function from zarr.core.indexing that fixes:
+class CollapsedDimensionError(BaseException):
+    pass
+
+
+# This is a (mostly) copy of the function from zarr.core.indexing that fixes:
 #   DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated
 # TODO: Upstream this fix
 def make_slice_selection(selection: tuple[np.ndarray | float]) -> list[slice]:
@@ -44,6 +49,14 @@ def make_slice_selection(selection: tuple[np.ndarray | float]) -> list[slice]:
                 ls.append(slice(dim_selection[0], dim_selection[-1] + 1, 1))
         else:
             ls.append(dim_selection)
+    if (
+        sum(isinstance(dim_selection, np.ndarray) for dim_selection in selection)
+        == sum(isinstance(dim_selection, slice) for dim_selection in selection)
+        == len(selection)
+    ):
+        raise DiscontiguousArrayError(
+            "Vindexing with multiple contiguous numpy arrays is not supported"
+        )
     return ls
 
 
@@ -66,11 +79,89 @@ def convert_chunk_to_primitive(
     )
 
 
+def resulting_shape_from_index(
+    array_shape: tuple[int, ...],
+    index_tuple: tuple[int | slice | EllipsisType | np.ndarray],
+    drop_axes: tuple[int, ...],
+    *,
+    pad: bool,
+) -> tuple[int, ...]:
+    result_shape = []
+    advanced_index_shapes = [
+        idx.shape for idx in index_tuple if isinstance(idx, np.ndarray)
+    ]
+    basic_shape_index = 0
+
+    # Broadcast all advanced indices, if any
+    if advanced_index_shapes:
+        broadcasted_shape = np.broadcast_shapes(*advanced_index_shapes)
+        result_shape.extend(broadcasted_shape)
+        basic_shape_index += len(
+            advanced_index_shapes
+        )  # Consume dimensions from array_shape
+
+    # Process each remaining index in index_tuple
+    for idx in index_tuple:
+        if isinstance(idx, int):
+            # Integer index reduces dimension, so skip this dimension in array_shape
+            basic_shape_index += 1
+        elif isinstance(idx, slice):
+            if idx.step is not None and idx.step > 1:
+                raise DiscontiguousArrayError(
+                    "Step size greater than 1 is not supported"
+                )
+            # Slice keeps dimension, adjust size accordingly
+            start, stop, _ = idx.indices(array_shape[basic_shape_index])
+            result_shape.append(stop - start)
+            basic_shape_index += 1
+        elif idx is Ellipsis:
+            # Calculate number of dimensions that Ellipsis should fill
+            num_to_fill = len(array_shape) - len(index_tuple) + 1
+            result_shape.extend(
+                array_shape[basic_shape_index : basic_shape_index + num_to_fill]
+            )
+            basic_shape_index += num_to_fill
+
+    # Step 4: Append remaining dimensions from array_shape if fewer indices were used
+    if basic_shape_index < len(array_shape) and pad:
+        result_shape.extend(array_shape[basic_shape_index:])
+
+    return tuple(size for idx, size in enumerate(result_shape) if idx not in drop_axes)
+
+
+def get_shape_for_selector(
+    selector_tuple: SelectorTuple,
+    shape: tuple[int, ...],
+    drop_axes: tuple[int, ...],
+    *,
+    pad: bool,
+) -> tuple[int, ...]:
+    if isinstance(selector_tuple, slice | np.ndarray):
+        return resulting_shape_from_index(
+            shape,
+            (selector_tuple,),
+            drop_axes,
+            pad=pad,
+        )
+    return resulting_shape_from_index(shape, selector_tuple, drop_axes, pad=pad)
+
+
 def make_chunk_info_for_rust_with_indices(
     batch_info: Iterable[
         tuple[ByteGetter | ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]
     ],
+    drop_axes: tuple,
 ) -> list[tuple[tuple[str, ChunkCoords, str, Any], list[slice], list[slice]]]:
+    # all?
+    for _, chunk_spec, chunk_selection, out_selection in batch_info:
+        shape_out_selection = get_shape_for_selector(
+            out_selection, chunk_spec.shape, (), pad=False
+        )
+        shape_chunk_selection = get_shape_for_selector(
+            chunk_selection, chunk_spec.shape, drop_axes, pad=True
+        )
+        if len(shape_chunk_selection) != len(shape_out_selection):
+            raise CollapsedDimensionError()
     return list(
         (
             convert_chunk_to_primitive(byte_getter, chunk_spec),

tests/conftest.py

@@ -116,6 +116,14 @@ def array_fixture(request: pytest.FixtureRequest) -> npt.NDArray[Any]:
     "test_roundtrip[vindex-discontinuous_in_chunk_array-contiguous_in_chunk_array]",
 ]
 
+# vindexing with two contiguous arrays would be converted to two slices but
+# in numpy indexing actually requires dropping a dimension, which in turn boils
+# down to integer indexing, which we can't do i.e., [np.array(1, 2), np.array(1, 2)] -> [slice(1, 2), slice(1, 2)]
+# is not a correct conversion, and thus we don't support the write operation
+zarrs_python_no_collapsed_dim = [
+    "test_roundtrip[vindex-contiguous_in_chunk_array-contiguous_in_chunk_array]"
+]
+
 
 def pytest_collection_modifyitems(
     config: pytest.Config, items: Iterable[pytest.Item]
@@ -128,5 +136,6 @@ def pytest_collection_modifyitems(
             item.name
             in zarr_python_default_codec_pipeline_failures
             + zarrs_python_no_discontinuous_writes
+            + zarrs_python_no_collapsed_dim
         ):
             item.add_marker(xfail_marker)