From fdf886d42abd1cf8a0a209f2131039581795d59a Mon Sep 17 00:00:00 2001 From: Hien Hoang Date: Sun, 1 Nov 2026 09:27:13 +0700 Subject: [PATCH] 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. --- python/python/lancedb/query.py | 87 +++++++- python/python/tests/test_pydantic.py | 300 +++++++++++++++++++++++++++ 2 files changed, 383 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index c9ca258e7..a572e7609 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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. diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 09af2f081..2a874b68e 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -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)."""