Files
lancedb/python/python
Mohammad Hijjawi 712a8b8699 fix(python): write plain Enum fields of a LanceModel as their values (#4357)
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.
2026-10-23 02:42:36 +08:00
..