mirror of
https://github.com/lancedb/lancedb.git
synced 2026-04-07 00:10:40 +00:00
feat(python): support Enum types in Pydantic to Arrow schema conversion (#3232)
## Summary Fixes #1846. Python `Enum` fields raised `TypeError: Converting Pydantic type to Arrow Type: unsupported type <enum 'SomethingTypes'>` when converting a Pydantic model to an Arrow schema. The fix adds Enum detection in `_pydantic_type_to_arrow_type`. When an Enum subclass is encountered, the value type of its members is inspected and mapped to the appropriate Arrow type: - `str`-valued enums (e.g. `class Status(str, Enum)`) → `pa.utf8()` - `int`-valued enums (e.g. `class Priority(int, Enum)`) → `pa.int64()` - Other homogeneous value types → the Arrow type for that Python type - Mixed-value or empty enums → `pa.utf8()` (safe fallback) This covers the common `(str, Enum)` and `(int, Enum)` mixin patterns used in practice. ## Changes - `python/python/lancedb/pydantic.py`: add Enum branch in `_pydantic_type_to_arrow_type` - `python/python/tests/test_pydantic.py`: add `test_enum_types` covering `str`, `int`, and `Optional` Enum fields ## Note on #2395 PR #2395 handles `StrEnum` (Python 3.11+) specifically, using a dictionary-encoded type. This PR handles the broader `(str, Enum)` / `(int, Enum)` mixin pattern that works across all Python versions and stores values as their natural Arrow type. AI assistance was used in developing this fix.
This commit is contained in:
@@ -3,6 +3,7 @@
|
||||
|
||||
import json
|
||||
from datetime import date, datetime
|
||||
from enum import Enum
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -673,3 +674,29 @@ async def test_aliases_in_lance_model_async(mem_db_async):
|
||||
assert hasattr(model, "name")
|
||||
assert hasattr(model, "distance")
|
||||
assert model.distance < 0.01
|
||||
|
||||
|
||||
def test_enum_types():
|
||||
"""Enum fields should map to the Arrow type of their value (issue #1846)."""
|
||||
|
||||
class StrStatus(str, Enum):
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
DONE = "done"
|
||||
|
||||
class IntPriority(int, Enum):
|
||||
LOW = 1
|
||||
MEDIUM = 2
|
||||
HIGH = 3
|
||||
|
||||
class TestModel(pydantic.BaseModel):
|
||||
status: StrStatus
|
||||
priority: IntPriority
|
||||
opt_status: Optional[StrStatus] = None
|
||||
|
||||
schema = pydantic_to_schema(TestModel)
|
||||
|
||||
assert schema.field("status").type == pa.dictionary(pa.int32(), pa.utf8())
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user