mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-31 02:18:27 +00:00
Compare commits
13 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8b14e2fe63 | |||
| d24b2dcacc | |||
| 2deccf21cf | |||
| ead4d27bfc | |||
| 5153e5a023 | |||
| 79f626b09e | |||
| d55446f71f | |||
| ae81d73563 | |||
| 8b7e13b0c6 | |||
| 676c5b7315 | |||
| 5093f37559 | |||
| 4f5c55888b | |||
| f95d4f583d |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.38.0-beta.10"
|
||||
current_version = "0.38.0-beta.11"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
Generated
+3
-3
@@ -5402,7 +5402,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5490,7 +5490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5515,7 +5515,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.38.0-beta.10</version>
|
||||
<version>0.38.0-beta.11</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -1292,6 +1292,18 @@ abstract updateFieldMetadata(updates): Promise<UpdateFieldMetadataResult>
|
||||
|
||||
Update per-field (column) metadata.
|
||||
|
||||
The following keys are treated specially, by convention, and should be
|
||||
used when appropriate:
|
||||
|
||||
- `lancedb:description`: for a human-readable description of a field.
|
||||
- `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
- `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
`feature_v2` might be in the same logical column.
|
||||
- `lancedb:status`: for status options (`production`, `candidate`,
|
||||
`deprecated`, `archived`) to designate the current life cycle state of
|
||||
this column.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **updates**: [`FieldMetadataUpdate`](../interfaces/FieldMetadataUpdate.md)[]
|
||||
|
||||
@@ -17,7 +17,8 @@ metadata: Record<string, null | string>;
|
||||
```
|
||||
|
||||
Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
default; a value of `null` deletes that key.
|
||||
default; a value of `null` deletes that key. See
|
||||
[Table.updateFieldMetadata](../classes/Table.md#updatefieldmetadata) for the conventional `lancedb:*` keys.
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.10</version>
|
||||
<version>0.38.0-beta.11</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+1
-1
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.10</version>
|
||||
<version>0.38.0-beta.11</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -6,6 +6,9 @@ import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
|
||||
import {
|
||||
Field as CurrentField,
|
||||
LargeBinary as CurrentLargeBinary,
|
||||
Schema as CurrentSchema,
|
||||
Vector as CurrentVector,
|
||||
convertToTable,
|
||||
tableFromIPC as currentTableFromIPC,
|
||||
@@ -36,6 +39,24 @@ function sampleRecords(): Array<Record<string, any>> {
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
it("preserves field metadata from a provided schema", async function () {
|
||||
const jsonMetadata = new Map([["ARROW:extension:name", "lance.json"]]);
|
||||
const schema = new CurrentSchema([
|
||||
new CurrentField("meta", new CurrentLargeBinary(), true, jsonMetadata),
|
||||
]);
|
||||
|
||||
const table = makeArrowTable(
|
||||
[{ meta: Buffer.from(JSON.stringify({ source: "test" })) }],
|
||||
{ schema },
|
||||
);
|
||||
|
||||
expect(table.schema.fields[0].metadata).toEqual(jsonMetadata);
|
||||
|
||||
const roundTripped = currentTableFromIPC(await fromTableToBuffer(table));
|
||||
expect(roundTripped.schema.fields[0].metadata).toEqual(jsonMetadata);
|
||||
});
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
"Arrow",
|
||||
(
|
||||
|
||||
@@ -187,6 +187,58 @@ describe("embedding functions", () => {
|
||||
const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
|
||||
expect(vector0).toEqual([1, 2, 3]);
|
||||
});
|
||||
it("should append multiple Python embeddings with the same alias", async () => {
|
||||
@register("python-mock")
|
||||
// biome-ignore lint/correctness/noUnusedVariables: the decorator registers this class
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((value) =>
|
||||
value === "hello world" ? [1, 2, 3] : [4, 5, 6],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const metadata = new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
'[{"source_column":"text1","vector_column":"vector1","name":"python-mock","model":{}},{"source_column":"text2","vector_column":"vector2","name":"python-mock","model":{}}]',
|
||||
],
|
||||
]);
|
||||
const schema = new Schema(
|
||||
[
|
||||
new Field("text1", new Utf8(), true),
|
||||
new Field("text2", new Utf8(), true),
|
||||
new Field(
|
||||
"vector1",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
new Field(
|
||||
"vector2",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
],
|
||||
metadata,
|
||||
);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = await db.createEmptyTable("test", schema);
|
||||
await table.add([{ text1: "hello world", text2: "goodbye world" }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector1))).toEqual([1, 2, 3]);
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector2))).toEqual([4, 5, 6]);
|
||||
});
|
||||
|
||||
it("should append generated vectors to a non-nullable schema", async () => {
|
||||
@register("non_nullable_schema_test")
|
||||
|
||||
@@ -3561,6 +3561,27 @@ describe("when creating an empty table", () => {
|
||||
expect((actualSchema.fields[1].type as Float64).precision).toBe(2);
|
||||
});
|
||||
|
||||
it("can add and query JSON data", async () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32(), true),
|
||||
new Field(
|
||||
"meta",
|
||||
new Utf8(),
|
||||
true,
|
||||
new Map([["ARROW:extension:name", "arrow.json"]]),
|
||||
),
|
||||
]);
|
||||
const table = await con.createEmptyTable("json", schema);
|
||||
const meta = JSON.stringify({ x: 1 });
|
||||
|
||||
await table.add([{ id: 1, meta }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(rows).toHaveLength(1);
|
||||
expect(rows[0].id).toBe(1);
|
||||
expect(rows[0].meta).toBe(meta);
|
||||
});
|
||||
|
||||
it("can create an empty table from schema that specifies field types by name", async () => {
|
||||
const schemaLike = {
|
||||
fields: [
|
||||
|
||||
@@ -170,7 +170,7 @@ test("basic table examples", async () => {
|
||||
// --8<-- [end:create_index]
|
||||
|
||||
// --8<-- [start:delete_rows]
|
||||
await tbl.delete('item = "fizz"');
|
||||
await tbl.delete("item = 'fizz'");
|
||||
// --8<-- [end:delete_rows]
|
||||
|
||||
// --8<-- [start:drop_table]
|
||||
|
||||
@@ -727,11 +727,11 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* Add a query vector to the search
|
||||
*
|
||||
* This method can be called multiple times to add multiple query vectors
|
||||
* to the search. If multiple query vectors are added, then they will be searched
|
||||
* in parallel, and the results will be concatenated. A column called `query_index`
|
||||
* will be added to indicate the index of the query vector that produced the result.
|
||||
*
|
||||
* Performance wise, this is equivalent to running multiple queries concurrently.
|
||||
* to the search. A column called `query_index` will be added to indicate the index
|
||||
* of the query vector that produced the result. Flat searches share one table scan
|
||||
* across the query vectors, avoiding the scan and memory amplification of running
|
||||
* multiple queries concurrently. Indexed searches may still perform per-vector
|
||||
* index work.
|
||||
*/
|
||||
addQueryVector(vector: IntoVector): VectorQuery {
|
||||
if (vector instanceof Promise) {
|
||||
|
||||
@@ -406,10 +406,11 @@ function matchingFields(fields: Field[], tree: FieldTree): Field[] {
|
||||
field.name,
|
||||
new Struct(matchingFields(struct.children, value)),
|
||||
field.nullable,
|
||||
field.metadata,
|
||||
),
|
||||
);
|
||||
} else {
|
||||
matches.push(new Field(field.name, value as DataType, field.nullable));
|
||||
matches.push(field);
|
||||
}
|
||||
}
|
||||
return matches;
|
||||
|
||||
+14
-1
@@ -630,6 +630,18 @@ export abstract class Table {
|
||||
|
||||
/**
|
||||
* Update per-field (column) metadata.
|
||||
*
|
||||
* The following keys are treated specially, by convention, and should be
|
||||
* used when appropriate:
|
||||
*
|
||||
* - `lancedb:description`: for a human-readable description of a field.
|
||||
* - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
* names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
* - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
* `feature_v2` might be in the same logical column.
|
||||
* - `lancedb:status`: for status options (`production`, `candidate`,
|
||||
* `deprecated`, `archived`) to designate the current life cycle state of
|
||||
* this column.
|
||||
* @param {FieldMetadataUpdate[]} updates One or more per-field updates. Each
|
||||
* update's metadata is merged into the field's existing metadata by default;
|
||||
* a value of `null` deletes that key, and `replace: true` swaps the whole map.
|
||||
@@ -1555,7 +1567,8 @@ export interface FieldMetadataUpdate {
|
||||
path: string;
|
||||
/**
|
||||
* Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
* default; a value of `null` deletes that key.
|
||||
* default; a value of `null` deletes that key. See
|
||||
* {@link Table.updateFieldMetadata} for the conventional `lancedb:*` keys.
|
||||
*/
|
||||
metadata: Record<string, string | null>;
|
||||
/** If true, replace the field's entire metadata map instead of merging. */
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.38.0-beta.10",
|
||||
"version": "0.38.0-beta.11",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
|
||||
@@ -3401,9 +3401,10 @@ class AsyncQuery(AsyncStandardQuery):
|
||||
pass in multiple vectors. When multiple vectors are passed in, if the vector
|
||||
column is with multivector type, then the vectors will be treated as a single
|
||||
query. Or the vectors will be treated as multiple queries, this can be useful
|
||||
if you want to find the nearest vectors to multiple query vectors.
|
||||
This is not expected to be faster than making multiple queries concurrently;
|
||||
it is just a convenience method. If multiple vectors are passed in then
|
||||
if you want to find the nearest vectors to multiple query vectors. Flat
|
||||
searches share one table scan across the query vectors, avoiding the scan
|
||||
and memory amplification of making multiple queries concurrently. If
|
||||
multiple vectors are passed in then
|
||||
an additional column `query_index` will be added to the results. This column
|
||||
will contain the index of the query vector that the result is nearest to.
|
||||
"""
|
||||
@@ -3532,8 +3533,8 @@ class AsyncFTSQuery(AsyncStandardQuery):
|
||||
|
||||
Typically, a single vector is passed in as the query. However, you can also
|
||||
pass in multiple vectors. This can be useful if you want to find the nearest
|
||||
vectors to multiple query vectors. This is not expected to be faster than
|
||||
making multiple queries concurrently; it is just a convenience method.
|
||||
vectors to multiple query vectors. Flat searches share one table scan across
|
||||
the query vectors instead of issuing concurrent full scans.
|
||||
If multiple vectors are passed in then an additional column `query_index`
|
||||
will be added to the results. This column will contain the index of the
|
||||
query vector that the result is nearest to.
|
||||
|
||||
@@ -2127,12 +2127,25 @@ class Table(ABC):
|
||||
----------
|
||||
updates : dict
|
||||
One or more dicts, each with:
|
||||
|
||||
- "path": str — dot-path to the field (e.g. "embedding" or "a.b.c").
|
||||
- "metadata": dict[str, str | None] — keys to set; a value of ``None``
|
||||
deletes that key.
|
||||
- "replace": bool, optional — replace the field's whole metadata map
|
||||
instead of merging (default False).
|
||||
|
||||
The following keys are treated specially, by convention, and should
|
||||
be used when appropriate:
|
||||
|
||||
- "lancedb:description": for a human-readable description of a field.
|
||||
- ``"lancedb:tag:<name>"`` for a user-defined key-value tag, where the
|
||||
suffix names the tag category; e.g. "lancedb:tag:model": "clip".
|
||||
- "lancedb:logical-column" for a column grouping; e.g. "feature_v1"
|
||||
and "feature_v2" might be in the same logical column.
|
||||
- "lancedb:status" for status options ("production", "candidate",
|
||||
"deprecated", "archived") to designate the current life cycle
|
||||
state of this column.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UpdateFieldMetadataResult
|
||||
|
||||
@@ -105,7 +105,7 @@ def test_quickstart(tmp_path):
|
||||
tbl.create_index(num_sub_vectors=1)
|
||||
# --8<-- [end:create_index]
|
||||
# --8<-- [start:delete_rows]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows]
|
||||
# --8<-- [start:drop_table]
|
||||
db.drop_table("my_table")
|
||||
@@ -201,7 +201,7 @@ async def test_quickstart_async(tmp_path):
|
||||
await tbl.create_index("vector")
|
||||
# --8<-- [end:create_index_async]
|
||||
# --8<-- [start:delete_rows_async]
|
||||
await tbl.delete('item = "fizz"')
|
||||
await tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows_async]
|
||||
# --8<-- [start:drop_table_async]
|
||||
await db.drop_table("my_table_async")
|
||||
|
||||
@@ -266,7 +266,7 @@ def test_table():
|
||||
tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_from_pydantic]
|
||||
# --8<-- [start:delete_row]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row]
|
||||
# --8<-- [start:delete_specific_row]
|
||||
data = [
|
||||
@@ -538,7 +538,7 @@ async def test_table_async():
|
||||
await async_tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_async_from_pydantic]
|
||||
# --8<-- [start:delete_row_async]
|
||||
await async_tbl.delete('item = "fizz"')
|
||||
await async_tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row_async]
|
||||
# --8<-- [start:delete_specific_row_async]
|
||||
data = [
|
||||
|
||||
@@ -897,6 +897,23 @@ def test_query_builder_batches(table):
|
||||
assert rs_list["id"][1] == 2
|
||||
|
||||
|
||||
def test_batch_vector_query_shares_filtered_flat_scan(table):
|
||||
query = (
|
||||
table.search([[1.0, 2.0], [3.0, 4.0]])
|
||||
.where("id > 0", prefilter=True)
|
||||
.limit(1)
|
||||
.select(["id"])
|
||||
)
|
||||
|
||||
plan = query.explain_plan(verbose=True)
|
||||
assert "KNNVectorDistance: queries=2" in plan
|
||||
assert "UnionExec" not in plan
|
||||
|
||||
results = query.to_arrow()
|
||||
assert len(results) == 2
|
||||
assert results["query_index"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_dynamic_projection(table):
|
||||
rs = (
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.10"
|
||||
version = "0.38.0-beta.11"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
|
||||
mod sql;
|
||||
|
||||
pub(crate) use sql::canonicalize_sql_predicate;
|
||||
pub use sql::expr_to_sql_string;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -1,10 +1,16 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::any::TypeId;
|
||||
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
|
||||
use datafusion_expr::Expr;
|
||||
use datafusion_sql::unparser::{self, dialect::Dialect};
|
||||
use datafusion_sql::sqlparser::{
|
||||
dialect::{Dialect as SqlParserDialect, GenericDialect},
|
||||
tokenizer::{Token, Tokenizer},
|
||||
};
|
||||
use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect};
|
||||
|
||||
/// Unparser dialect that matches the quoting style expected by the Lance SQL
|
||||
/// parser. Lance uses backtick (`` ` ``) as the only delimited-identifier
|
||||
@@ -19,7 +25,7 @@ use datafusion_sql::unparser::{self, dialect::Dialect};
|
||||
/// lower-case by the SQL parser, which would break case-sensitive schemas).
|
||||
struct LanceSqlDialect;
|
||||
|
||||
impl Dialect for LanceSqlDialect {
|
||||
impl UnparserDialect for LanceSqlDialect {
|
||||
fn identifier_quote_style(&self, identifier: &str) -> Option<char> {
|
||||
let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase())
|
||||
|| !identifier
|
||||
@@ -30,6 +36,61 @@ impl Dialect for LanceSqlDialect {
|
||||
}
|
||||
}
|
||||
|
||||
/// Lance's tokenizer dialect with SQL-standard double-quoted identifiers added.
|
||||
///
|
||||
/// Keep this deliberately small: Lance's parser wraps `GenericDialect` and
|
||||
/// delegates only identifier recognition, leaving every other dialect option at
|
||||
/// its default. In particular, `/*! ... */` remains an ordinary block comment.
|
||||
#[derive(Debug, Default)]
|
||||
struct PredicateDialect(GenericDialect);
|
||||
|
||||
impl SqlParserDialect for PredicateDialect {
|
||||
fn dialect(&self) -> TypeId {
|
||||
self.0.dialect()
|
||||
}
|
||||
|
||||
fn is_identifier_start(&self, ch: char) -> bool {
|
||||
self.0.is_identifier_start(ch)
|
||||
}
|
||||
|
||||
fn is_identifier_part(&self, ch: char) -> bool {
|
||||
self.0.is_identifier_part(ch)
|
||||
}
|
||||
|
||||
fn is_delimited_identifier_start(&self, ch: char) -> bool {
|
||||
ch == '"' || ch == '`'
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonicalize a raw SQL predicate for Lance's parser.
|
||||
///
|
||||
/// Lance wraps [`GenericDialect`] for identifier recognition while retaining the
|
||||
/// default dialect behavior for every other lexical option. [`PredicateDialect`]
|
||||
/// mirrors that contract and additionally recognizes `"` as an identifier
|
||||
/// delimiter, allowing this function to rewrite only those identifier tokens.
|
||||
pub fn canonicalize_sql_predicate(predicate: &str) -> crate::Result<String> {
|
||||
let dialect = PredicateDialect::default();
|
||||
let tokens = Tokenizer::new(&dialect, predicate)
|
||||
.with_unescape(false)
|
||||
.tokenize()
|
||||
.map_err(|err| crate::Error::InvalidInput {
|
||||
message: format!("invalid SQL predicate: {err}"),
|
||||
})?;
|
||||
|
||||
Ok(tokens
|
||||
.into_iter()
|
||||
.map(|token| match token {
|
||||
Token::Word(word) if word.quote_style == Some('"') => {
|
||||
// with_unescape(false) retains doubled double quotes. Decode
|
||||
// those before escaping any backticks for Lance's delimiter.
|
||||
let identifier = word.value.replace("\"\"", "\"").replace('`', "``");
|
||||
format!("`{identifier}`")
|
||||
}
|
||||
other => other.to_string(),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Prefix for placeholder strings inserted in place of binary literals. Chosen
|
||||
/// to be extremely unlikely to occur in user data.
|
||||
const BINARY_PLACEHOLDER_PREFIX: &str = "__lancedb_binary_placeholder_";
|
||||
@@ -113,3 +174,51 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
|
||||
}
|
||||
Ok(sql)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::canonicalize_sql_predicate;
|
||||
|
||||
#[test]
|
||||
fn normalizes_double_quoted_identifiers() {
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""PartyAbbrev" = 'D'"#).unwrap(),
|
||||
"`PartyAbbrev` = 'D'"
|
||||
);
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""MetaData"."userId" = 5"#).unwrap(),
|
||||
"`MetaData`.`userId` = 5"
|
||||
);
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""a""b" = 1"#).unwrap(),
|
||||
"`a\"b` = 1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_quotes_inside_literals_and_backticks() {
|
||||
let filter = r#"name = 'Alice "Ace"' AND `quoted"field` = 1"#;
|
||||
assert_eq!(canonicalize_sql_predicate(filter).unwrap(), filter);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_literals_and_comments_using_lance_dialect_rules() {
|
||||
let predicate = r#"path = '\' AND "PartyAbbrev" = 'D' -- unmatched " in comment"#;
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(predicate).unwrap(),
|
||||
r#"path = '\' AND `PartyAbbrev` = 'D' -- unmatched " in comment"#
|
||||
);
|
||||
|
||||
let predicate = r#"id = 1 /* unmatched " in block comment */"#;
|
||||
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
|
||||
|
||||
let predicate = r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#;
|
||||
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unterminated_double_quoted_identifier() {
|
||||
let error = canonicalize_sql_predicate(r#""PartyAbbrev = 'D'"#).unwrap_err();
|
||||
assert!(matches!(error, crate::Error::InvalidInput { .. }));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,6 +170,15 @@ pub(crate) fn plan(
|
||||
filter: Option<&str>,
|
||||
limit: Option<u64>,
|
||||
) -> Result<(MaterializedViewDefinition, Vec<ArrowField>, Lineage)> {
|
||||
let filter = filter
|
||||
.map(crate::expr::canonicalize_sql_predicate)
|
||||
.transpose()
|
||||
.map_err(|err| match err {
|
||||
Error::InvalidInput { message } => Error::InvalidInput {
|
||||
message: format!("invalid view filter: {message}"),
|
||||
},
|
||||
err => err,
|
||||
})?;
|
||||
let projections: Vec<(String, String)> = if projections.is_empty() {
|
||||
source_schema
|
||||
.fields()
|
||||
@@ -274,7 +283,7 @@ pub(crate) fn plan(
|
||||
declared.push(output);
|
||||
}
|
||||
|
||||
if let Some(filter) = filter {
|
||||
if let Some(filter) = filter.as_deref() {
|
||||
let expr = planner
|
||||
.parse_filter(filter)
|
||||
.map_err(|e| Error::InvalidInput {
|
||||
@@ -314,7 +323,7 @@ pub(crate) fn plan(
|
||||
.into_iter()
|
||||
.map(|(output, expression)| ViewProjection { output, expression })
|
||||
.collect(),
|
||||
filter: filter.map(String::from),
|
||||
filter,
|
||||
limit,
|
||||
inputs,
|
||||
};
|
||||
|
||||
@@ -46,8 +46,9 @@ use lance_table::format::Fragment;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{
|
||||
INCARNATION_META_KEY, MaterializedViewDefinition, REFRESHED_AT_MS_META_KEY,
|
||||
SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY,
|
||||
DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition,
|
||||
REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY,
|
||||
definition_to_metadata,
|
||||
};
|
||||
use crate::database::OpenTableRequest;
|
||||
use crate::table::{NativeTable, NativeTableExt, Table};
|
||||
@@ -197,8 +198,28 @@ pub(crate) async fn execute_refresh(
|
||||
),
|
||||
});
|
||||
}
|
||||
let definition_changed =
|
||||
definition.filter != replanned.filter || definition.inputs != replanned.inputs;
|
||||
let definition = &replanned;
|
||||
|
||||
// A watermark written for a legacy raw filter certifies the rows that
|
||||
// filter produced, not the canonical predicate above. Rebuild instead of
|
||||
// accepting or advancing it, and persist the migrated definition in the
|
||||
// same metadata commit that certifies the replacement rows.
|
||||
if definition_changed {
|
||||
return rebuild(
|
||||
view_native,
|
||||
&view_ds,
|
||||
&source_ds,
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
true,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let metadata = &view_ds.schema().metadata;
|
||||
let watermark: Option<u64> = metadata
|
||||
.get(SOURCE_VERSION_META_KEY)
|
||||
@@ -257,6 +278,7 @@ pub(crate) async fn execute_refresh(
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
false,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await
|
||||
@@ -271,6 +293,7 @@ pub(crate) async fn execute_refresh(
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
false,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await
|
||||
@@ -683,6 +706,7 @@ async fn incremental(
|
||||
view_ds.clone(),
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -704,6 +728,7 @@ async fn incremental(
|
||||
published,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -775,6 +800,7 @@ async fn incremental(
|
||||
published,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -824,12 +850,14 @@ async fn incremental(
|
||||
appended,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
Ok(Some(result))
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn rebuild(
|
||||
view_native: &NativeTable,
|
||||
view_ds: &Dataset,
|
||||
@@ -837,6 +865,7 @@ async fn rebuild(
|
||||
source_version: u64,
|
||||
source_ts: u128,
|
||||
definition: &MaterializedViewDefinition,
|
||||
persist_definition: bool,
|
||||
expected_incarnation: Option<&str>,
|
||||
) -> Result<RefreshMaterializedViewResult> {
|
||||
let rows_written = Arc::new(AtomicU64::new(0));
|
||||
@@ -867,6 +896,7 @@ async fn rebuild(
|
||||
replaced,
|
||||
source_version,
|
||||
source_ts,
|
||||
persist_definition.then_some(definition),
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -981,6 +1011,7 @@ async fn stamp_watermark(
|
||||
mut dataset: Dataset,
|
||||
source_version: u64,
|
||||
source_ts: u128,
|
||||
definition: Option<&MaterializedViewDefinition>,
|
||||
expected_incarnation: Option<&str>,
|
||||
) -> Result<u64> {
|
||||
ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?;
|
||||
@@ -993,27 +1024,32 @@ async fn stamp_watermark(
|
||||
.get(INCARNATION_META_KEY)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
dataset
|
||||
.update_schema_metadata([
|
||||
(INCARNATION_META_KEY.to_string(), Some(incarnation)),
|
||||
(
|
||||
SOURCE_VERSION_META_KEY.to_string(),
|
||||
Some(source_version.to_string()),
|
||||
),
|
||||
(
|
||||
SOURCE_VERSION_TS_META_KEY.to_string(),
|
||||
Some(source_ts.to_string()),
|
||||
),
|
||||
(
|
||||
REFRESHED_AT_MS_META_KEY.to_string(),
|
||||
Some(now_ms().to_string()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
])
|
||||
.await?;
|
||||
let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))];
|
||||
if let Some(definition) = definition {
|
||||
metadata.push((
|
||||
DEFINITION_META_KEY.to_string(),
|
||||
Some(definition_to_metadata(definition)?),
|
||||
));
|
||||
}
|
||||
metadata.extend([
|
||||
(
|
||||
SOURCE_VERSION_META_KEY.to_string(),
|
||||
Some(source_version.to_string()),
|
||||
),
|
||||
(
|
||||
SOURCE_VERSION_TS_META_KEY.to_string(),
|
||||
Some(source_ts.to_string()),
|
||||
),
|
||||
(
|
||||
REFRESHED_AT_MS_META_KEY.to_string(),
|
||||
Some(now_ms().to_string()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
]);
|
||||
dataset.update_schema_metadata(metadata).await?;
|
||||
let actual = dataset.version().version;
|
||||
if actual != predicted {
|
||||
return Err(Error::Runtime {
|
||||
@@ -1585,6 +1621,106 @@ mod tests {
|
||||
assert_eq!(read(view.table(), "x").await, vec![20, 40]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(
|
||||
("id", Int32, [1, 2, 3]),
|
||||
("PartyAbbrev", Utf8, ["D", "R", "D"])
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("src", batch)
|
||||
.write_options(crate::materialized_view::tests::stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
conn.create_materialized_view("democrats", "src")
|
||||
.select([("id", "id")])
|
||||
.only_if(r#""PartyAbbrev" = 'D'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Reopen from schema metadata so these assertions cover the stored
|
||||
// predicate and lineage, not only the declaration-time handle.
|
||||
let view = conn.open_materialized_view("democrats").await.unwrap();
|
||||
assert_eq!(
|
||||
view.definition().filter.as_deref(),
|
||||
Some("`PartyAbbrev` = 'D'")
|
||||
);
|
||||
assert_eq!(view.definition().inputs, ["PartyAbbrev", "id"]);
|
||||
|
||||
let result = view.refresh().execute().await.unwrap();
|
||||
assert_eq!(result.rows_written, 2);
|
||||
assert_eq!(read(view.table(), "id").await, vec![1, 3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(
|
||||
("id", Int32, [1, 2, 3]),
|
||||
("PartyAbbrev", Utf8, ["D", "R", "D"])
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("legacy_src", batch)
|
||||
.write_options(crate::materialized_view::tests::stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let view = conn
|
||||
.create_materialized_view("legacy_view", "legacy_src")
|
||||
.select([("id", "id")])
|
||||
.only_if(r#""PartyAbbrev" = 'X'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0);
|
||||
|
||||
// Model a definition and up-to-date watermark written before filter
|
||||
// canonicalization was applied to materialized views.
|
||||
let mut legacy = view.definition().clone();
|
||||
legacy.filter = Some(r#""PartyAbbrev" = 'D'"#.into());
|
||||
legacy.inputs = vec!["id".into()];
|
||||
let native = view.table().as_native().unwrap();
|
||||
let mut dataset = native.dataset.get().await.unwrap().as_ref().clone();
|
||||
let predicted = dataset.version().version + 1;
|
||||
dataset
|
||||
.update_schema_metadata([
|
||||
(
|
||||
DEFINITION_META_KEY.to_string(),
|
||||
Some(definition_to_metadata(&legacy).unwrap()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
native.dataset.update(dataset);
|
||||
|
||||
let reopened = conn.open_materialized_view("legacy_view").await.unwrap();
|
||||
let result = reopened.refresh().execute().await.unwrap();
|
||||
assert_eq!(result.mode, RefreshMode::Rebuild);
|
||||
assert_eq!(result.rows_written, 2);
|
||||
assert_eq!(read(reopened.table(), "id").await, vec![1, 3]);
|
||||
|
||||
// A fresh handle proves the migration was stored alongside the new
|
||||
// watermark and therefore happens only once.
|
||||
let migrated = conn.open_materialized_view("legacy_view").await.unwrap();
|
||||
assert_eq!(
|
||||
migrated.definition().filter.as_deref(),
|
||||
Some("`PartyAbbrev` = 'D'")
|
||||
);
|
||||
assert_eq!(migrated.definition().inputs, ["PartyAbbrev", "id"]);
|
||||
assert_eq!(
|
||||
migrated.refresh().execute().await.unwrap().mode,
|
||||
RefreshMode::NoOp
|
||||
);
|
||||
assert_eq!(read(migrated.table(), "id").await, vec![1, 3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_append_refreshes_incrementally() {
|
||||
let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await;
|
||||
@@ -2767,7 +2903,7 @@ mod tests {
|
||||
let stale = view_native.dataset.get().await.unwrap().as_ref().clone();
|
||||
view.table().delete("x = 1").await.unwrap();
|
||||
|
||||
let err = stamp_watermark(view_native, stale, 99, 99, None).await;
|
||||
let err = stamp_watermark(view_native, stale, 99, 99, None, None).await;
|
||||
assert!(err.is_err());
|
||||
|
||||
let result = view.refresh().execute().await.unwrap();
|
||||
|
||||
+273
-9
@@ -399,6 +399,9 @@ pub trait QueryBase {
|
||||
/// x > 5 OR y = 'test'
|
||||
/// ```
|
||||
///
|
||||
/// Identifiers may be delimited with SQL-standard double quotes or
|
||||
/// backticks. String literals must use single quotes.
|
||||
///
|
||||
/// Filtering performance can often be improved by creating a scalar index
|
||||
/// on the filter column(s).
|
||||
///
|
||||
@@ -913,6 +916,17 @@ impl QueryRequest {
|
||||
/// use different representations) the error is recorded and surfaced later
|
||||
/// by [`Self::check_filter`].
|
||||
pub(crate) fn add_filter(&mut self, new: QueryFilter) {
|
||||
let new = match new {
|
||||
QueryFilter::Sql(filter) => match crate::expr::canonicalize_sql_predicate(&filter) {
|
||||
Ok(filter) => QueryFilter::Sql(filter),
|
||||
Err(err) => {
|
||||
self.filter_error = Some(err.to_string());
|
||||
return;
|
||||
}
|
||||
},
|
||||
other => other,
|
||||
};
|
||||
|
||||
self.filter = Some(match self.filter.take() {
|
||||
None => new,
|
||||
Some(existing) => match and_filters(existing, new) {
|
||||
@@ -1174,12 +1188,12 @@ impl VectorQuery {
|
||||
|
||||
/// Add another query vector to the search.
|
||||
///
|
||||
/// Multiple searches will be dispatched as part of the query.
|
||||
/// This is a convenience method for adding multiple query vectors
|
||||
/// to the search. It is not expected to be faster than issuing
|
||||
/// multiple queries concurrently.
|
||||
/// Multiple searches will be dispatched as a batch. Flat searches share
|
||||
/// one table scan across the query vectors, avoiding the scan and memory
|
||||
/// amplification of issuing the searches concurrently. Indexed searches
|
||||
/// may still perform per-vector index work.
|
||||
///
|
||||
/// The output data will contain an additional columns `query_index` which
|
||||
/// The output data will contain an additional column `query_index` which
|
||||
/// will contain the index of the query vector that was used to generate the
|
||||
/// result.
|
||||
pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result<Self> {
|
||||
@@ -1646,10 +1660,14 @@ mod tests {
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use super::*;
|
||||
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
|
||||
use arrow::{
|
||||
array::downcast_array,
|
||||
compute::concat_batches,
|
||||
datatypes::{Int32Type, UInt8Type},
|
||||
};
|
||||
use arrow_array::{
|
||||
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
|
||||
types::Float32Type,
|
||||
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchIterator,
|
||||
StringArray, cast::AsArray, types::Float32Type,
|
||||
};
|
||||
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
@@ -1878,6 +1896,157 @@ mod tests {
|
||||
query.execute().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_double_quoted_predicates_across_table_operations() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let dataset_path = tmp_dir.path().join("test.lance");
|
||||
let uri = dataset_path.to_str().unwrap();
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
ArrowField::new("id", DataType::Int32, false),
|
||||
ArrowField::new("PartyAbbrev", DataType::Utf8, false),
|
||||
ArrowField::new("path", DataType::Utf8, false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
|
||||
Arc::new(StringArray::from(vec!["D", "R", "R", "D"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conn = connect(uri).execute().await.unwrap();
|
||||
let table = conn.create_table("parties", batch).execute().await.unwrap();
|
||||
let batches = table
|
||||
.query()
|
||||
.only_if(r#""PartyAbbrev" = 'D'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'D'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
// Public BaseTable dispatch cannot bypass canonicalization.
|
||||
let query = AnyQuery::Query(QueryRequest {
|
||||
filter: Some(QueryFilter::Sql(r#""PartyAbbrev" = 'D'"#.to_string())),
|
||||
..Default::default()
|
||||
});
|
||||
let batches = table
|
||||
.base_table()
|
||||
.query(&query, Default::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
assert_eq!(
|
||||
table
|
||||
.base_table()
|
||||
.count_rows(Some(crate::table::Filter::Sql(
|
||||
r#""PartyAbbrev" = 'D'"#.to_string(),
|
||||
)))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
for predicate in [
|
||||
r#"id = 1 -- unmatched " in a valid SQL comment"#,
|
||||
r#"id = 1 /* unmatched " in a valid SQL comment */"#,
|
||||
r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#,
|
||||
r#"path = '\' AND "PartyAbbrev" = 'D'"#,
|
||||
] {
|
||||
let batches = table
|
||||
.query()
|
||||
.only_if(predicate)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
|
||||
}
|
||||
|
||||
// The same canonical predicate contract applies to both merge filters.
|
||||
let source = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3])),
|
||||
Arc::new(StringArray::from(vec!["D", "R", "R"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge.when_not_matched_by_source_delete(Some(r#""PartyAbbrev" = 'D'"#.to_string()));
|
||||
let result = table
|
||||
.base_table()
|
||||
.merge_insert(
|
||||
merge,
|
||||
Box::new(RecordBatchIterator::new(vec![Ok(source)], schema.clone())),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.num_deleted_rows, 1);
|
||||
|
||||
let source = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3])),
|
||||
Arc::new(StringArray::from(vec!["U", "U", "U"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge.when_matched_update_all(Some(r#"target."PartyAbbrev" = 'D'"#.to_string()));
|
||||
merge
|
||||
.execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema)))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'U'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
|
||||
let update = table
|
||||
.update()
|
||||
.only_if(r#""PartyAbbrev" = 'R'"#)
|
||||
.column("PartyAbbrev", "'X'");
|
||||
table.base_table().update(update).await.unwrap();
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'X'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
let result = table
|
||||
.base_table()
|
||||
.delete(crate::table::Predicate::String(r#""PartyAbbrev" = 'X'"#))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.num_deleted_rows, 2);
|
||||
assert_eq!(table.count_rows(None).await.unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_with_transform() {
|
||||
let batches = make_non_empty_batches();
|
||||
@@ -2334,7 +2503,8 @@ mod tests {
|
||||
.limit(1);
|
||||
|
||||
let plan = query.explain_plan(true).await.unwrap();
|
||||
assert!(plan.contains("UnionExec"));
|
||||
assert!(plan.contains("KNNVectorDistance: queries=2"));
|
||||
assert!(!plan.contains("UnionExec"));
|
||||
|
||||
let results = query
|
||||
.execute()
|
||||
@@ -2349,6 +2519,100 @@ mod tests {
|
||||
// We don't guarantee order.
|
||||
assert!(query_index.values().contains(&0));
|
||||
assert!(query_index.values().contains(&1));
|
||||
|
||||
// Batch KNN does not support a per-query offset, so offset queries keep
|
||||
// the legacy per-vector plan to preserve their result semantics.
|
||||
let offset_query = table
|
||||
.query()
|
||||
.nearest_to(&[0.1, 0.2, 0.3, 0.4])
|
||||
.unwrap()
|
||||
.add_query_vector(&[0.5, 0.6, 0.7, 0.8])
|
||||
.unwrap()
|
||||
.limit(1)
|
||||
.offset(1);
|
||||
assert!(
|
||||
offset_query
|
||||
.explain_plan(true)
|
||||
.await
|
||||
.unwrap()
|
||||
.contains("UnionExec")
|
||||
);
|
||||
let offset_results = offset_query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
offset_results
|
||||
.iter()
|
||||
.map(RecordBatch::num_rows)
|
||||
.sum::<usize>(),
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_binary_query_vectors() {
|
||||
let vectors = FixedSizeListArray::from_iter_primitive::<UInt8Type, _, _>(
|
||||
vec![
|
||||
Some(vec![Some(0), Some(0)]),
|
||||
Some(vec![Some(255), Some(255)]),
|
||||
],
|
||||
2,
|
||||
);
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
ArrowField::new("id", DataType::Int32, false),
|
||||
ArrowField::new("vector", vectors.data_type().clone(), false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema,
|
||||
vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let table = conn
|
||||
.create_table("binary_batch", batch)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let query = table
|
||||
.query()
|
||||
.nearest_to(&[0.0, 0.0])
|
||||
.unwrap()
|
||||
.add_query_vector(&[255.0, 255.0])
|
||||
.unwrap()
|
||||
.distance_type(DistanceType::Hamming)
|
||||
.limit(1);
|
||||
|
||||
// Binary queries retain the per-vector plan because Lance's binary
|
||||
// nearest path requires primitive UInt8 query arrays.
|
||||
assert!(
|
||||
query
|
||||
.explain_plan(true)
|
||||
.await
|
||||
.unwrap()
|
||||
.contains("UnionExec")
|
||||
);
|
||||
|
||||
let results = query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let results = concat_batches(&results[0].schema(), &results).unwrap();
|
||||
assert_eq!(results.num_rows(), 2);
|
||||
|
||||
let ids = results["id"].as_primitive::<Int32Type>();
|
||||
assert!(ids.values().contains(&0));
|
||||
assert!(ids.values().contains(&1));
|
||||
let query_index = results["query_index"].as_primitive::<Int32Type>();
|
||||
assert!(query_index.values().contains(&0));
|
||||
assert!(query_index.values().contains(&1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -1379,10 +1379,11 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
query: &AnyQuery,
|
||||
version: Option<u64>,
|
||||
) -> Result<Vec<serde_json::Value>> {
|
||||
let query = query.canonicalized()?;
|
||||
let mut base_body = serde_json::json!({ "version": version });
|
||||
self.apply_branch_body(&mut base_body);
|
||||
|
||||
match query {
|
||||
match &query {
|
||||
AnyQuery::Query(query) => {
|
||||
let mut body = base_body.clone();
|
||||
self.apply_query_params(&mut body, query)?;
|
||||
@@ -2491,7 +2492,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
let mut body = if let Some(filter) = filter {
|
||||
let filter_sql = match filter {
|
||||
Filter::Sql(sql) => sql.clone(),
|
||||
Filter::Sql(sql) => crate::expr::canonicalize_sql_predicate(&sql)?,
|
||||
Filter::Datafusion(expr) => expr_to_sql_string(&expr)?,
|
||||
};
|
||||
serde_json::json!({ "predicate": filter_sql, "version": read_snapshot.version })
|
||||
@@ -2747,7 +2748,8 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(final_analyze)
|
||||
}
|
||||
|
||||
async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult> {
|
||||
async fn update(&self, mut update: UpdateBuilder) -> Result<UpdateResult> {
|
||||
update.canonicalize_filter()?;
|
||||
self.check_mutable().await?;
|
||||
let request = self
|
||||
.client
|
||||
@@ -2794,7 +2796,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult> {
|
||||
self.check_mutable().await?;
|
||||
let predicate_sql = match predicate {
|
||||
Predicate::String(s) => s.to_string(),
|
||||
Predicate::String(s) => crate::expr::canonicalize_sql_predicate(s)?,
|
||||
Predicate::Expr(expr) => expr_to_sql_string(expr)?,
|
||||
};
|
||||
let mut body = serde_json::json!({ "predicate": predicate_sql });
|
||||
@@ -2851,9 +2853,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
async fn merge_insert(
|
||||
&self,
|
||||
params: MergeInsertBuilder,
|
||||
mut params: MergeInsertBuilder,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
params.canonicalize_filters()?;
|
||||
self.check_mutable().await?;
|
||||
|
||||
let timeout = params.timeout;
|
||||
@@ -3864,13 +3867,17 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
request.body().unwrap().as_bytes().unwrap(),
|
||||
br#"{"predicate":"a > 10","version":null}"#
|
||||
br#"{"predicate":"`A` > 10","version":null}"#
|
||||
);
|
||||
|
||||
http::Response::builder().status(200).body("42").unwrap()
|
||||
});
|
||||
|
||||
let count = table.count_rows(Some("a > 10".into())).await.unwrap();
|
||||
let count = table
|
||||
.base_table()
|
||||
.count_rows(Some(Filter::Sql(r#""A" > 10"#.into())))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(count, 42);
|
||||
}
|
||||
|
||||
@@ -4353,7 +4360,7 @@ mod tests {
|
||||
assert_eq!(expression, "b - 1");
|
||||
|
||||
let only_if = value.get("predicate").unwrap().as_str().unwrap();
|
||||
assert_eq!(only_if, "b > 10");
|
||||
assert_eq!(only_if, "`B` > 10");
|
||||
}
|
||||
|
||||
if old_server {
|
||||
@@ -4369,14 +4376,12 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table
|
||||
let update = table
|
||||
.update()
|
||||
.column("a", "a + 1")
|
||||
.column("b", "b - 1")
|
||||
.only_if("b > 10")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
.only_if(r#""B" > 10"#);
|
||||
let result = table.base_table().update(update).await.unwrap();
|
||||
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
assert_eq!(result.rows_updated, if old_server { 0 } else { 5 });
|
||||
@@ -4463,10 +4468,10 @@ mod tests {
|
||||
|
||||
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
|
||||
assert_eq!(params["on"], "some_col");
|
||||
assert_eq!(params["when_matched_update_all"], "false");
|
||||
assert_eq!(params["when_matched_update_all"], "true");
|
||||
assert_eq!(params["when_not_matched_insert_all"], "false");
|
||||
assert_eq!(params["when_not_matched_by_source_delete"], "false");
|
||||
assert!(!params.contains_key("when_matched_update_all_filt"));
|
||||
assert_eq!(params["when_matched_update_all_filt"], "target.`A` > 0");
|
||||
assert!(!params.contains_key("when_not_matched_by_source_delete_filt"));
|
||||
assert!(!params.contains_key("use_index"));
|
||||
|
||||
@@ -4483,11 +4488,9 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table
|
||||
.merge_insert(&["some_col"])
|
||||
.execute(data)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["some_col"]);
|
||||
merge.when_matched_update_all(Some(r#"target."A" > 0"#.into()));
|
||||
let result = table.base_table().merge_insert(merge, data).await.unwrap();
|
||||
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
if !old_server {
|
||||
@@ -4549,7 +4552,7 @@ mod tests {
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
let predicate = body.get("predicate").unwrap().as_str().unwrap();
|
||||
assert_eq!(predicate, "id in (1, 2, 3)");
|
||||
assert_eq!(predicate, "`ID` in (1, 2, 3)");
|
||||
|
||||
if old_server {
|
||||
http::Response::builder()
|
||||
@@ -4567,7 +4570,11 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table.delete("id in (1, 2, 3)").await.unwrap();
|
||||
let result = table
|
||||
.base_table()
|
||||
.delete(Predicate::String(r#""ID" in (1, 2, 3)"#))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
}
|
||||
|
||||
@@ -4659,6 +4666,7 @@ mod tests {
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
let expected_body = serde_json::json!({
|
||||
"filter": "`A` > 0",
|
||||
"k": isize::MAX as usize,
|
||||
"prefilter": true,
|
||||
"vector": [], // Empty vector means no vector query.
|
||||
@@ -4674,9 +4682,13 @@ mod tests {
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let query = AnyQuery::Query(QueryRequest {
|
||||
filter: Some(QueryFilter::Sql(r#""A" > 0"#.into())),
|
||||
..Default::default()
|
||||
});
|
||||
let data = table
|
||||
.query()
|
||||
.execute()
|
||||
.base_table()
|
||||
.query(&query, Default::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.collect::<Vec<_>>()
|
||||
|
||||
@@ -1164,7 +1164,10 @@ impl Table {
|
||||
///
|
||||
/// * `filter` if present, only count rows matching the filter
|
||||
pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
|
||||
self.inner.count_rows(filter.map(Filter::Sql)).await
|
||||
let filter = filter
|
||||
.map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate).map(Filter::Sql))
|
||||
.transpose()?;
|
||||
self.inner.count_rows(filter).await
|
||||
}
|
||||
|
||||
/// Names of the blob v2 columns in this table, in declaration order.
|
||||
@@ -1364,7 +1367,13 @@ impl Table {
|
||||
/// # });
|
||||
/// ```
|
||||
pub async fn delete(&self, predicate: impl Into<Predicate<'_>>) -> Result<DeleteResult> {
|
||||
self.inner.delete(predicate.into()).await
|
||||
match predicate.into() {
|
||||
Predicate::String(predicate) => {
|
||||
let predicate = crate::expr::canonicalize_sql_predicate(predicate)?;
|
||||
self.inner.delete(Predicate::String(&predicate)).await
|
||||
}
|
||||
predicate @ Predicate::Expr(_) => self.inner.delete(predicate).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an index on the provided column(s).
|
||||
@@ -1777,7 +1786,23 @@ impl Table {
|
||||
self.inner.alter_columns(alterations).await
|
||||
}
|
||||
|
||||
/// Update per-field metadata (merges by default).
|
||||
/// Update per-field (column) metadata.
|
||||
///
|
||||
/// Each [`FieldMetadataUpdate`] is merged into the field's existing metadata
|
||||
/// by default; use [`FieldMetadataUpdate::remove`] to delete a key, or
|
||||
/// [`FieldMetadataUpdate::replace`] to swap the field's entire metadata map.
|
||||
///
|
||||
/// The following keys are treated specially, by convention, and should be
|
||||
/// used when appropriate:
|
||||
///
|
||||
/// - `lancedb:description`: for a human-readable description of a field.
|
||||
/// - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
/// names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
/// - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
/// `feature_v2` might be in the same logical column.
|
||||
/// - `lancedb:status`: for status options (`production`, `candidate`,
|
||||
/// `deprecated`, `archived`) to designate the current life cycle state of
|
||||
/// this column.
|
||||
pub async fn update_field_metadata(
|
||||
&self,
|
||||
updates: &[FieldMetadataUpdate],
|
||||
@@ -3223,7 +3248,10 @@ impl BaseTable for NativeTable {
|
||||
let dataset = self.dataset.get().await?;
|
||||
match filter {
|
||||
None => Ok(dataset.count_rows(None).await?),
|
||||
Some(Filter::Sql(sql)) => Ok(dataset.count_rows(Some(sql)).await?),
|
||||
Some(Filter::Sql(sql)) => {
|
||||
let sql = crate::expr::canonicalize_sql_predicate(&sql)?;
|
||||
Ok(dataset.count_rows(Some(sql)).await?)
|
||||
}
|
||||
Some(Filter::Datafusion(_)) => Err(Error::NotSupported {
|
||||
message: "Datafusion filters are not yet supported".to_string(),
|
||||
}),
|
||||
|
||||
@@ -31,8 +31,9 @@ pub(crate) async fn execute_delete(
|
||||
table.dataset.ensure_mutable()?;
|
||||
match predicate {
|
||||
Predicate::String(s) => {
|
||||
let predicate = crate::expr::canonicalize_sql_predicate(s)?;
|
||||
let mut dataset = (*table.dataset.get().await?).clone();
|
||||
let delete_result = dataset.delete(s).boxed().await?;
|
||||
let delete_result = dataset.delete(&predicate).boxed().await?;
|
||||
let num_deleted_rows = delete_result.num_deleted_rows;
|
||||
let version = dataset.version().version;
|
||||
table.dataset.update(dataset);
|
||||
|
||||
@@ -220,9 +220,32 @@ impl MergeInsertBuilder {
|
||||
///
|
||||
/// Returns version and statistics about the merge operation including the number of rows
|
||||
/// inserted, updated, and deleted.
|
||||
pub async fn execute(self, new_data: Box<dyn RecordBatchReader + Send>) -> Result<MergeResult> {
|
||||
pub async fn execute(
|
||||
mut self,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
self.canonicalize_filters()?;
|
||||
self.table.clone().merge_insert(self, new_data).await
|
||||
}
|
||||
|
||||
pub(crate) fn canonicalize_filters(&mut self) -> Result<()> {
|
||||
self.when_matched_update_all_filt =
|
||||
canonicalize_merge_filter(self.when_matched_update_all_filt.take())?;
|
||||
self.when_not_matched_by_source_delete_filt =
|
||||
canonicalize_merge_filter(self.when_not_matched_by_source_delete_filt.take())?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn canonicalize_merge_filter(filter: Option<MergeFilter>) -> Result<Option<MergeFilter>> {
|
||||
filter
|
||||
.map(|filter| match filter {
|
||||
MergeFilter::Sql(predicate) => {
|
||||
crate::expr::canonicalize_sql_predicate(&predicate).map(MergeFilter::Sql)
|
||||
}
|
||||
filter @ MergeFilter::Expr(_) => Ok(filter),
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Internal implementation of the merge insert logic
|
||||
@@ -230,9 +253,10 @@ impl MergeInsertBuilder {
|
||||
/// This logic was moved from NativeTable::merge_insert to keep table.rs clean.
|
||||
pub(crate) async fn execute_merge_insert(
|
||||
table: &NativeTable,
|
||||
params: MergeInsertBuilder,
|
||||
mut params: MergeInsertBuilder,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
params.canonicalize_filters()?;
|
||||
super::computed_columns::ensure_no_function_bindings_for_mutation(
|
||||
table.schema().await?.as_ref(),
|
||||
"merge_insert",
|
||||
@@ -1407,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);
|
||||
}
|
||||
}
|
||||
|
||||
+744
-32
@@ -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_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 futures::future::try_join_all;
|
||||
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,
|
||||
@@ -45,6 +56,22 @@ impl AnyQuery {
|
||||
Self::VectorQuery(query) => &query.base,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_mut(&mut self) -> &mut QueryRequest {
|
||||
match self {
|
||||
Self::Query(query) => query,
|
||||
Self::VectorQuery(query) => &mut query.base,
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonicalize any raw SQL filter immediately before backend dispatch.
|
||||
pub(crate) fn canonicalized(&self) -> Result<Self> {
|
||||
let mut query = self.clone();
|
||||
if let Some(QueryFilter::Sql(predicate)) = &mut query.base_mut().filter {
|
||||
*predicate = crate::expr::canonicalize_sql_predicate(predicate)?;
|
||||
}
|
||||
Ok(query)
|
||||
}
|
||||
}
|
||||
|
||||
//Decide between namespace or local
|
||||
@@ -53,15 +80,16 @@ pub async fn execute_query(
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<DatasetRecordBatchStream> {
|
||||
let query = query.canonicalized()?;
|
||||
// QueryTable pushdown runs the query server-side, but only on the main
|
||||
// branch: the namespace request carries no branch yet, so a branch handle
|
||||
// must fall through to local execution.
|
||||
if can_execute_namespace_query(table, query).await?
|
||||
if can_execute_namespace_query(table, &query).await?
|
||||
&& let Some(ref namespace_client) = table.namespace_client
|
||||
{
|
||||
return execute_namespace_query(table, namespace_client.clone(), query, options).await;
|
||||
return execute_namespace_query(table, namespace_client.clone(), &query, options).await;
|
||||
}
|
||||
execute_generic_query(table, query, options).await
|
||||
execute_generic_query(table, &query, options).await
|
||||
}
|
||||
|
||||
async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> Result<bool> {
|
||||
@@ -136,9 +164,10 @@ pub async fn create_plan(
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let query = query.canonicalized()?;
|
||||
let query = match query {
|
||||
AnyQuery::VectorQuery(query) => query.clone(),
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query.clone()),
|
||||
AnyQuery::VectorQuery(query) => query,
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query),
|
||||
};
|
||||
query.base.check_filter()?;
|
||||
|
||||
@@ -170,26 +199,48 @@ pub async fn create_plan(
|
||||
let mut column = query.column.clone();
|
||||
|
||||
let mut query_vector = query.query_vector.first().cloned();
|
||||
let mut is_batch_query = false;
|
||||
if query.query_vector.len() > 1 {
|
||||
if column.is_none() {
|
||||
// Infer a vector column with the same dimension of the query vector.
|
||||
let arrow_schema = Schema::from(ds_ref.schema());
|
||||
let arrow_schema = Schema::from(schema);
|
||||
column = Some(default_vector_column(
|
||||
&arrow_schema,
|
||||
Some(query.query_vector[0].len() as i32),
|
||||
)?);
|
||||
}
|
||||
let vector_field = schema.field(column.as_ref().unwrap()).unwrap();
|
||||
if let DataType::List(_) = vector_field.data_type() {
|
||||
// Multivector handling: concatenate into FixedSizeList<FixedSizeList<_>>
|
||||
let (_, element_type) =
|
||||
lance::index::vector::utils::get_vector_type(schema, column.as_ref().unwrap())?;
|
||||
let is_binary = matches!(element_type, DataType::UInt8);
|
||||
if matches!(vector_field.data_type(), DataType::List(_))
|
||||
|| (query.base.offset.unwrap_or(0) == 0 && !is_binary)
|
||||
{
|
||||
// Lance distinguishes these cases from the vector column type: a
|
||||
// list-like query against a List column is one multivector query,
|
||||
// while the same query against a FixedSizeList column is a batch of
|
||||
// independent queries. The batch path shares a single flat scan and
|
||||
// bounds retained candidate data instead of running one scan per
|
||||
// query vector.
|
||||
let vectors = query
|
||||
.query_vector
|
||||
.iter()
|
||||
.map(|arr| arr.as_ref())
|
||||
.collect::<Vec<_>>();
|
||||
let dim = vectors[0].len();
|
||||
if let Some((query_index, actual_dim)) = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find_map(|(index, vector)| (vector.len() != dim).then_some((index, vector.len())))
|
||||
{
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"query vector at index {query_index} has dimension {actual_dim}, expected {dim}"
|
||||
),
|
||||
});
|
||||
}
|
||||
let mut fsl_builder = FixedSizeListBuilder::with_capacity(
|
||||
Float32Builder::with_capacity(dim),
|
||||
Float32Builder::with_capacity(dim * vectors.len()),
|
||||
dim as i32,
|
||||
vectors.len(),
|
||||
);
|
||||
@@ -200,8 +251,12 @@ pub async fn create_plan(
|
||||
fsl_builder.append(true);
|
||||
}
|
||||
query_vector = Some(Arc::new(fsl_builder.finish()));
|
||||
is_batch_query = !matches!(vector_field.data_type(), DataType::List(_));
|
||||
} else {
|
||||
// Multiple query vectors: create a plan for each and union them
|
||||
// Lance's batch path has no per-query offset, and its binary path
|
||||
// requires primitive UInt8 queries rather than a fixed-size list.
|
||||
// Keep the prior plan shape for these cases so offsets are applied
|
||||
// per query and binary query vectors retain their primitive shape.
|
||||
let query_vecs = query.query_vector.clone();
|
||||
let plan_futures = query_vecs
|
||||
.into_iter()
|
||||
@@ -214,7 +269,7 @@ pub async fn create_plan(
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let plans = try_join_all(plan_futures).await?;
|
||||
let plans = futures::future::try_join_all(plan_futures).await?;
|
||||
return create_multi_vector_plan(plans);
|
||||
}
|
||||
}
|
||||
@@ -225,7 +280,7 @@ pub async fn create_plan(
|
||||
let column = if let Some(col) = column {
|
||||
col
|
||||
} else {
|
||||
let arrow_schema = Schema::from(ds_ref.schema());
|
||||
let arrow_schema = Schema::from(schema);
|
||||
default_vector_column(&arrow_schema, Some(query_vector.len() as i32))?
|
||||
};
|
||||
|
||||
@@ -251,10 +306,14 @@ pub async fn create_plan(
|
||||
}
|
||||
}
|
||||
|
||||
scanner.limit(
|
||||
query.base.limit.map(|limit| limit as i64),
|
||||
query.base.offset.map(|offset| offset as i64),
|
||||
)?;
|
||||
// For a batch query, `nearest` already applies k to each query vector.
|
||||
// Adding Scanner's global limit would truncate the combined result to k rows.
|
||||
if !is_batch_query {
|
||||
scanner.limit(
|
||||
query.base.limit.map(|limit| limit as i64),
|
||||
query.base.offset.map(|offset| offset as i64),
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(ef) = query.ef {
|
||||
scanner.ef(ef);
|
||||
@@ -327,11 +386,299 @@ pub async fn create_plan(
|
||||
scanner.order_by(Some(order_by.clone()))?;
|
||||
}
|
||||
|
||||
Ok(scanner.create_plan().await?)
|
||||
let mut plan = scanner
|
||||
.create_plan()
|
||||
.await
|
||||
.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.
|
||||
///
|
||||
/// DataFusion resolves nested fields but its `FieldNotFound` error only lists the
|
||||
/// top-level Arrow fields. This makes a missing leaf look unavailable even when it
|
||||
/// exists below a struct. Keep every other Lance/DataFusion error unchanged and
|
||||
/// enrich only this one schema error at the LanceDB query boundary.
|
||||
fn enrich_lance_field_not_found(
|
||||
error: lance::Error,
|
||||
schema: &lance_core::datatypes::Schema,
|
||||
) -> Error {
|
||||
let Some(field) = find_missing_field(&error) else {
|
||||
return error.into();
|
||||
};
|
||||
field_not_found_error(field, &Schema::from(schema))
|
||||
}
|
||||
|
||||
fn field_not_found_diagnostic(
|
||||
error: &(dyn std::error::Error + 'static),
|
||||
schema: &Schema,
|
||||
) -> Option<Error> {
|
||||
let field = find_missing_field(error)?;
|
||||
Some(field_not_found_error(field, schema))
|
||||
}
|
||||
|
||||
fn field_not_found_error(field: &Column, schema: &Schema) -> Error {
|
||||
let valid_fields = leaf_field_paths(schema);
|
||||
let mut message = format!("Schema error: No field named {}", field.quoted_flat_name());
|
||||
if !valid_fields.is_empty() {
|
||||
message.push_str(". Valid fields are ");
|
||||
message.push_str(&valid_fields.join(", "));
|
||||
}
|
||||
message.push('.');
|
||||
|
||||
Error::InvalidInput { message }
|
||||
}
|
||||
|
||||
fn find_missing_field<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a Column> {
|
||||
if let Some(DataFusionError::SchemaError(schema_error, _)) =
|
||||
error.downcast_ref::<DataFusionError>()
|
||||
&& let SchemaError::FieldNotFound { field, .. } = schema_error.as_ref()
|
||||
{
|
||||
return Some(field);
|
||||
}
|
||||
|
||||
error.source().and_then(find_missing_field)
|
||||
}
|
||||
|
||||
fn leaf_field_paths(schema: &Schema) -> Vec<String> {
|
||||
fn format_segment(segment: &str) -> String {
|
||||
// Quote every segment instead of maintaining a SQL keyword list. Bare
|
||||
// lowercase names such as `true` can be parsed as expressions rather
|
||||
// than identifiers, while backticks preserve all field names in both
|
||||
// local SQL parsers.
|
||||
format!("`{}`", segment.replace('`', "``"))
|
||||
}
|
||||
|
||||
fn visit(fields: &arrow_schema::Fields, path: &mut Vec<String>, paths: &mut Vec<String>) {
|
||||
for field in fields {
|
||||
// Neither local planner can address an empty field-path segment,
|
||||
// even when it is backtick-quoted. Do not advertise leaves beneath
|
||||
// such a segment as valid filter fields.
|
||||
if field.name().is_empty() {
|
||||
continue;
|
||||
}
|
||||
path.push(field.name().clone());
|
||||
match field.data_type() {
|
||||
DataType::Struct(children) if !children.is_empty() => {
|
||||
visit(children, path, paths);
|
||||
}
|
||||
_ => {
|
||||
paths.push(
|
||||
path.iter()
|
||||
.map(|segment| format_segment(segment))
|
||||
.collect::<Vec<_>>()
|
||||
.join("."),
|
||||
);
|
||||
}
|
||||
}
|
||||
path.pop();
|
||||
}
|
||||
}
|
||||
|
||||
let mut paths = Vec::new();
|
||||
visit(schema.fields(), &mut Vec::new(), &mut paths);
|
||||
paths
|
||||
}
|
||||
|
||||
//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(
|
||||
@@ -687,7 +1034,10 @@ async fn parse_arrow_ipc_response(bytes: bytes::Bytes) -> Result<DatasetRecordBa
|
||||
#[cfg(test)]
|
||||
#[allow(deprecated)]
|
||||
mod tests {
|
||||
use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array};
|
||||
use arrow_array::{
|
||||
ArrayRef, FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray,
|
||||
StructArray,
|
||||
};
|
||||
use futures::TryStreamExt;
|
||||
use lance_arrow::FixedSizeListArrayExt;
|
||||
use std::sync::{
|
||||
@@ -696,7 +1046,7 @@ mod tests {
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::query::{QueryExecutionOptions, QueryRequest};
|
||||
use crate::query::{ExecutableQuery, QueryBase, QueryExecutionOptions, QueryRequest};
|
||||
use crate::table::BaseTable;
|
||||
|
||||
fn fixed_size_list_array(values: Vec<f32>, dimension: i32) -> FixedSizeListArray {
|
||||
@@ -837,7 +1187,6 @@ mod tests {
|
||||
async fn test_execute_query_local_routing() {
|
||||
use crate::connect;
|
||||
use crate::table::query::execute_query;
|
||||
use arrow_array::{Int32Array, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
@@ -877,6 +1226,164 @@ mod tests {
|
||||
assert_eq!(count, 2); // 4 and 5
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_filter_field_lists_nested_fields_in_local_planners() {
|
||||
use crate::connect;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let metadata = Arc::new(StructArray::from(vec![
|
||||
(
|
||||
Arc::new(Field::new("year", DataType::Int32, false)),
|
||||
Arc::new(Int32Array::from(vec![2024])) as ArrayRef,
|
||||
),
|
||||
(
|
||||
Arc::new(Field::new("genre", DataType::Utf8, false)),
|
||||
Arc::new(StringArray::from(vec!["fiction"])) as ArrayRef,
|
||||
),
|
||||
(
|
||||
Arc::new(Field::new("Title", DataType::Int32, false)),
|
||||
Arc::new(Int32Array::from(vec![7])) as ArrayRef,
|
||||
),
|
||||
(
|
||||
Arc::new(Field::new("true", DataType::Int32, false)),
|
||||
Arc::new(Int32Array::from(vec![8])) as ArrayRef,
|
||||
),
|
||||
(
|
||||
Arc::new(Field::new("", DataType::Int32, false)),
|
||||
Arc::new(Int32Array::from(vec![10])) as ArrayRef,
|
||||
),
|
||||
]));
|
||||
let vector = Arc::new(fixed_size_list_array(vec![0.0, 1.0], 2));
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
Field::new("vector", vector.data_type().clone(), false),
|
||||
Field::new("content", DataType::Utf8, false),
|
||||
Field::new("metadata", metadata.data_type().clone(), false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema,
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1])),
|
||||
vector,
|
||||
Arc::new(StringArray::from(vec!["example"])),
|
||||
metadata,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let table = conn
|
||||
.create_table("nested_error", batch)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = table
|
||||
.query()
|
||||
.only_if("year = 2024")
|
||||
.execute()
|
||||
.await
|
||||
.err()
|
||||
.expect("query should reject the unqualified nested field");
|
||||
let case_sensitive_path = "`metadata`.`Title`";
|
||||
let keyword_path = "`metadata`.`true`";
|
||||
let expected = format!(
|
||||
"No field named year. Valid fields are `id`, `vector`, `content`, `metadata`.`year`, `metadata`.`genre`, {case_sensitive_path}, {keyword_path}."
|
||||
);
|
||||
|
||||
assert!(
|
||||
error.to_string().contains(&expected),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
for (path, value) in [(case_sensitive_path, 7), (keyword_path, 8)] {
|
||||
table
|
||||
.query()
|
||||
.only_if(format!("{path} = {value}"))
|
||||
.execute()
|
||||
.await
|
||||
.expect("the path advertised by the diagnostic should be reusable");
|
||||
}
|
||||
|
||||
table.set_unenforced_primary_key(["id"]).await.unwrap();
|
||||
table
|
||||
.set_lsm_write_spec(crate::table::LsmWriteSpec::unsharded())
|
||||
.await
|
||||
.unwrap();
|
||||
let lsm_error = table
|
||||
.query()
|
||||
.only_if("year = 2024")
|
||||
.execute()
|
||||
.await
|
||||
.err()
|
||||
.expect("LSM query should reject the unqualified nested field");
|
||||
|
||||
assert!(
|
||||
lsm_error.to_string().contains(&expected),
|
||||
"unexpected LSM error: {lsm_error}"
|
||||
);
|
||||
for (path, value) in [(case_sensitive_path, 7), (keyword_path, 8)] {
|
||||
table
|
||||
.query()
|
||||
.only_if(format!("{path} = {value}"))
|
||||
.execute()
|
||||
.await
|
||||
.expect("the path advertised by the diagnostic should be reusable in LSM queries");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_leaf_field_paths_preserve_arbitrary_depth() {
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
fn nested_field(path: &[&str]) -> Field {
|
||||
let mut segments = path.iter().rev();
|
||||
let mut field = Field::new(
|
||||
*segments.next().expect("path must have a leaf"),
|
||||
DataType::Int32,
|
||||
false,
|
||||
);
|
||||
for segment in segments {
|
||||
field = Field::new(*segment, DataType::Struct(vec![field].into()), false);
|
||||
}
|
||||
field
|
||||
}
|
||||
|
||||
let schema = Schema::new(vec![
|
||||
nested_field(&["a", "b", "c", "d", "e"]),
|
||||
nested_field(&["metadata", "child.with.dot"]),
|
||||
nested_field(&["metadata", "Title"]),
|
||||
nested_field(&["metadata", "123child"]),
|
||||
nested_field(&["metadata", "child`tick"]),
|
||||
nested_field(&["metadata", ""]),
|
||||
nested_field(&["", "child"]),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
leaf_field_paths(&schema),
|
||||
vec![
|
||||
"`a`.`b`.`c`.`d`.`e`",
|
||||
"`metadata`.`child.with.dot`",
|
||||
"`metadata`.`Title`",
|
||||
"`metadata`.`123child`",
|
||||
"`metadata`.`child``tick`",
|
||||
]
|
||||
);
|
||||
|
||||
let source = DataFusionError::SchemaError(
|
||||
Box::new(SchemaError::FieldNotFound {
|
||||
field: Box::new(Column::from_name("missing")),
|
||||
valid_fields: Vec::new(),
|
||||
}),
|
||||
Box::new(None),
|
||||
);
|
||||
let error = field_not_found_diagnostic(&source, &schema).unwrap();
|
||||
assert!(
|
||||
error.to_string().contains(
|
||||
"Valid fields are `a`.`b`.`c`.`d`.`e`, `metadata`.`child.with.dot`, `metadata`.`Title`, `metadata`.`123child`, `metadata`.`child``tick`"
|
||||
),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CountingNamespaceClient {
|
||||
query_table_calls: AtomicUsize,
|
||||
@@ -1088,7 +1595,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_plan_multivector_structure() {
|
||||
async fn test_create_plan_batch_vector_uses_shared_scan() {
|
||||
use arrow_array::{Float32Array, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
@@ -1115,11 +1622,18 @@ mod tests {
|
||||
.unwrap();
|
||||
let native_table = table.as_native().unwrap();
|
||||
|
||||
// This triggers the "create_multi_vector_plan" logic branch
|
||||
// A batch of vectors against a fixed-size vector column should use
|
||||
// Lance's native batch KNN path instead of independent scan plans.
|
||||
let q1 = Arc::new(Float32Array::from(vec![1.0, 2.0]));
|
||||
let q2 = Arc::new(Float32Array::from(vec![3.0, 4.0]));
|
||||
|
||||
let req = VectorQueryRequest {
|
||||
base: QueryRequest {
|
||||
filter: Some(QueryFilter::Sql("id >= 0".to_string())),
|
||||
limit: Some(1),
|
||||
select: Select::Columns(vec!["id".to_string()]),
|
||||
..Default::default()
|
||||
},
|
||||
column: Some("vector".to_string()),
|
||||
query_vector: vec![q1, q2],
|
||||
..Default::default()
|
||||
@@ -1136,22 +1650,220 @@ mod tests {
|
||||
.indent(true)
|
||||
.to_string();
|
||||
|
||||
// We expect a RepartitionExec wrapping a UnionExec
|
||||
assert!(
|
||||
display.contains("RepartitionExec"),
|
||||
"Plan should include Repartitioning"
|
||||
display.contains("KNNVectorDistance: queries=2"),
|
||||
"plan should use native batch KNN, got:\n{display}"
|
||||
);
|
||||
assert!(
|
||||
display.contains("UnionExec"),
|
||||
"Plan should include a Union of multiple searches"
|
||||
!display.contains("UnionExec"),
|
||||
"flat batch KNN should share one scan, got:\n{display}"
|
||||
);
|
||||
// We expect the projection to add the 'query_index' column (logic inside multi_vector_plan)
|
||||
assert!(
|
||||
display.contains("query_index"),
|
||||
"Plan should add query_index column"
|
||||
"plan should add query_index column, got:\n{display}"
|
||||
);
|
||||
}
|
||||
|
||||
#[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;
|
||||
|
||||
@@ -27,6 +27,8 @@ use std::sync::Arc;
|
||||
|
||||
use arrow_array::Array;
|
||||
use arrow_schema::{DataType, Schema as ArrowSchema};
|
||||
use datafusion::common::{DataFusionError, ToDFSchema};
|
||||
use datafusion::prelude::SessionContext;
|
||||
use datafusion_physical_plan::expressions::Column;
|
||||
use datafusion_physical_plan::projection::ProjectionExec;
|
||||
use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
|
||||
@@ -128,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)
|
||||
@@ -391,7 +397,21 @@ fn base_scanner(
|
||||
}
|
||||
if let Some(filter) = &query.base.filter {
|
||||
scanner = match filter {
|
||||
QueryFilter::Sql(sql) => scanner.filter(sql)?,
|
||||
QueryFilter::Sql(sql) => {
|
||||
// Parse here instead of inside `LsmScanner::filter` so the typed
|
||||
// DataFusion `FieldNotFound` error is still available for the
|
||||
// same nested-field enrichment used by the ordinary scanner.
|
||||
let schema = ArrowSchema::from(dataset.schema());
|
||||
let df_schema = schema.clone().to_dfschema().map_err(|error| {
|
||||
enrich_filter_error(error, &schema, "Failed to create DFSchema")
|
||||
})?;
|
||||
let expr = SessionContext::new()
|
||||
.parse_sql_expr(sql, &df_schema)
|
||||
.map_err(|error| {
|
||||
enrich_filter_error(error, &schema, "Failed to parse filter expression")
|
||||
})?;
|
||||
scanner.filter_expr(expr)
|
||||
}
|
||||
QueryFilter::Datafusion(expr) => scanner.filter_expr(expr.clone()),
|
||||
QueryFilter::Substrait(_) => {
|
||||
return Err(Error::NotSupported {
|
||||
@@ -403,6 +423,12 @@ fn base_scanner(
|
||||
Ok(scanner)
|
||||
}
|
||||
|
||||
fn enrich_filter_error(error: DataFusionError, schema: &ArrowSchema, context: &str) -> Error {
|
||||
super::field_not_found_diagnostic(&error, schema).unwrap_or_else(|| Error::InvalidInput {
|
||||
message: format!("{context}: {error}"),
|
||||
})
|
||||
}
|
||||
|
||||
/// Plain scan: filter / projection / limit over base ∪ SSTables ∪ in-memory.
|
||||
/// The plain scan applies limit and offset inside the planner.
|
||||
async fn plain_plan(
|
||||
|
||||
@@ -55,7 +55,9 @@ pub struct DropColumnsResult {
|
||||
pub struct FieldMetadataUpdate {
|
||||
/// Dot-separated path to the field (e.g. `"embedding"` or `"address.zip"`).
|
||||
pub path: String,
|
||||
/// Keys to set (`Some`) or delete (`None`).
|
||||
/// Keys to set (`Some`) or delete (`None`). See
|
||||
/// [`Table::update_field_metadata`](crate::Table::update_field_metadata) for
|
||||
/// the conventional `lancedb:*` keys.
|
||||
pub metadata: HashMap<String, Option<String>>,
|
||||
/// If `true`, replace the field's entire metadata map instead of merging.
|
||||
pub replace: bool,
|
||||
|
||||
@@ -62,22 +62,33 @@ impl UpdateBuilder {
|
||||
}
|
||||
|
||||
/// Executes the update operation.
|
||||
pub async fn execute(self) -> Result<UpdateResult> {
|
||||
pub async fn execute(mut self) -> Result<UpdateResult> {
|
||||
if self.columns.is_empty() {
|
||||
Err(Error::InvalidInput {
|
||||
message: "at least one column must be specified in an update operation".to_string(),
|
||||
})
|
||||
} else {
|
||||
self.canonicalize_filter()?;
|
||||
self.parent.clone().update(self).await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn canonicalize_filter(&mut self) -> Result<()> {
|
||||
self.filter = self
|
||||
.filter
|
||||
.take()
|
||||
.map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal implementation of the update logic
|
||||
pub(crate) async fn execute_update(
|
||||
table: &NativeTable,
|
||||
update: UpdateBuilder,
|
||||
mut update: UpdateBuilder,
|
||||
) -> Result<UpdateResult> {
|
||||
update.canonicalize_filter()?;
|
||||
table.dataset.ensure_mutable()?;
|
||||
|
||||
// 1. Snapshot the current dataset
|
||||
|
||||
Reference in New Issue
Block a user