Merge pull request #16 from alimanfoo/blosc_use_context

Implement support for using global blosc functions, resolves #15.

Author: alimanfoo

zarr/__init__.py

@@ -2,8 +2,19 @@
 from __future__ import absolute_import, print_function, division
 
 
+import atexit
+
+
 from zarr.core import empty, zeros, ones, full, array, open
-from zarr.ext import blosc_version
+from zarr.ext import blosc_version, init as _init, destroy as _destroy, \
+    set_blosc_options
 from zarr import defaults
 from zarr import constants
 from zarr.version import version as __version__
+
+
+import multiprocessing
+ncores = multiprocessing.cpu_count()
+_init()
+set_blosc_options(use_context=False, nthreads=ncores)
+atexit.register(_destroy)

zarr/ext.pyx

@@ -17,6 +17,7 @@ import shutil
 import tempfile
 from collections import namedtuple
 from glob import glob
+import multiprocessing
 import fasteners
 
 
@@ -52,6 +53,14 @@ cdef extern from "blosc.h":
         BLOSC_VERSION_STRING,
         BLOSC_VERSION_DATE
 
+    void blosc_init()
+    void blosc_destroy()
+    int blosc_set_nthreads(int nthreads)
+    int blosc_set_compressor(const char *compname)
+    int blosc_compress(int clevel, int doshuffle, size_t typesize,
+                       size_t nbytes, void *src, void *dest,
+                       size_t destsize) nogil
+    int blosc_decompress(void *src, void *dest, size_t destsize) nogil
     int blosc_compname_to_compcode(const char *compname)
     int blosc_compress_ctx(int clevel, int doshuffle, size_t typesize,
                            size_t nbytes, const void* src, void* dest,
@@ -64,7 +73,7 @@ cdef extern from "blosc.h":
 
 
 def blosc_version():
-    """Return the version of c-blosc that zarr was compiled with."""
+    """Return the version of blosc that zarr was compiled with."""
 
     # all the 'decode' contorsions are for Python 3 returning actual strings
     ver_str = <char *> BLOSC_VERSION_STRING
@@ -76,6 +85,43 @@ def blosc_version():
     return ver_str, ver_date
 
 
+def init():
+    blosc_init()
+
+
+def destroy():
+    blosc_destroy()
+
+
+_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:
+            nthreads = multiprocessing.cpu_count()
+        blosc_set_nthreads(nthreads)
+
+
 ###############################################################################
 # DEBUG LOGGING                                                               #
 ###############################################################################
@@ -122,7 +168,7 @@ def _normalize_cparams(cname=None, clevel=None, shuffle=None):
         cname = cname.encode()
     # check compressor is available
     if blosc_compname_to_compcode(cname) < 0:
-        raise ValueError("compressor not available: %s" % cname)
+        raise ValueError('compressor not available: %s' % cname)
 
     # determine compression level
     clevel = clevel if clevel is not None else defaults.clevel
@@ -389,8 +435,12 @@ cdef class Chunk(BaseChunk):
         cdef int ret
 
         # do decompression
-        with nogil:
-            ret = blosc_decompress_ctx(self._data, dest, self._nbytes, 1)
+        if _blosc_use_context:
+            with nogil:
+                ret = blosc_decompress_ctx(self._data, dest, self._nbytes, 1)
+        else:
+            with nogil:
+                ret = blosc_decompress(self._data, dest, self._nbytes)
 
         # handle errors
         if ret <= 0:
@@ -408,12 +458,20 @@ cdef class Chunk(BaseChunk):
         dest = <char *> malloc(self._nbytes + BLOSC_MAX_OVERHEAD)
 
         # perform compression
-        with nogil:
-            cbytes = blosc_compress_ctx(self._clevel, self._shuffle,
-                                        self._itemsize, self._nbytes,
-                                        source, dest,
-                                        self._nbytes + BLOSC_MAX_OVERHEAD,
-                                        self._cname, 0, 1)
+        if _blosc_use_context:
+            with nogil:
+                cbytes = blosc_compress_ctx(self._clevel, self._shuffle,
+                                            self._itemsize, self._nbytes,
+                                            source, dest,
+                                            self._nbytes + BLOSC_MAX_OVERHEAD,
+                                            self._cname, 0, 1)
+        else:
+            # compressor should have been checked already
+            assert blosc_set_compressor(self._cname) >= 0
+            with nogil:
+                cbytes = blosc_compress(self._clevel, self._shuffle,
+                                        self._itemsize, self._nbytes, source,
+                                        dest, self._nbytes + BLOSC_MAX_OVERHEAD)
 
         # check compression was successful
         if cbytes <= 0:
@@ -505,8 +563,12 @@ cdef class PersistentChunk(BaseChunk):
         source = PyBytes_AsString(data)
 
         # do decompression
-        with nogil:
-            ret = blosc_decompress_ctx(source, dest, self._nbytes, 1)
+        if _blosc_use_context:
+            with nogil:
+                ret = blosc_decompress_ctx(source, dest, self._nbytes, 1)
+        else:
+            with nogil:
+                ret = blosc_decompress(source, dest, self._nbytes)
 
         # handle errors
         if ret <= 0:
@@ -542,12 +604,20 @@ cdef class PersistentChunk(BaseChunk):
         dest = <char *> malloc(self._nbytes + BLOSC_MAX_OVERHEAD)
 
         # perform compression
-        with nogil:
-            cbytes = blosc_compress_ctx(self._clevel, self._shuffle,
-                                        self._itemsize, self._nbytes,
-                                        source, dest,
-                                        self._nbytes + BLOSC_MAX_OVERHEAD,
-                                        self._cname, 0, 1)
+        if _blosc_use_context:
+            with nogil:
+                cbytes = blosc_compress_ctx(self._clevel, self._shuffle,
+                                            self._itemsize, self._nbytes,
+                                            source, dest,
+                                            self._nbytes + BLOSC_MAX_OVERHEAD,
+                                            self._cname, 0, 1)
+        else:
+            # compressor should have been checked already
+            assert blosc_set_compressor(self._cname) >= 0
+            with nogil:
+                cbytes = blosc_compress(self._clevel, self._shuffle,
+                                        self._itemsize, self._nbytes, source,
+                                        dest, self._nbytes + BLOSC_MAX_OVERHEAD)
 
         # check compression was successful
         if cbytes <= 0:
@@ -1321,6 +1391,7 @@ cdef class SynchronizedPersistentArray(PersistentArray):
 ###############################################################################
 
 
+# noinspection PyUnresolvedReferences,PyProtectedMember
 cdef _lazy_get_chunk(BaseArray array, tuple cidx):
     try:
         chunk = array._cdata[cidx]

zarr/tests/test_array.py

@@ -13,7 +13,7 @@
 from zarr.ext import Array, SynchronizedArray, PersistentArray, \
     SynchronizedPersistentArray, LazyArray, SynchronizedLazyArray, \
     LazyPersistentArray, SynchronizedLazyPersistentArray
-from zarr import defaults
+from zarr import defaults, set_blosc_options
 
 
 class ArrayTests(object):
@@ -479,3 +479,14 @@ def create_array(self, **kwargs):
             lambda: shutil.rmtree(path) if os.path.exists(path) else None
         )
         return SynchronizedLazyPersistentArray(**kwargs)
+
+
+class TestArrayUsingContext(TestArray):
+
+    @classmethod
+    def setUpClass(cls):
+        set_blosc_options(use_context=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        set_blosc_options(use_context=False)

zarr/tests/test_chunk.py

@@ -11,7 +11,7 @@
 from numpy.testing import assert_array_equal
 from zarr.ext import Chunk, PersistentChunk, SynchronizedChunk, \
     SynchronizedPersistentChunk
-from zarr import defaults
+from zarr import defaults, set_blosc_options
 
 
 class ChunkTests(object):
@@ -252,3 +252,14 @@ def test_shuffles():
 
     # expect improvement from bitshuffle
     assert c2.cbytes < c1.cbytes
+
+
+class TestChunkUsingContext(TestChunk):
+
+    @classmethod
+    def setUpClass(cls):
+        set_blosc_options(use_context=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        set_blosc_options(use_context=False)

zarr/util.py

@@ -4,14 +4,14 @@
 
 def human_readable_size(size):
     if size < 2**10:
-        return "%s" % size
+        return '%s' % size
     elif size < 2**20:
-        return "%.1fK" % (size / float(2**10))
+        return '%.1fK' % (size / float(2**10))
     elif size < 2**30:
-        return "%.1fM" % (size / float(2**20))
+        return '%.1fM' % (size / float(2**20))
     elif size < 2**40:
-        return "%.1fG" % (size / float(2**30))
+        return '%.1fG' % (size / float(2**30))
     elif size < 2**50:
-        return "%.1fT" % (size / float(2**40))
+        return '%.1fT' % (size / float(2**40))
     else:
-        return "%.1fP" % (size / float(2**50))
+        return '%.1fP' % (size / float(2**50))