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 52a1ccee5a:

- 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 d864e7198f with its
verified matching native extension reused read-only. All 174
native/config paths are unchanged from that base. Latest main
bbc4aebe55 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.
This commit is contained in:
LeonSGP
2026-10-23 02:45:51 +08:00
committed by GitHub
parent e1e87e1cae
commit 8552fd2f5b
3 changed files with 98 additions and 2 deletions
+4 -2
View File
@@ -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.
@@ -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"
+6
View File
@@ -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))]