shuffle with cython

numcodecs/__init__.py

@@ -92,6 +92,9 @@
 from numcodecs.base64 import Base64
 register_codec(Base64)
 
+from numcodecs.shuffle import Shuffle
+register_codec(Shuffle)
+
 try:
     from numcodecs.msgpacks import MsgPack
     register_codec(MsgPack)

numcodecs/shuffle.c

@@ -0,0 +1 @@
+#error Do not use this file, it is the result of a failed Cython compilation.

numcodecs/shuffle.pyx

@@ -0,0 +1,93 @@
+import numpy as np
+from .compat import ensure_bytes, ensure_contiguous_ndarray, ndarray_copy, ensure_ndarray
+from .abc import Codec
+
+
+cdef _doShuffle(const unsigned char[:] src, unsigned char[:] des, int element_size):
+    cdef int count, i, j, offset, byte_index
+    count = len(src) // element_size
+    for i in range(count):
+        offset = i*element_size
+        e = src[offset:(offset+element_size)]
+        for byte_index in range(element_size):
+            j = byte_index*count + i
+            des[j] = e[byte_index]
+    return des
+
+
+cdef _doUnshuffle(const unsigned char[:] src, unsigned char[:] des, int element_size):
+    cdef int count, i, j, offset, byte_index
+    count = len(src) // element_size
+    for i in range(element_size):
+        offset = i*count
+        e = src[offset:(offset+count)]
+        for byte_index in range(count):
+            j = byte_index*element_size + i
+            des[j] = e[byte_index]
+    return des
+
+
+def _shuffle(element_size, buf, out):
+    if element_size <= 1:
+        out.view(buf.dtype)[:len(buf)] = buf[:]
+        return out  # no shuffling needed
+
+    buf_size = buf.nbytes
+    if buf_size % element_size != 0:
+        raise ValueError("Shuffle buffer is not an integer multiple of elementsize")
+
+    _doShuffle(buf.view("uint8"), out.view("uint8"), element_size)
+
+    return out
+
+
+def _unshuffle(element_size, buf, out):
+    if element_size <= 1:
+        out.view(buf.dtype)[:len(buf)] = buf[:]
+        return out  # no shuffling needed
+
+    buf_size = buf.nbytes
+    if buf_size % element_size != 0:
+        raise ValueError("Shuffle buffer is not an integer multiple of elementsize")
+
+    _doUnshuffle(buf.view("uint8"), out.view("uint8"), element_size)
+
+    return out
+
+
+class Shuffle(Codec):
+    """Codec providing shuffle
+
+    Parameters
+    ----------
+    elementsize : int
+        Size in bytes of the array elements.  Default = 4
+
+    """
+
+    codec_id = 'shuffle'
+
+    def __init__(self, elementsize=4):
+        self.elementsize = elementsize
+
+    def encode(self, buf, out=None):
+        buf = ensure_contiguous_ndarray(buf)
+        if out is None:
+            out = np.zeros(buf.nbytes, dtype='uint8')
+        else:
+            out = ensure_contiguous_ndarray(out)
+        return _shuffle(self.elementsize, buf, out)
+
+    def decode(self, buf, out=None):
+        buf = ensure_contiguous_ndarray(buf)
+        if out is None:
+            out = np.zeros(buf.nbytes, dtype='uint8')
+        else:
+            out = ensure_contiguous_ndarray(out)
+        return _unshuffle(self.elementsize, buf, out)
+
+    def __repr__(self):
+        r = '%s(elementsize=%s)' % \
+            (type(self).__name__,
+             self.elementsize)
+        return r

numcodecs/tests/test_shuffle.py

@@ -0,0 +1,135 @@
+from multiprocessing import Pool
+from multiprocessing.pool import ThreadPool
+
+
+import numpy as np
+import pytest
+
+
+try:
+    from numcodecs.shuffle import Shuffle
+except ImportError:  # pragma: no cover
+    pytest.skip(
+        "numcodecs.shuffle not available", allow_module_level=True
+    )
+
+
+from numcodecs.tests.common import (check_encode_decode,
+                                    check_config)
+
+
+codecs = [
+    Shuffle(),
+    Shuffle(elementsize=0),
+    Shuffle(elementsize=4),
+    Shuffle(elementsize=8)
+]
+
+
+# mix of dtypes: integer, float, bool, string
+# mix of shapes: 1D, 2D, 3D
+# mix of orders: C, F
+arrays = [
+    np.arange(1000, dtype='i4'),
+    np.linspace(1000, 1001, 1000, dtype='f8'),
+    np.random.normal(loc=1000, scale=1, size=(100, 10)),
+    np.random.randint(0, 2, size=1000, dtype=bool).reshape(100, 10, order='F'),
+    np.random.choice([b'a', b'bb', b'ccc'], size=1000).reshape(10, 10, 10),
+    np.random.randint(0, 2**60, size=1000, dtype='u8').view('M8[ns]'),
+    np.random.randint(0, 2**60, size=1000, dtype='u8').view('m8[ns]'),
+    np.random.randint(0, 2**25, size=1000, dtype='u8').view('M8[m]'),
+    np.random.randint(0, 2**25, size=1000, dtype='u8').view('m8[m]'),
+    np.random.randint(-2**63, -2**63 + 20, size=1000, dtype='i8').view('M8[ns]'),
+    np.random.randint(-2**63, -2**63 + 20, size=1000, dtype='i8').view('m8[ns]'),
+    np.random.randint(-2**63, -2**63 + 20, size=1000, dtype='i8').view('M8[m]'),
+    np.random.randint(-2**63, -2**63 + 20, size=1000, dtype='i8').view('m8[m]'),
+]
+
+
+@pytest.mark.parametrize('array', arrays)
+@pytest.mark.parametrize('codec', codecs)
+def test_encode_decode(array, codec):
+    check_encode_decode(array, codec)
+
+
+def test_config():
+    codec = Shuffle()
+    check_config(codec)
+    codec = Shuffle(elementsize=8)
+    check_config(codec)
+
+
+def test_repr():
+    expect = "Shuffle(elementsize=0)"
+    actual = repr(Shuffle(elementsize=0))
+    assert expect == actual
+    expect = "Shuffle(elementsize=4)"
+    actual = repr(Shuffle(elementsize=4))
+    assert expect == actual
+    expect = "Shuffle(elementsize=8)"
+    actual = repr(Shuffle(elementsize=8))
+    assert expect == actual
+    expect = "Shuffle(elementsize=16)"
+    actual = repr(Shuffle(elementsize=16))
+    assert expect == actual
+
+
+def test_eq():
+    assert Shuffle() == Shuffle()
+    assert Shuffle(elementsize=16) != Shuffle()
+
+
+def _encode_worker(data):
+    compressor = Shuffle()
+    enc = compressor.encode(data)
+    return enc
+
+
+def _decode_worker(enc):
+    compressor = Shuffle()
+    data = compressor.decode(enc)
+    return data
+
+
+@pytest.mark.parametrize('pool', (Pool, ThreadPool))
+def test_multiprocessing(pool):
+    data = np.arange(1000000)
+    enc = _encode_worker(data)
+
+    pool = pool(5)
+
+    # test with process pool and thread pool
+
+    # test encoding
+    enc_results = pool.map(_encode_worker, [data] * 5)
+    assert all([len(enc) == len(e) for e in enc_results])
+
+    # test decoding
+    dec_results = pool.map(_decode_worker, [enc] * 5)
+    assert all([data.nbytes == len(d) for d in dec_results])
+
+    # tidy up
+    pool.close()
+    pool.join()
+
+
+# def test_err_decode_object_buffer():
+#     check_err_decode_object_buffer(Shuffle())
+
+
+# def test_err_encode_object_buffer():
+#     check_err_encode_object_buffer(Shuffle())
+
+# def test_decompression_error_handling():
+#     for codec in codecs:
+#         with pytest.raises(RuntimeError):
+#             codec.decode(bytearray())
+#         with pytest.raises(RuntimeError):
+#             codec.decode(bytearray(0))
+
+
+def test_incompatible_elementsize():
+    with pytest.raises(ValueError):
+        arr = np.arange(1001, dtype='u1')
+        codec = Shuffle(elementsize=4)
+        codec.encode(arr)

setup.py

@@ -250,6 +250,29 @@ def compat_extension():
     return extensions
 
 
+def shuffle_extension():
+    info('setting up shuffle extension')
+
+    extra_compile_args = list(base_compile_args)
+
+    if have_cython:
+        sources = ['numcodecs/shuffle.pyx']
+    else:
+        sources = ['numcodecs/shuffle.c']
+
+    # define extension module
+    extensions = [
+        Extension('numcodecs.shuffle',
+                  sources=sources,
+                  extra_compile_args=extra_compile_args),
+    ]
+
+    if have_cython:
+        extensions = cythonize(extensions)
+
+    return extensions
+
+
 if sys.platform == 'win32':
     ext_errors = (CCompilerError, DistutilsExecError, DistutilsPlatformError,
                   IOError, ValueError)
@@ -291,7 +314,7 @@ def run_setup(with_extensions):
 
     if with_extensions:
         ext_modules = (blosc_extension() + zstd_extension() + lz4_extension() +
-                       compat_extension() + vlen_extension())
+                       compat_extension() + shuffle_extension() + vlen_extension())
         cmdclass = dict(build_ext=ve_build_ext)
     else:
         ext_modules = []