diff --git a/CHANGELOG.md b/CHANGELOG.md index f4cc444..1790204 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.60.0] - 2026-08-25 + +### Changed + +- `tilebox-datasets`: Dataset queries now return short enum names as strings (for example, `HH`) instead of numeric + values with a name mapping in the xarray variable attributes. Repeated enum fields can now be queried and + round-tripped through all supported ingestion inputs. + ## [0.59.0] - 2026-08-18 ### Added @@ -471,7 +479,8 @@ the first client that does not cache data (since it's already on the local file - Released under the [MIT](https://opensource.org/license/mit) license. - Released packages: `tilebox-datasets`, `tilebox-workflows`, `tilebox-storage`, `tilebox-grpc` -[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.59.0...HEAD +[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.60.0...HEAD +[0.60.0]: https://github.com/tilebox/tilebox-python/compare/v0.59.0...v0.60.0 [0.59.0]: https://github.com/tilebox/tilebox-python/compare/v0.58.0...v0.59.0 [0.58.0]: https://github.com/tilebox/tilebox-python/compare/v0.57.0...v0.58.0 [0.57.0]: https://github.com/tilebox/tilebox-python/compare/v0.56.0...v0.57.0 diff --git a/tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py b/tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py index 9a1a243..bb4bf47 100644 --- a/tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py +++ b/tilebox-datasets/tests/protobuf_conversion/test_protobuf_xarray.py @@ -1,3 +1,4 @@ +from enum import Enum from uuid import UUID import pandas as pd @@ -25,8 +26,14 @@ from tilebox.datasets.datasets.stac.v1.core_pb2 import Provider as ProviderPB2 from tilebox.datasets.datasets.stac.v1.processing_pb import ProcessingSoftware from tilebox.datasets.datasets.stac.v1.processing_pb2 import ProcessingSoftware as ProcessingSoftwarePB2 +from tilebox.datasets.datasets.stac.v1.sar_pb2 import ( + SAR_POLARIZATION_HH, + SAR_POLARIZATION_VV, + SARProperties, +) from tilebox.datasets.datasets.stac.v1.storage_pb import Storage from tilebox.datasets.datasets.stac.v1.storage_pb2 import Storage as StoragePB2 +from tilebox.datasets.datasets.v1.well_known_types_pb2 import ProcessingLevel from tilebox.datasets.protobuf_conversion.field_types import _AssetsDisplay from tilebox.datasets.protobuf_conversion.protobuf_xarray import MessageToXarrayConverter from tilebox.datasets.protobuf_conversion.to_protobuf import to_messages @@ -87,7 +94,7 @@ def test_convert_datapoint(datapoint: ExampleDatapoint) -> None: # noqa: PLR091 ) assert isinstance(dataset.some_geometry.item(), Polygon | MultiPolygon) - assert dataset.some_enum.item() == datapoint.some_enum + assert dataset.some_enum.item() == ProcessingLevel.Name(datapoint.some_enum).removeprefix("PROCESSING_LEVEL_") 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) @@ -213,6 +220,89 @@ def test_convert_stac_messages_to_protobuf_py() -> None: assert "access_profiles" not in html +def test_convert_scalar_and_repeated_enums_to_names_and_round_trip() -> None: + class Polarization(Enum): + HH = "HH" + VV = "VV" + + file_descriptor = descriptor_pb2.FileDescriptorProto( + name="tests/protobuf_conversion/enum_datapoint.proto", + package="tests.protobuf_conversion.enums", + ) + enum_descriptor = file_descriptor.enum_type.add(name="SARPolarization") + for name, number in (("SAR_POLARIZATION_UNSPECIFIED", 0), ("SAR_POLARIZATION_HH", 1), ("SAR_POLARIZATION_VV", 2)): + enum_descriptor.value.add(name=name, number=number) + message_descriptor = file_descriptor.message_type.add(name="EnumDatapoint") + message_descriptor.field.add( + name="primary_polarization", + number=1, + label=descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL, + type=descriptor_pb2.FieldDescriptorProto.TYPE_ENUM, + type_name=".tests.protobuf_conversion.enums.SARPolarization", + ) + message_descriptor.field.add( + name="polarizations", + number=2, + label=descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED, + type=descriptor_pb2.FieldDescriptorProto.TYPE_ENUM, + type_name=".tests.protobuf_conversion.enums.SARPolarization", + ) + descriptor = Default().AddSerializedFile(file_descriptor.SerializeToString()) + message_type = GetMessageClass(descriptor.message_types_by_name["EnumDatapoint"]) + messages = [ + message_type(primary_polarization=1, polarizations=[0, 1]), + message_type(primary_polarization=2, polarizations=[2]), + message_type(), + ] + + converter = MessageToXarrayConverter() + converter.convert_all(messages) + dataset = converter.finalize("time") + + assert dataset.primary_polarization.dtype == object + assert dataset.primary_polarization[:2].to_numpy().tolist() == ["HH", "VV"] + assert pd.isna(dataset.primary_polarization[2].item()) + assert dataset.polarizations.dtype == object + assert dataset.polarizations[0].to_numpy().tolist() == ["UNSPECIFIED", "HH"] + assert dataset.polarizations[1, 0].item() == "VV" + assert pd.isna(dataset.polarizations[1, 1].item()) + assert pd.isna(dataset.polarizations[2].to_numpy()).all() + assert dataset.polarizations.dims == ("time", "n_polarizations") + assert dataset.primary_polarization.attrs == {} + assert dataset.polarizations.attrs == {} + assert to_messages(dataset, message_type) == messages + + expected = message_type(primary_polarization=1, polarizations=[1, 2, 0]) + record = {"primary_polarization": "HH", "polarizations": ["HH", Polarization.VV, 0]} + assert to_messages([record], message_type) == [expected] + assert to_messages({name: [value] for name, value in record.items()}, message_type) == [expected] + assert to_messages(pd.DataFrame([record]), message_type) == [expected] + + with pytest.raises(ValueError, match="Record 0: Field 'polarizations': Invalid enum name 'INVALID'"): + to_messages([{"polarizations": ["INVALID"]}], message_type) + + +def test_convert_sar_polarizations_with_short_names_in_both_directions() -> None: + messages = [ + SARProperties(polarizations=[SAR_POLARIZATION_HH, SAR_POLARIZATION_VV]), + SARProperties(polarizations=[SAR_POLARIZATION_VV]), + SARProperties(), + ] + converter = MessageToXarrayConverter() + converter.convert_all(messages) + + dataset = converter.finalize("item") + + assert dataset.polarizations[0].to_numpy().tolist() == ["HH", "VV"] + assert dataset.polarizations[1, 0].item() == "VV" + assert pd.isna(dataset.polarizations[1, 1].item()) + assert pd.isna(dataset.polarizations[2].to_numpy()).all() + + other_fields = [field.name for field in SARProperties.DESCRIPTOR.fields if field.name != "polarizations"] + assert to_messages(dataset, SARProperties, ignore_fields=other_fields) == messages + assert to_messages([{"polarizations": ["HH", "VV"]}], SARProperties) == [messages[0]] + + @given(lists(example_datapoints(generated_fields=True, missing_fields=True), min_size=5, max_size=30)) def test_convert_datapoints(datapoints: list[ExampleDatapoint]) -> None: # noqa: C901, PLR0912 converter = MessageToXarrayConverter() diff --git a/tilebox-datasets/tests/test_client.py b/tilebox-datasets/tests/test_client.py index a411296..a154be1 100644 --- a/tilebox-datasets/tests/test_client.py +++ b/tilebox-datasets/tests/test_client.py @@ -145,8 +145,7 @@ def test_find_datapoint() -> None: if not skip_data: assert datapoint.granule_name.item() == "S2A_MSIL1C_20220713T002201_N0400_R102_T08XNS_20220713T015332.SAFE" - processing_level = datapoint.processing_level.item() - assert datapoint.processing_level.attrs["names"][processing_level] == "L1C" + assert datapoint.processing_level.item() == "L1C" assert datapoint.copernicus_id.item() == "65505f82-76dd-5e85-b947-a6c879e07446" assert isinstance(datapoint.geometry.item(), Polygon) else: diff --git a/tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py b/tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py index 250349d..07fcb82 100644 --- a/tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py +++ b/tilebox-datasets/tilebox/datasets/protobuf_conversion/field_types.py @@ -1,5 +1,7 @@ +import re from collections.abc import Sequence from datetime import timedelta +from enum import Enum from typing import Any from uuid import UUID @@ -120,21 +122,30 @@ def to_proto(self, value: Any) -> bool: 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 + super().__init__(object) 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: + def from_proto(self, value: ProtoFieldValue) -> str: 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: + try: + return self._values_to_name[value] + except KeyError as error: + raise ValueError(f"Invalid enum value {value}") from error + + def to_proto(self, value: str | int | Enum) -> int: + if isinstance(value, Enum): + value = value.name if isinstance(value, (str, np.str_)): - return self._names_to_value[value] - if int(value) not in self._values_to_name: + try: + return self._names_to_value[value] + except KeyError as error: + raise ValueError(f"Invalid enum name {value!r}") from error + integer_value = int(value) + if integer_value not in self._values_to_name: raise ValueError(f"Invalid enum value {value}") # during ingestion, we can raise an error here - return value + return integer_value class TimestampField(ProtobufFieldType): @@ -360,8 +371,11 @@ def _camel_to_uppercase(name: str) -> str: Examples: >>> _camel_to_uppercase("ProcessingLevel") 'PROCESSING_LEVEL' + >>> _camel_to_uppercase("SARPolarization") + 'SAR_POLARIZATION' """ - return "".join(["_" + c.lower() if c.isupper() else c for c in name]).lstrip("_").upper() + name = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", name) + return re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name).upper() def is_missing(value: Any) -> bool: diff --git a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py index 640bb5a..e6dda31 100644 --- a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py +++ b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py @@ -13,10 +13,8 @@ from numpy.typing import NDArray from tilebox.datasets.protobuf_conversion.field_types import ( - EnumField, ProtobufFieldType, ProtoFieldValue, - enum_mapping_from_field_descriptor, infer_field_type, ) @@ -313,26 +311,6 @@ def _resize(self) -> None: self._data = data -class _EnumFieldConverter(_SimpleFieldConverter): - def __init__(self, field_name: str, enum_names: dict[int, str]) -> None: - """ - A field converter for the enum type. - - Args: - field_name: The name of enum field in the protobuf message - """ - super().__init__(field_name, EnumField(enum_names)) - self._enum_names = enum_names - - def finalize( - self, dataset: xr.Dataset, count: int, dimension_names: tuple[str, ...], skip_if_empty: bool = False - ) -> str | None: - field_name = super().finalize(dataset, count, dimension_names, skip_if_empty) - if field_name is not None: - dataset[field_name].attrs["names"] = self._enum_names - return field_name - - def _create_field_converters(message: Message, buffer_size: int) -> dict[str, _FieldConverter]: """ Create a dictionary mapping from field names to field converters for the given protobuf message descriptor. @@ -369,13 +347,6 @@ def _create_field_converter(field: FieldDescriptor) -> _FieldConverter: Returns: A field converter for the given protobuf field descriptor """ - # special handling for enums: - if field.type == FieldDescriptor.TYPE_ENUM: - if field.is_repeated: - raise NotImplementedError("Repeated enum fields are not supported") - - return _EnumFieldConverter(field.name, enum_mapping_from_field_descriptor(field)) - field_type = infer_field_type(field) if field.is_repeated: return _ArrayFieldConverter(field.name, field_type)