mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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:
@@ -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
|
||||
--------
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user