fix(python)!: reject mixed-value enums in schemas

A mixed-value Pydantic Enum previously produced a UTF-8 Arrow schema even when some member values were non-string, causing writes to fail after schema construction. `pydantic_to_schema` now rejects actual member values with more than one type using `TypeError`, including direct, optional, list, and nested fields. Homogeneous string and integer enums and empty enums retain their mappings, and existing plain-Enum writes remain supported. This intentionally changes mixed-value enums from failing at write time to failing at schema construction.

Validation on the exact head: the Python suite passed with 3034 passed, 126 skipped, and 70 deselected, including 12 mixed-value container cases, the empty-enum case, homogeneous schema coverage, and real plain-Enum writes. All applicable exact-head checks and workflows completed successfully.
This commit is contained in:
Minh Vu
2026-11-02 20:04:45 +08:00
committed by GitHub
parent b5d3561066
commit 0b20965d09
2 changed files with 37 additions and 2 deletions
+7 -2
View File
@@ -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
--------
+30
View File
@@ -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."""