TFDS Issue #737: use md5 instead of siphash.

PiperOrigin-RevId: 257178144

Author: pierrot0

setup.py

@@ -42,7 +42,6 @@
     'absl-py',
     'attrs',
     'dill',  # TODO(tfds): move to TESTS_REQUIRE.
-    'siphash',
     'future',
     'numpy',
     'promise',
@@ -58,7 +57,6 @@
 
 TESTS_REQUIRE = [
     'apache-beam',
-    # 'csiphash',  # https://github.com/tensorflow/datasets/issues/737
     'jupyter',
     'pytest',
     'pytest-xdist',
@@ -124,8 +122,6 @@
 
 EXTRAS_REQUIRE = {
     'apache-beam': ['apache-beam'],
-    # https://github.com/tensorflow/datasets/issues/737
-    # 'siphash': ['csiphash'],
     'tensorflow': ['tensorflow>=1.13.0'],
     'tensorflow_gpu': ['tensorflow-gpu>=1.13.0'],
     'tests': TESTS_REQUIRE + all_dataset_extras,

tensorflow_datasets/core/dataset_builder.py

@@ -898,7 +898,8 @@ def _prepare_split(self, split_generator, max_examples_per_split):
       return self._prepare_split_legacy(generator, split_info)
     fname = "{}-{}.tfrecord".format(self.name, split_generator.name)
     fpath = os.path.join(self._data_dir, fname)
-    writer = tfrecords_writer.Writer(self._example_specs, fpath)
+    writer = tfrecords_writer.Writer(self._example_specs, fpath,
+                                     hash_salt=split_generator.name)
     for key, record in utils.tqdm(generator, unit=" examples",
                                   total=split_info.num_examples, leave=False):
       example = self.info.features.encode_example(record)

tensorflow_datasets/core/hashing.py

@@ -13,80 +13,75 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Stable hashing function using SipHash (https://131002.net/siphash/).
+"""Stable hashing function using md5.
 
 Note that the properties we are looking at here are:
- 1- Good distribution of hashes;
+ 1- Good distribution of hashes (random uniformity);
  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.
+A bit of history:
+ - CityHash was first used. However the C implementation used complex
+ instructions and was hard to compile on some platforms.
+ - Siphash was chosen as a replacement, because although being slower,
+ it has a simpler implementation and it has a pure Python implementation, making
+ it easier to distribute TFDS on Windows or MAC. However, the used library
+ (reference C implementation wrapped using cffi) crashed the python interpreter
+ on py3 with tf1.13.
+ - So md5, although being slower that the two above works everywhere and is
+ still faster than a pure python implementation of siphash.
 
 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.
+
+The split name is being used as salt to avoid having the same keys in two splits
+result in same order.
 """
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import struct
+import hashlib
 
-import siphash
 import six
+import tensorflow as tf
+
+
+def _to_bytes(data):
+  if not isinstance(data, (six.string_types, bytes)):
+    data = str(data)
+  return tf.compat.as_bytes(data)
+
+
+class Hasher(object):
+  """Hasher: to initialize a md5 with salt."""
+
+  def __init__(self, salt):
+    self._md5 = hashlib.md5(_to_bytes(salt))
+
+  def hash_key(self, key):
+    """Returns 128 bits hash of given key.
 
-_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)
+    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:
+      128 bits integer, hash of key.
+    """
+    md5 = self._md5.copy()
+    md5.update(_to_bytes(key))
+    return int(md5.hexdigest(), 16)

tensorflow_datasets/core/hashing_test.py

@@ -19,89 +19,23 @@
 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)
-    # https://github.com/tensorflow/datasets/issues/737
-    # hashing._CSIPHASH_AVAILABLE = True
-    # self._assert_hashes(ints, EXPECTED_INT_HASHES)
+    hasher = hashing.Hasher(salt='')
+    res = hasher.hash_key(0)
+    self.assertEqual(res, 276215275525073243129443018166533317850)
+    res = hasher.hash_key(123455678901234567890)
+    self.assertEqual(res, 6876359009333865997613257802033240610)
 
   def test_ascii(self):
-    letters = string.ascii_lowercase + string.ascii_uppercase
-    hashing._CSIPHASH_AVAILABLE = False
-    self._assert_hashes(letters, EXPECTED_ASCII_HASHES)
-    # https://github.com/tensorflow/datasets/issues/737
-    # hashing._CSIPHASH_AVAILABLE = True
-    # self._assert_hashes(letters, EXPECTED_ASCII_HASHES)
+    hasher = hashing.Hasher(salt='')
+    res = hasher.hash_key('foo')
+    self.assertEqual(res, 229609063533823256041787889330700985560)
 
 
 if __name__ == '__main__':

tensorflow_datasets/core/shuffle.py

@@ -29,8 +29,6 @@
 
 from tensorflow_datasets.core import hashing
 
-HKEY_SIZE = 64  # Hash of keys is 64 bits.
-
 # Approximately how much data to store in memory before writing to disk.
 # If the amount of data to shuffle is < MAX_MEM_BUFFER_SIZE, no intermediary
 # data is written to disk.
@@ -47,6 +45,21 @@
 # the number of buckets might warrant some changes in implementation.
 BUCKETS_NUMBER = 1000  # Number of buckets to pre-sort and hold generated data.
 
+HKEY_SIZE = 128  # Hash of keys is 128 bits (md5).
+HKEY_SIZE_BYTES = HKEY_SIZE // 8
+
+
+def _hkey_to_bytes(hkey):
+  """Converts 128 bits integer hkey to binary representation."""
+  max_int64 = 0xFFFFFFFFFFFFFFFF
+  return struct.pack('QQ', (hkey >> 64) & max_int64, hkey & max_int64)
+
+
+def _read_hkey(buff):
+  """Reads from fobj and returns hkey (128 bites integer)."""
+  a, b = struct.unpack('QQ', buff)
+  return (a << 64) | b
+
 
 def _get_shard(hkey, shards_number):
   """Returns shard number (int) for given hashed key (int)."""
@@ -93,7 +106,8 @@ def add(self, key, data):
     if not self._fobj:
       self._fobj = tf.io.gfile.GFile(self._path, mode='wb')
     data_size = len(data)
-    self._fobj.write(struct.pack('LL', key, data_size))
+    self._fobj.write(_hkey_to_bytes(key))
+    self._fobj.write(struct.pack('L', data_size))
     self._fobj.write(data)
     self._length += 1
     self._size += data_size
@@ -107,10 +121,12 @@ def read_values(self):
     res = []
     with tf.io.gfile.GFile(path, 'rb') as fobj:
       while True:
-        buff = fobj.read(16)
+        buff = fobj.read(HKEY_SIZE_BYTES)
         if not buff:
           break
-        hkey, size = struct.unpack('LL', buff)
+        hkey = _read_hkey(buff)
+        size_bytes = fobj.read(8)
+        size = struct.unpack('L', size_bytes)[0]
         data = fobj.read(size)
         res.append((hkey, data))
     return res
@@ -121,14 +137,17 @@ def del_file(self):
 
 
 class Shuffler(object):
-  """Stores data in temp buckets, restitute it shuffled.
+  """Stores data in temp buckets, restitute it shuffled."""
 
-  Args:
-    dirpath: directory in which to store temporary files.
-  """
+  def __init__(self, dirpath, hash_salt):
+    """Initialize Shuffler.
 
-  def __init__(self, dirpath):
+    Args:
+      dirpath (string): directory in which to store temporary files.
+      hash_salt (string or bytes): salt to hash keys.
+    """
     grp_name = uuid.uuid4()
+    self._hasher = hashing.Hasher(hash_salt)
     self._buckets = [
         _Bucket(os.path.join(dirpath, 'bucket_%s_%03d.tmp' % (grp_name, i)))
         for i in range(BUCKETS_NUMBER)]
@@ -162,7 +181,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 = hashing.hash_key(key)
+    hkey = self._hasher.hash_key(key)
     self._total_bytes += len(data)
     if self._in_memory:
       self._add_to_mem_buffer(hkey, data)