add block indexing capabilities

Author: Altay Sansal

docs/api/core.rst

@@ -10,6 +10,8 @@ The Array class (``zarr.core``)
     .. automethod:: set_basic_selection
     .. automethod:: get_mask_selection
     .. automethod:: set_mask_selection
+    .. automethod:: get_block_selection
+    .. automethod:: set_block_selection
     .. automethod:: get_coordinate_selection
     .. automethod:: set_coordinate_selection
     .. automethod:: get_orthogonal_selection

zarr/core.py

@@ -22,6 +22,8 @@
     OIndex,
     OrthogonalIndexer,
     VIndex,
+    BlockIndex,
+    BlockIndexer,
     PartialChunkIterator,
     check_fields,
     check_no_multi_fields,
@@ -139,6 +141,7 @@ class Array:
     info
     vindex
     oindex
+    blocks
     write_empty_chunks
     meta_array
 
@@ -154,6 +157,8 @@ class Array:
     set_mask_selection
     get_coordinate_selection
     set_coordinate_selection
+    get_block_selection
+    set_block_selection
     digest
     hexdigest
     resize
@@ -230,6 +235,7 @@ def __init__(
         # initialize indexing helpers
         self._oindex = OIndex(self)
         self._vindex = VIndex(self)
+        self._blocks = BlockIndex(self)
 
     def _load_metadata(self):
         """(Re)load metadata from store."""
@@ -577,6 +583,12 @@ def vindex(self):
         :func:`set_mask_selection` for documentation and examples."""
         return self._vindex
 
+    @property
+    def blocks(self):
+        """Shortcut for blocked chunked indexing, see :func:`get_block_selection` and
+        :func:`set_block_selection` for documentation and examples."""
+        return self._blocks
+
     @property
     def write_empty_chunks(self) -> bool:
         """A Boolean, True if chunks composed of the array's fill value
@@ -814,7 +826,8 @@ def __getitem__(self, selection):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, get_orthogonal_selection,
-        set_orthogonal_selection, vindex, oindex, __setitem__
+        set_orthogonal_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __setitem__
 
         """
         fields, pure_selection = pop_fields(selection)
@@ -933,7 +946,8 @@ def get_basic_selection(self, selection=Ellipsis, out=None, fields=None):
         --------
         set_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, get_orthogonal_selection,
-        set_orthogonal_selection, vindex, oindex, __getitem__, __setitem__
+        set_orthogonal_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1089,7 +1103,8 @@ def get_orthogonal_selection(self, selection, out=None, fields=None):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, set_orthogonal_selection,
-        vindex, oindex, __getitem__, __setitem__
+        get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1160,7 +1175,8 @@ def get_coordinate_selection(self, selection, out=None, fields=None):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_orthogonal_selection, set_orthogonal_selection, set_coordinate_selection,
-        vindex, oindex, __getitem__, __setitem__
+        get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1185,6 +1201,89 @@ def get_coordinate_selection(self, selection, out=None, fields=None):
 
         return out
 
+    def get_block_selection(self, selection, out=None, fields=None):
+        """Retrieve a selection of individual chunk blocks, by providing the indices
+        (coordinates) for each chunk block.
+
+        Parameters
+        ----------
+        selection : tuple
+            An integer (coordinate) or slice for each dimension of the array.
+        out : ndarray, optional
+            If given, load the selected data directly into this array.
+        fields : str or sequence of str, optional
+            For arrays with a structured dtype, one or more fields can be specified to
+            extract data for.
+
+        Returns
+        -------
+        out : ndarray
+            A NumPy array containing the data for the requested selection.
+
+        Examples
+        --------
+        Setup a 2-dimensional array::
+
+            >>> import zarr
+            >>> import numpy as np
+            >>> z = zarr.array(np.arange(100).reshape(10, 10), chunks=(3, 3))
+
+        Retrieve items by specifying their block coordinates::
+
+            >>> z.get_block_selection((1, slice(None)))
+            array([[30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
+                   [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
+                   [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]])
+
+        Which is equivalent to:
+            >>> z[3:6, :]
+            array([[30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
+                   [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
+                   [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]])
+
+        For convenience, the block selection functionality is also available via the
+        `blocks` property, e.g.::
+
+            >>> z.blocks[1]
+            array([[30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
+                   [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
+                   [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]])
+
+        Notes
+        -----
+        Block indexing is a convenience indexing method to work on individual chunks
+        with chunk index slicing. It has the same concept as Dask's `Array.blocks`
+        indexing.
+
+        Slices are supported. However, only with a step size of one.
+
+        Block index arrays may be multidimensional to index multidimensional arrays.
+        For example:
+
+            >>> z.blocks[0, 1:3]
+            array([[ 3,  4,  5,  6,  7,  8],
+                   [13, 14, 15, 16, 17, 18],
+                   [23, 24, 25, 26, 27, 28]])
+
+        See Also
+        --------
+        get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
+        get_orthogonal_selection, set_orthogonal_selection, get_coordinate_selection,
+        set_coordinate_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
+
+        """
+        if not self._cache_metadata:
+            self._load_metadata()
+
+        # check args
+        check_fields(fields, self._dtype)
+
+        # setup indexer
+        indexer = BlockIndexer(selection, self)
+
+        return self._get_selection(indexer=indexer, out=out, fields=fields)
+
     def get_mask_selection(self, selection, out=None, fields=None):
         """Retrieve a selection of individual items, by providing a Boolean array of the
         same shape as the array against which the selection is being made, where True
@@ -1238,8 +1337,8 @@ def get_mask_selection(self, selection, out=None, fields=None):
         --------
         get_basic_selection, set_basic_selection, set_mask_selection,
         get_orthogonal_selection, set_orthogonal_selection, get_coordinate_selection,
-        set_coordinate_selection, vindex, oindex, __getitem__, __setitem__
-
+        set_coordinate_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
         """
 
         # refresh metadata
@@ -1376,7 +1475,8 @@ def __setitem__(self, selection, value):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, get_orthogonal_selection,
-        set_orthogonal_selection, vindex, oindex, __getitem__
+        set_orthogonal_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__
 
         """
         fields, pure_selection = pop_fields(selection)
@@ -1464,7 +1564,8 @@ def set_basic_selection(self, selection, value, fields=None):
         --------
         get_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, get_orthogonal_selection,
-        set_orthogonal_selection, vindex, oindex, __getitem__, __setitem__
+        set_orthogonal_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1555,7 +1656,8 @@ def set_orthogonal_selection(self, selection, value, fields=None):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_coordinate_selection, set_coordinate_selection, get_orthogonal_selection,
-        vindex, oindex, __getitem__, __setitem__
+        get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1627,7 +1729,8 @@ def set_coordinate_selection(self, selection, value, fields=None):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
         get_orthogonal_selection, set_orthogonal_selection, get_coordinate_selection,
-        vindex, oindex, __getitem__, __setitem__
+        get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """
 
@@ -1654,6 +1757,89 @@ def set_coordinate_selection(self, selection, value, fields=None):
 
         self._set_selection(indexer, value, fields=fields)
 
+    def set_block_selection(self, selection, value, fields=None):
+        """Modify a selection of individual blocks, by providing the chunk indices
+        (coordinates) for each block to be modified.
+
+        Parameters
+        ----------
+        selection : tuple
+            An integer (coordinate) or slice for each dimension of the array.
+        value : scalar or array-like
+            Value to be stored into the array.
+        fields : str or sequence of str, optional
+            For arrays with a structured dtype, one or more fields can be specified to set
+            data for.
+
+        Examples
+        --------
+        Setup a 2-dimensional array::
+
+            >>> import zarr
+            >>> import numpy as np
+            >>> z = zarr.zeros((6, 6), dtype=int, chunks=2)
+
+        Set data for a selection of items::
+
+            >>> z.set_block_selection((1, 0), 1)
+            >>> z[...]
+            array([[0, 0, 0, 0, 0, 0],
+                   [0, 0, 0, 0, 0, 0],
+                   [1, 1, 0, 0, 0, 0],
+                   [1, 1, 0, 0, 0, 0],
+                   [0, 0, 0, 0, 0, 0],
+                   [0, 0, 0, 0, 0, 0]])
+
+        For convenience, this functionality is also available via the `blocks` property.
+        E.g.::
+
+            >>> z.blocks[2, 1] = 4
+            >>> z[...]
+            array([[0, 0, 0, 0, 0, 0],
+                   [0, 0, 0, 0, 0, 0],
+                   [1, 1, 0, 0, 0, 0],
+                   [1, 1, 0, 0, 0, 0],
+                   [0, 0, 4, 4, 0, 0],
+                   [0, 0, 4, 4, 0, 0]])
+
+            >>> z.blocks[:, 2] = 7
+            >>> z[...]
+            array([[0, 0, 0, 0, 7, 7],
+                   [0, 0, 0, 0, 7, 7],
+                   [1, 1, 0, 0, 7, 7],
+                   [1, 1, 0, 0, 7, 7],
+                   [0, 0, 4, 4, 7, 7],
+                   [0, 0, 4, 4, 7, 7]])
+
+        Notes
+        -----
+        Block indexing is a convenience indexing method to work on individual chunks
+        with chunk index slicing. It has the same concept as Dask's `Array.blocks`
+        indexing.
+
+        Slices are supported. However, only with a step size of one.
+
+        See Also
+        --------
+        get_basic_selection, set_basic_selection, get_mask_selection, set_mask_selection,
+        get_orthogonal_selection, set_orthogonal_selection, get_coordinate_selection,
+        get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
+
+        """
+        # guard conditions
+        if self._read_only:
+            raise ReadOnlyError()
+
+        # refresh metadata
+        if not self._cache_metadata:
+            self._load_metadata_nosync()
+
+        # setup indexer
+        indexer = BlockIndexer(selection, self)
+
+        self._set_selection(indexer, value, fields=fields)
+
     def set_mask_selection(self, selection, value, fields=None):
         """Modify a selection of individual items, by providing a Boolean array of the
         same shape as the array against which the selection is being made, where True
@@ -1712,7 +1898,8 @@ def set_mask_selection(self, selection, value, fields=None):
         --------
         get_basic_selection, set_basic_selection, get_mask_selection,
         get_orthogonal_selection, set_orthogonal_selection, get_coordinate_selection,
-        set_coordinate_selection, vindex, oindex, __getitem__, __setitem__
+        set_coordinate_selection, get_block_selection, set_block_selection,
+        vindex, oindex, blocks, __getitem__, __setitem__
 
         """