Fix incompatibility of tfds.Builder.as_data_source with MultiSplitInfo.

PiperOrigin-RevId: 857137643

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -890,18 +890,22 @@ def as_data_source(
         )
       chosen_format = file_format
     else:
-      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,
-      )
+      if info.file_format in suitable_formats:
+        chosen_format = info.file_format
+      else:
+        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
-    )
+    if info.file_format != 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
+      )
 
     # Create a dataset for each of the given splits
     def build_single_data_source(split: str) -> Sequence[Any]:

tensorflow_datasets/core/dataset_info.py

@@ -538,6 +538,9 @@ def set_file_format(
       updated_split_infos = []
       for split_info in self.splits.values():
         if split_info.filename_template is None:
+          logging.warning(
+              "Split %s has no filename template, skipping.", split_info.name
+          )
           continue
         updated_split_info = split_info.replace(
             filename_template=split_info.filename_template.replace(