diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index db8723a99..d2b5cb22d 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -259,6 +259,15 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType: f"{py_type}. Specify the element type, e.g. List[int] instead " "of a bare List." ) + # An Arrow list has a single element type, so a tuple such as + # Tuple[int, str] cannot be stored as one. Tuple[int, ...] is homogeneous. + element_types = {arg for arg in args if arg is not Ellipsis} + if len(element_types) > 1: + raise TypeError( + "Converting Pydantic type to Arrow Type: unsupported type " + f"{py_type}. All tuple elements must have the same type; use a " + "nested LanceModel for a fixed set of fields with different types." + ) child = args[0] return _pydantic_list_child_to_arrow(child, field) raise TypeError( @@ -419,6 +428,15 @@ def pydantic_to_schema(model: Type[pydantic.BaseModel]) -> pa.Schema: pyarrow.Schema The Arrow Schema + Raises + ------ + TypeError + If a field has a type with no Arrow equivalent. A ``Tuple`` field becomes + an Arrow list, so all its elements must have the same type, as in + ``Tuple[int, int]`` or ``Tuple[float, ...]``; for a fixed set of fields + with different types, such as ``Tuple[int, str]``, use a nested + ``LanceModel`` (an Arrow struct) instead. + Examples -------- diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 6717bbc32..09af2f081 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -200,6 +200,27 @@ def test_bare_generic_raises_type_error(): pydantic_to_schema(TestModel) +def test_tuple_with_mixed_element_types_raises_type_error(): + # An Arrow list has one element type, so Tuple[int, str] has no Arrow + # equivalent. It used to map to list and fail later, on write. + class TestModel(pydantic.BaseModel): + pair: Tuple[int, str] + + with pytest.raises(TypeError, match="same type"): + pydantic_to_schema(TestModel) + + +def test_homogeneous_tuples_map_to_lists(): + class TestModel(pydantic.BaseModel): + pair: Tuple[int, int] + many: Tuple[float, ...] + + schema = pydantic_to_schema(TestModel) + + assert schema.field("pair").type == pa.list_(pa.int64()) + assert schema.field("many").type == pa.list_(pa.float64()) + + def test_nested_struct_list_optional(): class SplitInfo(pydantic.BaseModel): start_frame: int