Add support for splits in Croissant TFDS builder.

PiperOrigin-RevId: 652936333

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builders/croissant_builder.py

@@ -149,6 +149,7 @@ def __init__(
       float_dtype: type_utils.TfdsDType | None = np.float32,
       mapping: Mapping[str, epath.PathLike] | None = None,
       overwrite_version: str | None = None,
+      filters: Mapping[str, Any] | None = None,
       **kwargs: Any,
   ):
     """Initializes a CroissantBuilder.
@@ -170,6 +171,10 @@ def __init__(
         it to `~/Downloads/document.csv`, you can specify
         `mapping={"document.csv": "~/Downloads/document.csv"}`.
       overwrite_version: Semantic version of the dataset to be set.
+      filters: A dict of filters to apply to the records at preparation time (in
+        the `_generate_examples` function). The keys should be field names and
+        the values should be the values to filter by. If a record matches all
+        the filters, it will be included in the dataset.
       **kwargs: kwargs to pass to GeneratorBasedBuilder directly.
     """
     if mapping is None:
@@ -201,6 +206,7 @@ def __init__(
 
     self._int_dtype = int_dtype
     self._float_dtype = float_dtype
+    self._filters = filters or {}
 
     super().__init__(
         **kwargs,
@@ -222,19 +228,11 @@ def _info(self) -> dataset_info.DatasetInfo:
         disable_shuffling=self._disable_shuffling,
     )
 
-  def get_record_set(self, record_set_id: str):
-    """Returns the desired record set from self.metadata."""
-    for record_set in self.dataset.metadata.record_sets:
-      if huggingface_utils.convert_hf_name(record_set.id) == record_set_id:
-        return record_set
-    raise ValueError(
-        f'Did not find any record set with the name {record_set_id}.'
-    )
-
   def get_features(self) -> Optional[feature_lib.FeatureConnector]:
     """Infers the features dict for the required record set."""
-    record_set = self.get_record_set(self.builder_config.name)
-
+    record_set = croissant_utils.get_record_set(
+        self.builder_config.name, metadata=self.metadata
+    )
     fields = record_set.fields
     features = {}
     for field in fields:
@@ -249,18 +247,53 @@ def get_features(self) -> Optional[feature_lib.FeatureConnector]:
   def _split_generators(
       self, dl_manager: download.DownloadManager
   ) -> Dict[splits_lib.Split, split_builder_lib.SplitGenerator]:
-    # This will be updated when partitions are implemented in Croissant, ref to:
-    # https://docs.google.com/document/d/1saz3usja6mk5ugJXNF64_uSXsOzIgbIV28_bu1QamVY
-    return {'default': self._generate_examples()}  # pylint: disable=unreachable
+    # If a split recordset is joined for the required record set, we generate
+    # splits accordingly. Otherwise, it generates a single `default` split with
+    # all the records.
+    record_set = croissant_utils.get_record_set(
+        self.builder_config.name, metadata=self.metadata
+    )
+    if split_reference := croissant_utils.get_split_recordset(
+        record_set, metadata=self.metadata
+    ):
+      return {
+          split['name']: self._generate_examples(
+              filters={
+                  **self._filters,
+                  split_reference.reference_field.id: split['name'].encode(),
+              }
+          )
+          for split in split_reference.split_record_set.data
+      }
+    else:
+      return {'default': self._generate_examples(filters=self._filters)}
 
   def _generate_examples(
       self,
+      filters: dict[str, Any],
   ) -> split_builder_lib.SplitGenerator:
-    record_set = self.get_record_set(self.builder_config.name)
+    """Generates the examples for the given record set.
+
+    Args:
+      filters: A dict of filters to apply to the records. The keys should be
+        field names and the values should be the values to filter by. If a
+        record matches all the filters, it will be included in the dataset.
+
+    Yields:
+      A tuple of (index, record) for each record in the dataset.
+    """
+    record_set = croissant_utils.get_record_set(
+        self.builder_config.name, metadata=self.metadata
+    )
     records = self.dataset.records(record_set.id)
     for i, record in enumerate(records):
       # Some samples might not be TFDS-compatible as-is, e.g. from croissant
       # describing HuggingFace datasets, so we convert them here. This shouldn't
       # impact datasets which are already TFDS-compatible.
       record = huggingface_utils.convert_hf_value(record, self.info.features)
-      yield i, record
+      # After partition implementation, the filters will be applied from
+      # mlcroissant `dataset.records` directly.
+      # `records = records.filter(f == v for f, v in filters.items())``
+      # For now, we apply them in TFDS.
+      if all(record[filter] == value for filter, value in filters.items()):
+        yield i, record

tensorflow_datasets/core/utils/croissant_utils.py

@@ -17,15 +17,27 @@
 
 from __future__ import annotations
 
+import dataclasses
 import typing
 
 from tensorflow_datasets.core.utils import huggingface_utils
+from tensorflow_datasets.core.utils.lazy_imports_utils import mlcroissant as mlc
 
 if typing.TYPE_CHECKING:
   # pylint: disable=g-bad-import-order
   import mlcroissant as mlc
 
-_HUGGINGFACE_URL_PREFIX = 'https://huggingface.co/datasets/'
+_HUGGINGFACE_URL_PREFIX = "https://huggingface.co/datasets/"
+
+
+@dataclasses.dataclass(frozen=True)
+class SplitReference:
+  """Information about a split reference in a Croissant dataset."""
+
+  # A split record set in a Croissant dataset.
+  split_record_set: mlc.RecordSet
+  # A field from another record set that references split_record_set.
+  reference_field: mlc.Field
 
 
 def get_dataset_name(dataset: mlc.Dataset) -> str:
@@ -41,6 +53,66 @@ def get_tfds_dataset_name(dataset: mlc.Dataset) -> str:
   return huggingface_utils.convert_hf_name(dataset_name)
 
 
+def get_record_set(record_set_id: str, metadata: mlc.Metadata) -> mlc.RecordSet:
+  """Returns the desired record set from a dataset's metadata."""
+  for record_set in metadata.record_sets:
+    if huggingface_utils.convert_hf_name(record_set.id) == record_set_id:
+      return record_set
+  raise ValueError(
+      f"Did not find any record set with the name {record_set_id}."
+  )
+
+
+def get_field(field_id: str, metadata: mlc.Metadata) -> mlc.Field:
+  """Returns the desired field from a dataset's metadata."""
+  for record_set in metadata.record_sets:
+    for field in record_set.fields:
+      if field.id == field_id:
+        return field
+  raise ValueError(f"Did not find any field with the name {field_id}.")
+
+
+def get_record_set_for_field(
+    field_id: str, metadata: mlc.Metadata
+) -> mlc.RecordSet:
+  """Given a field id, returns the record set it belongs to, if any."""
+  for record_set in metadata.record_sets:
+    for field in record_set.fields:
+      if field.id == field_id:
+        return record_set
+  raise ValueError(f"Did not find any record set with field {field_id}.")
+
+
+def get_split_recordset(
+    record_set: mlc.RecordSet, metadata: mlc.Metadata
+) -> SplitReference | None:
+  """If a given recordset references a split recordset, returns it.
+
+  Args:
+    record_set: The record set to check.
+    metadata: The metadata of the dataset.
+
+  Returns:
+    If found, a tuple containing: (the field referencing the split record set,
+    and the split record set), None otherwise.
+  """
+  for field in record_set.fields:
+    if field.references and field.references.field:
+      # Check that the referenced record set is of type `cr:Split`.
+      referenced_field = get_field(field.references.field, metadata)
+      record_sets = [
+          node
+          for node in referenced_field.predecessors
+          if isinstance(node, mlc.RecordSet)
+      ]
+      if not record_sets:
+        raise ValueError("field {field.id} has no RecordSet")
+      referenced_record_set = record_sets[0]
+      if str(mlc.DataType.SPLIT) in referenced_record_set.data_types:
+        return SplitReference(referenced_record_set, field)
+  return None
+
+
 def get_record_set_ids(metadata: mlc.Metadata) -> list[str]:
   """Returns record set ids of the given MLcroissant metadata.
 

tensorflow_datasets/core/utils/croissant_utils_test.py

@@ -55,3 +55,35 @@ def test_get_record_set_ids():
   )
   record_set_ids = croissant_utils.get_record_set_ids(metadata=metadata)
   assert record_set_ids == ['record_set_1']
+
+
+def test_get_split_recordset_with_no_split_recordset():
+  record_sets = [
+      mlc.RecordSet(
+          id='labels',
+          key='name',
+          fields=[
+              mlc.Field(
+                  id='labels/label',
+                  name='label',
+                  data_types=mlc.DataType.TEXT,
+              )
+          ],
+          data=[{'label': 'bird'}, {'label': 'bike'}],
+      ),
+      mlc.RecordSet(
+          id='samples',
+          fields=[
+              mlc.Field(
+                  id='samples/label',
+                  data_types=mlc.DataType.TEXT,
+                  references=mlc.Source(field='labels/label'),
+              )
+          ],
+      ),
+  ]
+  metadata = mlc.Metadata(name='dummy', url='dum.my', record_sets=record_sets)
+  split_recordset = croissant_utils.get_split_recordset(
+      record_set=metadata.record_sets[0], metadata=metadata
+  )
+  assert split_recordset is None