From 8552fd2f5b8213ce6be5fdd9d3cdfe7978f10a1d Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:45:51 +0800 Subject: [PATCH] fix(python): prevent embedding aliases from replacing registered functions (#4462) Registering a different embedding class under an existing alias silently replaced the original class. A deterministic real local SDK table reproduction changed the top result for the same query from alpha to beta after reopening the persisted table. Check the resolved key before inserting a class. Preserve repeated registration of the same class under the same nonempty explicit alias as idempotent, as existing SDK tests rely on it. Repeated implicit class-name registration, including an empty alias, still raises KeyError. Failed registration preserves the mapping and class metadata. Existing embedding tests receive isolated registry mappings; test_util.py remains unchanged. Eight regression cases cover different-class alias collisions, alias/class-name collisions, implicit and empty-alias duplicates, equal class names under distinct aliases, explicit-alias idempotence, and real local creation/reopening/search. This is independent of #4356, which fixes no-alias serialization. Validation on 52a1ccee5a5b8dbfff8568abafb57099c55383b9: - Original registry plus eight new cases: 4 failed, 4 passed; Python 3.11.14; exit 1; 2.902604 s. Candidate plus full utility module: 83 passed, 1 skipped; exit 0; 3.144414 s. - Python 3.11.14 full official non-slow/non-S3 suite: 1532 passed, 36 skipped, 68 deselected; exit 0; 152.631543 s. - Python 3.13.11 full official non-slow/non-S3 suite: 1529 passed, 37 skipped, 68 deselected; exit 0; 163.970571 s. - Official Python 3.13 doctest with existing compatible embeddings/torch/mlx dependencies: 71 passed, 6 skipped; exit 0; 9.635034 s. - Pyright with matching interpreter/source path: 0 errors, 0 warnings; exit 0; 4.299743 s. - Repository Ruff 0.9.9 check and format: exit 0; 157 files already formatted; 0.047418 s and 0.016815 s. Pre-commit changed files: four hooks passed, two inapplicable skipped; exit 0; 0.406136 s. git diff --check passed. Base is official d864e7198fe55a4e5fbe747fee574737fa102aab with its verified matching native extension reused read-only. All 174 native/config paths are unchanged from that base. Latest main bbc4aebe55b69471fa1bd7655599b8d2e3f9352c has byte-identical registry production code; its complete native build was not run. Verification was macOS arm64, with optional dependency/service skips above; Linux/Windows CI remains remote verification. Exact commands, cwd, identities, raw logs and intermediate failures are in the local manifest. The earlier strict same-class check caused 13 utility failures and was corrected in production; it is not described as a baseline failure. Reproduction commands after preparing the repository's declared environment (from repository root unless noted): ```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) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/registry.py python/python/tests/test_embedding_registry.py python/python/tests/test_embeddings.py git diff --check ``` These relative commands describe reproducible source/cwd equivalents. The recorded local runs used absolute paths to existing Python 3.11.14/3.13.11 environments, a verified matching reused native extension, and read-only compatible optional-dependency cache paths for the 3.13 doctest, rather than a newly installed `.venv`. Full tests used separate task-private working directories and `--basetemp` directories to avoid cross-run database collisions. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally by the collaborating agents. --- python/python/lancedb/embeddings/registry.py | 6 +- .../python/tests/test_embedding_registry.py | 88 +++++++++++++++++++ python/python/tests/test_embeddings.py | 6 ++ 3 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_registry.py diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 3b940fe2e..728ff36da 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -69,9 +69,11 @@ class EmbeddingFunctionRegistry: def decorator(cls): if not issubclass(cls, EmbeddingFunction): raise TypeError("Must be a subclass of EmbeddingFunction") - if cls.__name__ in self._functions: - raise KeyError(f"{cls.__name__} was already registered") key = alias or cls.__name__ + if key in self._functions and ( + not alias or self._functions[key] is not cls + ): + raise KeyError(f"{key} was already registered") self._functions[key] = cls # Store the name the class is registered under, so the table # metadata written for it can be looked up again when reading. diff --git a/python/python/tests/test_embedding_registry.py b/python/python/tests/test_embedding_registry.py new file mode 100644 index 000000000..07c10bd7e --- /dev/null +++ b/python/python/tests/test_embedding_registry.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import EmbeddingFunctionRegistry, get_registry +from lancedb.pydantic import LanceModel, Vector + + +class FirstEmbedding(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] + + +class SecondEmbedding(FirstEmbedding): + def generate_embeddings(self, texts): + return [[0.0, 1.0] if text == "alpha" else [1.0, 0.0] for text in texts] + + +@pytest.mark.parametrize( + "first_alias,second_alias,second_class", + [ + ("shared", "shared", SecondEmbedding), + (None, "FirstEmbedding", SecondEmbedding), + ("SecondEmbedding", None, SecondEmbedding), + (None, None, FirstEmbedding), + ("", "", FirstEmbedding), + ], +) +def test_duplicate_registration_preserves_original( + first_alias, second_alias, second_class +): + registry = EmbeddingFunctionRegistry() + registry.register(first_alias)(FirstEmbedding) + key = first_alias or "FirstEmbedding" + original_alias = FirstEmbedding.__embedding_function_registry_alias__ + second_attributes = dict(second_class.__dict__) + with pytest.raises(KeyError, match=f"{key} was already registered"): + registry.register(second_alias)(second_class) + assert registry.get(key) is FirstEmbedding + assert FirstEmbedding.__embedding_function_registry_alias__ == original_alias + assert dict(second_class.__dict__) == second_attributes + + +def test_class_name_does_not_collide_with_distinct_alias(): + registry = EmbeddingFunctionRegistry() + registry.register()(FirstEmbedding) + same_name = type("FirstEmbedding", (SecondEmbedding,), {}) + registry.register("separate")(same_name) + assert registry.get("FirstEmbedding") is FirstEmbedding + assert registry.get("separate") is same_name + + +def test_same_class_explicit_alias_registration_is_idempotent(): + registry = EmbeddingFunctionRegistry() + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + attributes = dict(FirstEmbedding.__dict__) + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + assert registry.get("shared") is FirstEmbedding + assert dict(FirstEmbedding.__dict__) == attributes + + +def test_alias_collision_does_not_change_persisted_table(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + alias = "test-alias-collision-persisted-table" + registry.register(alias)(FirstEmbedding) + function = FirstEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField() + vector: Vector(2) = function.VectorField() + + 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" + assert registry.register(alias)(FirstEmbedding) is FirstEmbedding + with pytest.raises(KeyError, match=f"{alias} was already registered"): + registry.register(alias)(SecondEmbedding) + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, FirstEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 51907f141..d6d128845 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -23,6 +23,12 @@ from lancedb.embeddings.utils import retry from lancedb.pydantic import LanceModel, Vector +@pytest.fixture(autouse=True) +def isolated_embedding_registry(monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + + def mock_embed_func(input_data): return [np.random.randn(128).tolist() for _ in range(len(input_data))]