CroissantBuilder: Add support for fields with subfields in datatype_converter.

PiperOrigin-RevId: 683118823

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builders/croissant_builder.py

@@ -61,7 +61,7 @@
 
 
 def datatype_converter(
-    field,
+    field: mlc.Field,
     int_dtype: Optional[type_utils.TfdsDType] = np.int64,
     float_dtype: Optional[type_utils.TfdsDType] = np.float32,
 ):
@@ -83,8 +83,15 @@ def datatype_converter(
     raise NotImplementedError('Not implemented yet.')
 
   field_data_type = field.data_type
-
   if not field_data_type:
+    # Fields with sub fields are of type None
+    if field.sub_fields:
+      return features_dict.FeaturesDict({
+          subfield.id: datatype_converter(
+              subfield, int_dtype=int_dtype, float_dtype=float_dtype
+          )
+          for subfield in field.sub_fields
+      })
     return None
   elif field_data_type == int:
     return int_dtype

tensorflow_datasets/core/dataset_builders/croissant_builder_test.py

@@ -20,7 +20,9 @@
 from tensorflow_datasets import testing
 from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core.dataset_builders import croissant_builder
+from tensorflow_datasets.core.features import features_dict
 from tensorflow_datasets.core.features import image_feature
+from tensorflow_datasets.core.features import tensor_feature
 from tensorflow_datasets.core.features import text_feature
 from tensorflow_datasets.core.utils.lazy_imports_utils import mlcroissant as mlc
 
@@ -95,28 +97,52 @@ def test_simple_datatype_converter(field, feature_type, int_dtype, float_dtype):
 
 
 @pytest.mark.parametrize(
-    ["field", "feature_type"],
+    ["field", "feature_type", "subfield_types"],
     [
         (
             mlc.Field(data_types=mlc.DataType.TEXT, description="Text feature"),
             text_feature.Text,
+            None,
         ),
         (
             mlc.Field(data_types=mlc.DataType.DATE, description="Date feature"),
             text_feature.Text,
+            None,
         ),
         (
             mlc.Field(
                 data_types=mlc.DataType.IMAGE_OBJECT,
                 description="Image feature",
             ),
             image_feature.Image,
+            None,
+        ),
+        (
+            mlc.Field(
+                id="person",
+                data_types=[],
+                description="A field with subfields",
+                sub_fields=[
+                    mlc.Field(id="person/name", data_types=mlc.DataType.TEXT),
+                    mlc.Field(id="person/age", data_types=mlc.DataType.INTEGER),
+                ],
+            ),
+            features_dict.FeaturesDict,
+            {
+                "person/name": text_feature.Text,
+                "person/age": tensor_feature.Tensor,
+            },
         ),
     ],
 )
-def test_complex_datatype_converter(field, feature_type):
+def test_complex_datatype_converter(field, feature_type, subfield_types):
   actual_feature = croissant_builder.datatype_converter(field)
   assert isinstance(actual_feature, feature_type)
+  if subfield_types:
+    for feature_name in actual_feature.keys():
+      assert isinstance(
+          actual_feature[feature_name], subfield_types[feature_name]
+      )
 
 
 @pytest.fixture(name="crs_builder")

tensorflow_datasets/core/utils/conversion_utils_test.py

@@ -140,6 +140,14 @@ def test_convert_value_raises(value, feature):
             ),
             {'foo': b''},
         ),
+        (
+            {'name': b'Name', 'age': 100},
+            feature_lib.FeaturesDict({
+                'name': feature_lib.Text(),
+                'age': feature_lib.Scalar(dtype=np.int32),
+            }),
+            {'name': b'Name', 'age': 100},
+        ),
         # nan, but the feature type is not float
         (
             np.nan,