Add helper function to check whether data has been generated in a specific file format

PiperOrigin-RevId: 664771687

Author: tomvdw

tensorflow_datasets/core/dataset_builder.py

@@ -836,32 +836,28 @@ def as_data_source(
         " choose another data_dir or delete the data."
     )
 
-    if info.file_format is None and not info.alternative_file_formats:
+    available_formats = info.available_file_formats()
+    if not available_formats:
       raise ValueError(
           "Dataset info file format is not set! For random access, one of the"
           f" following formats is required: {random_access_formats_msg}"
       )
 
-    if (
-        info.file_format is None
-        or info.file_format not in random_access_formats
-    ):
-      available_formats = set(info.alternative_file_formats)
-      suitable_formats = available_formats.intersection(random_access_formats)
-      if suitable_formats:
-        chosen_format = suitable_formats.pop()
-        logging.info(
-            "Found random access formats: %s. Chose to use %s. Overriding file"
-            " format in the dataset info.",
-            ", ".join([f.name for f in suitable_formats]),
-            chosen_format,
-        )
-        # Change the dataset info to read from a random access format.
-        info.set_file_format(
-            chosen_format, override=True, override_if_initialized=True
-        )
-      else:
-        raise NotImplementedError(unsupported_format_msg)
+    suitable_formats = available_formats.intersection(random_access_formats)
+    if suitable_formats:
+      chosen_format = suitable_formats.pop()
+      logging.info(
+          "Found random access formats: %s. Chose to use %s. Overriding file"
+          " format in the dataset info.",
+          ", ".join([f.name for f in suitable_formats]),
+          chosen_format,
+      )
+      # Change the dataset info to read from a random access format.
+      info.set_file_format(
+          chosen_format, override=True, override_if_initialized=True
+      )
+    else:
+      raise NotImplementedError(unsupported_format_msg)
 
     # Create a dataset for each of the given splits
     def build_single_data_source(split: str) -> Sequence[Any]:

tensorflow_datasets/core/dataset_info.py

@@ -532,6 +532,13 @@ def add_alternative_file_format(
     self._alternative_file_formats.append(file_format)
     self.as_proto.alternative_file_formats.append(file_format.value)
 
+  def available_file_formats(self) -> set[file_adapters.FileFormat]:
+    formats = set()
+    if self.file_format:
+      formats.add(self.file_format)
+    formats.update(self.alternative_file_formats)
+    return formats
+
   @property
   def splits(self) -> splits_lib.SplitDict:
     return self._splits
@@ -1265,6 +1272,19 @@ def update_info_proto_with_features(
   return completed_info_proto
 
 
+def supports_file_format(
+    dataset_info_proto: dataset_info_pb2.DatasetInfo,
+    file_format: str | file_adapters.FileFormat,
+) -> bool:
+  """Returns whether the given file format is supported by the dataset."""
+  if isinstance(file_format, file_adapters.FileFormat):
+    file_format = file_format.value
+  return (
+      file_format == dataset_info_proto.file_format
+      or file_format in dataset_info_proto.alternative_file_formats
+  )
+
+
 class MetadataDict(Metadata, dict):
   """A `tfds.core.Metadata` object that acts as a `dict`.
 

tensorflow_datasets/core/dataset_info_test.py

@@ -682,6 +682,22 @@ def test_create_redistribution_info_proto_unsupported_fields():
     )
 
 
+def test_supports_file_format():
+  dataset_info_proto = dataset_info_pb2.DatasetInfo(
+      file_format=file_adapters.FileFormat.TFRECORD.value,
+      alternative_file_formats=[file_adapters.FileFormat.RIEGELI.value],
+  )
+  assert dataset_info.supports_file_format(
+      dataset_info_proto, file_format=file_adapters.FileFormat.TFRECORD
+  )
+  assert dataset_info.supports_file_format(
+      dataset_info_proto, file_format=file_adapters.FileFormat.RIEGELI
+  )
+  assert not dataset_info.supports_file_format(
+      dataset_info_proto, file_format=file_adapters.FileFormat.PARQUET
+  )
+
+
 # pylint: disable=g-inconsistent-quotes
 _INFO_STR = '''tfds.core.DatasetInfo(
     name='mnist',