Expose tf.data.Dataset cardinality in TFDS

PiperOrigin-RevId: 296016166

Author: Conchylicultor

tensorflow_datasets/core/splits.py

@@ -73,7 +73,7 @@ def __init__(self, file_instructions):
     """Constructor.
 
     Args:
-      file_instructions: _FileInstructionOutput
+      file_instructions: FileInstructions
     """
     self._file_instructions = file_instructions
 

tensorflow_datasets/core/tfrecords_reader.py

@@ -165,7 +165,8 @@ def _read_files(
     files,
     parse_fn,
     read_config,
-    shuffle_files):
+    shuffle_files,
+    num_examples):
   """Returns tf.data.Dataset for given file instructions.
 
   Args:
@@ -176,6 +177,8 @@ def _read_files(
     read_config: `tfds.ReadConfig`, Additional options to configure the
       input pipeline (e.g. seed, num parallel reads,...).
     shuffle_files (bool): Defaults to False. True to shuffle input files.
+    num_examples: `int`, if defined, set the cardinality on the
+      tf.data.Dataset instance with `tf.data.experimental.with_cardinality`.
   """
   # Eventually apply a transformation to the instruction function.
   # This allow the user to have direct control over the interleave order.
@@ -211,7 +214,13 @@ def _read_files(
       cycle_length=parallel_reads,
       block_length=block_length,
       num_parallel_calls=tf.data.experimental.AUTOTUNE,
-      )
+  )
+
+  # If the number of examples read in the tf-record is known, we forward
+  # the information to the tf.data.Dataset object.
+  # Check the `tf.data.experimental` for backward compatibility with TF <= 2.1
+  if num_examples and hasattr(tf.data.experimental, 'assert_cardinality'):
+    ds = ds.apply(tf.data.experimental.assert_cardinality(num_examples))
 
   # TODO(tfds): Should merge the default options with read_config to allow users
   # to overwrite the default options.
@@ -265,28 +274,27 @@ def read(
        ReadInstruction instance. Otherwise a dict/list of tf.data.Dataset
        corresponding to given instructions param shape.
     """
-    def _read_instruction_to_file_instructions(instruction):
+    def _read_instruction_to_ds(instruction):
       file_instructions = make_file_instructions(name, split_infos, instruction)
       files = file_instructions.file_instructions
       if not files:
         msg = 'Instruction "%s" corresponds to no data!' % instruction
         raise AssertionError(msg)
-      return tuple(files)
+      return self.read_files(
+          files=tuple(files),
+          read_config=read_config,
+          shuffle_files=shuffle_files,
+          num_examples=file_instructions.num_examples,
+      )
 
-    files = utils.map_nested(
-        _read_instruction_to_file_instructions, instructions, map_tuple=False)
-    return utils.map_nested(
-        functools.partial(
-            self.read_files, read_config=read_config,
-            shuffle_files=shuffle_files),
-        files,
-        map_tuple=False)
+    return tf.nest.map_structure(_read_instruction_to_ds, instructions)
 
   def read_files(
       self,
       files,
       read_config,
-      shuffle_files
+      shuffle_files,
+      num_examples=None,
   ):
     """Returns single tf.data.Dataset instance for the set of file instructions.
 
@@ -296,6 +304,8 @@ def read_files(
         skip/take indicates which example read in the shard: `ds.skip().take()`
       read_config: `tfds.ReadConfig`, the input pipeline options
       shuffle_files (bool): If True, input files are shuffled before being read.
+      num_examples: `int`, if defined, set the cardinality on the
+        tf.data.Dataset instance with `tf.data.experimental.with_cardinality`.
 
     Returns:
        a tf.data.Dataset instance.
@@ -308,7 +318,9 @@ def read_files(
         files=files,
         read_config=read_config,
         parse_fn=self._parser.parse_example,
-        shuffle_files=shuffle_files)
+        shuffle_files=shuffle_files,
+        num_examples=num_examples,
+    )
     return dataset
 
 

tensorflow_datasets/core/tfrecords_reader_test.py

@@ -26,6 +26,8 @@
 from absl.testing import absltest
 import six
 
+import tensorflow as tf
+
 import tensorflow_datasets as tfds
 from tensorflow_datasets import testing
 from tensorflow_datasets.core import example_parser
@@ -35,6 +37,10 @@
 from tensorflow_datasets.core.utils import read_config as read_config_lib
 
 
+# Skip the cardinality test for backward compatibility with TF <= 2.1.
+_SKIP_CARDINALITY_TEST = not hasattr(tf.data.experimental, 'assert_cardinality')
+
+
 class GetDatasetFilesTest(testing.TestCase):
 
   NAME2SHARD_LENGTHS = {
@@ -308,19 +314,34 @@ def test_noskip_notake(self):
     read_data = list(tfds.as_numpy(ds))
     self.assertEqual(read_data, [six.b(l) for l in 'abcdefghijkl'])
 
+    if not _SKIP_CARDINALITY_TEST:
+      # Check that the cardinality is correctly set.
+      self.assertEqual(
+          tf.data.experimental.cardinality(ds).numpy(), len(read_data))
+
   def test_overlap(self):
     self._write_tfrecord('train', 5, 'abcdefghijkl')
     ds = self.reader.read('mnist', 'train+train[:2]', self.SPLIT_INFOS)
     read_data = list(tfds.as_numpy(ds))
     self.assertEqual(read_data, [six.b(l) for l in 'abcdefghijklab'])
 
+    if not _SKIP_CARDINALITY_TEST:
+      # Check that the cardinality is correctly set.
+      self.assertEqual(
+          tf.data.experimental.cardinality(ds).numpy(), len(read_data))
+
   def test_complex(self):
     self._write_tfrecord('train', 5, 'abcdefghijkl')
     self._write_tfrecord('test', 3, 'mnopqrs')
     ds = self.reader.read('mnist', 'train[1:-1]+test[:-50%]', self.SPLIT_INFOS)
     read_data = list(tfds.as_numpy(ds))
     self.assertEqual(read_data, [six.b(l) for l in 'bcdefghijkmno'])
 
+    if not _SKIP_CARDINALITY_TEST:
+      # Check that the cardinality is correctly set.
+      self.assertEqual(
+          tf.data.experimental.cardinality(ds).numpy(), len(read_data))
+
   def test_shuffle_files(self):
     self._write_tfrecord('train', 5, 'abcdefghijkl')
     ds = self.reader.read('mnist', 'train', self.SPLIT_INFOS,