Fix interleave support in TensorFlow 1.14.

Author: 2sin18

docs/conf.py

@@ -24,7 +24,7 @@
 project = 'HybridBackend'
 author = 'Alibaba Group Holding Limited'
 copyright = '2021 Alibaba Group Holding Limited'  # pylint: disable=redefined-builtin
-release = '0.1.0'
+release = 'latest'
 
 # -- General configuration ---------------------------------------------------
 

hybridbackend/tensorflow/data/parquet_dataset_v1.py

@@ -23,7 +23,6 @@
 from __future__ import print_function
 
 from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.ops import readers
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.util import nest
@@ -186,13 +185,8 @@ def _build_dataset(
     '''
     if num_parallel_reads is None:
       return filenames.flat_map(dataset_creator)
-    if num_parallel_reads == dataset_ops.AUTOTUNE:
-      return filenames.interleave(
-          dataset_creator, num_parallel_calls=num_parallel_reads)
-    return readers.ParallelInterleaveDataset(
-        filenames, dataset_creator,
-        cycle_length=num_parallel_reads,
+    return filenames.interleave(
+        dataset_creator,
+        cycle_length=num_parallel_reads if num_parallel_reads > 0 else 1,
         block_length=num_sequential_reads,
-        sloppy=True,
-        buffer_output_elements=None,
-        prefetch_input_elements=1)
+        num_parallel_calls=num_parallel_reads)

tests/tensorflow/data/parquet_dataset_test.py

@@ -26,6 +26,10 @@
 import tempfile
 
 from tensorflow.python.data.ops.dataset_ops import Dataset
+try:
+  from tensorflow.python.data.ops.dataset_ops import AUTOTUNE
+except ImportError:
+  from tensorflow.python.data.experimental.ops.optimization import AUTOTUNE
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
@@ -129,6 +133,56 @@ def gen_filenames():
       with self.assertRaises(errors.OutOfRangeError):
         sess.run(batch)
 
+  def test_read_from_generator_parallel(self):
+    num_epochs = 2
+    batch_size = 100
+    with ops.Graph().as_default() as graph:
+      def gen_filenames():
+        for i in range(num_epochs + 1):
+          if i == num_epochs:
+            return  # raise StopIteration
+          yield self._filename
+      filenames = Dataset.from_generator(
+          gen_filenames, dtypes.string, tensor_shape.TensorShape([]))
+      fields = [
+          DataFrame.Field('A', dtypes.int64),
+          DataFrame.Field('C', dtypes.int64)]
+      ds = filenames.apply(
+          read_parquet(batch_size, fields=fields, num_parallel_reads=3))
+      ds = ds.prefetch(4)
+      batch = make_one_shot_iterator(ds).get_next()
+
+    with self.test_session(use_gpu=False, graph=graph) as sess:
+      for _ in range(len(self._df) * num_epochs // batch_size):
+        sess.run(batch)
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(batch)
+
+  def test_read_from_generator_parallel_auto(self):
+    num_epochs = 2
+    batch_size = 100
+    with ops.Graph().as_default() as graph:
+      def gen_filenames():
+        for i in range(num_epochs + 1):
+          if i == num_epochs:
+            return  # raise StopIteration
+          yield self._filename
+      filenames = Dataset.from_generator(
+          gen_filenames, dtypes.string, tensor_shape.TensorShape([]))
+      fields = [
+          DataFrame.Field('A', dtypes.int64),
+          DataFrame.Field('C', dtypes.int64)]
+      ds = filenames.apply(
+          read_parquet(batch_size, fields=fields, num_parallel_reads=AUTOTUNE))
+      ds = ds.prefetch(4)
+      batch = make_one_shot_iterator(ds).get_next()
+
+    with self.test_session(use_gpu=False, graph=graph) as sess:
+      for _ in range(len(self._df) * num_epochs // batch_size):
+        sess.run(batch)
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(batch)
+
 
 if __name__ == '__main__':
   os.environ['CUDA_VISIBLE_DEVICES'] = ''