From 554dc01f4d653f95bf044059e1f88ad7e3e71457 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:06 +0800 Subject: [PATCH] fix(python): preserve custom embedding field metadata (#4463) `SourceField(json_schema_extra={...})` and `VectorField(json_schema_extra={...})` currently raise `TypeError` because the helper supplies the same `Field` keyword twice. Explicit `None` also fails. This prevents users from adding ordinary Pydantic JSON schema metadata to embedding fields. Copy and merge the supplied dictionary before attaching the helper's internal embedding identity. Caller dictionaries remain unchanged, and reserved embedding keys retain the real function. Other `Field` arguments keep their existing behavior. This change covers dictionaries and `None`; callable `json_schema_extra` remains unsupported by the existing embedding metadata reader and is outside this fix. The 11 regression cases exercise both helpers, empty/custom/None metadata, reserved-key collisions, other Field arguments, and a deterministic local embedding model through JSON schema generation, table creation, reopening, and search. On the original source, 9 fail and 2 controls pass; the candidate passes all 11 on Python 3.11 and 3.13. Validation uses read-only reused environments and the native extension built for d864e719. The production helper is byte-identical to current main bbc4aebe before this fix; no current-main native rebuild was performed. The branch contains only this helper fix and its regression test, independently based on d864e719. Reproducible checkout commands below assume the documented environment is installed as `.venv`; the actual local runs used existing interpreters through absolute paths and isolated generated-data working directories: ```sh 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) git diff --check .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/base.py python/python/tests/test_embedding_field_metadata.py ``` Doctests require the repository's tests/dev/embeddings dependencies, including mlx on this platform. All checks below ran against frozen head 3fa6820. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally. Frozen-head validation (all exit 0): - Python 3.11 full: 1535 passed, 36 skipped, 68 deselected; 190.272415 seconds wall time. - Python 3.13 full: 1532 passed, 37 skipped, 68 deselected; 249.289778 seconds wall time. - Python 3.13 doctests: 71 passed, 6 skipped; 10.391002 seconds wall time. - Pyright: 0 errors, 0 warnings; 4.884641 seconds wall time. - pre-commit: 4 passed, 2 skipped; 0.747764 seconds wall time. - Repository Ruff check: passed; 0.078177 seconds wall time. - Repository Ruff format: 157 files already formatted; 0.034561 seconds wall time. --- python/python/lancedb/embeddings/base.py | 8 +- .../tests/test_embedding_field_metadata.py | 84 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_field_metadata.py diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index f26a940df..96ca6556c 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -239,7 +239,9 @@ class EmbeddingFunction(BaseModel, ABC): Creates a pydantic Field that can automatically annotate the source column for this embedding function """ - return Field(json_schema_extra={"source_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["source_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def VectorField(self, **kwargs): """ @@ -250,7 +252,9 @@ class EmbeddingFunction(BaseModel, ABC): static type checkers. LanceDB will infer the fixed vector dimension from this embedding function. """ - return Field(json_schema_extra={"vector_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["vector_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def __eq__(self, __value: object) -> bool: if not hasattr(__value, "__dict__"): diff --git a/python/python/tests/test_embedding_field_metadata.py b/python/python/tests/test_embedding_field_metadata.py new file mode 100644 index 000000000..68754ee35 --- /dev/null +++ b/python/python/tests/test_embedding_field_metadata.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import copy + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import get_registry +from lancedb.pydantic import LanceModel, Vector + + +class LocalEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts] + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +@pytest.mark.parametrize("extra", [None, {}, {"x-purpose": "retrieval"}]) +def test_embedding_field_preserves_extra_metadata(method, tag, extra): + function = LocalEmbedding.create(max_retries=0) + original = copy.deepcopy(extra) + field = getattr(function, method)(json_schema_extra=extra) + assert extra == original + assert field.json_schema_extra[tag] is function + for key, value in (extra or {}).items(): + assert field.json_schema_extra[key] == value + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +def test_extra_metadata_cannot_replace_embedding_identity(method, tag): + function = LocalEmbedding.create(max_retries=0) + supplied = {tag: "caller value", "x-purpose": "retrieval"} + field = getattr(function, method)(json_schema_extra=supplied) + assert supplied[tag] == "caller value" + assert field.json_schema_extra[tag] is function + assert field.json_schema_extra["x-purpose"] == "retrieval" + + +@pytest.mark.parametrize("method", ["SourceField", "VectorField"]) +def test_embedding_field_preserves_other_field_arguments(method): + function = LocalEmbedding.create(max_retries=0) + field = getattr(function, method)(title="Custom title", description="Custom field") + assert field.title == "Custom title" + assert field.description == "Custom field" + + +def test_extra_metadata_preserves_schema_and_local_embedding(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + registry.register("test-embedding-field-extra")(LocalEmbedding) + function = LocalEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField(json_schema_extra={"x-purpose": "source"}) + vector: Vector(2) = function.VectorField( + json_schema_extra={"x-purpose": "embedding"} + ) + + json_schema = Row.model_json_schema() + assert json_schema["properties"]["text"]["x-purpose"] == "source" + assert json_schema["properties"]["vector"]["x-purpose"] == "embedding" + configs = Row.parse_embedding_functions() + assert len(configs) == 1 + assert configs[0].function is function + assert configs[0].source_column == "text" + assert configs[0].vector_column == "vector" + db = lancedb.connect(tmp_path) + table = db.create_table( + "rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}] + ) + assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, LocalEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"