Compare commits

..

6 Commits

Author SHA1 Message Date
Gatefixer 8b14e2fe63 Merge main into gatekeeper/fix-3530-1 2026-08-27 08:01:24 +00:00
Gatefixer d55446f71f Merge remote-tracking branch 'origin/main' into gatekeeper/fix-3530-1
# Conflicts:
#	rust/lancedb/src/table/query.rs
2026-08-26 20:47:49 +00:00
Gatefixer 676c5b7315 fix: normalize cosine scores in LSM plans 2026-08-25 21:06:27 +00:00
Gatefixer 5093f37559 Merge main into gatekeeper/fix-3530-1 2026-08-25 20:39:08 +00:00
Gatefixer 4f5c55888b fix: normalize cosine scores at ANN boundaries 2026-08-06 08:24:54 +00:00
Gatefixer f95d4f583d fix: return cosine-scaled ANN distances 2026-08-06 03:17:30 +00:00
31 changed files with 750 additions and 1611 deletions
+2 -4
View File
@@ -270,8 +270,7 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
yield name, expr.to_sql()
return
for column in projection:
if isinstance(column, str):
@@ -281,8 +280,7 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
yield name, expr.to_sql()
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
-2
View File
@@ -87,7 +87,6 @@ class PyExpr:
def contains(self, substr: "PyExpr") -> "PyExpr": ...
def isin(self, values: List["PyExpr"]) -> "PyExpr": ...
def cast(self, data_type: pa.DataType) -> "PyExpr": ...
def column_name(self) -> Optional[str]: ...
def to_sql(self) -> str: ...
def expr_col(name: str) -> PyExpr: ...
@@ -609,7 +608,6 @@ class PyQueryRequest:
filter: Optional[Union[str, bytes]]
full_text_search: Optional[FullTextQuery]
select: Optional[Union[str, List[str]]]
select_source_columns: Optional[Dict[str, str]]
fast_search: Optional[bool]
with_row_id: Optional[bool]
use_lsm: Optional[bool]
+7 -38
View File
@@ -16,7 +16,6 @@ from typing import (
Iterable,
List,
Literal,
Mapping,
Optional,
Union,
)
@@ -688,35 +687,17 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError("serialize is not supported for this connection type")
def create_function(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> FunctionVersion:
def create_function(self, definition: UdfDefinition) -> FunctionVersion:
"""Register a scalar Python UDF and wait for its immutable version.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
This is the blocking counterpart of :meth:`create_function_async`.
Local connections raise ``NotImplementedError``.
"""
return self.create_function_async(definition, secrets=secrets).wait()
return self.create_function_async(definition).wait()
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
Submission returns a typed job. The immutable Function version becomes
available only when :meth:`Job.wait` succeeds. Local connections raise
``NotImplementedError``.
@@ -1424,13 +1405,8 @@ class LanceDBConnection(DBConnection):
return Job(self._conn.job(job_id))
@override
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition))
return Job(job)
@override
@@ -2249,24 +2225,17 @@ class AsyncConnection(object):
return AsyncJob(self._inner.job(job_id))
async def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
self, definition: UdfDefinition
) -> AsyncJob[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
The returned typed job resolves to the immutable Function version.
Local connections raise ``NotImplementedError``.
"""
if not isinstance(definition, UdfDefinition):
raise TypeError("create_function_async requires a @udf definition")
inner = await self._inner.create_function_async(
definition._submission_json(secrets)
definition.registration_request.to_canonical_json()
)
return _typed_job(inner, FunctionVersion.from_json)
+1 -5
View File
@@ -249,10 +249,6 @@ class Expr:
# ── utilities ────────────────────────────────────────────────────────────
def _column_name(self) -> str | None:
"""Return the source name when this is a bare column expression."""
return self._inner.column_name()
def to_sql(self) -> str:
"""Render the expression as a SQL string (useful for debugging)."""
return self._inner.to_sql()
@@ -316,7 +312,7 @@ def func(name: str, *args: ExprLike) -> Expr:
--------
>>> from lancedb.expr import col, func
>>> func("lower", col("name"))
Expr(lower(`name`))
Expr(lower(name))
"""
inner_args = [_coerce(a)._inner for a in args]
return Expr(expr_func(name, inner_args))
+6 -106
View File
@@ -4,7 +4,7 @@
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, secret resolution, and execution are owned by Sophon.
environment bake, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
@@ -229,7 +229,7 @@ class PythonEnvironmentSpec(_RemoteValue):
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with non-secret environment values.
"""Remote runtime definition with environment values.
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
their unknown payload fields are intentionally not retained by the client.
@@ -268,7 +268,6 @@ class FunctionVersion(_RemoteValue):
runtime: PythonRuntimeSpec
runtime_digest: str
environment_digest: str
required_secrets: tuple[str, ...] = ()
created_at: str
def __call__(self, **inputs: Any) -> FunctionApplication:
@@ -330,17 +329,12 @@ class FunctionVersion(_RemoteValue):
class FunctionRegistrationRequest(_RemoteValue):
"""Stable remote registration envelope produced by :func:`udf`.
Only secret names are represented. Secret values are supplied separately
when the definition is submitted and are not part of this durable value.
"""
"""Stable remote registration envelope produced by :func:`udf`."""
name: str
artifact: FunctionArtifactRequest
signature: FunctionSignature
runtime: PythonRuntimeSpec
required_secrets: tuple[str, ...] = ()
class FunctionVersionRef(_OpenRemoteValue):
@@ -485,27 +479,6 @@ class RefreshColumnResult(_RemoteValue):
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
# Keep this byte limit aligned with Sophon's MAX_FUNCTION_SECRET_VALUE_BYTES.
_MAX_FUNCTION_SECRET_VALUE_BYTES = 64 * 1024
_MAX_FUNCTION_SECRET_VALUES_BYTES = 512 * 1024
def _validate_secret_value(name: str, value: Any) -> str:
"""Validate one secret value before building the create request."""
if not isinstance(value, str):
raise TypeError(f"Function secret {name!r} value must be a string")
if not value:
raise ValueError(f"Function secret {name!r} value must be non-empty")
if "\0" in value:
raise ValueError(f"Function secret {name!r} value must not contain NUL")
value_bytes = len(value.encode("utf-8"))
if value_bytes > _MAX_FUNCTION_SECRET_VALUE_BYTES:
raise ValueError(
f"Function secret {name!r} value exceeds the "
f"{_MAX_FUNCTION_SECRET_VALUE_BYTES}-byte limit"
)
return value
_GRAMMAR_PRIMITIVES = (
@@ -936,7 +909,6 @@ class UdfDefinition:
output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
pip: tuple[str, ...],
env: Mapping[str, str],
secrets: tuple[str, ...],
python_version: Optional[str],
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
@@ -963,17 +935,6 @@ class UdfDefinition:
for key, value in environment.items()
):
raise TypeError("Function env keys and values must be strings")
required_secrets = tuple(sorted(set(secrets)))
invalid_secrets = [
secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret)
]
if invalid_secrets:
raise ValueError(f"invalid Function secret names: {invalid_secrets!r}")
overlap = set(environment) & set(required_secrets)
if overlap:
raise ValueError(
f"Function env and secret names must be disjoint: {sorted(overlap)!r}"
)
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
@@ -1002,65 +963,14 @@ class UdfDefinition:
),
signature=signature,
runtime=runtime,
required_secrets=required_secrets,
)
functools.update_wrapper(self, function)
@property
def registration_request(self) -> FunctionRegistrationRequest:
"""The immutable, value-free client model for a Function submission."""
"""The immutable request sent by ``create_function_async``."""
return self._request
def _submission_json(self, secrets: Optional[Mapping[str, str]]) -> str:
"""Build one registration submission without retaining values on self."""
if secrets is None:
secret_values: Mapping[str, str] = {}
elif not isinstance(secrets, Mapping):
raise TypeError("Function secrets must be a mapping of names to strings")
else:
secret_values = secrets
if any(not isinstance(name, str) for name in secret_values):
raise TypeError("Function secret names must be strings")
expected = set(self._request.required_secrets)
provided = set(secret_values)
if provided != expected:
missing = sorted(expected - provided)
unexpected = sorted(provided - expected)
details = []
if missing:
details.append(f"missing: {missing!r}")
if unexpected:
details.append(f"unexpected: {unexpected!r}")
raise ValueError(
"Function secret values must exactly match the declared secrets ("
+ "; ".join(details)
+ ")"
)
canonical_values = {}
total_bytes = 0
for name in sorted(secret_values):
value = _validate_secret_value(name, secret_values[name])
total_bytes += len(value.encode("utf-8"))
if total_bytes > _MAX_FUNCTION_SECRET_VALUES_BYTES:
raise ValueError(
"Function secret values exceed the "
f"{_MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"
)
canonical_values[name] = value
submission = self._request._known_dict()
if canonical_values:
submission["secret_values"] = canonical_values
return json.dumps(
submission,
ensure_ascii=False,
allow_nan=False,
sort_keys=True,
separators=(",", ":"),
)
def __call__(self, *args, **kwargs):
return self._function(*args, **kwargs)
@@ -1078,7 +988,6 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
@@ -1093,7 +1002,6 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
@@ -1124,10 +1032,7 @@ def udf(
conda_channels : sequence of str, optional
Conda channels in priority order; requires ``conda``.
env : mapping of str to str, optional
Non-secret environment variables. Use ``secrets`` for credentials.
secrets : sequence of str, optional
Names of secrets required by the callable. Supply their values separately
to ``create_function`` or ``create_function_async``.
Environment variables included in the Function definition.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
@@ -1149,15 +1054,11 @@ def udf(
Examples
--------
>>> from lancedb import udf
>>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"])
>>> @udf(pip=["numpy==2.2.0"])
... def score(value: float) -> float:
... return value * 2
>>> score(1.5)
3.0
>>> db.create_function( # doctest: +SKIP
... score, secrets={"MODEL_TOKEN": "user-secret-value"}
... )
"""
def decorate(target: Callable[..., Any]) -> UdfDefinition:
@@ -1168,7 +1069,6 @@ def udf(
output_schema=output_schema,
pip=tuple(pip),
env={} if env is None else env,
secrets=tuple(secrets),
python_version=python_version,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
+4 -12
View File
@@ -167,12 +167,6 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
return {"columns": projection}
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
if req.select_source_columns is not None:
return req.select_source_columns
return req.select
def _scanner_kwargs_for_query(
query: Query,
blob_mode: BlobMode,
@@ -2805,16 +2799,15 @@ class AsyncQueryBase(object):
req = self._inner.to_query_request()
schema = await self._table.schema()
projection = _query_request_projection(req)
self._blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
projection,
req.select,
with_row_id=self._with_row_id,
)
if not self._blob_auto_row_id:
self._blob_paths = ()
return
self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
self._inner.with_row_id()
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
@@ -3901,15 +3894,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
blob_paths: tuple[str, ...] = ()
if self._table is not None:
schema = await self._table.schema()
projection = _query_request_projection(req)
blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
projection,
req.select,
with_row_id=self._with_row_id,
)
if blob_auto_row_id:
blob_paths = tuple(
blob_v2_projection_sources(schema, projection).keys()
blob_v2_projection_sources(schema, req.select).keys()
)
self._blob_auto_row_id = blob_auto_row_id
self._blob_paths = blob_paths
+3 -10
View File
@@ -7,7 +7,7 @@ import json
import logging
from concurrent.futures import ThreadPoolExecutor
import sys
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Mapping, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from urllib.parse import urlparse
import warnings
@@ -742,15 +742,8 @@ class RemoteDBConnection(DBConnection):
return Job(self._conn.job(job_id))
@override
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
return Job(
LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
)
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
return Job(LOOP.run(self._conn.create_function_async(definition)))
@override
def get_function(self, name: str, *, version: str) -> FunctionVersion:
+4 -7
View File
@@ -36,7 +36,6 @@ from lancedb._lancedb import (
UpdateResult,
)
from lancedb.embeddings.base import EmbeddingFunctionConfig
from lancedb.expr import Expr
from lancedb.index import (
FTS,
BTree,
@@ -864,7 +863,7 @@ class RemoteTable(Table):
def update(
self,
where: Optional[Union[str, Expr]] = None,
where: Optional[str] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -875,11 +874,9 @@ class RemoteTable(Table):
Parameters
----------
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
+16 -26
View File
@@ -1744,7 +1744,7 @@ class Table(ABC):
@abstractmethod
def update(
self,
where: Optional[Union[str, Expr]] = None,
where: Optional[str] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -1759,11 +1759,9 @@ class Table(ABC):
Parameters
----------
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -1781,7 +1779,6 @@ class Table(ABC):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -1791,7 +1788,7 @@ class Table(ABC):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -3844,7 +3841,7 @@ class LanceTable(Table):
def update(
self,
where: Optional[Union[str, Expr]] = None,
where: Optional[str] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -3855,11 +3852,9 @@ class LanceTable(Table):
Parameters
----------
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -3877,7 +3872,6 @@ class LanceTable(Table):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -3887,7 +3881,7 @@ class LanceTable(Table):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -6001,7 +5995,7 @@ class AsyncTable:
self,
updates: Optional[Dict[str, Any]] = None,
*,
where: Optional[Union[str, Expr]] = None,
where: Optional[str] = None,
updates_sql: Optional[Dict[str, str]] = None,
) -> UpdateResult:
"""
@@ -6016,11 +6010,9 @@ class AsyncTable:
The updates to apply. The keys should be the name of the column to
update. The values should be the new values to assign. This is
required unless updates_sql is supplied.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. Only rows that satisfy this filter will
be updated.
where: str, optional
An SQL filter that controls which rows are updated. For example, 'x = 2'
or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated.
updates_sql: dict, optional
The updates to apply, expressed as SQL expression strings. The keys should
be column names. The values should be SQL expressions. These can be SQL
@@ -6038,14 +6030,13 @@ class AsyncTable:
--------
>>> import asyncio
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> async def demo_update():
... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]})
... db = await lancedb.connect_async("./.lancedb")
... table = await db.create_table("my_table", data)
... # x is [1, 2], vector is [[1, 2], [3, 4]]
... await table.update({"vector": [10, 10]}, where=col("x") == 2)
... await table.update({"vector": [10, 10]}, where="x = 2")
... # x is [1, 2], vector is [[1, 2], [10, 10]]
... await table.update(updates_sql={"x": "x + 1"})
... # x is [2, 3], vector is [[1, 2], [10, 10]]
@@ -6059,8 +6050,7 @@ class AsyncTable:
if updates is not None:
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
predicate = where.to_sql() if isinstance(where, Expr) else where
return await self._inner.update(updates_sql, predicate)
return await self._inner.update(updates_sql, where)
async def add_columns(
self,
+1 -72
View File
@@ -8,12 +8,7 @@ import pyarrow.compute as pc
import pytest
import lancedb
from lancedb._blob import (
blob_v2_projection_sources,
read_row_ids_from_hits,
stash_auto_row_ids,
)
from lancedb.expr import col
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
from lancedb.index import FTS
from lancedb.schema import blob_column_paths, blob_v2_column_paths
@@ -75,14 +70,6 @@ def test_blob_v2_column_paths_include_list_children():
]
def test_blob_v2_projection_sources_use_typed_column_name():
schema = pa.schema([lancedb.blob("blob")])
assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
"blob_alias": "blob"
}
def _legacy_v1_table(name):
db = lancedb.connect("memory:///")
schema = pa.schema(
@@ -179,20 +166,6 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
@pytest.mark.asyncio
async def test_async_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///typed_blob_projection")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
table = await db.create_table("typed_blob_projection", schema=schema)
await table.add([{"id": 1, "blob": b"alpha"}])
hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert blobs.to_pylist() == [b"alpha"]
def test_fetch_blobs_round_trip():
table = _blob_table(
"round_trip",
@@ -430,50 +403,6 @@ async def test_blob_v2_hybrid_fetch_blobs_async():
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
@pytest.mark.asyncio
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///hybrid_typed_blob")
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("text", pa.utf8()),
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
lancedb.blob("blob"),
]
)
table = await db.create_table("hybrid_typed_blob", schema=schema)
await table.add(
[
{
"id": 1,
"text": "hello alpha",
"vector": [1.0, 0.0],
"blob": b"alpha",
},
{
"id": 2,
"text": "hello beta",
"vector": [0.9, 0.1],
"blob": b"beta",
},
]
)
await table.create_index("text", config=FTS(with_position=False))
hits = await (
table.query()
.nearest_to([1.0, 0.0])
.nearest_to_text("hello")
.select({"blob_alias": col("blob")})
.limit(2)
.to_arrow()
)
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
def test_blob_file_seek_read_and_read_range():
payload = _identifiable_payload(1024)
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
+21 -21
View File
@@ -52,7 +52,7 @@ class TestExprConstruction:
def test_func(self):
e = func("lower", col("name"))
assert isinstance(e, Expr)
assert e.to_sql() == "lower(`name`)"
assert e.to_sql() == "lower(name)"
def test_func_unknown_raises(self):
with pytest.raises(Exception):
@@ -115,7 +115,7 @@ class TestExprOperators:
def test_and_operator(self):
e = (col("age") > lit(18)) & (col("status") == lit("active"))
assert isinstance(e, Expr)
assert e.to_sql() == "((age > 18) AND (`status` = 'active'))"
assert e.to_sql() == "((age > 18) AND (status = 'active'))"
def test_or_operator(self):
e = (col("a") == lit(1)) | (col("b") == lit(2))
@@ -166,7 +166,7 @@ class TestExprOperators:
def test_coerce_plain_str(self):
e = col("name") == "alice"
assert isinstance(e, Expr)
assert e.to_sql() == "(`name` = 'alice')"
assert e.to_sql() == "(name = 'alice')"
def test_reflexive_comparisons(self):
# 10 < col("age") swaps to col("age") > 10
@@ -198,85 +198,85 @@ class TestExprBytesLiteral:
def test_bytes_equality_expr_sql(self):
e = col("data") == lit(b"\xca\xfe")
assert e.to_sql() == "(`data` = X'CAFE')"
assert e.to_sql() == "(data = X'CAFE')"
def test_bytes_ne_expr_sql(self):
e = col("data") != lit(b"\xff")
assert e.to_sql() == "(`data` <> X'FF')"
assert e.to_sql() == "(data <> X'FF')"
def test_bytes_compound_expr_sql(self):
e = (col("data") == lit(b"\x01")) & (col("id") > lit(5))
assert e.to_sql() == "((`data` = X'01') AND (id > 5))"
assert e.to_sql() == "((data = X'01') AND (id > 5))"
def test_bytes_in_function_call(self):
# Regression test: binary literals inside scalar function calls
# used to fail because DataFusion's unparser does not support Binary
# scalars. Now handled via a placeholder-substitution rewrite.
e = func("contains", col("data"), lit(b"\xff"))
assert e.to_sql() == "contains(`data`, X'FF')"
assert e.to_sql() == "contains(data, X'FF')"
def test_bytes_in_not(self):
e = ~(col("data") == lit(b"\xff"))
assert e.to_sql() == "NOT (`data` = X'FF')"
assert e.to_sql() == "NOT (data = X'FF')"
class TestExprStringMethods:
def test_lower(self):
e = col("name").lower()
assert isinstance(e, Expr)
assert e.to_sql() == "lower(`name`)"
assert e.to_sql() == "lower(name)"
def test_upper(self):
e = col("name").upper()
assert isinstance(e, Expr)
assert e.to_sql() == "upper(`name`)"
assert e.to_sql() == "upper(name)"
def test_contains(self):
e = col("text").contains(lit("hello"))
assert isinstance(e, Expr)
assert e.to_sql() == "contains(`text`, 'hello')"
assert e.to_sql() == "contains(text, 'hello')"
def test_contains_with_str_coerce(self):
e = col("text").contains("hello")
assert isinstance(e, Expr)
assert e.to_sql() == "contains(`text`, 'hello')"
assert e.to_sql() == "contains(text, 'hello')"
def test_chained_lower_eq(self):
e = col("name").lower() == lit("alice")
assert isinstance(e, Expr)
assert e.to_sql() == "(lower(`name`) = 'alice')"
assert e.to_sql() == "(lower(name) = 'alice')"
class TestExprCast:
def test_cast_string(self):
e = col("id").cast("string")
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
assert e.to_sql() == "CAST(id AS VARCHAR)"
def test_cast_int32(self):
e = col("score").cast("int32")
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(score, 'Int32')"
assert e.to_sql() == "CAST(score AS INTEGER)"
def test_cast_float64(self):
e = col("val").cast("float64")
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(val, 'Float64')"
assert e.to_sql() == "CAST(val AS DOUBLE)"
def test_cast_pyarrow_type(self):
e = col("score").cast(pa.int32())
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(score, 'Int32')"
assert e.to_sql() == "CAST(score AS INTEGER)"
def test_cast_pyarrow_float64(self):
e = col("val").cast(pa.float64())
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(val, 'Float64')"
assert e.to_sql() == "CAST(val AS DOUBLE)"
def test_cast_pyarrow_string(self):
e = col("id").cast(pa.string())
assert isinstance(e, Expr)
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
assert e.to_sql() == "CAST(id AS VARCHAR)"
def test_cast_pyarrow_and_string_equivalent(self):
# pa.int32() and "int32" should produce equivalent SQL
@@ -597,14 +597,14 @@ class TestExprIsin:
def test_isin_strs(self):
assert (
col("status").isin(["active", "pending"]).to_sql()
== "`status` IN ('active', 'pending')"
== "status IN ('active', 'pending')"
)
def test_isin_coerces_and_mixes(self):
assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)"
def test_isin_empty(self):
assert col("id").isin([]).to_sql() == "false"
assert col("id").isin([]).to_sql() == "id IN ()"
def test_isin_filter(self, simple_table):
result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow()
@@ -37,21 +37,6 @@ def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def assert_no_secret_values(value):
if isinstance(value, dict):
for key, child in value.items():
assert key not in {
"secret_value",
"secret_values",
"resolved_secret",
"resolved_secrets",
}
assert_no_secret_values(child)
elif isinstance(value, list):
for child in value:
assert_no_secret_values(child)
def test_public_function_values_are_in_api_reference():
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
rendered = docs.read_text()
@@ -109,7 +94,6 @@ def test_function_version_identity_is_immutable_and_exact():
version = FunctionVersion.from_json(json.dumps(value))
assert version.name == "embed"
assert version.version == "fv_01K3EXACT"
assert version.required_secrets == ("HF_TOKEN",)
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
@@ -292,15 +276,6 @@ def test_refresh_result_rejects_non_u64_values(field):
RefreshColumnResult.from_json(json.dumps(value))
def test_canonical_client_values_contain_secret_names_only():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
canonical = json.loads(version.to_canonical_json())
assert canonical["required_secrets"] == ["HF_TOKEN"]
assert_no_secret_values(canonical)
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
@@ -19,13 +19,7 @@ import pyarrow as pa
import pytest
import lancedb
from lancedb.functions import (
_MAX_FUNCTION_SECRET_VALUE_BYTES,
_MAX_FUNCTION_SECRET_VALUES_BYTES,
FunctionRegistrationRequest,
UdfDefinition,
udf,
)
from lancedb.functions import UdfDefinition, udf
THRESHOLD = 20
_CACHE = None
@@ -45,28 +39,12 @@ FIXTURES = (
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
secrets=["API_TOKEN"],
python_version="3.12",
)
def normalize_score(value: float) -> float:
return value / 100.0
def _assert_no_secret_values(value):
if isinstance(value, dict):
for key, child in value.items():
assert key not in {
"secret_value",
"secret_values",
"resolved_secret",
"resolved_secrets",
}
_assert_no_secret_values(child)
elif isinstance(value, list):
for child in value:
_assert_no_secret_values(child)
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
assert isinstance(normalize_score, UdfDefinition)
assert normalize_score(25.0) == 0.25
@@ -81,8 +59,6 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
"kind": "scalar_to_arrow_batch",
"version": 1,
}
assert request["required_secrets"] == ["API_TOKEN"]
_assert_no_secret_values(request)
def _run_packaged(definition, *args):
@@ -396,7 +372,6 @@ def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
output_schema=None,
pip=(),
env={},
secrets=(),
python_version=None,
)
with pytest.raises(ValueError, match="binds that name to another value"):
@@ -551,54 +526,13 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
return value
def test_secret_names_are_canonical_and_disjoint_from_environment():
@udf(secrets=["Z_TOKEN", "A_TOKEN", "Z_TOKEN"])
def canonical_secrets(value: int) -> int:
return value
assert canonical_secrets.registration_request.required_secrets == (
"A_TOKEN",
"Z_TOKEN",
)
with pytest.raises(ValueError, match="must be disjoint"):
@udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"])
def overlapping(value: int) -> int:
return value
def test_declared_secret_api_still_requires_explicit_create_values():
@udf(secrets=["API_TOKEN"])
def declared_secret(value: int) -> int:
return value
with pytest.raises(ValueError, match="missing"):
declared_secret._submission_json(None)
submission = json.loads(
declared_secret._submission_json({"API_TOKEN": "explicit-secret"})
)
assert submission["required_secrets"] == ["API_TOKEN"]
assert submission["secret_values"] == {"API_TOKEN": "explicit-secret"}
def test_no_secrets_preserve_canonical_registration_shape():
@udf
def no_secrets(value: int) -> int:
return value
canonical = json.loads(no_secrets.registration_request.to_canonical_json())
assert "required_secrets" not in canonical
assert json.loads(no_secrets._submission_json(None)) == canonical
def test_local_function_catalog_operations_are_not_supported(tmp_path):
db = lancedb.connect(tmp_path)
message = "Function catalog operations are not supported by this database"
with pytest.raises(NotImplementedError, match=message):
db.create_function(normalize_score, secrets={"API_TOKEN": "value"})
db.create_function(normalize_score)
with pytest.raises(NotImplementedError, match=message):
db.create_function_async(normalize_score, secrets={"API_TOKEN": "value"})
db.create_function_async(normalize_score)
with pytest.raises(NotImplementedError, match=message):
db.get_function("normalize_score", version="fv_exact")
@@ -628,7 +562,6 @@ def _mock_remote_function_catalog():
"runtime": body["runtime"],
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": body.get("required_secrets", []),
"created_at": "2026-08-21T00:00:00Z",
}
response = {"job_id": "job-register"}
@@ -675,9 +608,7 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
registration = db.create_function_async(
normalize_score, secrets={"API_TOKEN": "secret-value"}
)
registration = db.create_function_async(normalize_score)
assert registration.id == "job-register"
created = registration.wait()
reopened = db.get_function("normalize_score", version=created.version)
@@ -686,18 +617,9 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
assert reopened.name == "normalize_score"
assert reopened.version == "fv_exact"
create_request = state["requests"][0][1]
expected = json.loads(normalize_score.registration_request.to_canonical_json())
expected["secret_values"] = {"API_TOKEN": "secret-value"}
assert create_request == expected
durable_request = FunctionRegistrationRequest.from_json(json.dumps(create_request))
assert not hasattr(durable_request, "secret_values")
assert "secret_values" not in json.loads(durable_request.to_canonical_json())
assert "secret_values" not in json.loads(
assert create_request == json.loads(
normalize_score.registration_request.to_canonical_json()
)
assert "secret-value" not in repr(normalize_score)
assert "secret-value" not in repr(normalize_score.registration_request)
assert not hasattr(created, "secret_values")
def test_blocking_remote_registration_returns_function_version():
@@ -708,9 +630,7 @@ def test_blocking_remote_registration_returns_function_version():
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
created = db.create_function(
normalize_score, secrets={"API_TOKEN": "blocking-secret"}
)
created = db.create_function(normalize_score)
assert created.name == "normalize_score"
assert created.version == "fv_exact"
@@ -718,121 +638,3 @@ def test_blocking_remote_registration_returns_function_version():
"/v1/functions/create",
"/v1/jobs/describe",
]
assert state["requests"][0][1]["secret_values"] == {"API_TOKEN": "blocking-secret"}
@pytest.mark.parametrize(
("secret_values", "error_type", "message"),
[
(None, ValueError, "missing"),
({}, ValueError, "missing"),
({"OTHER": "value"}, ValueError, "missing.*unexpected"),
({"API_TOKEN": ""}, ValueError, "non-empty"),
({"API_TOKEN": "bad\0value"}, ValueError, "NUL"),
({"API_TOKEN": 123}, TypeError, "must be a string"),
([("API_TOKEN", "value")], TypeError, "must be a mapping"),
],
)
def test_secret_values_are_validated_before_remote_request(
secret_values, error_type, message
):
with _mock_remote_function_catalog() as (host, state):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
with pytest.raises(error_type, match=message):
db.create_function_async(normalize_score, secrets=secret_values)
assert state["requests"] == []
@pytest.mark.parametrize(
"value",
[
"x" * _MAX_FUNCTION_SECRET_VALUE_BYTES,
"é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8"))),
],
ids=["ascii", "multibyte"],
)
def test_secret_value_accepts_exact_utf8_byte_limit(value):
submission = json.loads(normalize_score._submission_json({"API_TOKEN": value}))
assert submission["secret_values"]["API_TOKEN"] == value
assert len(value.encode("utf-8")) == _MAX_FUNCTION_SECRET_VALUE_BYTES
@pytest.mark.parametrize(
"value",
[
"x" * (_MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
"é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8")) + 1),
],
ids=["ascii", "multibyte"],
)
def test_secret_value_rejects_over_utf8_byte_limit_before_json_construction(
monkeypatch, value
):
def fail_if_json_construction_starts(self):
pytest.fail("oversized secret reached JSON construction")
monkeypatch.setattr(
FunctionRegistrationRequest, "_known_dict", fail_if_json_construction_starts
)
with pytest.raises(ValueError, match=r"exceeds the 65536-byte limit"):
normalize_score._submission_json({"API_TOKEN": value})
def test_secret_values_accept_exact_aggregate_utf8_byte_limit(monkeypatch):
names = tuple(f"SECRET_{index}" for index in range(8))
value = "é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8")))
values = {name: value for name in names}
monkeypatch.setattr(
normalize_score,
"_request",
normalize_score._request._copy(update={"required_secrets": names}),
)
submission = json.loads(normalize_score._submission_json(values))
assert submission["secret_values"] == values
assert sum(len(item.encode("utf-8")) for item in values.values()) == (
_MAX_FUNCTION_SECRET_VALUES_BYTES
)
def test_secret_values_reject_aggregate_over_limit_before_construction(monkeypatch):
names = tuple(f"SECRET_{index}" for index in range(9))
values = {name: "x" * _MAX_FUNCTION_SECRET_VALUE_BYTES for name in names}
monkeypatch.setattr(
normalize_score,
"_request",
normalize_score._request._copy(update={"required_secrets": names}),
)
def fail_if_json_construction_starts(self):
pytest.fail("oversized aggregate reached JSON construction")
monkeypatch.setattr(
FunctionRegistrationRequest, "_known_dict", fail_if_json_construction_starts
)
with pytest.raises(ValueError, match=r"exceed.*524288-byte request limit"):
normalize_score._submission_json(values)
@pytest.mark.asyncio
async def test_async_remote_registration_submits_secret_values_only_once():
with _mock_remote_function_catalog() as (host, state):
db = await lancedb.connect_async(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
registration = await db.create_function_async(
normalize_score, secrets={"API_TOKEN": "async-secret"}
)
created = await registration.wait()
assert state["requests"][0][1]["secret_values"] == {"API_TOKEN": "async-secret"}
assert not hasattr(created, "secret_values")
-15
View File
@@ -675,21 +675,6 @@ def test_distance_range(table: lancedb.table.Table):
assert res["_distance"].to_pylist() == [min_dist, max_dist]
@pytest.mark.parametrize("expression", ["1 - _distance", "1.0 - _distance"])
def test_select_arithmetic_with_distance(table, expression):
result = (
table.search([10, 10])
.select({"similarity": expression, "_distance": "_distance"})
.distance_type("cosine")
.to_arrow()
)
assert result.schema.field("similarity").type == pa.float32()
assert result["similarity"].to_pylist() == pytest.approx(
[1 - distance for distance in result["_distance"].to_pylist()]
)
@pytest.mark.asyncio
async def test_distance_range_async(table_async: AsyncTable):
q = [0, 0]
-158
View File
@@ -11,7 +11,6 @@ import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal
from time import sleep
from typing import List
from unittest.mock import patch
@@ -337,21 +336,6 @@ async def test_update_async(mem_db_async: AsyncConnection):
assert await table.count_rows("id == 10") == 1
@pytest.mark.asyncio
async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection):
values = ["5", "4.66e-84", "it's"]
table = await mem_db_async.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = await table.update({"result": value}, where=col("field") == value)
assert update_res.rows_updated == 1
assert (await table.to_arrow())["result"].to_pylist() == values
def test_create_table(mem_db: DBConnection):
schema = pa.schema(
{
@@ -2359,148 +2343,6 @@ def test_update(mem_db: DBConnection):
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
def test_update_expr_filter_literals(mem_db: DBConnection):
values = ["5", "4.66e-84", "it's"]
table = mem_db.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = table.update(where=col("field") == value, values={"result": value})
assert update_res.rows_updated == 1
assert table.to_arrow()["result"].to_pylist() == values
def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
low = Decimal("1.234567890123456789")
high = Decimal("1.234567890123456790")
decimal_schema = pa.schema(
[("val", pa.decimal128(19, 18)), ("result", pa.string())]
)
decimal_table = mem_db.create_table(
"update_expr_decimal",
pa.table(
{"val": [low, high], "result": ["old", "old"]},
schema=decimal_schema,
),
)
predicate = col("val") < lit(high)
assert decimal_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
keyword_table = mem_db.create_table(
"update_expr_keyword", [{"null": 1, "result": "old"}]
)
predicate = col("null") == 1
assert keyword_table.search().where(predicate).to_arrow().num_rows == 1
result = keyword_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
empty_in_table = mem_db.create_table(
"update_expr_empty_in", [{"id": 1, "result": "old"}]
)
predicate = col("id").isin([])
assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0
result = empty_in_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
marker = "__lancedb_binary_placeholder_0__"
binary_schema = pa.schema(
[("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())]
)
binary_table = mem_db.create_table(
"update_expr_binary",
pa.table(
{
"payload": [b"\x01", b"\x02"],
"text": ["other", marker],
"result": ["old", "old"],
},
schema=binary_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker)
assert binary_table.search().where(predicate).to_arrow().num_rows == 2
result = binary_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
nonfinite_table = mem_db.create_table(
"update_expr_nonfinite",
[{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
)
predicate = col("x") < float("inf")
assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
result = nonfinite_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
float16_table = mem_db.create_table(
"update_expr_float16",
[{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
)
predicate = col("x").cast(pa.float16()) < 2.0
assert float16_table.search().where(predicate).to_arrow().num_rows == 1
result = float16_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
string_cast_table = mem_db.create_table(
"update_expr_string_cast",
[{"x": 1, "result": "old"}, {"x": 2, "result": "old"}],
)
predicate = col("x").cast("string") == "1"
assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1
result = string_cast_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
quoted_identifier_schema = pa.schema(
[("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())]
)
quoted_identifier_table = mem_db.create_table(
"update_expr_quoted_identifier",
pa.table(
{"payload": [b"\x01"], "odd'name": [1], "result": ["old"]},
schema=quoted_identifier_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1)
assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1
result = quoted_identifier_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
decimal256_schema = pa.schema(
[("val", pa.decimal256(40, 2)), ("result", pa.string())]
)
decimal256_table = mem_db.create_table(
"update_expr_decimal256",
pa.table(
{
"val": [Decimal("1.00"), Decimal("3.00")],
"result": ["old", "old"],
},
schema=decimal256_schema,
),
)
predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal256_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
binary_empty_table = mem_db.create_table(
"update_expr_binary_empty",
pa.table(
{"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
),
)
predicate = (col("payload") == lit(b"\x01")).isin([])
assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
assert predicate.to_sql() == "false"
result = binary_empty_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
def test_update_with_arrow_scalar(mem_db: DBConnection):
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
table = mem_db.create_table("my_table", schema=schema)
-8
View File
@@ -130,14 +130,6 @@ impl PyExpr {
// ── utilities ────────────────────────────────────────────────────────────
/// Return the referenced column name for a bare column expression.
fn column_name(&self) -> Option<String> {
match &self.0 {
DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()),
_ => None,
}
}
/// Render the expression as a SQL string (useful for debugging).
fn to_sql(&self) -> PyResult<String> {
lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string()))
-23
View File
@@ -1,7 +1,6 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
@@ -326,7 +325,6 @@ pub struct PyQueryRequest {
pub filter: Option<PyQueryFilter>,
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
pub select: PySelect,
pub select_source_columns: Option<HashMap<String, String>>,
pub fast_search: Option<bool>,
pub with_row_id: Option<bool>,
pub use_lsm: Option<bool>,
@@ -357,7 +355,6 @@ impl From<AnyQuery> for PyQueryRequest {
full_text_search: query_request
.full_text_search
.map(|fts| PyLanceDB(fts.query)),
select_source_columns: PySelect::source_columns(&query_request.select),
select: PySelect(query_request.select),
fast_search: Some(query_request.fast_search),
with_row_id: Some(query_request.with_row_id),
@@ -383,7 +380,6 @@ impl From<AnyQuery> for PyQueryRequest {
offset: vector_query.base.offset,
filter: vector_query.base.filter.map(PyQueryFilter),
full_text_search: None,
select_source_columns: PySelect::source_columns(&vector_query.base.select),
select: PySelect(vector_query.base.select),
fast_search: Some(vector_query.base.fast_search),
with_row_id: Some(vector_query.base.with_row_id),
@@ -416,25 +412,6 @@ impl From<AnyQuery> for PyQueryRequest {
#[derive(Clone)]
pub struct PySelect(Select);
impl PySelect {
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
match select {
Select::Expr(pairs) => Some(
pairs
.iter()
.filter_map(|(output, expr)| match expr {
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
Some((output.clone(), column.name.clone()))
}
_ => None,
})
.collect(),
),
_ => None,
}
}
}
impl<'py> IntoPyObject<'py> for PySelect {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;
+4 -120
View File
@@ -157,7 +157,7 @@ mod tests {
use datafusion_common::ScalarValue;
let expr = col("data").eq(lit(ScalarValue::Binary(Some(vec![0xca, 0xfe]))));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "(`data` = X'CAFE')");
assert_eq!(sql, "(data = X'CAFE')");
}
#[test]
@@ -167,7 +167,7 @@ mod tests {
let int_expr = col("id").gt(lit(5i64));
let combined = bin_expr.and(int_expr);
let sql = expr_to_sql_string(&combined).unwrap();
assert_eq!(sql, "((`data` = X'01') AND (id > 5))");
assert_eq!(sql, "((data = X'01') AND (id > 5))");
}
#[test]
@@ -185,7 +185,7 @@ mod tests {
// serialized correctly (regression test for placeholder rewrite path).
let expr = contains(col("data"), lit(ScalarValue::Binary(Some(vec![0xff]))));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "contains(`data`, X'FF')");
assert_eq!(sql, "contains(data, X'FF')");
}
#[test]
@@ -196,7 +196,7 @@ mod tests {
.eq(lit(ScalarValue::Binary(Some(vec![0xab, 0xcd]))))
.not();
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "NOT (`data` = X'ABCD')");
assert_eq!(sql, "NOT (data = X'ABCD')");
}
#[test]
@@ -206,122 +206,6 @@ mod tests {
assert!(sql.contains("IN"), "expected IN in: {}", sql);
}
#[test]
fn test_empty_is_in() {
let expr = is_in(col("id"), vec![]);
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
}
#[test]
fn test_empty_is_in_discards_binary_children() {
use datafusion_common::ScalarValue;
let expr = is_in(
col("payload").eq(lit(ScalarValue::Binary(Some(vec![0x01])))),
vec![],
);
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
}
#[test]
fn test_keyword_identifier() {
let expr = col("null").eq(lit(1i64));
assert_eq!(expr_to_sql_string(&expr).unwrap(), "(`null` = 1)");
}
#[test]
fn test_decimal_literal_preserves_type() {
use datafusion_common::ScalarValue;
let expr = col("val").lt(lit(ScalarValue::Decimal128(
Some(1_234_567_890_123_456_790),
19,
18,
)));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(
sql,
"(val < arrow_cast('1.234567890123456790', 'Decimal128(19, 18)'))"
);
}
#[test]
fn test_non_finite_float_literal_preserves_type() {
let expr = col("x").lt(lit(f64::INFINITY));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(x < arrow_cast('inf', 'Float64'))"
);
}
#[test]
fn test_cast_uses_arrow_type_name() {
let string = expr_cast(col("x"), DataType::Utf8);
assert_eq!(
expr_to_sql_string(&string).unwrap(),
"arrow_cast(x, 'Utf8')"
);
let int32 = expr_cast(col("x"), DataType::Int32);
assert_eq!(
expr_to_sql_string(&int32).unwrap(),
"arrow_cast(x, 'Int32')"
);
let expr = expr_cast(col("x"), DataType::Float16).lt(lit(2.0));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(arrow_cast(x, 'Float16') < 2.0)"
);
let decimal = expr_cast(lit("2.00"), DataType::Decimal256(40, 2));
assert_eq!(
expr_to_sql_string(&decimal).unwrap(),
"arrow_cast('2.00', 'Decimal256(40, 2)')"
);
}
#[test]
fn test_binary_placeholder_does_not_rewrite_user_string() {
use datafusion_common::ScalarValue;
let marker = "__lancedb_binary_placeholder_0__";
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.or(col("text").eq(lit(marker)));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"((payload = X'01') OR (`text` = '__lancedb_binary_placeholder_0__'))"
);
}
#[test]
fn test_binary_binding_skips_quoted_identifiers() {
use datafusion_common::ScalarValue;
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.and(col("odd'name").eq(lit(1i64)))
.and(col("odd`'name").eq(lit(2i64)));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(((payload = X'01') AND (`odd'name` = 1)) AND (`odd``'name` = 2))"
);
}
#[test]
fn test_binary_placeholder_collision_search_is_linear() {
use datafusion_common::ScalarValue;
let collision_shaped = format!("__lancedb_binary_placeholder_0__{}", "_".repeat(64_000));
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.and(col("text").eq(lit(collision_shaped.clone())));
let sql = expr_to_sql_string(&expr).unwrap();
assert!(sql.contains("X'01'"));
assert!(sql.contains(&format!("'{collision_shaped}'")));
}
#[test]
fn test_multiple_binary_literals() {
use datafusion_common::ScalarValue;
+42 -220
View File
@@ -1,24 +1,13 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::{
any::TypeId,
collections::{HashMap, HashSet},
};
use std::any::TypeId;
use arrow_array::types::{
Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
};
use arrow_schema::DataType;
use datafusion_common::ScalarValue;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_expr::Expr;
use datafusion_functions::core::expr_fn::{
arrow_cast as datafusion_arrow_cast, arrow_try_cast as datafusion_arrow_try_cast,
};
use datafusion_sql::sqlparser::{
dialect::{Dialect as SqlParserDialect, GenericDialect},
keywords::ALL_KEYWORDS,
tokenizer::{Token, Tokenizer},
};
use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect};
@@ -38,13 +27,11 @@ struct LanceSqlDialect;
impl UnparserDialect for LanceSqlDialect {
fn identifier_quote_style(&self, identifier: &str) -> Option<char> {
let identifier_upper = identifier.to_ascii_uppercase();
let needs_quote =
(identifier_upper != "ID" && ALL_KEYWORDS.contains(&identifier_upper.as_str()))
|| identifier.chars().any(|c| c.is_ascii_uppercase())
|| !identifier.chars().enumerate().all(|(i, c)| {
c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit())
});
let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase())
|| !identifier
.chars()
.enumerate()
.all(|(i, c)| c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit()));
if needs_quote { Some('`') } else { None }
}
}
@@ -113,128 +100,24 @@ fn bytes_to_hex_sql(bytes: &[u8]) -> String {
format!("X'{hex}'")
}
fn string_literals(expr: &Expr) -> HashSet<String> {
let mut literals = HashSet::new();
/// Returns true if *expr* contains a `Binary` or `LargeBinary` scalar literal
/// anywhere in its subtree. DataFusion's SQL unparser cannot serialize those
/// variants, so we route such expressions through a placeholder-substitution
/// path that emits SQL `X'...'` byte-string literals.
fn has_binary_literal(expr: &Expr) -> bool {
let mut found = false;
let _ = expr.apply(&mut |e: &Expr| {
if let Expr::Literal(
ScalarValue::Utf8(Some(value))
| ScalarValue::LargeUtf8(Some(value))
| ScalarValue::Utf8View(Some(value)),
_,
) = e
{
literals.insert(value.clone());
}
Ok(TreeNodeRecursion::Continue)
});
literals
}
fn typed_string_literal(value: String, data_type: DataType) -> Expr {
datafusion_arrow_cast(
Expr::Literal(ScalarValue::Utf8(Some(value)), None),
Expr::Literal(ScalarValue::Utf8(Some(data_type.to_string())), None),
)
}
fn next_binary_placeholder(user_strings: &HashSet<String>, next_id: &mut usize) -> String {
loop {
let placeholder = format!("{BINARY_PLACEHOLDER_PREFIX}{}__", *next_id);
*next_id += 1;
if !user_strings.contains(&placeholder) {
return placeholder;
}
}
}
fn bind_binary_literals(
sql: &str,
mut bindings: HashMap<String, Vec<u8>>,
) -> crate::Result<String> {
let bytes = sql.as_bytes();
let mut output = Vec::with_capacity(bytes.len());
let mut index = 0;
// Walk SQL string tokens once. Placeholders are plain, unescaped string
// literals, so this remains linear even when user strings are large or
// deliberately resemble the placeholder prefix.
while index < bytes.len() {
if bytes[index] == b'`' {
let identifier_start = index;
index += 1;
let mut identifier_end = None;
while index < bytes.len() {
if bytes[index] == b'`' {
if index + 1 < bytes.len() && bytes[index + 1] == b'`' {
index += 2;
} else {
index += 1;
identifier_end = Some(index);
break;
}
} else {
index += 1;
}
}
let Some(identifier_end) = identifier_end else {
return Err(crate::Error::InvalidInput {
message: "unterminated identifier while binding binary literal".to_string(),
});
};
output.extend_from_slice(&bytes[identifier_start..identifier_end]);
continue;
}
if bytes[index] != b'\'' {
output.push(bytes[index]);
index += 1;
continue;
}
let literal_start = index;
index += 1;
let content_start = index;
let mut escaped = false;
let mut content_end = None;
while index < bytes.len() {
if bytes[index] == b'\'' {
if index + 1 < bytes.len() && bytes[index + 1] == b'\'' {
escaped = true;
index += 2;
} else {
content_end = Some(index);
index += 1;
break;
}
} else {
index += 1;
}
}
let Some(content_end) = content_end else {
return Err(crate::Error::InvalidInput {
message: "unterminated string while binding binary literal".to_string(),
});
};
let placeholder = &sql[content_start..content_end];
if !escaped && let Some(value) = bindings.remove(placeholder) {
output.extend_from_slice(bytes_to_hex_sql(&value).as_bytes());
if matches!(
e,
Expr::Literal(ScalarValue::Binary(_) | ScalarValue::LargeBinary(_), _)
) {
found = true;
Ok(TreeNodeRecursion::Stop)
} else {
output.extend_from_slice(&bytes[literal_start..index]);
Ok(TreeNodeRecursion::Continue)
}
}
if !bindings.is_empty() {
return Err(crate::Error::InvalidInput {
message: "failed to bind binary literal while serializing expression".to_string(),
});
}
String::from_utf8(output).map_err(|e| crate::Error::InvalidInput {
message: format!("failed to bind binary literal: {e}"),
})
});
found
}
fn run_unparser(expr: &Expr) -> crate::Result<String> {
@@ -247,37 +130,25 @@ fn run_unparser(expr: &Expr) -> crate::Result<String> {
}
pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
// DataFusion's unparser needs a few adaptations before its SQL can be
// reparsed by Lance without changing the typed expression's semantics:
//
// * decimal literals need an explicit cast to preserve precision and scale;
// * casts need exact Arrow type names rather than SQL type aliases;
// * an empty IN list is valid in DataFusion but invalid SQL;
// * binary literals are unsupported by the unparser and need placeholders.
// Eliminate empty membership expressions before visiting their children.
// Otherwise a discarded binary child could leave behind a stale binding.
// Fast path: no binary literals — DataFusion's unparser handles everything.
if !has_binary_literal(expr) {
return run_unparser(expr);
}
// Slow path: DataFusion's unparser cannot serialize `Binary`/`LargeBinary`
// scalars, so we rewrite each one to a unique string-literal placeholder,
// let the unparser do the rest of the work, then substitute the SQL
// `X'...'` byte-string literal back in. This keeps the operator/function
// serialization logic centralized in DataFusion and works for every
// expression node type the unparser supports.
let mut bindings: Vec<Vec<u8>> = Vec::new();
let rewritten = expr
.clone()
.transform(|e: Expr| match e {
Expr::InList(in_list) if in_list.list.is_empty() => Ok(Transformed::yes(
Expr::Literal(ScalarValue::Boolean(Some(in_list.negated)), None),
)),
other => Ok(Transformed::no(other)),
})
.map_err(|e| crate::Error::InvalidInput {
message: format!("failed to rewrite expression: {e}"),
})?
.data;
let user_strings = string_literals(&rewritten);
let mut next_placeholder_id = 0;
let mut binary_bindings = HashMap::new();
let rewritten = rewritten
.transform(|e: Expr| match e {
Expr::Literal(ScalarValue::Binary(Some(bytes)), m)
| Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), m) => {
let placeholder = next_binary_placeholder(&user_strings, &mut next_placeholder_id);
binary_bindings.insert(placeholder.clone(), bytes);
let placeholder = format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, bindings.len());
bindings.push(bytes);
Ok(Transformed::yes(Expr::Literal(
ScalarValue::Utf8(Some(placeholder)),
m,
@@ -287,57 +158,6 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
| Expr::Literal(ScalarValue::LargeBinary(None), m) => {
Ok(Transformed::yes(Expr::Literal(ScalarValue::Null, m)))
}
Expr::Literal(ScalarValue::Decimal32(Some(value), precision, scale), _m) => {
let value = Decimal32Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal32(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), _m) => {
let value = Decimal64Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal64(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), _m) => {
let value = Decimal128Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal128(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), _m) => {
let value = Decimal256Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal256(precision, scale),
)))
}
Expr::Literal(ScalarValue::Float16(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float16)),
),
Expr::Literal(ScalarValue::Float32(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float32)),
),
Expr::Literal(ScalarValue::Float64(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float64)),
),
Expr::Cast(cast) => Ok(Transformed::yes(datafusion_arrow_cast(
*cast.expr,
Expr::Literal(
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
None,
),
))),
Expr::TryCast(cast) => Ok(Transformed::yes(datafusion_arrow_try_cast(
*cast.expr,
Expr::Literal(
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
None,
),
))),
other => Ok(Transformed::no(other)),
})
.map_err(|e| crate::Error::InvalidInput {
@@ -345,12 +165,14 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
})?
.data;
let sql = run_unparser(&rewritten)?;
if binary_bindings.is_empty() {
Ok(sql)
} else {
bind_binary_literals(&sql, binary_bindings)
let mut sql = run_unparser(&rewritten)?;
for (i, bytes) in bindings.iter().enumerate() {
// The unparser quotes string literals with single quotes, so the
// placeholder appears as `'__lancedb_binary_placeholder_<i>__'`.
let quoted = format!("'{}{}__'", BINARY_PLACEHOLDER_PREFIX, i);
sql = sql.replace(&quoted, &bytes_to_hex_sql(bytes));
}
Ok(sql)
}
#[cfg(test)]
+4 -309
View File
@@ -5,9 +5,9 @@
//! backend-neutral terminal result of a computed-column refresh.
//!
//! This module contains client/wire values only. Catalog persistence,
//! environment bake, secret resolution, and execution are owned by Sophon.
//! environment bake, and execution are owned by Sophon.
use std::collections::{BTreeMap, BTreeSet};
use std::collections::BTreeMap;
use serde::de::{self, DeserializeOwned};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -15,16 +15,6 @@ use serde_json::Value;
use crate::{Error, Result};
// Keep these byte limits aligned with Sophon's Function submission validation.
pub(crate) const MAX_FUNCTION_SECRET_VALUE_BYTES: usize = 64 * 1024;
const MAX_FUNCTION_SECRET_VALUES_BYTES: usize = 512 * 1024;
fn is_portable_environment_name(name: &str) -> bool {
let mut bytes = name.bytes();
matches!(bytes.next(), Some(b'A'..=b'Z' | b'a'..=b'z' | b'_'))
&& bytes.all(|byte| matches!(byte, b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_'))
}
fn invalid_json(error: impl std::fmt::Display) -> Error {
Error::InvalidInput {
message: format!("invalid remote Function JSON: {error}"),
@@ -208,11 +198,6 @@ pub struct PythonEnvironmentSpec {
}
/// Reproducible Python runtime definition understood by Sophon.
///
/// `env` contains non-secret values. Secret values are submission-only in the
/// client model and do not become part of this public runtime identity;
/// [`FunctionVersion::required_secrets`] contains names only. Sophon persists
/// submitted values separately in the private execution artifact.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum PythonRuntimeSpec {
@@ -254,7 +239,7 @@ impl PythonRuntimeSpec {
}
}
/// Non-secret environment variables, or `None` for an unknown kind.
/// Environment variables, or `None` for an unknown kind.
pub fn env(&self) -> Option<&BTreeMap<String, String>> {
match self {
Self::Python { env, .. } => Some(env),
@@ -339,8 +324,6 @@ pub struct FunctionVersion {
runtime: PythonRuntimeSpec,
runtime_digest: String,
environment_digest: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
required_secrets: Vec<String>,
created_at: String,
}
@@ -373,12 +356,6 @@ impl FunctionVersion {
&self.environment_digest
}
/// Required secret names. Resolved values exist only in Sophon's private
/// execution artifact and worker launch path.
pub fn required_secrets(&self) -> &[String] {
&self.required_secrets
}
pub fn created_at(&self) -> &str {
&self.created_at
}
@@ -420,115 +397,12 @@ pub struct FunctionArtifactRequest {
}
/// Stable request envelope for remote immutable Function registration.
///
/// Secret values are submission-only in the client model. Sophon persists them
/// in the database-scoped private execution artifact; returned
/// [`FunctionVersion`] and Job metadata contain only
/// [`Self::required_secrets`] names. Debug formatting always redacts values.
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FunctionRegistrationRequest {
pub name: String,
pub artifact: FunctionArtifactRequest,
pub signature: FunctionSignature,
pub runtime: PythonRuntimeSpec,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub required_secrets: Vec<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub secret_values: BTreeMap<String, String>,
}
impl FunctionRegistrationRequest {
pub(crate) fn validate_secret_values(&self) -> Result<()> {
let mut required = BTreeSet::new();
for name in &self.required_secrets {
if !is_portable_environment_name(name) {
return Err(Error::InvalidInput {
message: format!(
"Function secret name {name:?} must be a portable environment variable name"
),
});
}
if !required.insert(name) {
return Err(Error::InvalidInput {
message: format!("Function required_secrets contains duplicate name {name:?}"),
});
}
}
if let PythonRuntimeSpec::Python { env, .. } = &self.runtime
&& let Some(name) = required.iter().find(|name| env.contains_key(**name))
{
return Err(Error::InvalidInput {
message: format!(
"Function runtime env and secret names must be disjoint: {name:?}"
),
});
}
let provided = self.secret_values.keys().collect::<BTreeSet<_>>();
if required != provided {
return Err(Error::InvalidInput {
message: "Function secret_values keys must exactly match required_secrets"
.to_string(),
});
}
let mut total_bytes = 0usize;
for (name, value) in &self.secret_values {
if value.is_empty() {
return Err(Error::InvalidInput {
message: format!("Function secret {name:?} value must be non-empty"),
});
}
if value.contains('\0') {
return Err(Error::InvalidInput {
message: format!("Function secret {name:?} value must not contain NUL"),
});
}
if value.len() > MAX_FUNCTION_SECRET_VALUE_BYTES {
return Err(Error::InvalidInput {
message: format!(
"Function secret {name:?} value exceeds the \
{MAX_FUNCTION_SECRET_VALUE_BYTES}-byte limit"
),
});
}
total_bytes =
total_bytes
.checked_add(value.len())
.ok_or_else(|| Error::InvalidInput {
message: "Function secret values exceed the request byte limit".to_string(),
})?;
}
if total_bytes > MAX_FUNCTION_SECRET_VALUES_BYTES {
return Err(Error::InvalidInput {
message: format!(
"Function secret values exceed the \
{MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"
),
});
}
Ok(())
}
}
impl std::fmt::Debug for FunctionRegistrationRequest {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let secret_values = self
.secret_values
.keys()
.map(|name| (name, "[REDACTED]"))
.collect::<BTreeMap<_, _>>();
formatter
.debug_struct("FunctionRegistrationRequest")
.field("name", &self.name)
.field("artifact", &self.artifact)
.field("signature", &self.signature)
.field("runtime", &self.runtime)
.field("required_secrets", &self.required_secrets)
.field("secret_values", &secret_values)
.finish()
}
}
impl_json!(FunctionRegistrationRequest);
@@ -713,185 +587,6 @@ impl RefreshColumnResult {
impl_json!(RefreshColumnResult);
#[cfg(test)]
mod secret_value_tests {
use super::{
FunctionRegistrationRequest, MAX_FUNCTION_SECRET_VALUE_BYTES,
MAX_FUNCTION_SECRET_VALUES_BYTES, PythonRuntimeSpec,
};
use crate::Error;
fn request() -> FunctionRegistrationRequest {
FunctionRegistrationRequest::from_json(include_str!(
"../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
))
.unwrap()
}
#[test]
fn validates_secret_name_and_value_invariants() {
let missing = request();
assert!(matches!(
missing.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("exactly match")
));
let mut empty = request();
empty
.secret_values
.insert("API_TOKEN".to_string(), String::new());
assert!(matches!(
empty.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("non-empty")
));
let mut nul = request();
nul.secret_values
.insert("API_TOKEN".to_string(), "before\0after".to_string());
assert!(matches!(
nul.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("NUL")
));
let mut unexpected = request();
unexpected
.secret_values
.insert("OTHER".to_string(), "value".to_string());
assert!(matches!(
unexpected.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("exactly match")
));
}
#[test]
fn rejects_invalid_duplicate_and_overlapping_secret_declarations() {
let mut invalid_name = request();
invalid_name.required_secrets = vec!["BAD=NAME".to_string()];
invalid_name
.secret_values
.insert("BAD=NAME".to_string(), "secret".to_string());
let mut duplicate = request();
duplicate.required_secrets = vec!["API_TOKEN".to_string(), "API_TOKEN".to_string()];
duplicate
.secret_values
.insert("API_TOKEN".to_string(), "secret".to_string());
let mut overlap = request();
overlap
.secret_values
.insert("API_TOKEN".to_string(), "secret".to_string());
if let PythonRuntimeSpec::Python { env, .. } = &mut overlap.runtime {
env.insert("API_TOKEN".to_string(), "public".to_string());
}
assert!(matches!(
invalid_name.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("portable environment variable")
));
assert!(matches!(
duplicate.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("duplicate")
));
assert!(matches!(
overlap.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("must be disjoint")
));
}
#[test]
fn enforces_portable_secret_name_boundaries() {
for name in ["A", "_", "A0_"] {
let mut request = request();
request.required_secrets = vec![name.to_string()];
request
.secret_values
.insert(name.to_string(), "secret".to_string());
request.validate_secret_values().unwrap();
}
for name in ["", "0TOKEN", "BAD-NAME", "TÖKEN"] {
let mut request = request();
request.required_secrets = vec![name.to_string()];
request
.secret_values
.insert(name.to_string(), "secret".to_string());
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message })
if message.contains("portable environment variable")
));
}
}
#[test]
fn accepts_exact_secret_value_utf8_byte_limit() {
for value in [
"x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES),
"é".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES / "é".len()),
] {
assert_eq!(value.len(), MAX_FUNCTION_SECRET_VALUE_BYTES);
let mut request = request();
request.secret_values.insert("API_TOKEN".to_string(), value);
request.validate_secret_values().unwrap();
}
}
#[test]
fn rejects_secret_value_over_utf8_byte_limit() {
for value in [
"x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
"é".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES / "é".len() + 1),
] {
assert!(value.len() > MAX_FUNCTION_SECRET_VALUE_BYTES);
let mut request = request();
request.secret_values.insert("API_TOKEN".to_string(), value);
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("65536-byte limit")
));
}
}
#[test]
fn rejects_aggregate_secret_value_bytes_over_server_limit() {
let mut request = request();
request.required_secrets = (0..9).map(|index| format!("SECRET_{index}")).collect();
request.secret_values = request
.required_secrets
.iter()
.map(|name| (name.clone(), "x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES)))
.collect();
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message })
if message.contains(&format!("{MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"))
));
}
#[test]
fn accepts_exact_aggregate_secret_value_byte_limit() {
let mut request = request();
request.required_secrets = (0..8).map(|index| format!("SECRET_{index}")).collect();
request.secret_values = request
.required_secrets
.iter()
.map(|name| (name.clone(), "x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES)))
.collect();
assert_eq!(
request
.secret_values
.values()
.map(String::len)
.sum::<usize>(),
MAX_FUNCTION_SECRET_VALUES_BYTES
);
request.validate_secret_values().unwrap();
}
}
#[cfg(test)]
mod conda_environment_tests {
use super::PythonEnvironmentSpec;
+15 -100
View File
@@ -7,7 +7,6 @@ use reqwest::{
Body, Request, RequestBuilder, Response,
header::{HeaderMap, HeaderValue},
};
use serde_json::Value;
use std::{collections::HashMap, future::Future, str::FromStr, sync::Arc, time::Duration};
use crate::error::{Error, Result};
@@ -15,60 +14,6 @@ use crate::remote::db::RemoteOptions;
use crate::remote::retry::{ResolvedRetryConfig, RetryCounter};
const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
const REDACTED_JSON_VALUE: &str = "[REDACTED]";
const SUPPRESSED_JSON_BODY: &str = "[JSON BODY SUPPRESSED]";
fn is_sensitive_json_field(name: &str) -> bool {
name.to_ascii_lowercase().contains("secret")
}
fn redact_sensitive_json_fields(value: &mut Value) {
match value {
Value::Object(fields) => {
for (name, child) in fields {
if is_sensitive_json_field(name) {
*child = Value::String(REDACTED_JSON_VALUE.to_string());
} else {
redact_sensitive_json_fields(child);
}
}
}
Value::Array(values) => values.iter_mut().for_each(redact_sensitive_json_fields),
_ => {}
}
}
fn redacted_json_body(request: &Request) -> Option<String> {
let body = request.body()?.as_bytes()?;
let mut value = serde_json::from_slice(body).ok()?;
redact_sensitive_json_fields(&mut value);
serde_json::to_string(&value).ok()
}
fn request_log_message(request: &Request, request_id: &str) -> String {
let prefix = format!(
"Sending request_id={}: {} {}",
request_id,
request.method(),
request.url()
);
let content_type = request
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next());
if content_type.is_some_and(|value| value.eq_ignore_ascii_case("application/json")) {
// Never format the raw Request here: its Debug representation is not a
// redaction boundary and may include the original body. If the JSON body
// cannot be structurally parsed, suppress it instead of logging raw bytes.
let body = redacted_json_body(request).unwrap_or_else(|| SUPPRESSED_JSON_BODY.to_string());
format!("{prefix} with body {body}")
} else {
// Method and URL are sufficient request context. Raw Request formatting
// may expose headers or a non-JSON body, so it is never a logging fallback.
prefix
}
}
/// Configuration for TLS/mTLS settings.
#[derive(Clone, Debug)]
@@ -894,9 +839,22 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
}
}
pub(crate) fn log_request(&self, request: &Request, request_id: &str) {
pub(crate) fn log_request(&self, request: &Request, request_id: &String) {
if log::log_enabled!(log::Level::Debug) {
debug!("{}", request_log_message(request, request_id));
let content_type = request
.headers()
.get("content-type")
.map(|v| v.to_str().unwrap());
if content_type == Some("application/json") {
let body = request.body().as_ref().unwrap().as_bytes().unwrap();
let body = String::from_utf8_lossy(body);
debug!(
"Sending request_id={}: {:?} with body {}",
request_id, request, body
);
} else {
debug!("Sending request_id={}: {:?}", request_id, request);
}
}
}
@@ -1119,49 +1077,6 @@ mod tests {
ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
}
#[test]
fn test_request_log_message_redacts_secrets_and_never_formats_raw_requests() {
const SECRET_SENTINEL: &str = "udf-secret-log-sentinel-7e4e";
const MALFORMED_SENTINEL: &str = "malformed-secret-log-sentinel-b652";
const NON_JSON_SENTINEL: &str = "non-json-secret-log-sentinel-7fd1";
let request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.json(&serde_json::json!({
"name": "uses_secret",
"nested": {
"secret_values": {"OPENAI_API_KEY": SECRET_SENTINEL},
"safe": "visible-value"
}
}))
.build()
.unwrap();
let log_message = request_log_message(&request, "valid-json");
let malformed_request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.header("content-type", "application/json; charset=utf-8")
.body(format!(r#"{{"secret_values":"{MALFORMED_SENTINEL}""#))
.build()
.unwrap();
let malformed_log_message = request_log_message(&malformed_request, "malformed-json");
let non_json_request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.header("content-type", "text/plain")
.body(NON_JSON_SENTINEL)
.build()
.unwrap();
let non_json_log_message = request_log_message(&non_json_request, "non-json");
assert!(log_message.contains("visible-value"));
assert!(log_message.contains(REDACTED_JSON_VALUE));
assert!(!log_message.contains(SECRET_SENTINEL));
assert!(malformed_log_message.contains(SUPPRESSED_JSON_BODY));
assert!(!malformed_log_message.contains(MALFORMED_SENTINEL));
assert!(!non_json_log_message.contains(NON_JSON_SENTINEL));
}
#[test]
fn test_timeout_config_default() {
let config = TimeoutConfig::default();
+2 -33
View File
@@ -554,7 +554,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
&self,
request: FunctionRegistrationRequest,
) -> Result<Job<FunctionVersion>> {
request.validate_secret_values()?;
let req = self.client.post("/v1/functions/create").json(&request);
let (request_id, response) = self.client.send(req).await?;
let response = self.client.check_response(&request_id, response).await?;
@@ -2643,8 +2642,7 @@ mod tests {
);
const FUNCTION_JOB: &str =
include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json");
let mut expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
expected["secret_values"] = serde_json::json!({"API_TOKEN": "secret-value"});
let expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
let conn = Connection::new_with_handler(move |request| match request.url().path() {
"/v1/functions/create" => {
assert_eq!(request.method(), &reqwest::Method::POST);
@@ -2662,10 +2660,7 @@ mod tests {
.unwrap(),
path => panic!("unexpected path: {path}"),
});
let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
request
.secret_values
.insert("API_TOKEN".to_string(), "secret-value".to_string());
let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
let job = conn.create_function_async(request).await.unwrap();
assert_eq!(job.id(), Some("job-function-1"));
let version = job.wait().await.unwrap();
@@ -2673,32 +2668,6 @@ mod tests {
assert_eq!(version.version(), "fv_01K3EXACT");
}
#[tokio::test]
async fn test_create_function_async_validates_secrets_before_serialization_and_send() {
const REQUEST: &str = include_str!(
"../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
);
let sends = Arc::new(AtomicUsize::new(0));
let sends_ref = sends.clone();
let conn = Connection::new_with_handler(move |_| {
sends_ref.fetch_add(1, Ordering::SeqCst);
http::Response::builder().status(500).body("").unwrap()
});
let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
request.secret_values.insert(
"API_TOKEN".to_string(),
"x".repeat(crate::function::MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
);
let error = conn.create_function_async(request).await.unwrap_err();
assert!(matches!(
error,
Error::InvalidInput { message } if message.contains("65536-byte limit")
));
assert_eq!(sends.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_get_function_requires_and_sends_exact_version() {
const VERSION: &str = include_str!(
+191
View File
@@ -1431,4 +1431,195 @@ mod lsm_tests {
"LSM vector search must rank the memtable row first"
);
}
#[tokio::test]
async fn lsm_cosine_distance_scale_and_mixed_tier_ordering() {
use arrow::array::{FixedSizeListBuilder, Float32Builder};
use arrow::datatypes::Float32Type;
use crate::index::Index;
use crate::index::vector::IvfPqIndexBuilder;
const DIM: usize = 8;
const N: usize = 256;
fn normalized_vector(state: &mut u64) -> Vec<f32> {
let mut vector = (0..DIM)
.map(|_| {
*state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1);
((*state >> 32) as u32 as f32 / u32::MAX as f32) * 2.0 - 1.0
})
.collect::<Vec<_>>();
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
vector.iter_mut().for_each(|value| *value /= norm);
vector
}
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new(
"vec",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
DIM as i32,
),
false,
),
]));
let make_batch = |rows: Vec<(i64, Vec<f32>)>| {
let ids = rows.iter().map(|(id, _)| *id).collect::<Vec<_>>();
let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32);
for (_, vector) in &rows {
vectors.values().append_slice(vector);
vectors.append(true);
}
RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors.finish())],
)
.unwrap()
};
let first_result = |batches: &[RecordBatch]| {
let batch = &batches[0];
let id = batch["id"].as_primitive::<Int64Type>().value(0);
let distance = batch["_distance"].as_primitive::<Float32Type>().value(0);
(id, distance)
};
let mut state = 42;
let base_rows = (0..N)
.map(|id| (id as i64, normalized_vector(&mut state)))
.collect::<Vec<_>>();
let query = normalized_vector(&mut state);
let dir = tempdir().unwrap();
let conn = connect(dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let base = make_batch(base_rows);
let reader: Box<dyn RecordBatchReader + Send> =
Box::new(RecordBatchIterator::new(vec![Ok(base)], schema.clone()));
let table = conn
.create_table("cosine_lsm", reader)
.execute()
.await
.unwrap();
table.set_unenforced_primary_key(["id"]).await.unwrap();
table
.create_index(
&["vec"],
Index::IvfPq(
IvfPqIndexBuilder::default()
.distance_type(crate::DistanceType::Cosine)
.num_partitions(1)
.num_sub_vectors(1),
),
)
.name("vec_cosine".to_string())
.execute()
.await
.unwrap();
table
.set_lsm_write_spec(
LsmWriteSpec::unsharded().with_maintained_indexes(vec!["vec_cosine".to_string()]),
)
.await
.unwrap();
let base_only = table
.query()
.nearest_to(query.as_slice())
.unwrap()
.limit(1)
.use_lsm(false)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let (base_id, public_distance) = first_result(&base_only);
let lsm = table
.query()
.nearest_to(query.as_slice())
.unwrap()
.limit(1)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let (lsm_id, lsm_distance) = first_result(&lsm);
assert_eq!(lsm_id, base_id);
assert!(
(lsm_distance - public_distance).abs() < 1e-5,
"LSM cosine distance {lsm_distance} did not use the public scale {public_distance}"
);
// Add an exact memtable result whose distance lies between the public ANN
// score and its doubled internal score. Correctly normalized plans still
// rank the ANN row first; mixed units would incorrectly rank this row first.
assert!(public_distance > 0.0 && public_distance < 4.0 / 3.0);
let memtable_distance = public_distance * 1.5;
let cosine_similarity = 1.0 - memtable_distance;
let mut orthogonal = normalized_vector(&mut state);
let projection = orthogonal
.iter()
.zip(&query)
.map(|(left, right)| left * right)
.sum::<f32>();
for (value, query_value) in orthogonal.iter_mut().zip(&query) {
*value -= projection * query_value;
}
let norm = orthogonal
.iter()
.map(|value| value * value)
.sum::<f32>()
.sqrt();
orthogonal.iter_mut().for_each(|value| *value /= norm);
let sine = (1.0 - cosine_similarity * cosine_similarity).sqrt();
let memtable_vector = query
.iter()
.zip(&orthogonal)
.map(|(query_value, orthogonal_value)| {
cosine_similarity * query_value + sine * orthogonal_value
})
.collect::<Vec<_>>();
let mut merge = table.merge_insert(&[]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
let memtable = make_batch(vec![(N as i64, memtable_vector)]);
merge
.execute(Box::new(RecordBatchIterator::new(
vec![Ok(memtable)],
schema,
)))
.await
.unwrap();
let mixed = table
.query()
.nearest_to(query.as_slice())
.unwrap()
.limit(1)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let (mixed_id, mixed_distance) = first_result(&mixed);
assert_eq!(
mixed_id, base_id,
"mixed LSM tiers must compare ANN and exact distances in public units"
);
assert!((mixed_distance - public_distance).abs() < 1e-5);
}
}
+414 -5
View File
@@ -1,7 +1,10 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use std::{
collections::{HashSet, VecDeque},
sync::Arc,
};
mod lsm;
@@ -17,15 +20,23 @@ use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder};
use arrow::datatypes::{Float32Type, UInt8Type};
use arrow_array::Array;
use arrow_schema::{DataType, Schema};
use datafusion_common::{Column, DataFusionError, SchemaError};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_common::{Column, DataFusionError, ScalarValue, SchemaError};
use datafusion_expr::Operator;
use datafusion_physical_expr::expressions::{BinaryExpr, Column as PhysicalColumn, Literal};
use datafusion_physical_plan::PhysicalExpr;
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::union::UnionExec;
use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary};
use lance::dataset::mem_wal::DatasetMemWalExt;
use lance::dataset::scanner::DatasetRecordBatchStream;
use lance::dataset::scanner::Scanner;
use lance::index::DatasetIndexInternalExt;
use lance::io::exec::ANNIvfSubIndexExec;
use lance_datafusion::exec::{analyze_plan as lance_analyze_plan, execute_plan};
use lance_index::metrics::NoOpMetricsCollector;
use lance_index::vector::{DIST_COL, quantizer::QuantizationType};
use lance_linalg::distance::DistanceType as LanceDistanceType;
use lance_namespace::LanceNamespace;
use lance_namespace::models::{
QueryTableRequest as NsQueryTableRequest, QueryTableRequestColumns,
@@ -375,10 +386,30 @@ pub async fn create_plan(
scanner.order_by(Some(order_by.clone()))?;
}
scanner
let mut plan = scanner
.create_plan()
.await
.map_err(|error| enrich_lance_field_not_found(error, schema))
.map_err(|error| enrich_lance_field_not_found(error, schema))?;
let normalized_l2_indices = normalized_l2_ann_indices(plan.as_ref()).await?;
if !normalized_l2_indices.is_empty() {
// Rebuild only the affected ANN nodes with internal normalized squared-L2
// bounds. Exact branches keep the public cosine bounds from `plan`.
let internal_plan = if query.lower_bound.is_some() || query.upper_bound.is_some() {
scanner.distance_range(
query.lower_bound.map(|bound| bound / COSINE_ANN_SCALE),
query.upper_bound.map(|bound| bound / COSINE_ANN_SCALE),
);
scanner
.create_plan()
.await
.map_err(|error| enrich_lance_field_not_found(error, schema))?
} else {
plan.clone()
};
plan = normalize_ann_branches(plan, internal_plan, &normalized_l2_indices)?;
}
Ok(plan)
}
/// Replace DataFusion's top-level field candidates with qualified leaf paths.
@@ -470,6 +501,184 @@ fn leaf_field_paths(schema: &Schema) -> Vec<String> {
//Helper functions below
const COSINE_ANN_SCALE: f32 = 0.5;
/// Find ANN index segments whose scores use normalized squared L2 for cosine search.
///
/// Cosine PQ/SQ/RQ indices normalize their vectors and use squared L2 internally. This
/// preserves ranking, but squared L2 over unit vectors is twice the cosine distance. Flat
/// cosine indices calculate cosine directly, so they are not included.
async fn normalized_l2_ann_indices(plan: &dyn ExecutionPlan) -> Result<HashSet<String>> {
let mut ann_plans = Vec::new();
find_ann_plans(plan, &mut ann_plans);
let mut checked = HashSet::new();
let mut normalized_l2 = HashSet::new();
for ann in ann_plans {
if ann.query().metric_type != Some(LanceDistanceType::Cosine) {
continue;
}
for index in ann.indices() {
let uuid = index.uuid.to_string();
if !checked.insert(uuid.clone()) {
continue;
}
let vector_index = ann
.dataset()
.open_vector_index(&ann.query().column, &index.uuid, &NoOpMetricsCollector)
.await?;
let (_, quantization_type) = vector_index.sub_index_type();
if matches!(
quantization_type,
QuantizationType::Product | QuantizationType::Scalar | QuantizationType::Rabit
) {
normalized_l2.insert(uuid);
}
}
}
Ok(normalized_l2)
}
/// Normalize affected ANN outputs before their parent plan nodes consume them.
///
/// This is used by planners that do not support distance ranges, such as the MemWAL
/// LSM planner. The standard scanner path rebuilds a second plan when it also needs
/// to translate range bounds, then calls [`normalize_ann_branches`] directly.
pub(super) async fn normalize_cosine_ann_branches(
plan: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
let normalized_l2_indices = normalized_l2_ann_indices(plan.as_ref()).await?;
if normalized_l2_indices.is_empty() {
return Ok(plan);
}
normalize_ann_branches(plan.clone(), plan, &normalized_l2_indices)
}
fn find_ann_plans<'a>(plan: &'a dyn ExecutionPlan, ann_plans: &mut Vec<&'a ANNIvfSubIndexExec>) {
if let Some(ann) = plan.downcast_ref::<ANNIvfSubIndexExec>() {
ann_plans.push(ann);
}
for child in plan.children() {
find_ann_plans(child.as_ref(), ann_plans);
}
}
fn collect_ann_plans(
plan: &Arc<dyn ExecutionPlan>,
ann_plans: &mut VecDeque<Arc<dyn ExecutionPlan>>,
) {
if plan.downcast_ref::<ANNIvfSubIndexExec>().is_some() {
ann_plans.push_back(plan.clone());
return;
}
for child in plan.children() {
collect_ann_plans(child, ann_plans);
}
}
/// Replace normalized-L2 ANN nodes with equivalent nodes that use internal bounds, then
/// convert their output to the public cosine scale before any generic plan node consumes it.
fn normalize_ann_branches(
public_plan: Arc<dyn ExecutionPlan>,
internal_plan: Arc<dyn ExecutionPlan>,
normalized_l2_indices: &HashSet<String>,
) -> Result<Arc<dyn ExecutionPlan>> {
let mut internal_ann_plans = VecDeque::new();
collect_ann_plans(&internal_plan, &mut internal_ann_plans);
let normalized =
replace_ann_branches(public_plan, &mut internal_ann_plans, normalized_l2_indices)?;
if !internal_ann_plans.is_empty() {
return Err(Error::Runtime {
message: "internal and public vector plans contained different ANN branches"
.to_string(),
});
}
Ok(normalized)
}
fn replace_ann_branches(
public_plan: Arc<dyn ExecutionPlan>,
internal_ann_plans: &mut VecDeque<Arc<dyn ExecutionPlan>>,
normalized_l2_indices: &HashSet<String>,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(public_ann) = public_plan.downcast_ref::<ANNIvfSubIndexExec>() {
let internal_plan = internal_ann_plans
.pop_front()
.ok_or_else(|| Error::Runtime {
message: "internal vector plan was missing an ANN branch".to_string(),
})?;
let internal_ann = internal_plan
.downcast_ref::<ANNIvfSubIndexExec>()
.expect("collected only ANN plans");
let same_indices = public_ann
.indices()
.iter()
.map(|index| &index.uuid)
.eq(internal_ann.indices().iter().map(|index| &index.uuid));
if public_ann.query().column != internal_ann.query().column
|| public_ann.query().metric_type != internal_ann.query().metric_type
|| !same_indices
{
return Err(Error::Runtime {
message: "internal and public vector plans had mismatched ANN branches".to_string(),
});
}
let normalized_count = public_ann
.indices()
.iter()
.filter(|index| normalized_l2_indices.contains(&index.uuid.to_string()))
.count();
if normalized_count == 0 {
return Ok(public_plan);
}
if normalized_count != public_ann.indices().len() {
return Err(Error::Runtime {
message: "one ANN branch mixed public and normalized-L2 distance scales"
.to_string(),
});
}
return scale_distance_column(internal_plan, COSINE_ANN_SCALE);
}
let children = public_plan
.children()
.into_iter()
.cloned()
.map(|child| replace_ann_branches(child, internal_ann_plans, normalized_l2_indices))
.collect::<Result<Vec<_>>>()?;
Ok(with_new_children_if_necessary(public_plan, children)?)
}
fn scale_distance_column(
plan: Arc<dyn ExecutionPlan>,
scale: f32,
) -> Result<Arc<dyn ExecutionPlan>> {
let schema = plan.schema();
if schema.column_with_name(DIST_COL).is_none() {
return Ok(plan);
}
let expressions: Vec<(Arc<dyn PhysicalExpr>, String)> = schema
.fields()
.iter()
.enumerate()
.map(|(index, field)| {
let column: Arc<dyn PhysicalExpr> = Arc::new(PhysicalColumn::new(field.name(), index));
let expression = if field.name() == DIST_COL {
let scale: Arc<dyn PhysicalExpr> =
Arc::new(Literal::new(ScalarValue::Float32(Some(scale))));
Arc::new(BinaryExpr::new(column, Operator::Multiply, scale))
as Arc<dyn PhysicalExpr>
} else {
column
};
(expression, field.name().clone())
})
.collect();
Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?))
}
// Take many execution plans and map them into a single plan that adds
// a query_index column and unions them.
pub(crate) fn create_multi_vector_plan(
@@ -1455,6 +1664,206 @@ mod tests {
);
}
#[tokio::test]
async fn test_cosine_pq_distance_uses_public_cosine_scale() {
use arrow_array::{Int32Array, RecordBatch, types::Float32Type};
use arrow_schema::{DataType, Field, Schema};
use crate::connect;
use crate::index::{Index, vector::IvfPqIndexBuilder};
fn normalized_vector(state: &mut u64, dimension: usize) -> Vec<f32> {
let mut vector = (0..dimension)
.map(|_| {
*state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1);
((*state >> 32) as u32 as f32 / u32::MAX as f32) * 2.0 - 1.0
})
.collect::<Vec<_>>();
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
vector.iter_mut().for_each(|value| *value /= norm);
vector
}
fn distances(batches: &[RecordBatch]) -> Vec<f32> {
batches
.iter()
.flat_map(|batch| {
batch[DIST_COL]
.as_primitive::<Float32Type>()
.values()
.to_vec()
})
.collect()
}
let conn = connect("memory://").execute().await.unwrap();
let dimension = 8;
let num_rows = 256;
let mut state = 42;
let values = (0..num_rows)
.flat_map(|_| normalized_vector(&mut state, dimension))
.collect::<Vec<_>>();
let query_vector = normalized_vector(&mut state, dimension);
let vectors = Arc::new(fixed_size_list_array(values, dimension as i32));
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("vector", vectors.data_type().clone(), false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from_iter_values(0..num_rows)), vectors],
)
.unwrap();
let table = conn
.create_table("test_cosine_pq_distance", batch)
.execute()
.await
.unwrap();
table
.create_index(
&["vector"],
Index::IvfPq(
IvfPqIndexBuilder::default()
.distance_type(crate::DistanceType::Cosine)
.num_partitions(1)
.num_sub_vectors(1),
),
)
.execute()
.await
.unwrap();
let approximate = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(5)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let refined = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(5)
.refine_factor(1)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let approximate_distances = distances(&approximate);
let refined_distances = distances(&refined);
assert_eq!(approximate_distances.len(), refined_distances.len());
for (approximate, refined) in approximate_distances.iter().zip(&refined_distances) {
assert!(
(approximate - refined).abs() < 1e-5,
"approximate cosine distance {approximate} did not use the public scale; refined distance was {refined}"
);
}
// Distance range bounds are public cosine distances too. Lance applies them to
// internal ANN scores, so the planner must translate the bounds before execution.
let nearest = approximate_distances[0];
let ranged = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(1)
.distance_range(Some(nearest - 1e-5), Some(nearest + 1e-5))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let ranged_distances = distances(&ranged);
assert_eq!(ranged_distances.len(), 1);
assert!((ranged_distances[0] - nearest).abs() < 1e-5);
let refined_ranged = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(1)
.refine_factor(1)
.distance_range(None, Some(nearest + 1e-5))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
distances(&refined_ranged).len(),
1,
"refinement must not apply public cosine bounds to internal ANN scores"
);
let aliased = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(1)
.select(Select::dynamic(&[("aliased_distance", "_distance")]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let batch = &aliased[0];
let aliased_distance = batch["aliased_distance"]
.as_primitive::<Float32Type>()
.value(0);
let public_distance = batch[DIST_COL].as_primitive::<Float32Type>().value(0);
assert!(
(aliased_distance - public_distance).abs() < 1e-5,
"distance aliases and auto-projected distances must use the same public scale"
);
// Appended rows take an exact fallback branch. Its public range filter must stay
// independent of the translated ANN bounds before both branches are merged.
let mut orthogonal = normalized_vector(&mut state, dimension);
let projection = orthogonal
.iter()
.zip(&query_vector)
.map(|(left, right)| left * right)
.sum::<f32>();
for (value, query_value) in orthogonal.iter_mut().zip(&query_vector) {
*value -= projection * query_value;
}
let norm = orthogonal
.iter()
.map(|value| value * value)
.sum::<f32>()
.sqrt();
orthogonal.iter_mut().for_each(|value| *value /= norm);
let appended_vectors = Arc::new(fixed_size_list_array(orthogonal, dimension as i32));
let appended = RecordBatch::try_new(
schema,
vec![Arc::new(Int32Array::from(vec![num_rows])), appended_vectors],
)
.unwrap();
table.add(appended).execute().await.unwrap();
let mixed = table
.vector_search(query_vector.as_slice())
.unwrap()
.limit(5)
.distance_range(None, Some(nearest + 1e-5))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let mixed_distances = distances(&mixed);
assert_eq!(mixed_distances.len(), 1);
assert!((mixed_distances[0] - nearest).abs() < 1e-5);
}
#[tokio::test]
async fn test_create_plan_applies_approx_mode_to_ann_query() {
use arrow_array::RecordBatch;
+4
View File
@@ -130,6 +130,10 @@ pub(super) async fn create_lsm_plan(
.await?
};
// Normalize cosine ANN arms before LSM merge and sort nodes compare their
// distances with exact SSTable and memtable arms.
let plan = super::normalize_cosine_ann_branches(plan).await?;
// Lance appends the primary-key columns internally for dedup and keeps them in
// the output; drop the ones the user did not request so the projection matches.
restore_projection(plan, &query, &pk_columns)
@@ -20,25 +20,6 @@ fn job_result(name: &str) -> Value {
serde_json::from_str::<Value>(&fixture(name)).expect("remote Job fixture")["result"].clone()
}
fn assert_no_secret_values(value: &Value) {
match value {
Value::Object(values) => {
for (key, value) in values {
assert!(
!matches!(
key.as_str(),
"secret_value" | "secret_values" | "resolved_secret" | "resolved_secrets"
),
"client canonical value must not model resolved secret material"
);
assert_no_secret_values(value);
}
}
Value::Array(values) => values.iter().for_each(assert_no_secret_values),
_ => {}
}
}
#[test]
fn function_version_job_result_matches_shared_canonical_golden() {
let result = job_result("remote_function_job.json");
@@ -47,7 +28,6 @@ fn function_version_job_result_matches_shared_canonical_golden() {
assert_eq!(version.name(), "embed");
assert_eq!(version.version(), "fv_01K3EXACT");
assert_eq!(version.runtime_digest(), "sha256:runtime");
assert_eq!(version.required_secrets(), &["HF_TOKEN"]);
assert_eq!(
version.to_canonical_json().expect("canonical JSON"),
fixture("remote_function_version.canonical.json").trim()
@@ -162,21 +142,3 @@ fn floating_point_application_literals_are_rejected_consistently() {
.contains("floating-point Function literals")
);
}
#[test]
fn canonical_client_values_contain_secret_names_only() {
let result = job_result("remote_function_job.json");
let version = FunctionVersion::from_json(&result.to_string()).expect("FunctionVersion result");
let canonical: Value = serde_json::from_str(
&version
.to_canonical_json()
.expect("canonical FunctionVersion"),
)
.expect("canonical JSON");
assert_eq!(
canonical["required_secrets"],
serde_json::json!(["HF_TOKEN"])
);
assert_no_secret_values(&canonical);
}
@@ -6,7 +6,6 @@ use std::path::PathBuf;
use lancedb::Error;
use lancedb::function::FunctionRegistrationRequest;
use serde_json::Value;
fn fixture(name: &str) -> String {
let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
@@ -15,25 +14,6 @@ fn fixture(name: &str) -> String {
fs::read_to_string(path).expect("fixture must be readable")
}
fn assert_no_secret_values(value: &Value) {
match value {
Value::Object(values) => {
for (key, value) in values {
assert!(
!matches!(
key.as_str(),
"secret_value" | "secret_values" | "resolved_secret" | "resolved_secrets"
),
"registration requests must not model resolved secret material"
);
assert_no_secret_values(value);
}
}
Value::Array(values) => values.iter().for_each(assert_no_secret_values),
_ => {}
}
}
#[test]
fn registration_request_matches_shared_canonical_golden() {
let request = FunctionRegistrationRequest::from_json(&fixture(
@@ -42,33 +22,10 @@ fn registration_request_matches_shared_canonical_golden() {
.expect("registration request");
assert_eq!(request.name, "normalize_score");
assert_eq!(request.artifact.adapter.kind, "scalar_to_arrow_batch");
assert_eq!(request.required_secrets, ["API_TOKEN"]);
assert!(request.secret_values.is_empty());
assert_eq!(
request.to_canonical_json().expect("canonical request"),
fixture("remote_function_registration_request.canonical.json").trim()
);
let value: Value =
serde_json::from_str(&request.to_canonical_json().expect("canonical request"))
.expect("request JSON");
assert_no_secret_values(&value);
}
#[test]
fn registration_request_serializes_secret_values_but_redacts_debug_output() {
let mut value: Value =
serde_json::from_str(&fixture("remote_function_registration_request.json")).unwrap();
value["secret_values"] = serde_json::json!({"API_TOKEN": "secret-plaintext"});
let request = FunctionRegistrationRequest::from_json(&value.to_string()).unwrap();
assert_eq!(request.secret_values["API_TOKEN"], "secret-plaintext");
let canonical = request.to_canonical_json().unwrap();
assert!(canonical.contains("secret-plaintext"));
let debug = format!("{request:?}");
assert!(debug.contains("API_TOKEN"));
assert!(debug.contains("[REDACTED]"));
assert!(!debug.contains("secret-plaintext"));
}
#[tokio::test]
@@ -24,7 +24,6 @@
},
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": ["HF_TOKEN"],
"created_at": "2026-08-21T00:00:00Z"
},
"future_job": {"trace_id": "trace-1"}
@@ -1 +1 @@
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"name":"normalize_score","required_secrets":["API_TOKEN"],"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"name":"normalize_score","runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
@@ -39,8 +39,5 @@
"env": {
"MODE": "test"
}
},
"required_secrets": [
"API_TOKEN"
]
}
}
@@ -1 +1 @@
{"artifact":{"digest":"sha256:code","entrypoint":"embed","kind":"python_callable"},"created_at":"2026-08-21T00:00:00Z","environment_digest":"sha256:environment","name":"embed","required_secrets":["HF_TOKEN"],"runtime":{"env":{"TOKENIZERS_PARALLELISM":"false"},"environment":{"kind":"pip","packages":["sentence-transformers>=3"]},"kind":"python","python_version":"3.12"},"runtime_digest":"sha256:runtime","signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list<float32>","kind":"scalar","nullable":false}},"version":"fv_01K3EXACT"}
{"artifact":{"digest":"sha256:code","entrypoint":"embed","kind":"python_callable"},"created_at":"2026-08-21T00:00:00Z","environment_digest":"sha256:environment","name":"embed","runtime":{"env":{"TOKENIZERS_PARALLELISM":"false"},"environment":{"kind":"pip","packages":["sentence-transformers>=3"]},"kind":"python","python_version":"3.12"},"runtime_digest":"sha256:runtime","signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list<float32>","kind":"scalar","nullable":false}},"version":"fv_01K3EXACT"}