Start stop for iterator (#621)

* added option to iterate over part of array

* fixed a small error in islice

* added test for array.islice

* mimick itertools.islice input param edge case handling

* add docstring for islice
make input argument error handling more verbose

* add missing input param edge case to tests

* fix islice docstring

* add islice to API docs

* API doc entry in wrong file fixed

Co-authored-by: jmoore <[email protected]>
Co-authored-by: Josh Moore <[email protected]>

Author: yetyetanotherusername

docs/api/hierarchy.rst

@@ -38,4 +38,4 @@ Groups (``zarr.hierarchy``)
     .. automethod:: zeros_like
     .. automethod:: ones_like
     .. automethod:: full_like
-    .. automethod:: move
+    .. automethod:: move

docs/release.rst

@@ -116,6 +116,10 @@ This release will be the last to support Python 3.5, next version of Zarr will b
   applies to reading.
   By :user:`Martin Durant <martindurant>`; :issue:`606`
 
+* Efficient iteration expanded with option to pass start and stop index via
+  ``array.islice``.
+  By :user:`Sebastian Grill <yetyetanotherusername>`, :issue:`615`.
+
 .. _release_2.4.0:
 
 2.4.0

zarr/core.py

@@ -458,18 +458,72 @@ def __array__(self, *args):
             a = a.astype(args[0])
         return a
 
-    def __iter__(self):
+    def islice(self, start=None, end=None):
+        """
+        Yield a generator for iterating over the entire or parts of the
+        array. Uses a cache so chunks only have to be decompressed once.
+
+        Parameters
+        ----------
+        start : int, optional
+            Start index for the generator to start at. Defaults to 0.
+        end : int, optional
+            End index for the generator to stop at. Defaults to self.shape[0].
+
+        Yields
+        ------
+        out : generator
+            A generator that can be used to iterate over the requested region
+            the array.
+
+        Examples
+        --------
+        Setup a 1-dimensional array::
+
+            >>> import zarr
+            >>> import numpy as np
+            >>> z = zarr.array(np.arange(100))
+
+        Iterate over part of the array:
+            >>> for value in z.islice(25, 30): value;
+            25
+            26
+            27
+            28
+            29
+        """
+
         if len(self.shape) == 0:
             # Same error as numpy
             raise TypeError("iteration over a 0-d array")
+        if start is None:
+            start = 0
+        if end is None or end > self.shape[0]:
+            end = self.shape[0]
+
+        if not isinstance(start, int) or start < 0:
+            raise ValueError('start must be a nonnegative integer')
+
+        if not isinstance(end, int) or end < 0:
+            raise ValueError('end must be a nonnegative integer')
+
         # Avoid repeatedly decompressing chunks by iterating over the chunks
         # in the first dimension.
         chunk_size = self.chunks[0]
-        for j in range(self.shape[0]):
+        chunk = None
+        for j in range(start, end):
             if j % chunk_size == 0:
                 chunk = self[j: j + chunk_size]
+            # init chunk if we start offset of chunk borders
+            elif chunk is None:
+                chunk_start = j - j % chunk_size
+                chunk_end = chunk_start + chunk_size
+                chunk = self[chunk_start:chunk_end]
             yield chunk[j % chunk_size]
 
+    def __iter__(self):
+        return self.islice()
+
     def __len__(self):
         if self.shape:
             return self.shape[0]

zarr/tests/test_core.py

@@ -1347,7 +1347,8 @@ def test_object_codec_warnings(self):
         if hasattr(z.store, 'close'):
             z.store.close()
 
-    def test_zero_d_iter(self):
+    def test_iteration_exceptions(self):
+        # zero d array
         a = np.array(1, dtype=int)
         z = self.create_array(shape=a.shape, dtype=int)
         z[...] = a
@@ -1357,6 +1358,27 @@ def test_zero_d_iter(self):
         with pytest.raises(TypeError):
             # noinspection PyStatementEffect
             list(z)
+
+        # input argument error handling
+        a = np.array((10, 10), dtype=int)
+        z = self.create_array(shape=a.shape, dtype=int)
+        z[...] = a
+
+        params = (
+            (-1, 0),
+            (0, -1),
+            (0.5, 1),
+            (0, 0.5)
+        )
+
+        for start, end in params:
+            with pytest.raises(ValueError):
+                # noinspection PyStatementEffect
+                list(z.islice(start, end))
+
+        # check behavior for start > end
+        assert [] == list(z.islice(6, 5))
+
         if hasattr(z.store, 'close'):
             z.store.close()
 
@@ -1385,6 +1407,30 @@ def test_iter(self):
             if hasattr(z.store, 'close'):
                 z.store.close()
 
+    def test_islice(self):
+        params = (
+            ((1,), (1,), 0, 1),
+            ((2,), (1,), 0, 1),
+            ((1,), (2,), 0, 1),
+            ((3,), (3,), 1, 2),
+            ((1000,), (100,), 150, 1050),
+            ((100,), (1000,), 25, 75),
+            ((1, 100), (1, 1), 0, 1),
+            ((100, 1), (3, 1), 56, 100),
+            ((100, 100), (10, 10), 13, 99),
+            ((10, 10, 10), (3, 3, 3), 2, 4),
+        )
+        for shape, chunks, start, end in params:
+            z = self.create_array(shape=shape, chunks=chunks, dtype=int)
+            a = np.arange(np.product(shape)).reshape(shape)
+            z[:] = a
+            end_array = min(end, a.shape[0])
+            for expect, actual in zip_longest(a[start:end_array],
+                                              z.islice(start, end)):
+                assert_array_equal(expect, actual)
+            if hasattr(z.store, 'close'):
+                z.store.close()
+
     def test_compressors(self):
         compressors = [
             None, BZ2(), Blosc(), LZ4(), Zlib(), GZip()