Merge pull request #11070 from lgeiger:fix-no-shuffle-beam-writer

PiperOrigin-RevId: 777974077

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/naming.py

@@ -666,7 +666,7 @@ def sharded_filepaths_pattern(
     `/path/dataset_name-split.fileformat@num_shards` or
     `/path/dataset_name-split@num_shards.fileformat` depending on the format.
     If `num_shards` is not given, then it returns
-    `/path/dataset_name-split.fileformat*`.
+    `/path/dataset_name-split.fileformat-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]`.
 
     Args:
       num_shards: optional specification of the number of shards.
@@ -681,7 +681,7 @@ def sharded_filepaths_pattern(
     elif use_at_notation:
       replacement = '@*'
     else:
-      replacement = '*'
+      replacement = '-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]'
     return _replace_shard_pattern(os.fspath(a_filepath), replacement)
 
   def glob_pattern(self, num_shards: int | None = None) -> str:

tensorflow_datasets/core/naming_test.py

@@ -459,7 +459,7 @@ def test_sharded_file_template_shard_index():
   )
   assert (
       os.fspath(template.sharded_filepaths_pattern())
-      == '/my/path/data/mnist-train.tfrecord*'
+      == '/my/path/data/mnist-train.tfrecord-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]'
   )
   assert (
       os.fspath(template.sharded_filepaths_pattern(num_shards=100))
@@ -474,7 +474,10 @@ def test_glob_pattern():
       filetype_suffix='tfrecord',
       data_dir=epath.Path('/data'),
   )
-  assert '/data/ds-train.tfrecord*' == template.glob_pattern()
+  assert (
+      '/data/ds-train.tfrecord-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]'
+      == template.glob_pattern()
+  )
   assert '/data/ds-train.tfrecord-*-of-00042' == template.glob_pattern(
       num_shards=42
   )

tensorflow_datasets/core/writer.py

@@ -816,8 +816,9 @@ def finalize(self) -> tuple[list[int], int]:
     logging.info("Finalizing writer for %s", self._filename_template.split)
     # We don't know the number of shards, the length of each shard, nor the
     # total size, so we compute them here.
-    prefix = epath.Path(self._filename_template.filepath_prefix())
-    shards = self._filename_template.data_dir.glob(f"{prefix.name}*")
+    shards = self._filename_template.data_dir.glob(
+        self._filename_template.glob_pattern()
+    )
 
     def _get_length_and_size(shard: epath.Path) -> tuple[epath.Path, int, int]:
       length = self._file_adapter.num_examples(shard)

tensorflow_datasets/core/writer_test.py

@@ -592,39 +592,47 @@ def test_write_beam(self, file_format: file_adapters.FileFormat):
 
     with tempfile.TemporaryDirectory() as tmp_dir:
       tmp_dir = epath.Path(tmp_dir)
-      filename_template = naming.ShardedFileTemplate(
-          dataset_name='foo',
-          split='train',
-          filetype_suffix=file_format.file_suffix,
-          data_dir=tmp_dir,
-      )
-      writer = writer_lib.NoShuffleBeamWriter(
-          serializer=testing.DummySerializer('dummy specs'),
-          filename_template=filename_template,
-          file_format=file_format,
-      )
+
+      def get_writer(split):
+        filename_template = naming.ShardedFileTemplate(
+            dataset_name='foo',
+            split=split,
+            filetype_suffix=file_format.file_suffix,
+            data_dir=tmp_dir,
+        )
+        return writer_lib.NoShuffleBeamWriter(
+            serializer=testing.DummySerializer('dummy specs'),
+            filename_template=filename_template,
+            file_format=file_format,
+        )
+
       to_write = [(i, str(i).encode('utf-8')) for i in range(10)]
       # Here we need to disable type check as `beam.Create` is not capable of
       # inferring the type of the PCollection elements.
       options = beam.options.pipeline_options.PipelineOptions(
           pipeline_type_check=False
       )
-      with beam.Pipeline(options=options, runner=_get_runner()) as pipeline:
-
-        @beam.ptransform_fn
-        def _build_pcollection(pipeline):
-          pcollection = pipeline | 'Start' >> beam.Create(to_write)
-          return writer.write_from_pcollection(pcollection)
-
-        _ = pipeline | 'test' >> _build_pcollection()  # pylint: disable=no-value-for-parameter
-      shard_lengths, total_size = writer.finalize()
-      self.assertNotEmpty(shard_lengths)
-      self.assertEqual(sum(shard_lengths), 10)
-      self.assertGreater(total_size, 10)
+      writers = [get_writer(split) for split in ('train-b', 'train')]
+
+      for writer in writers:
+        with beam.Pipeline(options=options, runner=_get_runner()) as pipeline:
+
+          @beam.ptransform_fn
+          def _build_pcollection(pipeline, writer):
+            pcollection = pipeline | 'Start' >> beam.Create(to_write)
+            return writer.write_from_pcollection(pcollection)
+
+          _ = pipeline | 'test' >> _build_pcollection(writer)
+
       files = list(tmp_dir.iterdir())
-      self.assertGreaterEqual(len(files), 1)
+      self.assertGreaterEqual(len(files), 2)
       for f in files:
         self.assertIn(file_format.file_suffix, f.name)
+      for writer in writers:
+        shard_lengths, total_size = writer.finalize()
+        self.assertNotEmpty(shard_lengths)
+        self.assertEqual(sum(shard_lengths), 10)
+        self.assertGreater(total_size, 10)
 
 
 class CustomExampleWriter(writer_lib.ExampleWriter):