mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
fix(python): reject embedding source updates without vectors
Fixes #4761 `Table.update` does not compute embeddings, so changing an embedding source column could leave the previous embedding in place. Source-column updates now raise `ValueError` unless every associated vector column is supplied in the same call. The guard reads Python embedding metadata and native `lancedb::column_definitions` without loading embedding models. It uses Lance's field-path parser for quoted targets and covers synchronous and asynchronous local and remote tables. Supplying every associated vector column, vector-only updates, and unrelated updates remain supported. This is a breaking change for source-only updates. The API documentation and Update rows guide explain the requirement and the `merge_insert` alternative for regenerating embeddings. Validation at the reviewed head covered 44 targeted public update cases and the full Python suite (3107 passed, 162 skipped, 70 deselected). All applicable exact-head checks and workflows completed successfully.
This commit is contained in:
lancedb-gatefixer[bot]
@@ -29,6 +29,8 @@ from .types import BaseTokenizerType
|
||||
|
||||
_experimental_functions_enabled: bool
|
||||
|
||||
def _parse_field_path(path: str) -> List[str]: ...
|
||||
|
||||
IvfHnswPq: type[HnswPq] = HnswPq
|
||||
IvfHnswSq: type[HnswSq] = HnswSq
|
||||
IvfHnswFlat: type[HnswFlat] = HnswFlat
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import json
|
||||
from typing import Dict, Optional, Set, Tuple, Type
|
||||
from typing import Dict, Mapping, Optional, Set, Tuple, Type, Union
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
@@ -13,6 +13,51 @@ from .base import (
|
||||
)
|
||||
|
||||
|
||||
def _embedding_column_pairs(
|
||||
metadata: Optional[Mapping[Union[bytes, str], Union[bytes, str]]],
|
||||
) -> list[Tuple[str, str]]:
|
||||
"""Return source/vector pairs without restoring embedding functions."""
|
||||
pairs = []
|
||||
seen = set()
|
||||
|
||||
def add_pair(source_column: str, vector_column: str) -> None:
|
||||
pair = (source_column, vector_column)
|
||||
if pair not in seen:
|
||||
seen.add(pair)
|
||||
pairs.append(pair)
|
||||
|
||||
if metadata is None:
|
||||
return pairs
|
||||
|
||||
serialized = metadata.get(
|
||||
b"embedding_functions", metadata.get("embedding_functions")
|
||||
)
|
||||
if serialized is not None:
|
||||
if isinstance(serialized, bytes):
|
||||
serialized = serialized.decode("utf-8")
|
||||
for config in json.loads(serialized):
|
||||
add_pair(config["source_column"], config["vector_column"])
|
||||
|
||||
serialized = metadata.get(
|
||||
b"lancedb::column_definitions", metadata.get("lancedb::column_definitions")
|
||||
)
|
||||
if serialized is not None:
|
||||
if isinstance(serialized, bytes):
|
||||
serialized = serialized.decode("utf-8")
|
||||
for column in json.loads(serialized):
|
||||
kind = column.get("kind")
|
||||
if not isinstance(kind, dict) or "Embedding" not in kind:
|
||||
continue
|
||||
embedding = kind["Embedding"]
|
||||
source_column = embedding["source_column"]
|
||||
vector_column = embedding.get("dest_column")
|
||||
if vector_column is None:
|
||||
vector_column = f"{source_column}_embedding"
|
||||
add_pair(source_column, vector_column)
|
||||
|
||||
return pairs
|
||||
|
||||
|
||||
class EmbeddingFunctionRegistry:
|
||||
"""
|
||||
This is a singleton class used to register embedding functions
|
||||
|
||||
@@ -992,6 +992,11 @@ class RemoteTable(Table):
|
||||
This can be used to update zero to all rows depending on how many
|
||||
rows match the where clause.
|
||||
|
||||
Embeddings are not computed by this method. Updating an embedding source
|
||||
column requires providing all associated vector columns in the same call;
|
||||
otherwise a ValueError is raised. Use `merge_insert` to recompute embeddings
|
||||
automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
|
||||
@@ -34,6 +34,7 @@ from urllib.parse import urlparse
|
||||
from lancedb.scannable import _hf_batches, _register_optional_converters, to_scannable
|
||||
|
||||
from . import __version__
|
||||
from ._lancedb import _parse_field_path
|
||||
from ._blob import (
|
||||
BlobFile,
|
||||
_normalize_blob_descriptor,
|
||||
@@ -62,6 +63,7 @@ import numpy as np
|
||||
|
||||
from .common import DATA, VEC, VECTOR_COLUMN_NAME
|
||||
from .embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry
|
||||
from .embeddings.registry import _embedding_column_pairs
|
||||
from .index import (
|
||||
BTree,
|
||||
IvfFlat,
|
||||
@@ -2767,6 +2769,11 @@ class Table(ABC):
|
||||
Either `values` or `values_sql` must be provided. You cannot provide
|
||||
both.
|
||||
|
||||
Embeddings are not computed by this method. Updating an embedding source
|
||||
column requires providing all associated vector columns in the same call;
|
||||
otherwise a ValueError is raised. Use `merge_insert` to recompute embeddings
|
||||
automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
@@ -4929,6 +4936,11 @@ class LanceTable(Table):
|
||||
This can be used to update zero to all rows depending on how many
|
||||
rows match the where clause.
|
||||
|
||||
Embeddings are not computed by this method. Updating an embedding source
|
||||
column requires providing all associated vector columns in the same call;
|
||||
otherwise a ValueError is raised. Use `merge_insert` to recompute embeddings
|
||||
automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
@@ -7329,6 +7341,11 @@ class AsyncTable:
|
||||
If a filter is provided with `where` then only rows matching the
|
||||
filter will be updated. Otherwise all rows will be updated.
|
||||
|
||||
Embeddings are not computed by this method. Updating an embedding source
|
||||
column requires providing all associated vector columns in the same call;
|
||||
otherwise a ValueError is raised. Use `merge_insert` to recompute embeddings
|
||||
automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
updates: dict, optional
|
||||
@@ -7378,6 +7395,24 @@ class AsyncTable:
|
||||
if updates is not None:
|
||||
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
|
||||
|
||||
# Only the column mappings are needed here. Restoring embedding functions
|
||||
# could load models even for updates that provide their own vectors.
|
||||
schema = await self.schema()
|
||||
column_pairs = _embedding_column_pairs(schema.metadata)
|
||||
if column_pairs:
|
||||
# Metadata names are literal top-level fields; targets may be quoted.
|
||||
updated_columns = {tuple(_parse_field_path(name)) for name in updates_sql}
|
||||
for source_column, vector_column in column_pairs:
|
||||
source_updated = (source_column,) in updated_columns
|
||||
vector_updated = (vector_column,) in updated_columns
|
||||
if source_updated and not vector_updated:
|
||||
raise ValueError(
|
||||
f"Updating embedding source column {source_column!r} requires "
|
||||
f"updating vector column {vector_column!r} in the same call. "
|
||||
"Provide both columns or use merge_insert "
|
||||
"to recompute embeddings."
|
||||
)
|
||||
|
||||
predicate = where.to_sql() if isinstance(where, Expr) else where
|
||||
return await self._inner.update(updates_sql, predicate)
|
||||
|
||||
|
||||
@@ -956,7 +956,8 @@ def test_remote_permutation_is_picklable():
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_schema", [True, False])
|
||||
def test_create_table_embedding_functions(with_schema):
|
||||
@pytest.mark.parametrize("source_target", ["text", "`text`"])
|
||||
def test_create_table_embedding_functions(with_schema, source_target):
|
||||
func = MockNonNormTextEmbeddingFunction.create()
|
||||
config = EmbeddingFunctionConfig(
|
||||
source_column="text", vector_column="vector", function=func
|
||||
@@ -1054,6 +1055,11 @@ def test_create_table_embedding_functions(with_schema):
|
||||
func.compute_source_embeddings(["goodbye world"])[0].tolist()
|
||||
]
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="source column 'text'.*vector column 'vector'"
|
||||
):
|
||||
table.update(values={source_target: "updated text"})
|
||||
|
||||
|
||||
def test_create_table_positional_parameters_match_base_contract():
|
||||
from lancedb.db import DBConnection
|
||||
|
||||
@@ -2834,6 +2834,260 @@ def test_update(mem_db: DBConnection):
|
||||
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("async_api", [False, True])
|
||||
@pytest.mark.parametrize("use_sql", [False, True])
|
||||
@pytest.mark.parametrize("where", [None, "id = 1", col("id") == 1])
|
||||
@pytest.mark.parametrize("source_target", ["text", "`text`"])
|
||||
async def test_update_embedding_source_requires_vector(
|
||||
tmp_path, async_api, use_sql, where, source_target
|
||||
):
|
||||
func = MockTextEmbeddingFunction.create()
|
||||
|
||||
class Doc(LanceModel):
|
||||
id: int
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(func.ndims()) = func.VectorField()
|
||||
|
||||
data = [{"id": 1, "text": "hello world"}, {"id": 2, "text": "goodbye world"}]
|
||||
if async_api:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table("docs", data, schema=Doc)
|
||||
before = await table.to_arrow()
|
||||
version = await table.version()
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("docs", data, schema=Doc)
|
||||
before = table.to_arrow()
|
||||
version = table.version
|
||||
|
||||
# Rejection must happen before any rows or the table version change.
|
||||
updates = {source_target: "text || ' updated'" if use_sql else "updated text"}
|
||||
with pytest.raises(
|
||||
ValueError, match="source column 'text'.*vector column 'vector'"
|
||||
):
|
||||
if async_api:
|
||||
if use_sql:
|
||||
await table.update(where=where, updates_sql=updates)
|
||||
else:
|
||||
await table.update(updates, where=where)
|
||||
elif use_sql:
|
||||
table.update(where=where, values_sql=updates)
|
||||
else:
|
||||
table.update(where=where, values=updates)
|
||||
|
||||
if async_api:
|
||||
assert (await table.to_arrow()).equals(before)
|
||||
assert await table.version() == version
|
||||
else:
|
||||
assert table.to_arrow().equals(before)
|
||||
assert table.version == version
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("async_api", [False, True])
|
||||
@pytest.mark.parametrize("use_sql", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"source_target,vector_target,other_vector_target",
|
||||
[
|
||||
("text", "vector", "other_vector"),
|
||||
("`text`", "vector", "`other_vector`"),
|
||||
("text", "`vector`", "other_vector"),
|
||||
("`text`", "`vector`", "`other_vector`"),
|
||||
],
|
||||
)
|
||||
async def test_update_embedding_source_with_multiple_vectors(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
async_api,
|
||||
use_sql,
|
||||
source_target,
|
||||
vector_target,
|
||||
other_vector_target,
|
||||
):
|
||||
func = MockTextEmbeddingFunction.create()
|
||||
|
||||
class Doc(LanceModel):
|
||||
id: int
|
||||
text: str
|
||||
vector: Vector(func.ndims())
|
||||
other_vector: Vector(func.ndims())
|
||||
|
||||
configs = [
|
||||
EmbeddingFunctionConfig(
|
||||
source_column="text", vector_column=column, function=func
|
||||
)
|
||||
for column in ["vector", "other_vector"]
|
||||
]
|
||||
data = [{"id": 1, "text": "hello world"}, {"id": 2, "text": "goodbye world"}]
|
||||
if async_api:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table(
|
||||
"docs", data, schema=Doc, embedding_functions=configs
|
||||
)
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("docs", data, schema=Doc, embedding_functions=configs)
|
||||
|
||||
async def update(values):
|
||||
if async_api:
|
||||
if use_sql:
|
||||
return await table.update(updates_sql=values, where=col("id") == 1)
|
||||
return await table.update(values, where=col("id") == 1)
|
||||
if use_sql:
|
||||
return table.update(values_sql=values, where=col("id") == 1)
|
||||
return table.update(values=values, where=col("id") == 1)
|
||||
|
||||
vector = func.compute_source_embeddings(["updated text"])[0]
|
||||
updates = {source_target: "updated text", vector_target: vector}
|
||||
if use_sql:
|
||||
updates = {
|
||||
key: table_module.value_to_sql(value) for key, value in updates.items()
|
||||
}
|
||||
|
||||
# Checking column dependencies must not require loading an embedding model.
|
||||
monkeypatch.delitem(EmbeddingFunctionRegistry.get_instance()._functions, "test")
|
||||
with pytest.raises(ValueError, match="vector column 'other_vector'"):
|
||||
await update(updates)
|
||||
|
||||
updates[other_vector_target] = updates[vector_target]
|
||||
result = await update(updates)
|
||||
assert result.rows_updated == 1
|
||||
actual = await table.to_arrow() if async_api else table.to_arrow()
|
||||
rows = actual.sort_by("id").to_pylist()
|
||||
assert rows[0]["text"] == "updated text"
|
||||
assert np.allclose(rows[0]["vector"], vector)
|
||||
assert np.allclose(rows[0]["other_vector"], vector)
|
||||
assert rows[1]["text"] == "goodbye world"
|
||||
|
||||
# Vector-only and unrelated updates are still allowed on embedding tables.
|
||||
assert (await update({vector_target: updates[vector_target]})).rows_updated == 1
|
||||
assert (await update({"id": "3" if use_sql else 3})).rows_updated == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("async_api", [False, True])
|
||||
@pytest.mark.parametrize("source_target", ["text", "`text`"])
|
||||
async def test_update_embedding_source_requires_native_vector_metadata(
|
||||
tmp_path, async_api, source_target
|
||||
):
|
||||
vector_type = pa.list_(pa.float32(), 2)
|
||||
column_definitions = [
|
||||
{"kind": "Physical"},
|
||||
{"kind": "Physical"},
|
||||
{
|
||||
"kind": {
|
||||
"Embedding": {
|
||||
"source_column": "text",
|
||||
"dest_column": "vector",
|
||||
"embedding_name": "native-explicit",
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"kind": {
|
||||
"Embedding": {
|
||||
"source_column": "text",
|
||||
"dest_column": None,
|
||||
"embedding_name": "native-default",
|
||||
}
|
||||
}
|
||||
},
|
||||
]
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("text", pa.string()),
|
||||
pa.field("vector", vector_type),
|
||||
pa.field("text_embedding", vector_type),
|
||||
],
|
||||
metadata={
|
||||
b"lancedb::column_definitions": json.dumps(column_definitions).encode()
|
||||
},
|
||||
)
|
||||
data = pa.Table.from_pylist(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"text": "hello world",
|
||||
"vector": [1.0, 0.0],
|
||||
"text_embedding": [0.0, 1.0],
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"text": "goodbye world",
|
||||
"vector": [0.5, 0.5],
|
||||
"text_embedding": [0.25, 0.75],
|
||||
},
|
||||
],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
if async_api:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table("docs", data)
|
||||
table_schema = await table.schema()
|
||||
before = await table.to_arrow()
|
||||
version = await table.version()
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("docs", data)
|
||||
table_schema = table.schema
|
||||
before = table.to_arrow()
|
||||
version = table.version
|
||||
|
||||
metadata = table_schema.metadata or {}
|
||||
assert b"lancedb::column_definitions" in metadata
|
||||
assert b"embedding_functions" not in metadata
|
||||
|
||||
async def update(values):
|
||||
if async_api:
|
||||
return await table.update(values, where="id = 1")
|
||||
return table.update(values=values, where="id = 1")
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="source column 'text'.*vector column 'vector'"
|
||||
):
|
||||
await update({source_target: "updated text"})
|
||||
|
||||
if async_api:
|
||||
assert (await table.to_arrow()).equals(before)
|
||||
assert await table.version() == version
|
||||
else:
|
||||
assert table.to_arrow().equals(before)
|
||||
assert table.version == version
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="source column 'text'.*vector column 'text_embedding'"
|
||||
):
|
||||
await update({source_target: "updated text", "vector": [2.0, 3.0]})
|
||||
|
||||
if async_api:
|
||||
assert (await table.to_arrow()).equals(before)
|
||||
assert await table.version() == version
|
||||
else:
|
||||
assert table.to_arrow().equals(before)
|
||||
assert table.version == version
|
||||
|
||||
result = await update(
|
||||
{
|
||||
source_target: "updated text",
|
||||
"vector": [2.0, 3.0],
|
||||
"text_embedding": [4.0, 5.0],
|
||||
}
|
||||
)
|
||||
assert result.rows_updated == 1
|
||||
assert (await update({"vector": [6.0, 7.0]})).rows_updated == 1
|
||||
assert (await update({"id": 99})).rows_updated == 1
|
||||
|
||||
actual = await table.to_arrow() if async_api else table.to_arrow()
|
||||
rows = {row["id"]: row for row in actual.to_pylist()}
|
||||
assert rows[99]["text"] == "updated text"
|
||||
assert np.allclose(rows[99]["vector"], [6.0, 7.0])
|
||||
assert np.allclose(rows[99]["text_embedding"], [4.0, 5.0])
|
||||
assert rows[2]["text"] == "goodbye world"
|
||||
|
||||
|
||||
def test_update_expr_filter_literals(mem_db: DBConnection):
|
||||
values = ["5", "4.66e-84", "it's"]
|
||||
table = mem_db.create_table(
|
||||
|
||||
@@ -128,6 +128,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(table::tokenize, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(util::parse_field_path, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(expr_col, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(expr_lit, m)?)?;
|
||||
|
||||
@@ -61,6 +61,12 @@ pub fn validate_table_name(table_name: &str) -> PyResult<()> {
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
}
|
||||
|
||||
/// Parse field references with the same identifier syntax as Lance writes.
|
||||
#[pyfunction(name = "_parse_field_path")]
|
||||
pub fn parse_field_path(path: &str) -> PyResult<Vec<String>> {
|
||||
lance_core::datatypes::parse_field_path(path).map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
}
|
||||
|
||||
/// A wrapper around a LanceDB type to allow it to be used in Python
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PyLanceDB<T>(pub T);
|
||||
|
||||
Reference in New Issue
Block a user