Dataset Schema¶

The output schema lives in the autoware_ml/databases/schemas/ package. It is split into a top-level table definition and reusable nested data models that can be shared across dataset families. Every database implementation emits Sequence[DatasetRecord] from process_scenario_records() and persists rows as a Polars DataFrame (typically Parquet) using DatasetTableSchema.

Base building blocks (base_schemas.py)¶

• DatasetTableColumn - a NamedTuple pairing a column name with a Polars data type.
• BaseFieldSchema - base class for nested struct schemas. Subclasses declare DatasetTableColumn attributes and expose to_polars_field_schema() to produce pl.Field definitions for struct columns.
• DataModelInterface - abstract interface requiring to_dictionary() and load_from_dictionary() so every Pydantic data model can round-trip through a Polars DataFrame.

Top-level table (dataset_schemas.py)¶

• DatasetTableSchema - a frozen dataclass whose class-level attributes are DatasetTableColumn entries. Call DatasetTableSchema.to_polars_schema() to get a pl.Schema for constructing or validating a Polars DataFrame.
• DatasetRecord - a frozen Pydantic model (implementing DataModelInterface) representing a single row. One record is emitted per sample/frame by process_scenario_records().

class DatasetTableSchema:
    # Basic metadata
    SCENARIO_ID = DatasetTableColumn("scenario_id", pl.String)
    SAMPLE_ID = DatasetTableColumn("sample_id", pl.String)
    SAMPLE_INDEX = DatasetTableColumn("sample_index", pl.Int32)
    TIMESTAMP_SECONDS = DatasetTableColumn("timestamp_seconds", pl.Float64)
    LOCATION = DatasetTableColumn("location", pl.String)
    VEHICLE_TYPE = DatasetTableColumn("vehicle_type", pl.String)
    SCENARIO_NAME = DatasetTableColumn("scenario_name", pl.String)

    # Nested sensor data columns
    LIDAR_FRAMES = DatasetTableColumn("lidar_frames", pl.List(pl.Struct(...)))
    LIDAR_SOURCES = DatasetTableColumn("lidar_sources", pl.List(pl.Struct(...)))
    # Annotation fields
    CATEGORY_MAPPING = DatasetTableColumn("category_mapping", pl.Struct(...))
    BOXES_3D = DatasetTableColumn("boxes_3d", pl.List(pl.Struct(...)))

    @classmethod
    def to_polars_schema(cls) -> pl.Schema: ...

class DatasetRecord(BaseModel, DataModelInterface):
    scenario_id: str
    sample_id: str
    sample_index: int
    timestamp_seconds: float
    location: str | None
    vehicle_type: str | None
    scenario_name: str
    lidar_frames: Sequence[LidarFrameDataModel]
    lidar_sources: Sequence[LidarSourceDataModel] | None
    category_mapping: CategoryMappingDataModel | None
    boxes_3d: Sequence[Box3DDataModel] | None

    def to_dictionary(self) -> Mapping[str, Any]: ...
    @classmethod
    def load_from_dictionary(cls, data_model: Mapping[str, Any]) -> DatasetRecord: ...

Top-level columns¶

lidar_frames struct fields¶

Each list entry is a LidarFrameDataModel covering one keyframe or sweep:

lidar_sources struct fields¶

Each list entry is a LidarSourceDataModel describing one LiDAR sensor in the scene:

category_mapping struct fields¶

Nested data models¶

boxes_3d struct fields¶

Each list entry is a Box3DDataModel with the following struct fields:

Each nested module follows the same pattern: a *DatasetSchema class defines the Polars struct layout, and a matching *DataModel Pydantic class implements DataModelInterface for serialization. DatasetRecord.to_dictionary() delegates to these nested models when writing Parquet; DatasetRecord.load_from_dictionary() reconstructs them when reading back.

DatasetTableSchema, DatasetRecord, and every nested schema/data-model pair are kept in sync. When adding new columns/fields (e.g. 3D bounding boxes), add entries to the relevant *DatasetSchema and *DataModel, then wire the new column into DatasetTableSchema and DatasetRecord.

Extending the schema¶

See T4Dataset for a concrete example of how these schemas are populated from T4 annotations.

Implementation¶