Include a beam_sink class method to the file adapters.

PiperOrigin-RevId: 693445053

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/file_adapters.py

@@ -26,14 +26,17 @@
 from typing import Any, ClassVar, Type, TypeVar
 
 from etils import epy
+from tensorflow_datasets.core.utils.lazy_imports_utils import apache_beam as beam
 from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_module
 from tensorflow_datasets.core.utils.lazy_imports_utils import parquet as pq
 from tensorflow_datasets.core.utils.lazy_imports_utils import pyarrow as pa
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
+
 with epy.lazy_imports():
   # pylint: disable=g-import-not-at-top
   from etils import epath
+  from tensorflow_datasets.core import naming
   from tensorflow_datasets.core.utils import file_utils
   from tensorflow_datasets.core.utils import type_utils
 
@@ -167,6 +170,23 @@ def deserialize(cls, raw_example: bytes) -> Any:
     """
     return tf.train.Example.FromString(raw_example)
 
+  @classmethod
+  def beam_sink(
+      cls,
+      filename_template: naming.ShardedFileTemplate,
+      num_shards: int | None = None,
+  ) -> beam.PTransform:
+    """Returns a Beam sink for writing examples in the given file format."""
+    raise NotImplementedError()
+
+  @classmethod
+  def num_examples(cls, filename: epath.PathLike) -> int:
+    """Returns the number of examples in the given file."""
+    n = 0
+    for _ in cls.make_tf_data(filename):
+      n += 1
+    return n
+
 
 class TfRecordFileAdapter(FileAdapter):
   """File adapter for TFRecord file format."""
@@ -205,6 +225,20 @@ def write_examples(
         writer.write(serialized_example)
       writer.flush()
 
+  @classmethod
+  def beam_sink(
+      cls,
+      filename_template: naming.ShardedFileTemplate,
+      num_shards: int | None = None,
+  ) -> beam.PTransform:
+    """Returns a Beam sink for writing examples in the given file format."""
+    file_path_prefix = filename_template.sharded_filepaths_pattern(
+        num_shards=num_shards, use_at_notation=True
+    ).removesuffix('@*')
+    return beam.io.WriteToTFRecord(
+        file_path_prefix=file_path_prefix, num_shards=num_shards
+    )
+
 
 class RiegeliFileAdapter(FileAdapter):
   """File adapter for Riegeli file format."""

tensorflow_datasets/core/naming.py

@@ -657,6 +657,7 @@ def sharded_filepaths_pattern(
       self,
       *,
       num_shards: int | None = None,
+      use_at_notation: bool = False,
   ) -> str:
     """Returns a pattern describing all the file paths captured by this template.
 
@@ -668,13 +669,16 @@ def sharded_filepaths_pattern(
 
     Args:
       num_shards: optional specification of the number of shards.
+      use_at_notation: whether to return @* in case `num_shards` is `None`.
 
     Returns:
       the pattern describing all shards captured by this template.
     """
     a_filepath = self.sharded_filepath(shard_index=0, num_shards=1)
     if num_shards:
       replacement = f'@{num_shards}'
+    elif use_at_notation:
+      replacement = '@*'
     else:
       replacement = '*'
     return _replace_shard_pattern(os.fspath(a_filepath), replacement)

tensorflow_datasets/core/utils/lazy_imports_utils_untyped.py

@@ -0,0 +1,17 @@
+# coding=utf-8
+# Copyright 2024 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Lazy import utils untyped. Please, use lazy_imports_utils.py instead.
+"""