implement support for using global blosc functions

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
+_cpu_count = multiprocessing.cpu_count()
+_init()
+set_blosc_options(use_context=False, nthreads=_cpu_count)
+atexit.register(_destroy)

zarr/ext.pyx

@@ -18,6 +18,7 @@ import shutil
 import tempfile
 from collections import namedtuple
 from glob import glob
+import multiprocessing
 import fasteners
 
 
@@ -53,6 +54,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,
@@ -77,6 +86,26 @@ 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, nthreads=None):
+    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                                                               #
 ###############################################################################
@@ -123,7 +152,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
@@ -390,8 +419,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:
@@ -409,12 +442,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:
@@ -506,8 +547,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:
@@ -543,12 +588,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:
@@ -1322,6 +1375,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))