tested array getitem and setitem; cc #21

Author: alimanfoo

zarr/array.py

@@ -5,14 +5,45 @@
 from functools import reduce  # TODO PY2 compatibility
 import operator
 import itertools
+import multiprocessing
 import numpy as np
 
 
-from zarr.blosc import compress, decompress
+from zarr import blosc
 from zarr.util import is_total_slice, normalize_array_selection, \
     get_chunk_range
 
 
+_blosc_use_context = False
+
+
+def set_blosc_options(use_context=False, nthreads=None):
+    """Set options for how the blosc compressor is used.
+
+    Parameters
+    ----------
+    use_context : bool, optional
+        If False, blosc will be used in non-contextual mode, which is best
+        when using zarr in a single-threaded environment because it allows
+        blosc to use multiple threads internally. If True, blosc will be used
+        in contextual mode, which is better when using zarr in a
+        multi-threaded environment like dask.array because it avoids the blosc
+        global lock and so multiple blosc operations can be running
+        concurrently.
+    nthreads : int, optional
+        Number of internal threads to use when running blosc in non-contextual
+        mode.
+
+    """
+    global _blosc_use_context
+    _blosc_use_context = use_context
+    if not use_context:
+        if nthreads is None:
+            # diminishing returns beyond 4 threads?
+            nthreads = max(4, multiprocessing.cpu_count())
+        blosc.set_nthreads(nthreads)
+
+
 class Array(object):
 
     def __init__(self, store):
@@ -38,6 +69,10 @@ def __init__(self, store):
         # user-defined attributes
         self._attrs = store.attrs
 
+    @property
+    def store(self):
+        return self._store
+
     @property
     def shape(self):
         return self._shape
@@ -170,11 +205,13 @@ def __setitem__(self, key, value):
                 # determine index within value
                 value_selection = tuple(
                     slice(max(0, o - start), min(o + c - start, stop - start))
-                    for (start, stop), o, c, in zip(selection, offset, self._chunks)
+                    for (start, stop), o, c, in zip(selection, offset,
+                                                    self._chunks)
                 )
 
                 # put data
-                self._chunk_setitem(cidx, chunk_selection, value[value_selection])
+                self._chunk_setitem(cidx, chunk_selection,
+                                    value[value_selection])
 
     def _chunk_getitem(self, cidx, item, dest):
         """Obtain part or whole of a chunk.
@@ -192,26 +229,36 @@ def _chunk_getitem(self, cidx, item, dest):
 
         # override this in sub-classes, e.g., if need to use a lock
 
-        # obtain compressed data for chunk
-        cdata = self._store.data[cidx]
+        try:
+
+            # obtain compressed data for chunk
+            cdata = self._store.data[cidx]
 
-        if is_total_slice(item, self._chunks) and dest.flags.c_contiguous:
+        except KeyError:
 
-            # optimisation: we want the whole chunk, and the destination is
-            # C contiguous, so we can decompress directly from the chunk
-            # into the destination array
-            decompress(cdata, dest)
+            # chunk not initialized
+            if self._fill_value is not None:
+                dest.fill(self._fill_value)
 
         else:
 
-            # decompress chunk
-            chunk = np.empty(self._chunks, dtype=self._dtype)
-            decompress(cdata, chunk)
+            if is_total_slice(item, self._chunks) and dest.flags.c_contiguous:
 
-            # set data in output array
-            # (split into two lines for profiling)
-            tmp = chunk[item]
-            dest[:] = tmp
+                # optimisation: we want the whole chunk, and the destination is
+                # C contiguous, so we can decompress directly from the chunk
+                # into the destination array
+                blosc.decompress(cdata, dest, _blosc_use_context)
+
+            else:
+
+                # decompress chunk
+                chunk = np.empty(self._chunks, dtype=self._dtype)
+                blosc.decompress(cdata, chunk, _blosc_use_context)
+
+                # set data in output array
+                # (split into two lines for profiling)
+                tmp = chunk[item]
+                dest[:] = tmp
 
     def _chunk_setitem(self, cidx, key, value):
         """Replace part or whole of a chunk.
@@ -248,18 +295,30 @@ def _chunk_setitem(self, cidx, key, value):
         else:
             # partially replace the contents of this chunk
 
-            # obtain compressed data for chunk
-            cdata = self._store.data[cidx]
+            try:
+
+                # obtain compressed data for chunk
+                cdata = self._store.data[cidx]
+
+            except KeyError:
+
+                # chunk not initialized
+                chunk = np.empty(self.chunks, dtype=self._dtype)
+                if self._fill_value is not None:
+                    chunk.fill(self._fill_value)
+
+            else:
 
-            # decompress
-            chunk = np.empty(self._chunks, dtype=self._dtype)
-            decompress(cdata, chunk)
+                # decompress chunk
+                chunk = np.empty(self.chunks, dtype=self._dtype)
+                blosc.decompress(cdata, chunk, _blosc_use_context)
 
             # modify
             chunk[key] = value
 
         # compress
-        cdata = compress(chunk, self._cname, self._clevel, self._shuffle)
+        cdata = blosc.compress(chunk, self._cname, self._clevel,
+                               self._shuffle, _blosc_use_context)
 
         # store
         self._store.data[cidx] = cdata

zarr/blosc.pyx

@@ -37,8 +37,6 @@ cdef extern from "blosc.h":
 				           size_t blocksize, int numinternalthreads) nogil
     int blosc_decompress_ctx(const void *src, void *dest, size_t destsize,
                              int numinternalthreads) nogil
-    void blosc_cbuffer_sizes(void *cbuffer, size_t *nbytes,
-                             size_t *cbytes, size_t *blocksize)
 
 
 def version():
@@ -64,6 +62,10 @@ def compname_to_compcode(bytes cname):
     return blosc_compname_to_compcode(cname)
 
 
+def set_nthreads(int nthreads):
+    blosc_set_nthreads(nthreads)
+
+
 def decompress(bytes cdata, np.ndarray array, use_context):
     """Decompress data into a numpy array.
 

zarr/tests/test_array.py

@@ -3,19 +3,24 @@
 
 
 import numpy as np
-from nose.tools import eq_ as eq
+from numpy.testing import assert_array_equal
+from nose.tools import eq_ as eq, assert_is_instance
 
 
 from zarr.store.memory import MemoryStore
 from zarr.array import Array
 from zarr import defaults
 
 
+def create_array(shape, chunks, **kwargs):
+    store = MemoryStore(shape, chunks, **kwargs)
+    return Array(store)
+
+
 def test_1d():
 
     a = np.arange(1050)
-    store = MemoryStore(a.shape, chunks=100, dtype=a.dtype)
-    z = Array(store)
+    z = create_array(a.shape, chunks=100, dtype=a.dtype)
 
     # check properties
     eq(a.shape, z.shape)
@@ -28,4 +33,128 @@ def test_1d():
     eq(0, z.cbytes)
     eq(0, z.initialized)
 
-    # TODO
+    # check empty
+    b = z[:]
+    assert_is_instance(b, np.ndarray)
+    eq(a.shape, b.shape)
+    eq(a.dtype, b.dtype)
+
+    # set data
+    z[:] = a
+
+    # check properties
+    eq(a.nbytes, z.nbytes)
+    assert z.cbytes > 0
+    eq(z.store.cbytes, z.cbytes)
+    eq(11, z.initialized)
+
+    # check slicing
+    assert_array_equal(a, z[:])
+    assert_array_equal(a, z[...])
+    # noinspection PyTypeChecker
+    assert_array_equal(a, z[slice(None)])
+    assert_array_equal(a[:10], z[:10])
+    assert_array_equal(a[10:20], z[10:20])
+    assert_array_equal(a[-10:], z[-10:])
+    # ...across chunk boundaries...
+    assert_array_equal(a[:110], z[:110])
+    assert_array_equal(a[190:310], z[190:310])
+    assert_array_equal(a[-110:], z[-110:])
+
+    # check partial assignment
+    b = np.arange(1e5, 2e5)
+    z[190:310] = b[190:310]
+    assert_array_equal(a[:190], z[:190])
+    assert_array_equal(b[190:310], z[190:310])
+    assert_array_equal(a[310:], z[310:])
+
+
+def test_array_1d_fill_value():
+
+    for fill_value in -1, 0, 1, 10:
+
+        a = np.arange(1050)
+        f = np.empty_like(a)
+        f.fill(fill_value)
+        z = create_array(shape=a.shape, chunks=100, dtype=a.dtype,
+                         fill_value=fill_value)
+        z[190:310] = a[190:310]
+
+        assert_array_equal(f[:190], z[:190])
+        assert_array_equal(a[190:310], z[190:310])
+        assert_array_equal(f[310:], z[310:])
+
+
+def test_array_1d_set_scalar():
+
+    # setup
+    a = np.empty(100)
+    z = create_array(shape=a.shape, chunks=10, dtype=a.dtype)
+    z[:] = a
+    assert_array_equal(a, z[:])
+
+    for value in -1, 0, 1, 10:
+        a[15:35] = value
+        z[15:35] = value
+        assert_array_equal(a, z[:])
+        a[:] = value
+        z[:] = value
+        assert_array_equal(a, z[:])
+
+
+def test_array_2d():
+
+    a = np.arange(10000).reshape((1000, 10))
+    z = create_array(shape=a.shape, chunks=(100, 2), dtype=a.dtype)
+
+    # check properties
+    eq(a.shape, z.shape)
+    eq(a.dtype, z.dtype)
+    eq((100, 2), z.chunks)
+    eq(defaults.cname, z.cname)
+    eq(defaults.clevel, z.clevel)
+    eq(defaults.shuffle, z.shuffle)
+    eq(0, z.cbytes)
+    eq(0, z.initialized)
+
+    # set data
+    z[:] = a
+
+    # check properties
+    eq(a.nbytes, z.nbytes)
+    assert z.cbytes > 0
+    eq(50, z.initialized)
+
+    # check slicing
+    assert_array_equal(a, z[:])
+    assert_array_equal(a, z[...])
+    # noinspection PyTypeChecker
+    assert_array_equal(a, z[slice(None)])
+    assert_array_equal(a[:10], z[:10])
+    assert_array_equal(a[10:20], z[10:20])
+    assert_array_equal(a[-10:], z[-10:])
+    assert_array_equal(a[:, :2], z[:, :2])
+    assert_array_equal(a[:, 2:4], z[:, 2:4])
+    assert_array_equal(a[:, -2:], z[:, -2:])
+    assert_array_equal(a[:10, :2], z[:10, :2])
+    assert_array_equal(a[10:20, 2:4], z[10:20, 2:4])
+    assert_array_equal(a[-10:, -2:], z[-10:, -2:])
+    # ...across chunk boundaries...
+    assert_array_equal(a[:110], z[:110])
+    assert_array_equal(a[190:310], z[190:310])
+    assert_array_equal(a[-110:], z[-110:])
+    assert_array_equal(a[:, :3], z[:, :3])
+    assert_array_equal(a[:, 3:7], z[:, 3:7])
+    assert_array_equal(a[:, -3:], z[:, -3:])
+    assert_array_equal(a[:110, :3], z[:110, :3])
+    assert_array_equal(a[190:310, 3:7], z[190:310, 3:7])
+    assert_array_equal(a[-110:, -3:], z[-110:, -3:])
+
+    # check partial assignment
+    b = np.arange(10000, 20000).reshape((1000, 10))
+    z[190:310, 3:7] = b[190:310, 3:7]
+    assert_array_equal(a[:190], z[:190])
+    assert_array_equal(a[:, :3], z[:, :3])
+    assert_array_equal(b[190:310, 3:7], z[190:310, 3:7])
+    assert_array_equal(a[310:], z[310:])
+    assert_array_equal(a[:, 7:], z[:, 7:])