Pass num shards to the beam sink from the no shuffle beam writer

tomvdw · The TensorFlow Datasets Authors · commit 498ba3f62c36 · 2025-10-01T07:52:40.000-07:00
PiperOrigin-RevId: 813751246
diff --git a/tensorflow_datasets/core/writer.py b/tensorflow_datasets/core/writer.py
@@ -810,7 +810,7 @@ def write_from_pcollection(self, examples_pcollection):
           self._num_shards
       )
     return serialized_examples | "Write" >> self._file_adapter.beam_sink(
-        filename_template=self._filename_template
+        filename_template=self._filename_template, num_shards=self._num_shards
     )
 
   def finalize(self) -> tuple[list[int], int]:

Original file line number	Diff line number	Diff line change
`@@ -810,7 +810,7 @@ def write_from_pcollection(self, examples_pcollection):`
`810`	`810`	`self._num_shards`
`811`	`811`	`)`
`812`	`812`	`return serialized_examples \| "Write" >> self._file_adapter.beam_sink(`
`813`		`- filename_template=self._filename_template`
	`813`	`+ filename_template=self._filename_template, num_shards=self._num_shards`
`814`	`814`	`)`
`815`	`815`
`816`	`816`	`def finalize(self) -> tuple[list[int], int]:`