Dataset field annotations for roles, queryables and json path mapping (#44)

* Update protos

* Dataset field annotations for roles, queryables and json path mapping

Author: lukasbindreiter

CHANGELOG.md

@@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.56.0] - 2026-07-16
 
+### Added
+
+- `tilebox-datasets`: Added dataset field annotations for source JSON pointers, queryable fields, JSON Schema references,
+  and semantic roles, including STAC mappings for well-known dataset fields.
+
 ### Changed
 
 - `tilebox-workflows`: Reduced import-time overhead for release runners by lazily loading optional heavy dependencies
@@ -412,7 +417,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.55.1...HEAD
+[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.56.0...HEAD
+[0.56.0]: https://github.com/tilebox/tilebox-python/compare/v0.55.1...v0.56.0
 [0.55.1]: https://github.com/tilebox/tilebox-python/compare/v0.55.0...v0.55.1
 [0.55.0]: https://github.com/tilebox/tilebox-python/compare/v0.54.0...v0.55.0
 [0.54.0]: https://github.com/tilebox/tilebox-python/compare/v0.53.0...v0.54.0

pyproject.toml

@@ -69,9 +69,9 @@ exclude = [
     "*/.venv/*",
     # it's auto generated, don't lint it
     "*/datasets/v1/*",
+    "*/datasets/stac/v1/*",
     "*/workflows/v1/*",
     "*/tilebox/v1/*",
-    "*/buf/validate/*",
 ]
 
 [tool.ruff.lint]
@@ -85,7 +85,7 @@ ignore = [
     "DTZ", # flake8-datetimez: utc datetimes are useful
     "DJ",  # flake8-django: not needed
     "EM",  # flake8-errmsg: str directly in Exception constructor is accetable
-    "TCH", # flake8-type-checking: type checking blocks are weird
+    "TC",  # flake8-type-checking: type checking blocks are weird
     # specific rules
     "ANN401",  # any-type: allow Any in *args and **kwargs
     "S101",    # assert: allow usage of assert

tilebox-datasets/tests/data/datasets.py

@@ -31,6 +31,7 @@
     Field,
     FieldAnnotation,
     FieldDict,
+    FieldRole,
     ListDatasetsResponse,
 )
 from tilebox.datasets.message_pool import register_once
@@ -41,7 +42,11 @@ def field_annotations(draw: DrawFn) -> FieldAnnotation:
     """A hypothesis strategy for generating random field annotations"""
     description = draw(text(alphabet=string.ascii_letters, min_size=3, max_size=25))
     example_value = draw(text(alphabet=string.ascii_letters + string.digits + "-_", min_size=1, max_size=10))
-    return FieldAnnotation(description, example_value)
+    source_json_pointer = draw(text(alphabet=string.ascii_letters + string.digits + "/_-", max_size=25) | none())
+    queryable = draw(booleans())
+    json_schema_ref = draw(text(alphabet=string.ascii_letters + string.digits + "#/_-", max_size=25) | none())
+    roles = draw(lists(sampled_from(FieldRole), unique=True))
+    return FieldAnnotation(description, example_value, source_json_pointer, queryable, json_schema_ref, roles)
 
 
 @composite
@@ -73,12 +78,22 @@ def field_dicts(draw: DrawFn) -> FieldDict:
         )
     )
     annotation = draw(field_annotations())
+    roles = draw(
+        one_of(
+            lists(sampled_from(FieldRole), unique=True),
+            lists(just("primary_title"), unique=True),
+        )
+    )
 
     return {
         "name": name,
         "type": field_type,
         "description": annotation.description,
         "example_value": annotation.example_value,
+        "source_json_pointer": annotation.source_json_pointer,
+        "queryable": annotation.queryable,
+        "json_schema_ref": annotation.json_schema_ref,
+        "roles": roles,
     }
 
 
@@ -104,8 +119,7 @@ def fields(draw: DrawFn) -> Field:
     descriptor = FieldDescriptorProto(name=name, type=field_type, type_name=type_name, label=label)
 
     annotation = draw(field_annotations())
-    queryable = draw(booleans())
-    return Field(descriptor, annotation, queryable)
+    return Field(descriptor, annotation)
 
 
 @composite

tilebox-datasets/tests/data/test_datasets.py

@@ -1,3 +1,7 @@
+from unittest.mock import MagicMock, patch
+from uuid import uuid4
+
+import pytest
 from hypothesis import given
 
 from tests.data.datasets import (
@@ -14,26 +18,47 @@
     AnnotatedType,
     Dataset,
     DatasetGroup,
+    DatasetKind,
     DatasetType,
     Field,
     FieldAnnotation,
     FieldDict,
+    FieldRole,
     ListDatasetsResponse,
 )
+from tilebox.datasets.service import TileboxDatasetService
 
 
 @given(field_annotations())
 def test_field_annotations_to_message_and_back(annotation: FieldAnnotation) -> None:
     assert FieldAnnotation.from_message(annotation.to_message()) == annotation
 
 
+@pytest.mark.parametrize("attribute", ["source_json_pointer", "json_schema_ref"])
+def test_optional_field_annotation_presence(attribute: str) -> None:
+    absent = FieldAnnotation("description", "example").to_message()
+    assert not absent.HasField(attribute)
+
+    if attribute == "source_json_pointer":
+        present = FieldAnnotation("description", "example", source_json_pointer="").to_message()
+    else:
+        present = FieldAnnotation("description", "example", json_schema_ref="").to_message()
+    assert present.HasField(attribute)
+
+
 @given(field_dicts())
 def test_field_from_dict(field_dict: FieldDict) -> None:
     field = Field.from_dict(field_dict)
     assert field.descriptor.name == field_dict["name"]
     assert field.descriptor.type is not None
     assert field.annotation.description == field_dict.get("description", "")
     assert field.annotation.example_value == field_dict.get("example_value", "")
+    assert field.annotation.source_json_pointer == field_dict.get("source_json_pointer")
+    assert field.annotation.queryable == field_dict.get("queryable", False)
+    assert field.annotation.json_schema_ref == field_dict.get("json_schema_ref")
+    assert field.annotation.roles == [
+        FieldRole[role.upper()] if isinstance(role, str) else role for role in field_dict.get("roles", [])
+    ]
 
 
 @given(fields())
@@ -64,3 +89,44 @@ def test_dataset_groups_to_message_and_back(group: DatasetGroup) -> None:
 @given(list_datasets_responses())
 def test_list_datasets_responses_to_message_and_back(response: ListDatasetsResponse) -> None:
     assert ListDatasetsResponse.from_message(response.to_message()) == response
+
+
+@pytest.mark.parametrize("operation", ["create", "update"])
+def test_create_and_update_dataset_include_field_annotations(operation: str) -> None:
+    dataset_service = MagicMock()
+    service = TileboxDatasetService(dataset_service, MagicMock(), MagicMock(), MagicMock())
+    custom_fields: list[FieldDict] = [
+        {
+            "name": "title",
+            "type": str,
+            "source_json_pointer": "/properties/title",
+            "queryable": True,
+            "json_schema_ref": "https://example.com/schema.json#/title",
+            "roles": [FieldRole.PRIMARY_TITLE],
+        }
+    ]
+
+    with patch.object(Dataset, "from_message", side_effect=lambda message: message):
+        if operation == "create":
+            service.create_dataset(DatasetKind.SPATIOTEMPORAL, "example", "Example", custom_fields).get()
+            request = dataset_service.CreateDataset.call_args.args[0]
+        else:
+            service.update_dataset(DatasetKind.SPATIOTEMPORAL, uuid4(), "Example", custom_fields).get()
+            request = dataset_service.UpdateDataset.call_args.args[0]
+
+    annotations = {field.descriptor.name: field.annotation for field in request.type.fields}
+    assert annotations["time"].source_json_pointer == "/properties/datetime"
+    assert annotations["time"].json_schema_ref.endswith("datetime.json#/properties/datetime")
+    assert annotations["id"].source_json_pointer == "/properties/tilebox_id"
+    assert not annotations["id"].HasField("json_schema_ref")
+    assert annotations["ingestion_time"].source_json_pointer == "/properties/created"
+    assert annotations["ingestion_time"].json_schema_ref.endswith("datetime.json#/properties/created")
+    assert annotations["geometry"].source_json_pointer == "/geometry"
+    assert annotations["geometry"].json_schema_ref == "https://geojson.org/schema/Geometry.json"
+    assert all(
+        not annotation.queryable and not annotation.roles for name, annotation in annotations.items() if name != "title"
+    )
+    assert annotations["title"].source_json_pointer == "/properties/title"
+    assert annotations["title"].queryable
+    assert annotations["title"].json_schema_ref == "https://example.com/schema.json#/title"
+    assert list(annotations["title"].roles) == [FieldRole.PRIMARY_TITLE.value]

tilebox-datasets/tilebox/datasets/data/datasets.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
+from dataclasses import field as dataclass_field
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import TypedDict, get_args, get_origin
+from typing import Literal, TypeAlias, TypedDict, get_args, get_origin
 from uuid import UUID
 
 import numpy as np
@@ -24,17 +25,56 @@ class DatasetKind(Enum):
 _dataset_kind_int_to_enum = {kind.value: kind for kind in DatasetKind}
 
 
+class FieldRole(Enum):
+    PRIMARY_TITLE = dataset_type_pb2.FIELD_ROLE_PRIMARY_TITLE
+    """The primary human-readable title or name used when displaying a datapoint."""
+
+
+_field_roles_from_string = {role.name.lower(): role for role in FieldRole}
+_field_role_int_to_enum = {role.value: role for role in FieldRole}
+FieldRolesLike: TypeAlias = list[FieldRole] | list[Literal["primary_title"]]
+
+
 @dataclass(frozen=True)
 class FieldAnnotation:
     description: str
+    """A human-readable description of the field."""
     example_value: str
+    """An example value illustrating the field's expected contents."""
+    source_json_pointer: str | None = None
+    """An RFC 6901 JSON Pointer locating the field value in the source document."""
+    queryable: bool = False
+    """Whether the field is available for use in server-side filter expressions."""
+    json_schema_ref: str | None = None
+    """A JSON Schema reference for the field when it corresponds to a well-known schema."""
+    roles: list[FieldRole] = dataclass_field(default_factory=list)
+    """Semantic display roles fulfilled by the field."""
 
     @classmethod
     def from_message(cls, annotation: dataset_type_pb2.FieldAnnotation) -> "FieldAnnotation":
-        return cls(description=annotation.description, example_value=annotation.example_value)
+        return cls(
+            description=annotation.description,
+            example_value=annotation.example_value,
+            source_json_pointer=(
+                annotation.source_json_pointer if annotation.HasField("source_json_pointer") else None
+            ),
+            queryable=annotation.queryable,
+            json_schema_ref=annotation.json_schema_ref if annotation.HasField("json_schema_ref") else None,
+            roles=[_field_role_int_to_enum[role] for role in annotation.roles],
+        )
 
     def to_message(self) -> dataset_type_pb2.FieldAnnotation:
-        return dataset_type_pb2.FieldAnnotation(description=self.description, example_value=self.example_value)
+        annotation = dataset_type_pb2.FieldAnnotation(
+            description=self.description,
+            example_value=self.example_value,
+            queryable=self.queryable,
+            roles=[role.value for role in self.roles],
+        )
+        if self.source_json_pointer is not None:
+            annotation.source_json_pointer = self.source_json_pointer
+        if self.json_schema_ref is not None:
+            annotation.json_schema_ref = self.json_schema_ref
+        return annotation
 
 
 class FieldDict(TypedDict):
@@ -63,6 +103,10 @@ class FieldDict(TypedDict):
     ]
     description: NotRequired[str]
     example_value: NotRequired[str]
+    source_json_pointer: NotRequired[str | None]
+    queryable: NotRequired[bool]
+    json_schema_ref: NotRequired[str | None]
+    roles: NotRequired[FieldRolesLike]
 
 
 _TYPE_INFO: dict[type, tuple[FieldDescriptorProto.Type.ValueType, str | None]] = {
@@ -83,14 +127,12 @@ class FieldDict(TypedDict):
 class Field:
     descriptor: FieldDescriptorProto
     annotation: FieldAnnotation
-    queryable: bool
 
     @classmethod
     def from_message(cls, field: dataset_type_pb2.Field) -> "Field":
         return cls(
             descriptor=field.descriptor,
             annotation=FieldAnnotation.from_message(field.annotation),
-            queryable=field.annotation.queryable,
         )
 
     @classmethod
@@ -116,16 +158,19 @@ def from_dict(cls, field: FieldDict) -> "Field":
             annotation=FieldAnnotation(
                 description=field.get("description", ""),
                 example_value=field.get("example_value", ""),
+                source_json_pointer=field.get("source_json_pointer"),
+                queryable=field.get("queryable", False),
+                json_schema_ref=field.get("json_schema_ref"),
+                roles=[
+                    _field_roles_from_string[role] if isinstance(role, str) else role for role in field.get("roles", [])
+                ],
             ),
-            queryable=False,
         )
 
     def to_message(self) -> dataset_type_pb2.Field:
-        annotation = self.annotation.to_message()
-        annotation.queryable = self.queryable
         return dataset_type_pb2.Field(
             descriptor=self.descriptor,
-            annotation=annotation,
+            annotation=self.annotation.to_message(),
         )