Add get_file_spec method to DatasetBuilder.

PiperOrigin-RevId: 671362394

Author: fineguy

tensorflow_datasets/core/dataset_builder.py

@@ -553,6 +553,11 @@ def get_reference(
         data_dir=self.data_dir_root,
     )
 
+  def get_file_spec(self, split: str) -> str:
+    """Returns the file spec of the split."""
+    split_info: splits_lib.SplitInfo = self.info.splits[split]
+    return split_info.file_spec(self.info.file_format)
+
   def is_prepared(self) -> bool:
     """Returns whether this dataset is already downloaded and prepared."""
     return self.data_path.exists()
@@ -1082,7 +1087,7 @@ def _should_cache_ds(self, split, shuffle_files, read_config) -> bool:
     # shuffling is enabled, as this would effectively disable shuffling.
     # An exception is for single shard (as shuffling is a no-op).
     # Another exception is if reshuffle is disabled (shuffling already cached)
-    num_shards = len(self.info.splits[split].file_instructions)
+    num_shards = self.info.splits[split].num_shards
     if (
         shuffle_files
         and
@@ -1658,9 +1663,7 @@ def _get_filename_template(
         split=split_name,
         dataset_name=self.name,
         data_dir=self.data_path,
-        filetype_suffix=file_adapters.ADAPTER_FOR_FORMAT[
-            self.info.file_format
-        ].FILE_SUFFIX,
+        filetype_suffix=self.info.file_format.file_suffix,  # pytype: disable=attribute-error
     )
 
   def _generate_splits(

tensorflow_datasets/core/dataset_builder_test.py

@@ -634,6 +634,20 @@ def test_get_reference(self):
     )
     self.assertEqual(reference, expected_reference)
 
+  def test_get_file_spec(self):
+    builder = DummyDatasetWithConfigs(
+        config="plus1", data_dir=self.get_temp_dir()
+    )
+    builder.download_and_prepare()
+    self.assertEndsWith(
+        builder.get_file_spec("train"),
+        "dummy_dataset_with_configs/plus1/0.0.1/dummy_dataset_with_configs-train.tfrecord@1",
+    )
+    self.assertEndsWith(
+        builder.get_file_spec("test"),
+        "dummy_dataset_with_configs/plus1/0.0.1/dummy_dataset_with_configs-test.tfrecord@1",
+    )
+
   def test_load_as_data_source(self):
     data_dir = self.get_temp_dir()
     builder = DummyDatasetWithConfigs(

tensorflow_datasets/core/splits.py

@@ -205,18 +205,14 @@ def filenames(self) -> list[str]:
     """Returns the list of filenames."""
     if not self.filename_template:
       raise ValueError('No filename templates available.')
-    return sorted(
-        self.filename_template.sharded_filenames(len(self.shard_lengths))
-    )
+    return sorted(self.filename_template.sharded_filenames(self.num_shards))
 
   @property
   def filepaths(self) -> list[epath.Path]:
     """All the paths for all the files that are part of this split."""
     if not self.filename_template:
       raise ValueError('No filename templates available.')
-    return sorted(
-        self.filename_template.sharded_filepaths(len(self.shard_lengths))
-    )
+    return sorted(self.filename_template.sharded_filepaths(self.num_shards))
 
   def replace(self, **kwargs: Any) -> SplitInfo:
     """Returns a copy of the `SplitInfo` with updated attributes."""
@@ -421,7 +417,7 @@ def __init__(
       )
     self._dataset_name = dataset_name  # deprecated, please don't use
 
-  def __getitem__(self, key):
+  def __getitem__(self, key) -> SplitInfo | SubSplitInfo:
     if not self:
       raise KeyError(
           f'Trying to access `splits[{key!r}]` but `splits` is empty. '