adds partial_decompress capabilites

zarr/core.py

@@ -12,9 +12,10 @@
 
 from zarr.attrs import Attributes
 from zarr.codecs import AsType, get_codec
-from zarr.errors import err_array_not_found, err_read_only
+from zarr.errors import ArrayIndexError, err_array_not_found, err_read_only
 from zarr.indexing import (BasicIndexer, CoordinateIndexer, MaskIndexer,
-                           OIndex, OrthogonalIndexer, VIndex, check_fields,
+                           OIndex, OrthogonalIndexer, VIndex, PartialChunkIterator,
+                           check_fields,
                            check_no_multi_fields, ensure_tuple,
                            err_too_many_indices, is_contiguous_selection,
                            is_scalar, pop_fields)
@@ -1626,6 +1627,24 @@ def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
                     return
 
             # decode chunk
+            try:
+                if self._compressor and self._compressor.codec_id == 'blosc' \
+                   and not fields and self.dtype != object:
+                    tmp = np.empty(self._chunks, dtype=self.dtype)
+                    index_selection = PartialChunkIterator(chunk_selection, self.chunks)
+                    for start, nitems, partial_out_selection in index_selection:
+                        expected_shape = [
+                            len(range(*partial_out_selection[i].indices(self.chunks[0]+1)))
+                            if i < len(partial_out_selection) else dim
+                            for i, dim in enumerate(self.chunks)]
+                        chunk_partial = self._decode_chunk(
+                            cdata, start=start, nitems=nitems,
+                            expected_shape=expected_shape)
+                        tmp[partial_out_selection] = chunk_partial
+                    out[out_selection] = tmp[chunk_selection]
+                    return
+            except ArrayIndexError:
+                pass
             chunk = self._decode_chunk(cdata)
 
             # select data from chunk
@@ -1731,11 +1750,16 @@ def _chunk_setitem_nosync(self, chunk_coords, chunk_selection, value, fields=Non
     def _chunk_key(self, chunk_coords):
         return self._key_prefix + '.'.join(map(str, chunk_coords))
 
-    def _decode_chunk(self, cdata):
+    def _decode_chunk(self, cdata, start=None, nitems=None, expected_shape=None):    
 
         # decompress
         if self._compressor:
-            chunk = self._compressor.decode(cdata)
+            # only decode requested items
+            if (all([x is not None for x in [start, nitems]])
+                and self._compressor.codec_id == 'blosc'):
+                chunk = self._compressor.decode_partial(cdata, start, nitems)
+            else:
+                chunk = self._compressor.decode(cdata)
         else:
             chunk = cdata
 
@@ -1760,7 +1784,7 @@ def _decode_chunk(self, cdata):
 
         # ensure correct chunk shape
         chunk = chunk.reshape(-1, order='A')
-        chunk = chunk.reshape(self._chunks, order=self._order)
+        chunk = chunk.reshape(expected_shape or self._chunks, order=self._order)
 
         return chunk
 

zarr/errors.py

@@ -9,6 +9,10 @@ class CopyError(RuntimeError):
     pass
 
 
+class ArrayIndexError(IndexError):
+    pass
+
+
 def err_contains_group(path):
     raise ValueError('path %r contains a group' % path)
 

zarr/indexing.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from zarr.errors import (err_boundscheck, err_negative_step,
+from zarr.errors import (ArrayIndexError, err_boundscheck, err_negative_step,
                          err_too_many_indices, err_vindex_invalid_selection)
 
 
@@ -822,3 +822,83 @@ def pop_fields(selection):
         selection = tuple(s for s in selection if not isinstance(s, str))
         selection = selection[0] if len(selection) == 1 else selection
     return fields, selection
+
+
+def int_to_slice(dim_selection):
+    return slice(dim_selection, dim_selection+1, 1)
+
+def make_slice_selection(selection):
+    ls = []
+    for dim_selection in selection:
+        if is_integer(dim_selection):
+            ls.append(int_to_slice(dim_selection))
+        elif isinstance(dim_selection, np.ndarray):
+            if len(dim_selection) == 1:
+                ls.append(int_to_slice(dim_selection[0]))
+            else:
+                raise ArrayIndexError()
+        else:
+            ls.append(dim_selection)
+    return ls
+
+
+class PartialChunkIterator(object):
+    """Iterator tp retrieve the specific coordinates of requested data
+    from within a compressed chunk.
+
+    Parameters
+    -----------
+    selection : tuple
+        tuple of slice objects to take from the chunk
+    arr_shape : shape of chunk to select data from
+
+    Attributes
+    -----------
+    arr_shape
+    selection
+    """
+
+    def __init__(self, selection, arr_shape):
+        self.selection = make_slice_selection(selection)
+        self.arr_shape = arr_shape
+
+        # number of selection dimensions can't be greater than the number of chunk dimensions
+        if len(self.selection) > len(self.arr_shape):
+            raise ValueError('Selection has more dimensions then the array:\n'
+                             'selection dimensions = {len(self.selection)\n'
+                             'array dimensions = {len(self.arr_shape)}')
+
+        # any selection can not be out of the range of the chunk
+        self.selection_shape = np.empty(self.arr_shape)[self.selection].shape
+        if any([selection_dim < 0 or selection_dim > arr_dim for selection_dim, arr_dim
+                in zip(self.selection_shape, self.arr_shape)]):
+            raise IndexError('a selection index is out of range for the dimension')
+
+        for i, dim_size in enumerate(self.arr_shape[::-1]):
+            index = len(self.arr_shape) - (i+1)
+            if index <= len(self.selection)-1:
+                slice_size = self.selection_shape[index]
+                if slice_size == dim_size and index > 0:
+                    self.selection.pop()
+                else:
+                    break
+
+        chunk_loc_slices = []
+        last_dim_slice = None if self.selection[-1].step > 1 else self.selection.pop()
+        for i, sl in enumerate(self.selection):
+            dim_chunk_loc_slices = []
+            for i, x in enumerate(slice_to_range(sl, arr_shape[i])):
+                dim_chunk_loc_slices.append(slice(x, x+1, 1))
+            chunk_loc_slices.append(dim_chunk_loc_slices)
+        if last_dim_slice:
+            chunk_loc_slices.append([last_dim_slice])
+        self.chunk_loc_slices = list(itertools.product(*chunk_loc_slices))
+
+    def __iter__(self):
+        chunk1 = self.chunk_loc_slices[0]
+        nitems = (chunk1[-1].stop - chunk1[-1].start) * np.prod(self.arr_shape[len(chunk1):])
+        for chunk_selection in self.chunk_loc_slices:
+            start = 0
+            for i, sl in enumerate(chunk_selection):
+                start += sl.start * np.prod(self.arr_shape[i+1:])
+            yield int(start), int(nitems), chunk_selection

zarr/tests/test_core.py

@@ -278,6 +278,7 @@ def test_array_1d_selections(self):
         ix = [99, 100, 101]
         bix = np.zeros_like(a, dtype=bool)
         bix[ix] = True
+        print(a[ix])
         assert_array_equal(a[ix], z.get_orthogonal_selection(ix))
         assert_array_equal(a[ix], z.oindex[ix])
         assert_array_equal(a[ix], z.get_coordinate_selection(ix))

zarr/tests/test_indexing.py

@@ -5,7 +5,7 @@
 
 import zarr
 from zarr.indexing import (normalize_integer_selection, oindex, oindex_set,
-                           replace_ellipsis)
+                           replace_ellipsis, PartialChunkIterator)
 
 
 def test_normalize_integer_selection():
@@ -1289,3 +1289,61 @@ def test_set_selections_with_fields():
             a[key][ix] = v[key][ix]
             z.set_mask_selection(ix, v[key][ix], fields=fields)
             assert_array_equal(a, z[:])
+
+
+@pytest.mark.parametrize('selection, arr, expected', [
+    ((slice(5, 8, 1), slice(2, 4, 1), slice(0, 100, 1)),
+     np.arange(2, 100_002).reshape((100, 10, 100)),
+     [(5200, 200, (slice(5, 6, 1), slice(2, 4, 1))),
+      (6200, 200, (slice(6, 7, 1), slice(2, 4, 1))),
+      (7200, 200, (slice(7, 8, 1), slice(2, 4, 1)))]),
+    ((slice(5, 8, 1), slice(2, 4, 1), slice(0, 5, 1)),
+     np.arange(2, 100_002).reshape((100, 10, 100)),
+     [(5200.0, 5.0, (slice(5, 6, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (5300.0, 5.0, (slice(5, 6, 1), slice(3, 4, 1), slice(0, 5, 1))),
+      (6200.0, 5.0, (slice(6, 7, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (6300.0, 5.0, (slice(6, 7, 1), slice(3, 4, 1), slice(0, 5, 1))),
+      (7200.0, 5.0, (slice(7, 8, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (7300.0, 5.0, (slice(7, 8, 1), slice(3, 4, 1), slice(0, 5, 1)))]),
+    ((slice(5, 8, 1), slice(2, 4, 1), slice(0, 5, 1)),
+     np.asfortranarray(np.arange(2, 100_002).reshape((100, 10, 100))),
+     [(5200.0, 5.0, (slice(5, 6, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (5300.0, 5.0, (slice(5, 6, 1), slice(3, 4, 1), slice(0, 5, 1))),
+      (6200.0, 5.0, (slice(6, 7, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (6300.0, 5.0, (slice(6, 7, 1), slice(3, 4, 1), slice(0, 5, 1))),
+      (7200.0, 5.0, (slice(7, 8, 1), slice(2, 3, 1), slice(0, 5, 1))),
+      (7300.0, 5.0, (slice(7, 8, 1), slice(3, 4, 1), slice(0, 5, 1)))]),
+    ((slice(5, 8, 1), slice(2, 4, 1)),
+     np.arange(2, 100_002).reshape((100, 10, 100)),
+     [(5200, 200, (slice(5, 6, 1), slice(2, 4, 1))),
+      (6200, 200, (slice(6, 7, 1), slice(2, 4, 1))),
+      (7200, 200, (slice(7, 8, 1), slice(2, 4, 1)))]),
+    ((slice(0, 10, 1),),
+     np.arange(0, 10).reshape((10)),
+     [(0, 10, (slice(0, 10, 1),))]),
+    ((0,),
+     np.arange(0, 100).reshape((10, 10)),
+     [(0, 10, (slice(0, 1, 1),))]),
+    ((0,0,),
+     np.arange(0, 100).reshape((10, 10)),
+     [(0, 1, (slice(0, 1, 1), slice(0, 1, 1)))]),
+    ((0,),
+     np.arange(0, 10).reshape((10)),
+     [(0, 1, (slice(0, 1, 1),))]),
+    pytest.param((slice(5, 8, 1), slice(2, 4, 1), slice(0, 5, 1)),
+                 np.arange(2, 100002).reshape((10, 1, 10000)),
+                 None,
+                 marks=[pytest.mark.xfail(reason='slice 2 is out of range')]
+    ),
+    pytest.param((slice(5, 8, 1), slice(2, 4, 1), slice(0, 5, 1)),
+                 np.arange(2, 100_002).reshape((10, 10_000)),
+                 None,
+                 marks=[pytest.mark.xfail(reason='slice 2 is out of range')]
+    ),
+])
+def test_PartialChunkIterator(selection, arr, expected):
+    print(selection)
+    PCI = PartialChunkIterator(selection, arr.shape)
+    results = list(PCI)
+    assert(results == expected)
+