TFDS shuffling: swap cityhash for siphash (Issue #653, Issue #690).

PiperOrigin-RevId: 255375385

Author: pierrot0

setup.py

@@ -42,7 +42,7 @@
     'absl-py',
     'attrs',
     'dill',  # TODO(tfds): move to TESTS_REQUIRE.
-    'cityhash',
+    'siphash',
     'future',
     'numpy',
     'promise',
@@ -58,6 +58,7 @@
 
 TESTS_REQUIRE = [
     'apache-beam',
+    'csiphash',
     'jupyter',
     'pytest',
 ]
@@ -122,6 +123,7 @@
 
 EXTRAS_REQUIRE = {
     'apache-beam': ['apache-beam'],
+    'siphash': ['csiphash'],
     'tensorflow': ['tensorflow>=1.13.0'],
     'tensorflow_gpu': ['tensorflow-gpu>=1.13.0'],
     'tests': TESTS_REQUIRE + all_dataset_extras,

tensorflow_datasets/core/hashing.py

@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2019 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Stable hashing function using SipHash (https://131002.net/siphash/).
+
+Note that the properties we are looking at here are:
+ 1- Good distribution of hashes;
+ 2- Speed;
+ 3- Availability as portable library, giving the same hash independently of
+ platform.
+
+Crypto level hashing is not a requirement.
+
+Although SipHash being slower than alternatives (eg: CityHash), its
+implementation is simpler and it has a pure Python implementation, making it
+easier to distribute TFDS on Windows or MAC. The speed cost being only paid at
+data preparation time, portability wins.
+
+Python3 uses SipHash internally, but since it is not exposed and `hash` function
+is not guaranteed to use SipHash in the future, we have to import libraries.
+
+Changing the hash function should be done thoughfully, as it would change the
+order of datasets (and thus sets of records when using slicing API). If done,
+all datasets would need to have their major version bumped.
+
+Note that if we were to find a dataset for which two different keys give the
+same hash (collision), a solution could be to append the key to its hash.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import struct
+
+import siphash
+import six
+
+_CSIPHASH_AVAILABLE = False
+try:
+  import csiphash  # pylint: disable=g-import-not-at-top
+  _CSIPHASH_AVAILABLE = True
+except ImportError:
+  pass
+
+# SipHash needs a 16 bits key.
+_SECRET = b'\0' * 16
+
+
+def _siphash(data):
+  if _CSIPHASH_AVAILABLE:
+    hash_bytes = csiphash.siphash24(_SECRET, data)
+    # Equivalent to `int.from_bytes(hash_bytes, sys.byteorder)` in py3,
+    # but py2 compatible:
+    return struct.unpack('<Q', hash_bytes)[0]
+  else:
+    return siphash.SipHash24(_SECRET, data).hash()
+
+
+def hash_key(key):
+  """Returns 64 bits hash of given key.
+
+  Args:
+    key (bytes, string or anything convertible to a string): key to be hashed.
+      If the key is a string, it will be encoded to bytes using utf-8.
+      If the key is neither a string nor bytes, it will be converted to a str,
+        then to bytes.
+      This means that `"1"` (str) and `1` (int) will have the same hash. The
+      intent of the hash being to shuffle keys, it is recommended that all keys
+      of a given set to shuffle use a single type.
+
+  Returns:
+    64 bits integer, hash of key.
+  """
+  if not isinstance(key, (six.string_types, bytes)):
+    key = str(key)
+  if not isinstance(key, bytes):
+    key = key.encode('utf-8')
+  return _siphash(key)
+

tensorflow_datasets/core/hashing_test.py

@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2019 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for tensorflow_datasets.core.hashing."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import string
+
+from tensorflow_datasets import testing
+from tensorflow_datasets.core import hashing
+
+EXPECTED_ASCII_HASHES = [
+    10863254463029944905, 17270894748891556580, 1421958803217889556,
+    2755936516345535118, 10743397483099227609, 14847904966537228119,
+    8418086691718335568, 11915783202300667430, 9716548853924804360,
+    1754241953508342506, 3694815635762763843, 3596675173038026136,
+    8070480893871644201, 10828895317094483299, 6744895540894632903,
+    14841125270424448140, 17688157251436176611, 233329666533924104,
+    12405195962649360880, 6834955004197680363, 13969673605019971927,
+    7535964442537143799, 4233129741512941149, 6982149393400456965,
+    6390938501586896152, 10691964138414231865, 2507792285634992701,
+    8595505165718044977, 12668683514780208634, 4572719121647279127,
+    3793178347273963799, 12176803637319179781, 8203705581228860887,
+    6621546828295374121, 10489668860986262979, 2912571526136585784,
+    7938998531316856084, 9873236742883981124, 11452724586647173125,
+    17125311730203373552, 14019180516172084896, 4703666581826301487,
+    855367633508639125, 14196452709219207778, 12230774050365420318,
+    2537184170178946005, 7687283993009238581, 8820003036235826491,
+    17283458755135480410, 17416811303727010451, 5922709015603821319,
+    14428075881971433291]
+
+EXPECTED_INT_HASHES = [
+    13163753087252323914, 5003827105613308882, 11449545338359147399,
+    3672830208859661989, 5406800756778728304, 14481025309921804611,
+    1609946449970207933, 8255655750251093705, 8491335656787965458,
+    614618433269745394, 13284406229155908416, 15288099744915406683,
+    15843726386018726000, 5119555216682145351, 18142174963634300418,
+    9628501536262262909, 6944376202371927192, 14434159142587098114,
+    9198994799148619777, 10976198626365556062, 12201794547984161272,
+    3793153257673850654, 17671034425519684114, 1052793885524652273,
+    6624479012289984384, 8863054234635171593, 4965745346604460244,
+    9391177234155550519, 4717670777148145883, 17524121804784260174,
+    11037926627716248709, 6960985957824457329, 12195906204051370437,
+    10975328135781691390, 5730073803446725122, 13712792850809427923,
+    4455483863044151629, 3518672581294300691, 15747605586304671771,
+    13668305533495453291, 4654232860820002596, 10574044313476412487,
+    11212237458977771261, 15365614270461858889, 13872585532843456912,
+    17241372482471269826, 3462651383276316179, 12647419365920702661,
+    1995464140987078702, 1972561720831881829, 3955328643960597520,
+    1027369990565220197, 11322815499836299743, 5956980248780520012,
+    18278096046037631584, 5241067853637136955, 7630275157338448032,
+    1913046367367276708, 7440126551217879907, 7220748048216444121,
+    5064805925883037540, 1318705738537093225, 2730963110225791297,
+    6920195161195209846, 4682001368639242156, 9166607120404080113,
+    11268721706256216334, 5379201623047735445, 15999685243572303930,
+    13046608731927560566, 16276928149450612660, 16298571539550440629,
+    17035045450101282343, 14240119263724925078, 9965315260615748500,
+    14921974451741066715, 3620887669157180415, 14499246414755411500,
+    188410546870139183, 14101909720529780551, 1623775225152586541,
+    1826999275929156985, 5289921295512723016, 151075781360207052,
+    17598366794955569210, 1171265316432012145, 104641658363814304,
+    9264688353594391671, 458105873437653640, 1830791798008018143,
+    15150529348822956655, 7610023982731034430, 12031109555968877759,
+    2814999125315447998, 14537302745610214253, 14150033554738292901,
+    1316273336836242886, 4973610113424219627, 11435740729903845598,
+    4598673630928554393]
+
+
+class HashingTest(testing.TestCase):
+
+  def _assert_hashes(self, keys, expected_hashes):
+    hashes = [hashing.hash_key(k) for k in keys]
+    self.assertEqual(hashes, expected_hashes)
+
+  def test_ints(self):
+    ints = list(range(100))
+    hashing._CSIPHASH_AVAILABLE = False
+    self._assert_hashes(ints, EXPECTED_INT_HASHES)
+    hashing._CSIPHASH_AVAILABLE = True
+    self._assert_hashes(ints, EXPECTED_INT_HASHES)
+
+  def test_ascii(self):
+    letters = string.ascii_lowercase + string.ascii_uppercase
+    hashing._CSIPHASH_AVAILABLE = False
+    self._assert_hashes(letters, EXPECTED_ASCII_HASHES)
+    hashing._CSIPHASH_AVAILABLE = True
+    self._assert_hashes(letters, EXPECTED_ASCII_HASHES)
+
+
+if __name__ == '__main__':
+  testing.test_main()

tensorflow_datasets/core/shuffle.py

@@ -24,10 +24,10 @@
 import struct
 import uuid
 
-import cityhash
 import six
 import tensorflow as tf
 
+from tensorflow_datasets.core import hashing
 
 HKEY_SIZE = 64  # Hash of keys is 64 bits.
 
@@ -48,13 +48,6 @@
 BUCKETS_NUMBER = 1000  # Number of buckets to pre-sort and hold generated data.
 
 
-def _get_hashed_key(key):
-  """Returns hash (int) for given key."""
-  if not isinstance(key, (six.string_types, bytes)):
-    key = str(key)
-  return cityhash.CityHash64(key)
-
-
 def _get_shard(hkey, shards_number):
   """Returns shard number (int) for given hashed key (int)."""
   # We purposely do not use modulo (%) to keep global order across shards.
@@ -73,12 +66,14 @@ class _Bucket(object):
     key1 (8 bytes) | size1 (8 bytes) | data1 (size1 bytes) |
     key2 (8 bytes) | size2 (8 bytes) | data2 (size2 bytes) |
     ...
-
-  Args:
-    path (str): where to write the bucket file.
   """
 
   def __init__(self, path):
+    """Initialize a _Bucket instance.
+
+    Args:
+      path (str): where to write the bucket file.
+    """
     self._path = path
     self._fobj = None
     self._length = 0
@@ -167,7 +162,7 @@ def add(self, key, data):
     if not isinstance(data, six.binary_type):
       raise AssertionError('Only bytes (not %s) can be stored in Shuffler!' % (
           type(data)))
-    hkey = _get_hashed_key(key)
+    hkey = hashing.hash_key(key)
     self._total_bytes += len(data)
     if self._in_memory:
       self._add_to_mem_buffer(hkey, data)

tensorflow_datasets/core/shuffle_test.py

@@ -38,15 +38,15 @@
 ]
 
 _ORDERED_ITEMS = [
-    b'over',
-    b'jumps',
-    b'quick ',
     b' dog.',
+    b'over',
+    b'brown',
     b'The',
+    b' fox ',
     b' the ',
     b'lazy',
-    b'brown',
-    b' fox ',
+    b'quick ',
+    b'jumps',
 ]
 
 _TOTAL_SIZE = sum(len(rec) for rec in _ORDERED_ITEMS)
@@ -97,7 +97,6 @@ def test_duplicate_key(self):
     shuffler.add(2, b'b')
     shuffler.add(1, b'c')
     iterator = iter(shuffler)
-    self.assertEqual(next(iterator), b'b')
     self.assertEqual(next(iterator), b'a')
     with self.assertRaisesWithPredicateMatch(
         AssertionError, 'Two records share the same hashed key!'):

tensorflow_datasets/core/tfrecords_writer_test.py

@@ -107,7 +107,7 @@ def test_write(self):
     self.assertEqual(written_files,
                      ['foo.tfrecord-0000%s-of-00005' % i for i in range(5)])
     self.assertEqual(all_recs, [
-        [b'f', b'e'], [b'b'], [b'a', b'g'], [b'h'], [b'c', b'd'],
+        [b'f', b'c'], [b'a'], [b'd', b'g'], [b'h'], [b'b', b'e'],
     ])
 
 if __name__ == '__main__':