PydanticLogicalTypeFactory: typing.Literal fields unsupported as pipeline columns (Unsupported annotation)

[brian-arnold]

orcapod: typing.Literal fields in a pydantic/dataclass model aren't supported as pipeline columns

Repo / commit: nauticalab/orcapod-python @ main (fef6add24d82f66257461f36f3c416fe89a80509) — i.e. with PR #185 / ITL-432 (the fix for #184).
Environment: Python 3.12.10 · polars 1.41.2 · pyarrow 24.0.0 · pydantic 2.13.4 · starfix 0.3.1

Summary

PR #185 made pydantic models flow through pipelines as Arrow extension-type columns — thank you, that works. But registering a model whose fields use typing.Literal[...] fails:

ValueError: Unsupported annotation: typing.Literal['a', 'b']

Literal is one of the most common pydantic config patterns (enumerated string/int choices), so this blocks real configs. Our SpikeSortingConfig has 11 Literal fields (e.g. method_chosen: Literal["dredge_ap", "iterative_template", "medicine"]), and it can't be used as a broadcast config column until Literal is handled.

Minimal reproduction

import pydantic
from typing import Literal
from orcapod.contexts import get_default_context


class LiteralModel(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid", frozen=True)
    method: Literal["a", "b"]          # <-- the only thing that matters


# Fails. (A model with plain `method: str` registers fine; so do int/float/bool,
# list[int], list[str]. Only Literal is rejected.)
get_default_context().type_converter.register_python_class(LiteralModel)

The same failure occurs via the intended path — building a source with such a model as a column:

import orcapod as op
from orcapod.types import Schema
op.sources.DictSource(
    [{"config": LiteralModel(method="a")}],
    tag_columns=[],
    data_schema=Schema({"config": LiteralModel}),
)

Error / traceback (key frames)

File ".../orcapod/extension_types/pydantic_logical_type_factory.py", line 271, in create_for_python_type
    arrow_type = converter.register_python_class(annotation)
File ".../orcapod/semantic_types/universal_converter.py", line 266, in register_python_class
    return self._register_python_class_impl(annotation, in_progress)
File ".../orcapod/semantic_types/universal_converter.py", line 403, in _register_python_class_impl
    raise ValueError(f"Unsupported annotation: {annotation!r}")
ValueError: Unsupported annotation: typing.Literal['a', 'b']

Root cause

UniversalTypeConverter._register_python_class_impl (orcapod/semantic_types/universal_converter.py) walks a pydantic/dataclass model's field annotations to build the Arrow struct. It has branches for the scalar type_map, Optional/Union, list[T], set[T], dict[K,V], and the factory-backed (dataclass/BaseModel) cases — but no branch for typing.Literal — so a Literal[...] field falls through to the terminal raise ValueError("Unsupported annotation: ...").

Suggested fix

Add a Literal branch that maps to the Arrow type of the literal values' type. Literal values are restricted to hashable constants (str/int/bool/bytes/enum/None), so:

import typing
...
origin = get_origin(annotation)
args = get_args(annotation)
...
if origin is typing.Literal:
    value_types = {type(a) for a in args if a is not None}
    if len(value_types) == 1:
        return self.register_python_class(next(iter(value_types)))   # e.g. Literal[str,...] -> large_string
    raise ValueError(
        f"Mixed-type Literal is not supported: {annotation!r}. "
        f"All members must share one type (e.g. Literal['a','b'])."
    )

(Literal[str, ...] → large_string, Literal[int, ...] → int64, etc.) On the read path the value is just the underlying scalar, so no special deserialization is needed. Optionally preserve the allowed-value set as metadata, but mapping to the value type is enough to unblock.

Context

This is a follow-on to #184 / PR #185 (ITL-432). The generic pydantic-column support works; this is specifically the Literal field-type gap, which is common enough in real pydantic configs that it's worth supporting. Until then we stay pinned to the older arnoldb/pydantic commit c23fe264 (its struct-based PydanticModelConverter handled arbitrary BaseModel subclasses generically).