Faster array tokenize() (dask#2377)

* Faster array tokenize()

Optionally uses fast hashing functions if available (the packages
'mmh3' or 'cityhash' need to be installed), falls back on SHA1.

* Appease flake8

* Disable cityhash

* Reenable cityhash if post-0.2.1

* Also catch OverflowError (for mmh3)

* Add xxHash

Author: pitrou

continuous_integration/travis/install.sh

@@ -79,9 +79,12 @@ pip install -q \
     --upgrade --no-deps
 
 pip install -q \
+    cityhash \
     flake8 \
+    mmh3 \
     pandas_datareader \
-    pytest-xdist
+    pytest-xdist \
+    xxhash
 
 # Install dask
 pip install -q --no-deps -e .[complete]

dask/base.py

@@ -16,6 +16,7 @@
 from .compatibility import bind_method, unicode, PY3
 from .context import _globals
 from .core import flatten
+from .hashing import hash_buffer_hex
 from .utils import Dispatch, ensure_dict
 from .sharedict import ShareDict
 
@@ -387,15 +388,15 @@ def normalize_array(x):
                     x.shape, x.strides, offset)
         if x.dtype.hasobject:
             try:
-                data = md5('-'.join(x.flat).encode('utf-8')).hexdigest()
+                data = hash_buffer_hex('-'.join(x.flat).encode('utf-8'))
             except TypeError:
-                data = md5(b'-'.join([unicode(item).encode('utf-8') for item in
-                                      x.flat])).hexdigest()
+                data = hash_buffer_hex(b'-'.join([unicode(item).encode('utf-8') for item in
+                                                  x.flat]))
         else:
             try:
-                data = md5(x.ravel().view('i1').data).hexdigest()
+                data = hash_buffer_hex(x.ravel(order='K').view('i1'))
             except (BufferError, AttributeError, ValueError):
-                data = md5(x.copy().ravel().view('i1').data).hexdigest()
+                data = hash_buffer_hex(x.copy().ravel(order='K').view('i1'))
         return (data, x.dtype, x.shape, x.strides)
 
     normalize_token.register(np.dtype, repr)

dask/hashing.py

@@ -0,0 +1,102 @@
+from __future__ import absolute_import, division, print_function
+
+import binascii
+import hashlib
+import sys
+
+
+hashers = []  # In decreasing performance order
+
+
+# Timings on a largish array:
+# - CityHash is 2x faster than MurmurHash
+# - xxHash is slightly slower than CityHash
+# - MurmurHash is 8x faster than SHA1
+# - SHA1 is significantly faster than all other hashlib algorithms
+
+try:
+    import cityhash  # `pip install cityhash`
+except ImportError:
+    pass
+else:
+    # CityHash disabled unless the reference leak in
+    # https://github.com/escherba/python-cityhash/pull/16
+    # is fixed.
+    if cityhash.__version__ >= '0.2.2':
+        def _hash_cityhash(buf):
+            """
+            Produce a 16-bytes hash of *buf* using CityHash.
+            """
+            h = cityhash.CityHash128(buf)
+            if sys.version_info >= (3,):
+                return h.to_bytes(16, 'little')
+            else:
+                return binascii.a2b_hex('%032x' % h)
+
+        hashers.append(_hash_cityhash)
+
+try:
+    import xxhash  # `pip install xxhash`
+except ImportError:
+    pass
+else:
+    def _hash_xxhash(buf):
+        """
+        Produce a 8-bytes hash of *buf* using xxHash.
+        """
+        return xxhash.xxh64(buf).digest()
+
+    hashers.append(_hash_xxhash)
+
+try:
+    import mmh3  # `pip install mmh3`
+except ImportError:
+    pass
+else:
+    def _hash_murmurhash(buf):
+        """
+        Produce a 16-bytes hash of *buf* using MurmurHash.
+        """
+        return mmh3.hash_bytes(buf)
+
+    hashers.append(_hash_murmurhash)
+
+
+def _hash_sha1(buf):
+    """
+    Produce a 20-bytes hash of *buf* using SHA1.
+    """
+    return hashlib.sha1(buf).digest()
+
+
+hashers.append(_hash_sha1)
+
+
+def hash_buffer(buf, hasher=None):
+    """
+    Hash a bytes-like (buffer-compatible) object.  This function returns
+    a good quality hash but is not cryptographically secure.  The fastest
+    available algorithm is selected.  A fixed-length bytes object is returned.
+    """
+    if hasher is not None:
+        try:
+            return hasher(buf)
+        except (TypeError, OverflowError):
+            # Some hash libraries may have overly-strict type checking,
+            # not accepting all buffers
+            pass
+    for hasher in hashers:
+        try:
+            return hasher(buf)
+        except (TypeError, OverflowError):
+            pass
+    raise TypeError("unsupported type for hashing: %s" % (type(buf),))
+
+
+def hash_buffer_hex(buf, hasher=None):
+    """
+    Same as hash_buffer, but returns its result in hex-encoded form.
+    """
+    h = hash_buffer(buf, hasher)
+    s = binascii.b2a_hex(h)
+    return s.decode() if sys.version_info >= (3,) else s

dask/tests/test_delayed.py

@@ -403,7 +403,7 @@ def __call__(self):
     assert f().compute() == 2
 
 
-def test_name_consitent_across_instances():
+def test_name_consistent_across_instances():
     func = delayed(identity, pure=True)
 
     data = {'x': 1, 'y': 25, 'z': [1, 2, 3]}

dask/tests/test_hashing.py

@@ -0,0 +1,45 @@
+from __future__ import absolute_import, division, print_function
+
+import pytest
+
+from dask.hashing import hashers, hash_buffer, hash_buffer_hex
+
+
+np = pytest.importorskip('numpy')
+
+buffers = [
+    b'abc',
+    bytearray(b'123'),
+    memoryview(b'456'),
+    np.array(42),
+    np.ones((100, 100)),
+    np.zeros((100, 100), dtype=[('a', 'i4'), ('b', 'i2')]),
+    np.ones(10000, dtype=np.int8)[1:],  # unaligned
+]
+
+
+@pytest.mark.parametrize('x', buffers)
+def test_hash_buffer(x):
+    for hasher in [None] + hashers:
+        h = hash_buffer(x, hasher=hasher)
+        assert isinstance(h, bytes)
+        assert 8 <= len(h) < 32
+        assert h == hash_buffer(x, hasher=hasher)
+
+
+@pytest.mark.parametrize('x', buffers)
+def test_hash_buffer_hex(x):
+    for hasher in [None] + hashers:
+        h = hash_buffer_hex(x, hasher=hasher)
+        assert isinstance(h, str)
+        assert 16 <= len(h) < 64
+        assert h == hash_buffer_hex(x, hasher=hasher)
+
+
+@pytest.mark.parametrize('hasher', hashers)
+def test_hashers(hasher):
+    # Sanity check
+    x = b'x'
+    h = hasher(x)
+    assert isinstance(h, bytes)
+    assert 8 <= len(h) < 32