Allow ingestion of enums (#187)

Author: lukasbindreiter

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## Added
+
+- Support for ingestion of enum values into datasets
+
 ## [0.34.0] - 2025-04-15
 
 ## Added

tilebox-datasets/tests/data/datapoint.py

@@ -75,10 +75,10 @@ def example_datapoints(draw: DrawFn, generated_fields: bool = False, missing_fie
     maybe_none = none() if missing_fields else one_of()
 
     return ExampleDatapoint(
-        time=(draw(datetime_messages()) if generated_fields else None),
+        time=draw(datetime_messages()),
         id=(draw(uuid_messages()) if generated_fields else None),
         ingestion_time=(draw(datetime_messages()) if generated_fields else None),
-        # geometry=,  # skip for now
+        geometry=draw(geometry_messages()),
         some_string=draw(text(alphabet=string.ascii_letters + string.digits, min_size=1, max_size=10) | maybe_none),
         some_int=draw(integers(min_value=1, max_value=100) | maybe_none),
         some_double=draw(floats(min_value=1.0, max_value=100.0) | maybe_none),

tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py

@@ -5,14 +5,13 @@
 from hypothesis.strategies import lists
 from numpy.testing import assert_array_almost_equal, assert_array_equal
 from pandas import to_datetime
-from shapely import MultiPolygon, Polygon
+from shapely import MultiPolygon, Polygon, from_wkb
 from xarray.testing import assert_equal
 
 from tests.data.datapoint import datapoint_pages, datapoints, example_datapoints
 from tests.example_dataset.example_dataset_pb2 import ExampleDatapoint
 from tilebox.datasets.data.datapoint import Datapoint, DatapointPage
 from tilebox.datasets.data.time_interval import timestamp_to_datetime, us_to_datetime
-from tilebox.datasets.datasetsv1.well_known_types_pb2 import ProcessingLevel
 from tilebox.datasets.protobuf_conversion.protobuf_xarray import (
     MessageToXarrayConverter,
     TimeseriesToXarrayConverter,
@@ -45,6 +44,7 @@ def test_convert_datapoint(datapoint: ExampleDatapoint) -> None:  # noqa: PLR091
     assert us_to_datetime(to_datetime(ingestion_time, utc=True).value // 1000) == timestamp_to_datetime(
         datapoint.ingestion_time
     )
+    assert dataset.geometry.item() == from_wkb(datapoint.geometry.wkb)
 
     dataset = dataset.isel(time=0)  # select the only datapoint in the dataset
     assert dataset.some_string.item() == datapoint.some_string
@@ -72,8 +72,7 @@ def test_convert_datapoint(datapoint: ExampleDatapoint) -> None:  # noqa: PLR091
     )
 
     assert isinstance(dataset.some_geometry.item(), Polygon | MultiPolygon)
-    expected_level = {v: k for k, v in ProcessingLevel.items()}[datapoint.some_enum].removeprefix("PROCESSING_LEVEL_")
-    assert dataset.some_enum.item() == expected_level
+    assert dataset.some_enum.item() == datapoint.some_enum
 
     assert list(dataset.some_repeated_string.to_numpy()) == list(datapoint.some_repeated_string)
     assert_array_equal(dataset.some_repeated_int.to_numpy(), datapoint.some_repeated_int)

tilebox-datasets/tests/protobuf_conversion/test_to_protobuf.py

@@ -23,9 +23,6 @@ def test_xarray_dataset_to_protobuf_messages(messages: list[ExampleDatapoint]) -
                 assert not converted_message.HasField(field.name)
                 continue
 
-            if field.name == "some_enum":  # enum ingestion not implemented yet
-                continue
-
             assert getattr(expected_message, field.name) == getattr(converted_message, field.name), (
                 f"Field {field.name} mismatch"
             )

tilebox-datasets/tests/test_client.py

@@ -110,7 +110,8 @@ def test_find_datapoint() -> None:
 
         if not skip_data:
             assert datapoint.granule_name.item() == "S2A_MSIL1C_20220713T002201_N0400_R102_T08XNS_20220713T015332.SAFE"
-            assert datapoint.processing_level.item() == "L1C"
+            processing_level = datapoint.processing_level.item()
+            assert datapoint.processing_level.attrs["names"][processing_level] == "L1C"
             assert datapoint.copernicus_id.item() == "65505f82-76dd-5e85-b947-a6c879e07446"
             assert isinstance(datapoint.geometry.item(), Polygon)
         else:

tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py

@@ -78,6 +78,25 @@ def to_proto(self, value: Any) -> bool:
         return bool(value)
 
 
+class EnumField(ProtobufFieldType):
+    def __init__(self, name_lookup: dict[int, str]) -> None:
+        super().__init__(np.uint8)  # we support up to 256 different enum values for now
+        self._values_to_name = name_lookup
+        self._names_to_value = {name: value for value, name in name_lookup.items()}
+
+    def from_proto(self, value: ProtoFieldValue) -> int:
+        if not isinstance(value, int):
+            raise TypeError(f"Expected int message but got {type(value)}")
+        return value  # we don't parse the value when loading, to avoid having huge arrays of strings
+
+    def to_proto(self, value: str | int) -> int:
+        if isinstance(value, (str, np.str_)):
+            return self._names_to_value[value]
+        if int(value) not in self._values_to_name:
+            raise ValueError(f"Invalid enum value {value}")  # during ingestion, we can raise an error here
+        return value
+
+
 class TimestampField(ProtobufFieldType):
     def __init__(self) -> None:
         super().__init__("datetime64[ns]")
@@ -252,10 +271,47 @@ def infer_field_type(field: FieldDescriptor) -> ProtobufFieldType:
 
         return _MESSAGE_NAMES_TO_FIELDS[message_name]
 
+    if field.type == FieldDescriptor.TYPE_ENUM:
+        return EnumField(enum_mapping_from_field_descriptor(field))
+
     if field.type == FieldDescriptor.TYPE_BOOL:
         return BoolField()  # special handling, since we need to convert numpy bools to python bools
 
     if field.type not in _PROTOBUF_TYPE_TO_NUMPY_TYPE:
         raise ValueError(f"Unsupported field type {field.type}")
 
     return ProtobufFieldType(_PROTOBUF_TYPE_TO_NUMPY_TYPE[field.type])
+
+
+def enum_mapping_from_field_descriptor(field: FieldDescriptor) -> dict[int, str]:
+    """Create a mapping from enum values to their names.
+
+    Args:
+        field: The field descriptor to create the mapping for. Must be of type FieldDescriptor.TYPE_ENUM.
+    """
+    if field.type != FieldDescriptor.TYPE_ENUM:
+        raise ValueError("Expected field to be of type FieldDescriptor.TYPE_ENUM")
+
+    # remove the enum type prefix from the enum values
+    # e.g. FLIGHT_DIRECTION_ASCENDING of the FlightDirection enum will result in a value of ASCENDING
+    enum_type_prefix = _camel_to_uppercase(field.enum_type.name) + "_"
+    return {
+        v.number: str(v.name).removeprefix(enum_type_prefix)
+        for v in field.enum_type.values  # noqa: PD011  # enum_type is not a numpy array, even though ruff thinks it is
+    }
+
+
+def _camel_to_uppercase(name: str) -> str:
+    """Convert a camelCase name to an UPPER_CASE name.
+
+    Args:
+        name: The name to convert.
+
+    Returns:
+        The converted name.
+
+    Examples:
+        >>> _camel_to_uppercase("ProcessingLevel")
+        'PROCESSING_LEVEL'
+    """
+    return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_").upper()

tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py

@@ -14,7 +14,13 @@
 
 from tilebox.datasets.data.datapoint import Datapoint, DatapointPage
 from tilebox.datasets.message_pool import get_message_type
-from tilebox.datasets.protobuf_conversion.field_types import ProtobufFieldType, ProtoFieldValue, infer_field_type
+from tilebox.datasets.protobuf_conversion.field_types import (
+    EnumField,
+    ProtobufFieldType,
+    ProtoFieldValue,
+    enum_mapping_from_field_descriptor,
+    infer_field_type,
+)
 
 AnyMessage = TypeVar("AnyMessage", bound=Message)
 
@@ -375,21 +381,13 @@ def __init__(self, field_name: str, enum_names: dict[int, str]) -> None:
         Args:
             field_name: The name of enum field in the protobuf message
         """
-        super().__init__(field_name, ProtobufFieldType(np.uint8))  # we support up to 256 different enum values for now
+        super().__init__(field_name, EnumField(enum_names))
         self._enum_names = enum_names
 
     def finalize(self, dataset: xr.Dataset, count: int, dimension_names: tuple[str, ...]) -> str | None:
         field_name = super().finalize(dataset, count, dimension_names)
         if field_name is not None:
-            # convert the numeric enum values to the corresponding string names
-            int_values = dataset[field_name]
-            dataset[field_name] = (
-                int_values.dims,
-                np.vectorize(lambda i: self._enum_names.get(i, ""))(int_values.values),
-            )
-            dataset[field_name].attrs["enum_dict"] = ";".join(
-                f"{name}: {value}" for value, name in self._enum_names.items()
-            )
+            dataset[field_name].attrs["names"] = self._enum_names
         return field_name
 
 
@@ -418,22 +416,6 @@ def _create_field_converters(message: Message, buffer_size: int) -> dict[str, _F
     return converters
 
 
-def _camel_to_uppercase(name: str) -> str:
-    """Convert a camelCase name to an UPPER_CASE name.
-
-    Args:
-        name: The name to convert.
-
-    Returns:
-        The converted name.
-
-    Examples:
-        >>> _camel_to_uppercase("ProcessingLevel")
-        'PROCESSING_LEVEL'
-    """
-    return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_").upper()
-
-
 def _create_field_converter(field: FieldDescriptor) -> _FieldConverter:
     """
     Create a field converter for the given protobuf field descriptor.
@@ -450,12 +432,7 @@ def _create_field_converter(field: FieldDescriptor) -> _FieldConverter:
         if field.label == FieldDescriptor.LABEL_REPEATED:
             raise NotImplementedError("Repeated enum fields are not supported")
 
-        # remove the enum type prefix from the enum values
-        enum_type_prefix = _camel_to_uppercase(field.enum_type.name) + "_"
-        return _EnumFieldConverter(
-            field.name,
-            {v.number: str(v.name).removeprefix(enum_type_prefix) for v in field.enum_type.values},
-        )
+        return _EnumFieldConverter(field.name, enum_mapping_from_field_descriptor(field))
 
     field_type = infer_field_type(field)
     if field.label == FieldDescriptor.LABEL_OPTIONAL:  # simple fields (in proto3 every simple field is optional)

tilebox-datasets/tilebox/datasets/protobuf_conversion/to_protobuf.py

@@ -79,9 +79,6 @@ def to_messages(  # noqa: C901, PLR0912
         field_lengths[len(values)].append(field_name)  # to validate all fields have the same length
 
         descriptor = field_descriptors_by_name[field_name]
-        if descriptor.type == FieldDescriptor.TYPE_ENUM:
-            continue  # skip enums, not supported for now in ingestion
-
         field_type = infer_field_type(descriptor)
         if isinstance(field_type, GeobufField):
             continue  # legacy geometry type, ingestion is only supported for the new geometry type

tilebox-datasets/tilebox/datasets/sync/timeseries.py

@@ -139,7 +139,7 @@ def __init__(self, dataset: TimeseriesDataset, info: CollectionInfo) -> None:
 
     def __repr__(self) -> str:
         """Human readable representation of the collection."""
-        return repr(self._info)
+        return repr(self.info())
 
     @property
     def name(self) -> str:

tilebox-storage/tests/storage_data.py

@@ -34,35 +34,7 @@ def ers_granules(draw: DrawFn, ensure_quicklook: bool = False) -> ASFStorageGran
         quicklook = f"{_ASF_URL}/BROWSE/E{platform}/{granule_name}.jpg"
     urls = StorageURLs(f"{_ASF_URL}/{level}/E{platform}/{file_name}.zip", quicklook)
 
-    return ASFStorageGranule(time, granule_name, level, "ASF", file_size, md5sum, urls)
-
-
-@composite
-def s1_granules(draw: DrawFn) -> ASFStorageGranule:
-    """Generate a realistic-looking random Sentinel 1 granule."""
-    level = "RAW"
-    platform = draw(one_of(just("A"), just("B")))
-    acquisition_mode = draw(one_of(*(just(am) for am in ["IW", "EW", *[f"S{i + 1}" for i in range(6)]])))
-    start = draw(datetimes(min_value=datetime(2014, 6, 1), max_value=datetime(2050, 1, 1), timezones=just(None)))
-    stop = draw(datetimes(min_value=datetime(2014, 6, 1), max_value=datetime(2050, 1, 1), timezones=just(None)))
-    orbit = draw(integers(min_value=1, max_value=999_999))
-    random_1 = draw(text(alphabet=string.hexdigits, min_size=6, max_size=6)).upper()
-    random_2 = draw(text(alphabet=string.hexdigits, min_size=4, max_size=4)).upper()
-    file_size = draw(integers(min_value=10_000, max_value=999_999_999))
-    md5sum = draw(text(alphabet=string.hexdigits, min_size=32, max_size=32))
-
-    granule_name = (
-        f"S1{platform}_{acquisition_mode}_{level}__0SDV_{start:%Y%m%dT%H%M%S}_{stop:%Y%m%dT%H%M%S}_"
-        f"{orbit:06d}_S{random_1}_{random_2}"
-    )
-
-    urls = StorageURLs(f"{_ASF_URL}/{level}/S{platform}/{granule_name}.zip", None)
-    return ASFStorageGranule(start, granule_name, level, "ASF", file_size, md5sum, urls)
-
-
-@composite
-def asf_granules(draw: DrawFn) -> ASFStorageGranule:
-    return draw(one_of(ers_granules(), s1_granules()))
+    return ASFStorageGranule(time, granule_name, file_size, md5sum, urls)
 
 
 @composite