diff --git a/numcodecs/__init__.py b/numcodecs/__init__.py index be86f6da..0ba2c15e 100644 --- a/numcodecs/__init__.py +++ b/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) diff --git a/numcodecs/shuffle.c b/numcodecs/shuffle.c new file mode 100644 index 00000000..06f2230d --- /dev/null +++ b/numcodecs/shuffle.c @@ -0,0 +1 @@ +#error Do not use this file, it is the result of a failed Cython compilation. diff --git a/numcodecs/shuffle.pyx b/numcodecs/shuffle.pyx new file mode 100644 index 00000000..c7798d0b --- /dev/null +++ b/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 diff --git a/numcodecs/tests/test_shuffle.py b/numcodecs/tests/test_shuffle.py new file mode 100644 index 00000000..e9de1331 --- /dev/null +++ b/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) diff --git a/setup.py b/setup.py index 61748497..f2448fae 100644 --- a/setup.py +++ b/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 = []