Use the new NoShuffleBeamWriter when download_config.nondeterministic_order is True and save it to dataset_info proto for documentation.

PiperOrigin-RevId: 693628830

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -622,6 +622,10 @@ def download_and_prepare(
     data_path = self.data_path
     data_exists = data_path.exists()
 
+    # Saving nondeterministic_order in the DatasetInfo for documentation.
+    if download_config.nondeterministic_order:
+      self.info.set_nondeterministic_order(True)
+
     if download_config.download_mode == UPDATE_DATASET_INFO:
       self._update_dataset_info()
       return
@@ -1426,6 +1430,8 @@ def _get_filename_template(
       self, split_name: str
   ) -> naming.ShardedFileTemplate:
     """Returns a filename template for the given split."""
+    if self.info.file_format is None:
+      raise ValueError("File format is not set!")
     return naming.ShardedFileTemplate(
         split=split_name,
         dataset_name=self.name,
@@ -1728,6 +1734,7 @@ def _generate_splits(
             generator=generator,
             filename_template=filename_template,
             disable_shuffling=self.info.disable_shuffling,
+            nondeterministic_order=download_config.nondeterministic_order,
         )
         split_info_futures.append(future)
 

tensorflow_datasets/core/dataset_builder_beam_test.py

@@ -39,14 +39,16 @@ class DummyBeamDataset(dataset_builder.GeneratorBasedBuilder):
       'valid_725': 725,
   }
 
+  FEATURE_DICT = features.FeaturesDict({
+      'image': features.Image(shape=(16, 16, 1)),
+      'label': features.ClassLabel(names=['dog', 'cat']),
+      'id': tf.int32,
+  })
+
   def _info(self):
     return dataset_info.DatasetInfo(
         builder=self,
-        features=features.FeaturesDict({
-            'image': features.Image(shape=(16, 16, 1)),
-            'label': features.ClassLabel(names=['dog', 'cat']),
-            'id': tf.int32,
-        }),
+        features=self.FEATURE_DICT,
         supervised_keys=('x', 'x'),
         metadata=dataset_info.BeamMetadataDict(),
     )
@@ -71,6 +73,18 @@ def _generate_examples(self, num_examples):
     return examples
 
 
+class UnshuffledDummyBeamDataset(DummyBeamDataset):
+
+  def _info(self) -> dataset_info.DatasetInfo:
+    return dataset_info.DatasetInfo(
+        builder=self,
+        features=self.FEATURE_DICT,
+        supervised_keys=('x', 'x'),
+        metadata=dataset_info.BeamMetadataDict(),
+        disable_shuffling=True,
+    )
+
+
 class CommonPipelineDummyBeamDataset(DummyBeamDataset):
   EXPECTED_METADATA = {
       'label_sum_1000': 500,
@@ -151,12 +165,21 @@ def _compute_mean(examples):
   )
 
 
+def get_id(ex):
+  return ex['id']
+
+
 def make_default_config():
   return download.DownloadConfig()
 
 
 @pytest.mark.parametrize(
-    'dataset_cls', [DummyBeamDataset, CommonPipelineDummyBeamDataset]
+    'dataset_cls',
+    [
+        DummyBeamDataset,
+        CommonPipelineDummyBeamDataset,
+        UnshuffledDummyBeamDataset,
+    ],
 )
 @pytest.mark.parametrize(
     'make_dl_config',
@@ -178,29 +201,23 @@ def test_beam_datasets(
   assert data_path.exists()  # Dataset has been generated
 
   # Check number of shards/generated files
-  _test_shards(
-      data_path,
-      pattern='%s-test.tfrecord-{:05}-of-{:05}' % dataset_name,
-      # Liquid sharding is not guaranteed to always use the same number.
-      num_shards=builder.info.splits['test'].num_shards,
-  )
-  _test_shards(
-      data_path,
-      pattern='%s-train.tfrecord-{:05}-of-{:05}' % dataset_name,
-      num_shards=1,
-  )
+  for split in ['test', 'train']:
+    _test_shards(
+        data_path,
+        pattern='%s-%s.tfrecord-{:05}-of-{:05}' % (dataset_name, split),
+        num_shards=builder.info.splits[split].num_shards,
+    )
 
   ds = dataset_utils.as_numpy(builder.as_dataset())
 
-  def get_id(ex):
-    return ex['id']
-
+  test_examples = list(ds['test'])
+  train_examples = list(ds['train'])
   _assert_values_equal(
-      sorted(list(ds['test']), key=get_id),
+      sorted(test_examples, key=get_id),
       sorted([_gen_example(i)[1] for i in range(725)], key=get_id),
   )
   _assert_values_equal(
-      sorted(list(ds['train']), key=get_id),
+      sorted(train_examples, key=get_id),
       sorted([_gen_example(i)[1] for i in range(1000)], key=get_id),
   )
 

tensorflow_datasets/core/split_builder.py

@@ -34,6 +34,7 @@
   # pylint: disable=g-import-not-at-top
   from tensorflow_datasets.core import example_serializer
   from tensorflow_datasets.core import features as features_lib
+  from tensorflow_datasets.core import file_adapters
   from tensorflow_datasets.core import naming
   from tensorflow_datasets.core import splits as splits_lib
   from tensorflow_datasets.core import utils
@@ -410,14 +411,18 @@ def submit_split_generation(
       generator: SplitGenerator,
       filename_template: naming.ShardedFileTemplate,
       disable_shuffling: bool,
+      nondeterministic_order: bool,
   ) -> _SplitInfoFuture:
     """Start the split generation.
 
     Args:
-      split_name: Name of the split to generate
-      generator: Generator, beam.PTransform,... yielding the examples
+      split_name: Name of the split to generate.
+      generator: Generator, beam.PTransform,... yielding the examples.
       filename_template: Template to format the filename for a shard.
-      disable_shuffling: Specifies whether to shuffle the examples
+      disable_shuffling: Specifies whether to shuffle the examples.
+      nondeterministic_order: If True, it will not assure deterministic ordering
+        when writing' examples to disk. This might result in quicker dataset
+        preparation
 
     Returns:
       split_info_future: Future containing the `split_info`, once generation
@@ -433,13 +438,19 @@ def submit_split_generation(
     # Depending on the type of generator, we use the corresponding
     # `_build_from_xyz` method.
     if isinstance(generator, Iterable):
+      if nondeterministic_order:
+        logging.warning(
+            'Enabling `nondeterministic_order` for a dataset that does not use'
+            ' beam has no effect.'
+        )
       return self._build_from_generator(**build_kwargs)
     else:  # Otherwise, beam required
       unknown_generator_type = TypeError(
           f'Invalid split generator value for split `{split_name}`. '
           'Expected generator or apache_beam object. Got: '
           f'{type(generator)}'
       )
+      build_kwargs['nondeterministic_order'] = nondeterministic_order
       if isinstance(generator, beam.PTransform):
         # Generate the beam.PCollection
         pcollection = self.beam_pipeline | split_name >> generator
@@ -527,20 +538,35 @@ def _build_from_pcollection(
       generator: 'beam.PCollection[KeyExample]',
       filename_template: naming.ShardedFileTemplate,
       disable_shuffling: bool,
+      nondeterministic_order: bool,
   ) -> _SplitInfoFuture:
     """Split generator for `beam.PCollection`."""
     # TODO(tfds): Should try to add support to `max_examples_per_split`
-    beam_writer = writer_lib.BeamWriter(
-        serializer=example_serializer.ExampleSerializer(
-            self._features.get_serialized_info()
-        ),
-        filename_template=filename_template,
-        hash_salt=split_name,
-        disable_shuffling=disable_shuffling,
-        shard_config=self._shard_config,
-        example_writer=self._example_writer,
-        ignore_duplicates=self._ignore_duplicates,
+    serializer = example_serializer.ExampleSerializer(
+        self._features.get_serialized_info()
     )
+    if nondeterministic_order:
+      logging.info(
+          'Order of examples does not matter, using NoShuffleBeamWriter'
+      )
+      beam_writer = writer_lib.NoShuffleBeamWriter(
+          serializer=serializer,
+          file_format=file_adapters.FileFormat.from_value(
+              filename_template.filetype_suffix
+          ),
+          filename_template=filename_template,
+      )
+    else:
+      logging.info('Deterministic ordering is enabled, using BeamWriter')
+      beam_writer = writer_lib.BeamWriter(
+          serializer=serializer,
+          filename_template=filename_template,
+          hash_salt=split_name,
+          disable_shuffling=disable_shuffling,
+          shard_config=self._shard_config,
+          example_writer=self._example_writer,
+          ignore_duplicates=self._ignore_duplicates,
+      )
 
     def _encode_example(key_ex, encode_fn=self._features.encode_example):
       # We do not access self._features in this function to avoid pickling the