mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
fix(python): let to_pydantic read tables written by an aliased LanceModel
`LanceModel` instances with validation aliases write field-name keys to Arrow, while `to_pydantic` previously reconstructed models by alias. This normalizes field-name rows to the applicable single-component validation alias for both sync and async reads, while preserving existing alias keys, alias precedence, name-only validation, and collision behavior. Adds sync and async regression coverage for string aliases, `AliasChoices`, single-component `AliasPath`, `populate_by_name`, and explicit `validate_by_name` precedence.
This commit is contained in:
Hien Hoang
@@ -81,6 +81,85 @@ if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
T = TypeVar("T", bound="LanceModel")
|
||||
|
||||
|
||||
def _top_level_alias_key(alias: Any) -> str | None:
|
||||
if isinstance(alias, str):
|
||||
return alias
|
||||
if isinstance(alias, pydantic.AliasPath) and alias.path:
|
||||
key = alias.path[0]
|
||||
return key if isinstance(key, str) else None
|
||||
return None
|
||||
|
||||
|
||||
def _top_level_alias_keys(alias: Any) -> set[str]:
|
||||
options = alias.choices if isinstance(alias, pydantic.AliasChoices) else (alias,)
|
||||
return {
|
||||
key for option in options if (key := _top_level_alias_key(option)) is not None
|
||||
}
|
||||
|
||||
|
||||
def _single_component_alias(alias: Any) -> str | None:
|
||||
# AliasChoices validates against the first matching source, so only use
|
||||
# its first option as a field-name fallback.
|
||||
option = alias.choices[0] if isinstance(alias, pydantic.AliasChoices) else alias
|
||||
if isinstance(option, pydantic.AliasPath) and len(option.path) != 1:
|
||||
# A nested path needs a structured payload; do not infer one here.
|
||||
return None
|
||||
return _top_level_alias_key(option)
|
||||
|
||||
|
||||
def _rows_to_models(model: type[T], rows: List[dict]) -> List[T]:
|
||||
"""Build ``model`` instances from result rows.
|
||||
|
||||
A table created from a LanceModel stores each field under its name, while
|
||||
pydantic validates an aliased field by its alias. Pass a row's value under
|
||||
the alias when the row has the field name but none of the aliases, so a
|
||||
model with aliases can read the tables it wrote; rows keyed by any alias
|
||||
still work.
|
||||
"""
|
||||
if model.model_config.get("validate_by_alias", True) is False:
|
||||
# Name-only validation: pydantic reads field names and ignores aliases.
|
||||
return [model(**row) for row in rows]
|
||||
aliases = {
|
||||
name: field.validation_alias or field.alias
|
||||
for name, field in model.model_fields.items()
|
||||
}
|
||||
# populate_by_name is the pre-2.11 spelling; an explicit new flag wins.
|
||||
accepts_by_name = model.model_config.get(
|
||||
"validate_by_name",
|
||||
model.model_config.get("populate_by_name", False),
|
||||
)
|
||||
# Every top-level key some field reads, including the first step of an
|
||||
# AliasPath and each choice of an AliasChoices.
|
||||
alias_keys = {}
|
||||
for name, alias in aliases.items():
|
||||
keys = _top_level_alias_keys(alias)
|
||||
if not keys or accepts_by_name:
|
||||
keys.add(name)
|
||||
alias_keys[name] = keys
|
||||
readers = {}
|
||||
for name, keys in alias_keys.items():
|
||||
for key in keys:
|
||||
readers.setdefault(key, set()).add(name)
|
||||
renames = {
|
||||
name: target
|
||||
for name, alias in aliases.items()
|
||||
if (target := _single_component_alias(alias)) is not None
|
||||
# Do not move a source that another field reads.
|
||||
and target != name
|
||||
and not (readers.get(name, set()) - {name})
|
||||
# Do not add a source that another field accepts. Introducing a shared
|
||||
# alias key could change that field's Pydantic selection.
|
||||
and not (readers.get(target, set()) - {name})
|
||||
}
|
||||
for row in rows:
|
||||
for name, target in renames.items():
|
||||
if name in row and not any(key in row for key in alias_keys[name]):
|
||||
row[target] = row.pop(name)
|
||||
return [model(**row) for row in rows]
|
||||
|
||||
|
||||
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
|
||||
_SYNC_VECTOR_QUERY_TYPES = (list, tuple, np.ndarray, pa.Array, pa.ChunkedArray)
|
||||
|
||||
@@ -1340,7 +1419,7 @@ class LanceQueryBuilder(ABC):
|
||||
-------
|
||||
List[LanceModel]
|
||||
"""
|
||||
return [model(**row) for row in self.to_arrow(timeout=timeout).to_pylist()]
|
||||
return _rows_to_models(model, self.to_arrow(timeout=timeout).to_pylist())
|
||||
|
||||
def to_polars(self, *, timeout: Optional[timedelta] = None) -> "pl.DataFrame":
|
||||
"""
|
||||
@@ -3770,9 +3849,9 @@ class AsyncQueryBase(object):
|
||||
-------
|
||||
list[LanceModel]
|
||||
"""
|
||||
return [
|
||||
model(**row) for row in (await self.to_arrow(timeout=timeout)).to_pylist()
|
||||
]
|
||||
return _rows_to_models(
|
||||
model, (await self.to_arrow(timeout=timeout)).to_pylist()
|
||||
)
|
||||
|
||||
async def explain_plan(self, verbose: Optional[bool] = False):
|
||||
"""Return the execution plan for this query.
|
||||
|
||||
@@ -16,6 +16,7 @@ from lancedb.pydantic import (
|
||||
Vector,
|
||||
pydantic_to_schema,
|
||||
)
|
||||
from packaging.version import Version
|
||||
from pydantic import BaseModel
|
||||
from pydantic import Field
|
||||
|
||||
@@ -791,6 +792,305 @@ async def test_aliases_in_lance_model_async(mem_db_async):
|
||||
assert model.distance < 0.01
|
||||
|
||||
|
||||
def test_aliased_lance_model_reads_back_its_own_table(mem_db):
|
||||
# The schema and add() use field names, so to_pydantic() must accept rows
|
||||
# keyed by field name even though pydantic validates by alias.
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(alias="userId")
|
||||
vector: Vector(2)
|
||||
|
||||
tbl = mem_db.create_table("aliased_items", schema=Item)
|
||||
tbl.add([Item(userId=1, vector=[0.1, 0.2])])
|
||||
assert tbl.schema.names == ["user_id", "vector"]
|
||||
|
||||
[item] = tbl.search([0.1, 0.2]).limit(1).to_pydantic(Item)
|
||||
assert item.user_id == 1
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"validation_alias",
|
||||
[
|
||||
pydantic.AliasChoices("userId"),
|
||||
pydantic.AliasPath("userId"),
|
||||
],
|
||||
ids=["alias-choices", "alias-path"],
|
||||
)
|
||||
def test_aliased_lance_model_with_complex_validation_alias_reads_back_its_own_table(
|
||||
mem_db, validation_alias
|
||||
):
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(validation_alias=validation_alias)
|
||||
vector: Vector(2)
|
||||
|
||||
tbl = mem_db.create_table("aliased_items", schema=Item)
|
||||
tbl.add([Item(userId=1, vector=[0.1, 0.2])])
|
||||
assert tbl.schema.names == ["user_id", "vector"]
|
||||
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 1
|
||||
|
||||
|
||||
def test_lower_priority_alias_takes_precedence_over_field_name(mem_db):
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(
|
||||
validation_alias=pydantic.AliasChoices("primaryId", "secondaryId")
|
||||
)
|
||||
|
||||
row = {"user_id": 1, "secondaryId": 2}
|
||||
tbl = mem_db.create_table("aliased_items", data=[row])
|
||||
|
||||
[arrow_row] = tbl.search().to_arrow().to_pylist()
|
||||
assert Item(**arrow_row).user_id == 2
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 2
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
Version(pydantic.VERSION) < Version("2.11"),
|
||||
reason="validate_by_alias/validate_by_name require pydantic>=2.11",
|
||||
)
|
||||
def test_aliased_lance_model_with_name_only_validation(mem_db):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(
|
||||
validate_by_alias=False, validate_by_name=True
|
||||
)
|
||||
user_id: int = Field(default=0, alias="userId")
|
||||
|
||||
tbl = mem_db.create_table("name_only_items", schema=Item)
|
||||
item = Item(user_id=42)
|
||||
assert item.user_id == 42
|
||||
tbl.add([item])
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 42
|
||||
|
||||
|
||||
def test_aliased_lance_model_with_populate_by_name(mem_db):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(populate_by_name=True)
|
||||
user_id: int = Field(default=0, alias="userId")
|
||||
|
||||
tbl = mem_db.create_table("populate_by_name_items", schema=Item)
|
||||
item = Item(user_id=42)
|
||||
assert item.user_id == 42
|
||||
tbl.add([item])
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 42
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
Version(pydantic.VERSION) < Version("2.11"),
|
||||
reason="validate_by_name requires pydantic>=2.11",
|
||||
)
|
||||
def test_explicit_validate_by_name_false_overrides_populate_by_name(mem_db):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(
|
||||
populate_by_name=True,
|
||||
validate_by_name=False,
|
||||
)
|
||||
user_id: int = Field(default=0, alias="userId")
|
||||
|
||||
tbl = mem_db.create_table("validate_by_name_precedence_items", schema=Item)
|
||||
item = Item(userId=42)
|
||||
assert item.user_id == 42
|
||||
tbl.add([item])
|
||||
assert tbl.schema.names == ["user_id"]
|
||||
assert tbl.search().to_arrow().to_pylist() == [{"user_id": 42}]
|
||||
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.user_id == 42
|
||||
|
||||
|
||||
def test_field_name_fallback_preserves_populate_by_name_reader(mem_db):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(populate_by_name=True)
|
||||
value: int = Field(default=0, validation_alias="id")
|
||||
id: int = Field(default=-1, alias="recordId")
|
||||
|
||||
row = {"value": 1}
|
||||
assert Item(**row).model_dump() == {"value": 1, "id": -1}
|
||||
tbl = mem_db.create_table("populate_by_name_readers", data=[row])
|
||||
|
||||
[arrow_row] = tbl.search().to_arrow().to_pylist()
|
||||
assert arrow_row == row
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.model_dump() == {"value": 1, "id": -1}
|
||||
|
||||
|
||||
def test_field_name_that_is_another_fields_alias_keeps_its_meaning(mem_db):
|
||||
class Item(LanceModel):
|
||||
record_id: int = Field(default=0, alias="id")
|
||||
id: int = Field(default=-1, alias="legacy_id")
|
||||
|
||||
tbl = mem_db.create_table("alias_collision_items", data=[{"id": 42}])
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.model_dump() == {"record_id": 42, "id": -1}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"validation_alias",
|
||||
[
|
||||
pydantic.AliasChoices("id", "recordId"),
|
||||
pydantic.AliasPath("id"),
|
||||
pydantic.AliasChoices(pydantic.AliasPath("id"), "recordId"),
|
||||
],
|
||||
)
|
||||
def test_field_name_read_by_another_fields_alias_choices_or_path(
|
||||
mem_db, validation_alias
|
||||
):
|
||||
class Item(LanceModel):
|
||||
record_id: int = Field(default=0, validation_alias=validation_alias)
|
||||
id: int = Field(default=-1, alias="legacy_id")
|
||||
|
||||
tbl = mem_db.create_table("alias_choices_items", data=[{"id": 42}])
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.model_dump() == {"record_id": 42, "id": -1}
|
||||
|
||||
|
||||
def test_field_name_fallback_does_not_change_another_fields_alias_choice(mem_db):
|
||||
class Item(LanceModel):
|
||||
value: int = Field(
|
||||
default=0, validation_alias=pydantic.AliasChoices("id", "valueId")
|
||||
)
|
||||
record_id: int = Field(
|
||||
default=-1, validation_alias=pydantic.AliasChoices("id", "recordId")
|
||||
)
|
||||
|
||||
row = {"value": 1, "recordId": 42}
|
||||
assert Item(**row).model_dump() == {"value": 0, "record_id": 42}
|
||||
tbl = mem_db.create_table("shared_alias_items", data=[row])
|
||||
|
||||
[arrow_row] = tbl.search().to_arrow().to_pylist()
|
||||
assert arrow_row == row
|
||||
[item] = tbl.search().to_pydantic(Item)
|
||||
assert item.model_dump() == {"value": 0, "record_id": 42}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aliased_lance_model_reads_back_its_own_table_async(mem_db_async):
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(alias="userId")
|
||||
vector: Vector(2)
|
||||
|
||||
tbl = await mem_db_async.create_table("aliased_items", schema=Item)
|
||||
await tbl.add([Item(userId=1, vector=[0.1, 0.2])])
|
||||
|
||||
[item] = await tbl.vector_search([0.1, 0.2]).limit(1).to_pydantic(Item)
|
||||
assert item.user_id == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"validation_alias",
|
||||
[
|
||||
pydantic.AliasChoices("userId"),
|
||||
pydantic.AliasPath("userId"),
|
||||
],
|
||||
ids=["alias-choices", "alias-path"],
|
||||
)
|
||||
async def test_aliased_lance_model_with_complex_validation_alias_reads_back_async(
|
||||
mem_db_async, validation_alias
|
||||
):
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(validation_alias=validation_alias)
|
||||
vector: Vector(2)
|
||||
|
||||
tbl = await mem_db_async.create_table("aliased_items", schema=Item)
|
||||
await tbl.add([Item(userId=1, vector=[0.1, 0.2])])
|
||||
assert (await tbl.schema()).names == ["user_id", "vector"]
|
||||
|
||||
[item] = await tbl.vector_search([0.1, 0.2]).limit(1).to_pydantic(Item)
|
||||
assert item.user_id == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lower_priority_alias_takes_precedence_over_field_name_async(
|
||||
mem_db_async,
|
||||
):
|
||||
class Item(LanceModel):
|
||||
user_id: int = Field(
|
||||
validation_alias=pydantic.AliasChoices("primaryId", "secondaryId")
|
||||
)
|
||||
|
||||
row = {"user_id": 1, "secondaryId": 2}
|
||||
tbl = await mem_db_async.create_table("aliased_items", data=[row])
|
||||
|
||||
[arrow_row] = (await (await tbl.search()).to_arrow()).to_pylist()
|
||||
assert Item(**arrow_row).user_id == 2
|
||||
[item] = await (await tbl.search()).to_pydantic(Item)
|
||||
assert item.user_id == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_field_name_fallback_does_not_change_another_fields_alias_choice_async(
|
||||
mem_db_async,
|
||||
):
|
||||
class Item(LanceModel):
|
||||
value: int = Field(
|
||||
default=0, validation_alias=pydantic.AliasChoices("id", "valueId")
|
||||
)
|
||||
record_id: int = Field(
|
||||
default=-1, validation_alias=pydantic.AliasChoices("id", "recordId")
|
||||
)
|
||||
|
||||
row = {"value": 1, "recordId": 42}
|
||||
assert Item(**row).model_dump() == {"value": 0, "record_id": 42}
|
||||
tbl = await mem_db_async.create_table("shared_alias_items", data=[row])
|
||||
|
||||
[arrow_row] = (await (await tbl.search()).to_arrow()).to_pylist()
|
||||
assert arrow_row == row
|
||||
[item] = await (await tbl.search()).to_pydantic(Item)
|
||||
assert item.model_dump() == {"value": 0, "record_id": 42}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_field_name_fallback_preserves_populate_by_name_reader_async(
|
||||
mem_db_async,
|
||||
):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(populate_by_name=True)
|
||||
value: int = Field(default=0, validation_alias="id")
|
||||
id: int = Field(default=-1, alias="recordId")
|
||||
|
||||
row = {"value": 1}
|
||||
assert Item(**row).model_dump() == {"value": 1, "id": -1}
|
||||
tbl = await mem_db_async.create_table("populate_by_name_readers", data=[row])
|
||||
|
||||
[arrow_row] = (await (await tbl.search()).to_arrow()).to_pylist()
|
||||
assert arrow_row == row
|
||||
[item] = await (await tbl.search()).to_pydantic(Item)
|
||||
assert item.model_dump() == {"value": 1, "id": -1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
Version(pydantic.VERSION) < Version("2.11"),
|
||||
reason="validate_by_name requires pydantic>=2.11",
|
||||
)
|
||||
async def test_explicit_validate_by_name_false_overrides_populate_by_name_async(
|
||||
mem_db_async,
|
||||
):
|
||||
class Item(LanceModel):
|
||||
model_config = pydantic.ConfigDict(
|
||||
populate_by_name=True,
|
||||
validate_by_name=False,
|
||||
)
|
||||
user_id: int = Field(default=0, alias="userId")
|
||||
|
||||
tbl = await mem_db_async.create_table(
|
||||
"validate_by_name_precedence_items", schema=Item
|
||||
)
|
||||
item = Item(userId=42)
|
||||
assert item.user_id == 42
|
||||
await tbl.add([item])
|
||||
assert (await tbl.schema()).names == ["user_id"]
|
||||
assert (await (await tbl.search()).to_arrow()).to_pylist() == [{"user_id": 42}]
|
||||
|
||||
[item] = await (await tbl.search()).to_pydantic(Item)
|
||||
assert item.user_id == 42
|
||||
|
||||
|
||||
def test_enum_types():
|
||||
"""Enum fields should map to the Arrow type of their value (issue #1846)."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user