Add a blocks property to Cubed arrays to allow block-level indexing (#737)

* Add a `blocks` property to Cubed arrays to allow block-level indexing

This is similar to Dask and Zarr.

* Disallow multiple integer array indexes for `blocks`
Fix typo
Handle missing else case

Author: tomwhite

cubed/core/array.py

@@ -57,6 +57,13 @@ def zarray(self):
         """The underlying Zarr array. May only be used during the computation once the array has been created."""
         return open_if_lazy_zarr_array(self._zarray)
 
+    @property
+    def blocks(self):
+        """An array-like interface to the blocks of an array."""
+        from cubed.core.indexing import BlockView
+
+        return BlockView(self)
+
     @property
     def chunkmem(self):
         """Amount of memory in bytes that a single chunk uses."""

cubed/core/indexing.py

@@ -0,0 +1,94 @@
+import math
+from typing import TYPE_CHECKING
+
+import ndindex
+import numpy as np
+from toolz import map
+
+from cubed.core.ops import general_blockwise
+
+if TYPE_CHECKING:
+    from cubed.array_api.array_object import Array
+
+
+class BlockView:
+    """An array-like interface to the blocks of an array."""
+
+    def __init__(self, array: "Array"):
+        self.array = array
+
+    def __getitem__(self, key) -> "Array":
+        if not isinstance(key, tuple):
+            key = (key,)
+
+        # Canonicalize index
+        idx = ndindex.ndindex(key)
+        idx = idx.expand(self.array.numblocks)
+
+        if any(isinstance(ia, ndindex.Newaxis) for ia in idx.args):
+            raise ValueError("Slicing with xp.newaxis is not supported")
+
+        if sum(1 for ia in idx.args if isinstance(ia, ndindex.IntegerArray)) > 1:
+            raise NotImplementedError("Only one integer array index is allowed.")
+
+        # convert Integer to Slice so we don't lose dimensions
+        def convert_integer_index_to_slice(ia):
+            if isinstance(ia, ndindex.Integer):
+                return ndindex.Slice(ia.raw, ia.raw + 1)
+            return ia
+
+        idx = ndindex.Tuple(*(convert_integer_index_to_slice(ia) for ia in idx.args))
+
+        chunks = tuple(
+            tuple(np.array(ch)[ia].tolist())
+            for ia, ch in zip(idx.raw, self.array.chunks)
+        )
+        shape = tuple(map(sum, chunks))
+
+        identity = lambda a: a
+
+        def get_dim_index(ia, i):
+            if isinstance(ia, ndindex.Slice):
+                step = ia.step or 1
+                return ia.start + (step * i)
+            elif isinstance(ia, ndindex.IntegerArray):
+                return ia.raw[i]
+            else:
+                raise NotImplementedError(
+                    "Only integer, slice, or int array indexes are supported."
+                )
+
+        def key_function(out_key):
+            out_coords = out_key[1:]
+            in_coords = tuple(
+                get_dim_index(ia, bi) for ia, bi in zip(idx.args, out_coords)
+            )
+            return ((self.array.name, *in_coords),)
+
+        out = general_blockwise(
+            identity,
+            key_function,
+            self.array,
+            shapes=[shape],
+            dtypes=[self.array.dtype],
+            chunkss=[chunks],
+        )
+
+        from cubed import Array
+
+        assert isinstance(out, Array)  # single output
+        return out
+
+    @property
+    def size(self) -> int:
+        """
+        The total number of blocks in the array.
+        """
+        return math.prod(self.shape)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """
+        The number of blocks per axis.
+        """
+        return self.array.numblocks

cubed/tests/test_indexing.py

@@ -78,3 +78,34 @@ def test_multiple_int_array_indexes(spec):
     )
     with pytest.raises(NotImplementedError):
         a[[1, 2, 1], [2, 1, 0]]
+
+
+def test_blocks():
+    # based on dask tests
+    x = xp.arange(10, chunks=2)
+    assert x.blocks.shape == (5,)
+    assert x.blocks.size == 5
+
+    assert_array_equal(x.blocks[0], x[:2])
+    assert_array_equal(x.blocks[-1], x[-2:])
+    assert_array_equal(x.blocks[:3], x[:6])
+    assert_array_equal(x.blocks[[0, 1, 2]], x[:6])
+    assert_array_equal(x.blocks[[3, 0, 2]], np.array([6, 7, 0, 1, 4, 5]))
+
+    x = cubed.random.random((20, 20), chunks=(4, 5))
+    assert x.blocks.shape == (5, 4)
+    assert x.blocks.size == 20
+    assert_array_equal(x.blocks[0], x[:4])
+    assert_array_equal(x.blocks[0, :3], x[:4, :15])
+    assert_array_equal(x.blocks[:, :3], x[:, :15])
+
+    x = xp.ones((40, 40, 40), chunks=(10, 10, 10))
+    assert_array_equal(x.blocks[0, :, 0], np.ones((10, 40, 10)))
+
+    x = xp.ones((2, 2), chunks=1)
+    with pytest.raises(ValueError, match="newaxis is not supported"):
+        x.blocks[xp.newaxis, :, :]
+    with pytest.raises(NotImplementedError):
+        x.blocks[[0, 1], [0, 1]]
+    with pytest.raises(IndexError, match="out of bounds"):
+        x.blocks[100, 100]