fix(python): keep embedding functions registered without an alias readable (#4356)

`EmbeddingFunctionRegistry.register()` documents that the class name is
used when no alias is given, and it does register the class under that
name. But it stored the alias itself (`None`) as
`__embedding_function_registry_alias__`, which is what
`function_to_metadata()` writes into the table schema. A table using
such a function gets `"name": null` in its `embedding_functions`
metadata, and any later read of the functions fails:

```python
@get_registry().register()
class MyEmbed(TextEmbeddingFunction): ...

f = MyEmbed.create()
class Schema(LanceModel):
    text: str = f.SourceField()
    vector: Vector(4) = f.VectorField()

tbl = db.create_table("t", schema=Schema)
tbl.add([{"text": "hi"}])   # KeyError: None
tbl.search("hi")            # KeyError: None
```

This stores the key the class is registered under instead, so the class
name round-trips. For aliased classes the key is the alias, so their
metadata is unchanged. The TypeScript registry already falls back to
`ctor.name` in the same way.

Added `test_embedding_function_registered_without_alias`, which creates
a table with such a function, reopens it, and adds a row. It fails with
`KeyError` before the change. `test_embeddings.py`, `test_pydantic.py`,
`test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff
check` are clean.
This commit is contained in:
Mohammad Hijjawi
2026-10-23 02:42:21 +08:00
committed by GitHub
parent a0946282ca
commit 5f9a8af27c
2 changed files with 34 additions and 1 deletions
+3 -1
View File
@@ -73,7 +73,9 @@ class EmbeddingFunctionRegistry:
raise KeyError(f"{cls.__name__} was already registered")
key = alias or cls.__name__
self._functions[key] = cls
cls.__embedding_function_registry_alias__ = alias
# Store the name the class is registered under, so the table
# metadata written for it can be looked up again when reading.
cls.__embedding_function_registry_alias__ = key
return cls
return decorator
+31
View File
@@ -64,6 +64,37 @@ def test_embedding_function(tmp_path):
assert np.allclose(actual, expected)
def test_embedding_function_registered_without_alias(tmp_path):
registry = EmbeddingFunctionRegistry.get_instance()
@registry.register()
class NoAliasEmbeddingFunction(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
return [np.array([1.0, 2.0]) for _ in texts]
assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction
func = NoAliasEmbeddingFunction.create()
class Schema(LanceModel):
text: str = func.SourceField()
vector: Vector(2) = func.VectorField()
db = lancedb.connect(tmp_path)
db.create_table("test", schema=Schema)
# The class name is written to the table metadata, so the function can be
# resolved again when the table is reopened.
table = lancedb.connect(tmp_path).open_table("test")
table.add([{"text": "hello"}])
conf = table.embedding_functions["vector"]
assert isinstance(conf.function, NoAliasEmbeddingFunction)
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]]
def test_instructor_ndims_uses_instruction():
instructor = get_registry().get("instructor").create()
model = MagicMock()