From 712a8b86995397de03ae1c0389f68824505f80a5 Mon Sep 17 00:00:00 2001
From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com>
Date: Thu, 22 Oct 2026 19:42:36 +0100
Subject: [PATCH] fix(python): write plain Enum fields of a LanceModel as their
 values (#4357)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Since #3232, `pydantic_to_schema` maps an `Enum` field to the Arrow type
of its values, including enums without a `str`/`int` mixin. Writing such
models still failed: `model_to_dict()` returns `model.model_dump()`,
which keeps the Enum members, and Arrow can only convert members that
are also `str`/`int` instances.

```python
class Color(Enum):
    RED = "red"

class Item(LanceModel):
    color: Color

tbl = db.create_table("items", schema=Item)
tbl.add([Item(color=Color.RED)])
# pyarrow.lib.ArrowTypeError: Expected bytes, got a 'Color' object
```

`create_table(data=[...])` fails the same way, and int-valued enums
raise `ArrowInvalid`.

`model_to_dict()` now replaces Enum members with their values, including
in nested models and lists. Lists of plain scalars (vectors in
particular) are returned without being walked, so the per-row cost stays
about the same as `model_dump()` alone (21.3µs vs 20.3µs for a 1536-dim
model). Reading back is unchanged: Pydantic already validates the stored
value into the Enum member.

Added `test_plain_enum_values_are_written`, which covers str and int
valued enums, `Optional` and `List` fields, and an Enum inside a nested
model, and checks that `to_pydantic` round-trips the rows. It fails with
the `ArrowTypeError` above before the change. `test_pydantic.py`,
`test_embeddings.py`, `test_table.py` and `test_db.py` pass, and `ruff
format --check` / `ruff check` are clean.
---
 python/python/lancedb/pydantic.py    | 18 +++++++++-
 python/python/tests/test_pydantic.py | 49 ++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

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