diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 313c514c3..8fe3c1c93 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -524,4 +524,20 @@ 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 model.model_dump() + return _enum_members_to_values(model.model_dump()) + + +def _enum_members_to_values(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 + if isinstance(value, dict): + return {k: _enum_members_to_values(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 value diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 81674d735..24fb7ee89 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -733,3 +733,52 @@ def test_enum_types(): assert schema.field("priority").type == pa.int64() assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8()) assert schema.field("opt_status").nullable + + +def test_plain_enum_values_are_written(mem_db): + """Members of an Enum without a str/int mixin are written as their values.""" + + class Color(Enum): + RED = "red" + BLUE = "blue" + + class Priority(Enum): + LOW = 1 + HIGH = 2 + + class Tag(LanceModel): + color: Color + + class TestModel(LanceModel): + color: Color + priority: Priority + opt_color: Optional[Color] = None + colors: List[Color] + tag: Tag + vector: Vector(2) + + rows = [ + TestModel( + color=Color.RED, + priority=Priority.HIGH, + opt_color=Color.BLUE, + colors=[Color.BLUE, Color.RED], + tag=Tag(color=Color.BLUE), + vector=[1.0, 2.0], + ), + TestModel( + color=Color.BLUE, + priority=Priority.LOW, + colors=[], + tag=Tag(color=Color.RED), + vector=[3.0, 4.0], + ), + ] + tbl = mem_db.create_table("test", schema=TestModel) + tbl.add(rows) + + assert tbl.to_arrow().select(["color", "priority", "opt_color"]).to_pylist() == [ + {"color": "red", "priority": 2, "opt_color": "blue"}, + {"color": "blue", "priority": 1, "opt_color": None}, + ] + assert tbl.search().to_pydantic(TestModel) == rows