mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
fix(python): omit excluded model fields from persistence (#4465)
LanceDB serializes Pydantic rows with `model_dump()`, which omits `Field(exclude=True)`, but currently constructs an Arrow schema containing those same fields. Creating a local table therefore fails when an excluded field is non-nullable, including defaulted transient fields and nested models. Filter unconditionally excluded fields when converting model schemas, including nested structs. Skip excluded source/vector fields when deriving embedding configuration so persistent metadata does not refer to absent columns. Keep `field_names()` as model-field reflection. Tests exercise implicit and explicit schemas, synchronous and asynchronous create/add/reopen/readback, nested fields, defaults and default factories, retained fields, and local deterministic embeddings. Excluded values are not persisted; `to_pydantic` restores declared defaults. Conditional exclusion and readback of required excluded fields without defaults are outside this change. Implemented with assistance from gpt-6.1-sol; final diff reviewed locally. Local validation reuses the native extension matching official based864e7198f. The current-main Pydantic module is byte-identical to that base, but current-main native code has not been rebuilt. Slow/S3 tests, remote services, and other operating systems were not validated. Reproducible commands after preparing the documented development environment (relative `.venv` examples are equivalent commands; actual local runs use absolute read-only environments and fresh task-private working directories): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests/test_pydantic_excluded_fields.py -q -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/pydantic.py python/python/tests/test_pydantic_excluded_fields.py git diff --checkd864e7198fHEAD ``` Final frozen-source results: - EXCLUDE_TARGET_COMPLETE_ORIGINAL_RED: Python 3.11.14, exit 1, 1.900076000019908 seconds, {"failed"=>7}. - EXCLUDE_TARGET_EMBEDDING: Python 3.11.14, exit 0, 1.3686820000002626 seconds, {"passed"=>7}. - EXCLUDE_TARGET313: Python 3.13.11, exit 0, 1.7214069999754429 seconds, {"passed"=>7}. - EXCLUDE_FULL311: Python 3.11.14, exit 0, 217.56785399999353 seconds, {"passed"=>1531, "skipped"=>36, "deselected"=>68, "warnings"=>251}. - EXCLUDE_FULL313: Python 3.13.11, exit 0, 197.93847399999504 seconds, {"passed"=>1528, "skipped"=>37, "deselected"=>68, "warnings"=>251}. - EXCLUDE_DOCTEST313: Python 3.13.11, exit 0, 11.8951390000002 seconds, {"passed"=>71, "skipped"=>6, "warnings"=>2}. - EXCLUDE_PYRIGHT: Python 3.13.11, exit 0, 5.808195000019623 seconds, {"warnings"=>0, "errors"=>0}. - EXCLUDE_PRECOMMIT: Python 3.11.14, exit 0, 0.5624079999979585 seconds, {"passed"=>4, "skipped"=>2}. - EXCLUDE_RUFF_FINAL: Ruff 0.9.9, exit 0, 0.06722099997568876 seconds, {}. - EXCLUDE_FORMAT_FINAL: Ruff 0.9.9, exit 0, 0.033322000002954155 seconds, {"files_formatted"=>157}. - EXCLUDE_SYNTAX: Python 3.11.14, exit 0, 0.031629999983124435 seconds, {"files_compiled"=>2}. - EXCLUDE_DIFF_CHECK: Git, exit 0, 0.024460999993607402 seconds, {}.
This commit is contained in:
LeonSGP
@@ -268,7 +268,9 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType:
|
||||
|
||||
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
|
||||
return [
|
||||
_pydantic_to_field(name, field) for name, field in model.model_fields.items()
|
||||
_pydantic_to_field(name, field)
|
||||
for name, field in model.model_fields.items()
|
||||
if field.exclude is not True
|
||||
]
|
||||
|
||||
|
||||
@@ -501,6 +503,9 @@ class LanceModel(pydantic.BaseModel):
|
||||
|
||||
vec_and_function = []
|
||||
for name, field_info in cls.safe_get_fields().items():
|
||||
if field_info.exclude is True:
|
||||
continue
|
||||
|
||||
func = get_extras(field_info, "vector_column_for")
|
||||
if func is not None:
|
||||
vec_and_function.append([name, func])
|
||||
@@ -508,6 +513,8 @@ class LanceModel(pydantic.BaseModel):
|
||||
configs = []
|
||||
for vec, func in vec_and_function:
|
||||
for source, field_info in cls.safe_get_fields().items():
|
||||
if field_info.exclude is True:
|
||||
continue
|
||||
src_func = get_extras(field_info, "source_column_for")
|
||||
if src_func is func:
|
||||
# note we can't use == here since the function is a pydantic
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import lancedb
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
from lancedb.pydantic import LanceModel
|
||||
|
||||
|
||||
class Details(BaseModel):
|
||||
label: str
|
||||
transient: str = Field(default="nested default", exclude=True)
|
||||
|
||||
|
||||
class Row(LanceModel):
|
||||
id: int
|
||||
details: Details
|
||||
transient: str = Field(default="local default", exclude=True)
|
||||
generated: list[str] = Field(default_factory=lambda: ["local"], exclude=True)
|
||||
retained: str = Field(default="persisted", exclude=False)
|
||||
|
||||
|
||||
def rows():
|
||||
return [
|
||||
Row(
|
||||
id=i,
|
||||
details=Details(label=str(i), transient="private"),
|
||||
transient="private",
|
||||
generated=["private"],
|
||||
)
|
||||
for i in (1, 2)
|
||||
]
|
||||
|
||||
|
||||
def assert_schema(schema):
|
||||
assert schema.names == ["id", "details", "retained"]
|
||||
assert [field.name for field in schema.field("details").type] == ["label"]
|
||||
assert Row.field_names() == ["id", "details", "transient", "generated", "retained"]
|
||||
|
||||
|
||||
def assert_restored(restored):
|
||||
assert sorted(row.id for row in restored) == [1, 2]
|
||||
for row in restored:
|
||||
assert row.transient == "local default"
|
||||
assert row.generated == ["local"]
|
||||
assert row.details.transient == "nested default"
|
||||
assert row.details.label == str(row.id)
|
||||
assert row.retained == "persisted"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit_schema", [False, True])
|
||||
def test_excluded_fields_round_trip(tmp_path, explicit_schema):
|
||||
data = rows()
|
||||
kwargs = {"schema": Row} if explicit_schema else {}
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("rows", data=data[:1], **kwargs)
|
||||
table.add(data[1:])
|
||||
assert_schema(table.schema)
|
||||
reopened = lancedb.connect(tmp_path).open_table("rows")
|
||||
assert_restored(reopened.search().to_pydantic(Row))
|
||||
assert data[0].transient == "private"
|
||||
assert data[0].generated == ["private"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("explicit_schema", [False, True])
|
||||
async def test_excluded_fields_round_trip_async(tmp_path, explicit_schema):
|
||||
data = rows()
|
||||
kwargs = {"schema": Row} if explicit_schema else {}
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table("rows", data=data[:1], **kwargs)
|
||||
await table.add(data[1:])
|
||||
assert_schema(await table.schema())
|
||||
reopened_db = await lancedb.connect_async(tmp_path)
|
||||
reopened = await reopened_db.open_table("rows")
|
||||
assert_restored(await reopened.query().to_pydantic(Row))
|
||||
assert data[0].details.transient == "private"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("excluded", ["transient", "source", "vector"])
|
||||
def test_excluded_fields_keep_embedding_metadata_consistent(
|
||||
tmp_path, monkeypatch, excluded
|
||||
):
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import get_registry
|
||||
from lancedb.pydantic import Vector
|
||||
|
||||
calls = []
|
||||
|
||||
class LocalEmbedding(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 2
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
calls.extend(texts)
|
||||
return [[1.0, 0.0] for _ in texts]
|
||||
|
||||
registry = get_registry()
|
||||
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
|
||||
registry.register("test-excluded-fields")(LocalEmbedding)
|
||||
function = LocalEmbedding.create(max_retries=0)
|
||||
|
||||
class EmbeddedRow(LanceModel):
|
||||
text: str = function.SourceField(default="alpha", exclude=excluded == "source")
|
||||
vector: Vector(2) = function.VectorField(
|
||||
default_factory=lambda: [1.0, 0.0], exclude=excluded == "vector"
|
||||
)
|
||||
transient: str = Field(default="local", exclude=True)
|
||||
|
||||
configs = EmbeddedRow.parse_embedding_functions()
|
||||
assert len(configs) == (1 if excluded == "transient" else 0)
|
||||
db = lancedb.connect(tmp_path)
|
||||
data = [{"text": "alpha"}] if excluded == "transient" else [EmbeddedRow()]
|
||||
table = db.create_table("embeddings", schema=EmbeddedRow, data=data)
|
||||
assert "transient" not in table.schema.names
|
||||
assert ("text" in table.schema.names) == (excluded != "source")
|
||||
assert ("vector" in table.schema.names) == (excluded != "vector")
|
||||
reopened = lancedb.connect(tmp_path).open_table("embeddings")
|
||||
assert len(reopened.embedding_functions) == len(configs)
|
||||
restored = reopened.search().to_pydantic(EmbeddedRow)
|
||||
assert restored[0].text == "alpha"
|
||||
assert restored[0].vector == [1.0, 0.0]
|
||||
assert restored[0].transient == "local"
|
||||
if excluded == "transient":
|
||||
assert calls == ["alpha"]
|
||||
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
|
||||
else:
|
||||
assert calls == []
|
||||
Reference in New Issue
Block a user