diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index cc3070b72..4d77ecd15 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -9,16 +9,19 @@ import inspect import sys import types from abc import ABC, abstractmethod -from datetime import date, datetime +from datetime import date, datetime, time, timedelta from enum import Enum from typing import ( TYPE_CHECKING, Any, List, + Literal, Type, Union, _GenericAlias, GenericAlias, + get_args, + get_origin, ) import pyarrow as pa @@ -247,6 +250,19 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType: elif py_type is datetime: tz = get_extras(field, "tz") return pa.timestamp("us", tz=tz) + elif py_type is time: + return pa.time64("us") + elif py_type is timedelta: + return pa.duration("us") + elif get_origin(py_type) is Literal: + # Store a Literal as the type of its values, e.g. Literal["a", "b"] as utf8. + value_types = {type(value) for value in get_args(py_type)} + if len(value_types) != 1: + raise TypeError( + "Converting Pydantic type to Arrow Type: unsupported type " + f"{py_type}. All values of a Literal must have the same type." + ) + return _pydantic_type_to_arrow_type(value_types.pop(), field) elif getattr(py_type, "__origin__", None) in (list, tuple): # A bare, unparameterised ``typing.List`` / ``typing.Tuple`` matches this # branch (its ``__origin__`` is ``list`` / ``tuple``) but has no @@ -561,21 +577,36 @@ def get_extras(field_info: FieldInfo, key: str) -> Any: def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]: - """Convert a Pydantic model to a dictionary.""" - return _enum_members_to_values(model.model_dump()) + """Convert a Pydantic model to a dictionary. + + Raises + ------ + ValueError + If a ``datetime.time`` value has a non-null ``tzinfo``. LanceDB stores + times as timezone-naive wall-clock values and cannot preserve the offset. + """ + return _model_value_to_arrow_compatible(model.model_dump()) -def _enum_members_to_values(value: Any) -> Any: +def _model_value_to_arrow_compatible(value: Any) -> Any: # model_dump() keeps Enum members, and Arrow cannot convert a member of a # plain Enum (one without a str or int mixin). Write its value instead, # which is what the Arrow type of the field is derived from. if isinstance(value, Enum): - return value.value + value = value.value + if isinstance(value, time): + if value.tzinfo is not None: + raise ValueError( + "LanceDB supports timezone-naive wall-clock datetime.time values " + "only; time values with tzinfo cannot be written losslessly. Remove " + "the timezone offset or store it in a separate field." + ) + return value if isinstance(value, dict): - return {k: _enum_members_to_values(v) for k, v in value.items()} + return {k: _model_value_to_arrow_compatible(v) for k, v in value.items()} if isinstance(value, (list, tuple)) and value: # Skip lists of plain scalars, such as vectors, without walking them. if isinstance(value[0], (int, float, str, bytes)): return value - return [_enum_members_to_values(v) for v in value] + return [_model_value_to_arrow_compatible(v) for v in value] return value diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index eb6962131..93c42e6b7 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import json -from datetime import date, datetime +from datetime import date, datetime, time, timedelta, timezone from enum import Enum -from typing import List, Optional, Tuple +from typing import List, Literal, Optional, Tuple import pyarrow as pa import pydantic @@ -1194,3 +1194,111 @@ def test_plain_enum_values_are_written(mem_db): {"color": "blue", "priority": 1, "opt_color": None}, ] assert tbl.search().to_pydantic(TestModel) == rows + + +def test_literal_time_and_timedelta_types(): + class TestModel(pydantic.BaseModel): + status: Literal["todo", "done"] + priority: Literal[1, 2, 3] + flag: Literal[True] + tags: List[Literal["a", "b"]] + opt_status: Optional[Literal["todo", "done"]] = None + starts: time + starts_list: List[time] + opt_starts: Optional[time] = None + duration: timedelta + durations: List[timedelta] + + schema = pydantic_to_schema(TestModel) + + assert schema.field("status").type == pa.utf8() + assert schema.field("priority").type == pa.int64() + assert schema.field("flag").type == pa.bool_() + assert schema.field("tags").type == pa.list_(pa.utf8()) + assert schema.field("opt_status").type == pa.utf8() + assert schema.field("opt_status").nullable + assert schema.field("starts").type == pa.time64("us") + assert schema.field("starts_list").type == pa.list_(pa.time64("us")) + assert schema.field("opt_starts").type == pa.time64("us") + assert schema.field("opt_starts").nullable + assert schema.field("duration").type == pa.duration("us") + assert schema.field("durations").type == pa.list_(pa.duration("us")) + + +@pytest.mark.parametrize( + "annotation", + [ + Literal["a", 1], + Optional[Literal["a", 1]], + List[Literal["a", 1]], + ], + ids=["field", "optional", "list"], +) +def test_literal_with_mixed_value_types_raises_type_error(annotation): + TestModel = pydantic.create_model("TestModel", value=(annotation, ...)) + + with pytest.raises(TypeError, match="same type"): + pydantic_to_schema(TestModel) + + +def test_literal_time_and_timedelta_round_trip(mem_db): + class TestModel(LanceModel): + status: Literal["todo", "done"] + starts: time + duration: timedelta + vector: Vector(2) + + rows = [ + TestModel( + status="done", + starts=time(9, 30), + duration=timedelta(minutes=45), + vector=[1.0, 2.0], + ), + TestModel( + status="todo", + starts=time(17, 0, 0, 250), + duration=timedelta(days=1, microseconds=5), + vector=[3.0, 4.0], + ), + ] + tbl = mem_db.create_table("test", schema=TestModel) + tbl.add(rows) + + assert tbl.search().to_pydantic(TestModel) == rows + + +@pytest.mark.parametrize("container", ["field", "list", "nested", "enum"]) +def test_tz_aware_time_is_rejected_before_write(mem_db, container): + class Nested(LanceModel): + starts: time + + aware = time(9, 30, tzinfo=timezone(timedelta(hours=5, minutes=30))) + AwareTime = Enum("AwareTime", {"MORNING": aware}) + Model = pydantic.create_model( + "TestModel", + __base__=LanceModel, + **{ + container: { + "field": (time, ...), + "list": (List[time], ...), + "nested": (Nested, ...), + "enum": (AwareTime, ...), + }[container] + }, + ) + row = Model( + **{ + container: { + "field": aware, + "list": [aware], + "nested": Nested(starts=aware), + "enum": AwareTime.MORNING, + }[container] + } + ) + + table_name = f"timezone_aware_{container}" + with pytest.raises(ValueError, match="timezone-naive wall-clock.*tzinfo"): + mem_db.create_table(table_name, data=[row]) + assert table_name not in mem_db.table_names()