Include a nondeterministic_order attribute to download_config, and add a flag to TFDS build CLI to set it accordingly. The new flag/attribute defaults to False.

PiperOrigin-RevId: 691787693

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/download/download_manager.py

@@ -108,6 +108,9 @@ class DownloadConfig:
       used.
     ignore_duplicates: whether to ignore duplicated examples with the same key.
       If there are multiple examples with the same key, the first one is kept.
+    nondeterministic_order: If True, it will not assure deterministic ordering
+      when writing' examples to disk in the case of beam datasets. This might
+      result in quicker dataset preparation.
   """
 
   extract_dir: epath.PathLike | None = None
@@ -126,6 +129,7 @@ class DownloadConfig:
   min_shard_size: int = shard_utils.DEFAULT_MIN_SHARD_SIZE
   max_shard_size: int = shard_utils.DEFAULT_MAX_SHARD_SIZE
   ignore_duplicates: bool = False
+  nondeterministic_order: bool = False
 
   def get_shard_config(self) -> shard_utils.ShardConfig:
     return shard_utils.ShardConfig(

tensorflow_datasets/scripts/cli/build.py

@@ -32,7 +32,13 @@
 
 
 def register_subparser(parsers: argparse._SubParsersAction) -> None:  # pylint: disable=protected-access
-  """Add subparser for `build` command."""
+  """Add subparser for `build` command.
+
+  New flags should be added to `cli_utils` module.
+
+  Args:
+    parsers: The subparsers object to add the parser to.
+  """
   build_parser = parsers.add_parser(
       'build', help='Commands for downloading and preparing datasets.'
   )
@@ -357,6 +363,7 @@ def _download_and_prepare(
       skip_if_published=args.skip_if_published,
       overwrite=args.overwrite,
       beam_pipeline_options=args.beam_pipeline_options,
+      nondeterministic_order=args.nondeterministic_order,
   )
 
 

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -261,6 +261,14 @@ def add_generation_argument_group(parser: argparse.ArgumentParser):
       default=1,
       help='Number of parallel build processes.',
   )
+  generation_group.add_argument(
+      '--nondeterministic_order',
+      action='store_false',
+      help=(
+          'If True, it will not assure deterministic ordering when writing'
+          ' examples to disk. This might result in quicker dataset preparation.'
+      ),
+  )
 
 
 def add_publish_argument_group(parser: argparse.ArgumentParser):
@@ -300,6 +308,7 @@ def download_and_prepare(
     skip_if_published: bool,
     overwrite: bool,
     beam_pipeline_options: str | None,
+    nondeterministic_order: bool = False,
 ) -> None:
   """Generate a single builder."""
   dataset = builder.info.full_name
@@ -317,6 +326,8 @@ def download_and_prepare(
     download_config = download.DownloadConfig()
   if overwrite and not download_config.download_mode.overwrite_dataset:
     download_config.download_mode = download.GenerateMode.REUSE_CACHE_IF_EXISTS
+  if nondeterministic_order:
+    download_config.nondeterministic_order = True
 
   # Add Apache Beam options to download config
   try: