diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index d2b5cb22d..cc3070b72 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -304,8 +304,12 @@ def _pydantic_type_to_arrow_type(tp: Any, field: FieldInfo) -> pa.DataType: # Map Enum to the Arrow type of its value. # For string-valued enums, use dictionary encoding for efficiency. # For integer enums, use the native type. - # Fall back to utf8 for mixed-type or empty enums. value_types = {type(m.value) for m in tp} + if len(value_types) > 1: + raise TypeError( + f"Converting Pydantic Enum {tp.__name__} to Arrow Type: " + "all enum values must have the same type." + ) if len(value_types) == 1: value_type = value_types.pop() if value_type is str: @@ -435,7 +439,8 @@ def pydantic_to_schema(model: Type[pydantic.BaseModel]) -> pa.Schema: 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. + ``LanceModel`` (an Arrow struct) instead. Enum values must also all + have the same type. Examples -------- diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 2a874b68e..eb6962131 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -1117,6 +1117,36 @@ def test_enum_types(): assert schema.field("opt_status").nullable +@pytest.mark.parametrize("container", ["field", "optional", "list", "nested"]) +@pytest.mark.parametrize("values", [("text", 7), (1, 2.5), (False, "no")]) +def test_mixed_enum_values_are_rejected(container, values): + Mixed = Enum("Mixed", {"FIRST": values[0], "SECOND": values[1]}) + + class Nested(pydantic.BaseModel): + value: Mixed + + annotation = { + "field": Mixed, + "optional": Optional[Mixed], + "list": List[Mixed], + "nested": Nested, + }[container] + model = pydantic.create_model("Model", value=(annotation, ...)) + + with pytest.raises(TypeError, match="Enum.*Mixed.*same type"): + pydantic_to_schema(model) + + +def test_empty_enum_schema_is_unchanged(): + class Empty(Enum): + pass + + class Model(pydantic.BaseModel): + value: Empty + + assert pydantic_to_schema(Model).field("value").type == pa.utf8() + + def test_plain_enum_values_are_written(mem_db): """Members of an Enum without a str/int mixin are written as their values."""