rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 14:52:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge main into gatekeeper/fix-4260-1

Browse Source
This commit is contained in:
Gatefixer
2026-09-29 16:13:23 +00:00
parent d8b0ab5f83 ea223c4ec4
commit 19e373e751
35 changed files with 1975 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+44 -43
View File
@@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arc-swap",
"arrow",
@@ -5217,7 +5217,7 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "13.0.0-beta.15"
version = "13.0.0-beta.18"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5239,7 +5239,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5253,7 +5253,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5262,8 +5262,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrayref",
"crunchy",
@@ -5273,8 +5273,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5311,8 +5311,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow",
"arrow-array",
@@ -5343,8 +5343,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow",
"arrow-array",
@@ -5361,8 +5361,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"proc-macro2",
"quote",
@@ -5371,8 +5371,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5405,8 +5405,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5438,8 +5438,8 @@ dependencies = [
[[package]]
name = "lance-geo"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"datafusion",
"geo-traits",
@@ -5453,8 +5453,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arc-swap",
"arrow",
@@ -5524,8 +5524,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5548,8 +5548,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow",
"arrow-array",
@@ -5591,8 +5591,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5606,8 +5606,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow",
"async-trait",
@@ -5621,9 +5621,10 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arc-swap",
"arrow",
"arrow-ipc",
"arrow-schema",
@@ -5675,8 +5676,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5690,8 +5691,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow",
"arrow-array",
@@ -5731,8 +5732,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5745,8 +5746,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
version = "13.0.0-beta.18"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b"
dependencies = [
"frostem",
"icu_segmenter",
Cargo.toml
+14 -14
View File
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { version = "=13.0.0-beta.15", path = "python/vendor/lance-arrow" }
lance = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { version = "=13.0.0-beta.18", path = "python/vendor/lance-arrow" }
lancedb = { path = "rust/lancedb", default-features = false }
ahash = "0.8"
# Note that this one does not include pyarrow
docs/src/js/classes/Table.md
+4
View File
@@ -548,6 +548,10 @@ Bytes for `column` at row IDs from [Query.withRowId](Query.md#withrowid).
Reads the table's current checkout. IDs from another version can fail after
compaction unless stable row ids are enabled. Results keep input order and
duplicates. Null blobs are `null`. Empty blobs are empty buffers.
Remote servers limit each request to 1024 row IDs and 64 MiB of blob bytes.
The client splits requests automatically and reads an individual larger
blob through the Range route. This method still materializes all bytes in
memory; use [Table.fetchBlobFiles](Table.md#fetchblobfiles) for large values.
#### Parameters
docs/src/js/interfaces/AddDataOptions.md
+16
View File
@@ -10,6 +10,22 @@ Options for adding data to a table.
## Properties
### allowExternalBlobOutsideBases?
```ts
optional allowExternalBlobOutsideBases: boolean;
```
Whether blob URIs outside registered external bases may be written.
Defaults to `false`. This option is supported only for local/native
tables; remote tables return an error when it is enabled. An enabled write
stores the absolute URI reference without registering a base or copying
the external object into the database. The object must remain accessible
when the blob is read later.
***
### mode
```ts
java/pom.xml
+1 -1
View File
@@ -28,7 +28,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<arrow.version>15.0.0</arrow.version>
<lance-core.version>13.0.0-beta.14</lance-core.version>
<lance-core.version>13.0.0-beta.18</lance-core.version>
<spotless.skip>false</spotless.skip>
<spotless.version>2.30.0</spotless.version>
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
nodejs/__test__/remote.test.ts
+63
View File
@@ -85,6 +85,69 @@ async function withMockDatabase(
}
describe("remote connection", () => {
it("rejects the external blob opt-in without blocking regular adds", async () => {
let insertRequests = 0;
const describeRequests: string[] = [];
await withMockDatabase(
(req, res) => {
const requestPath = req.url ?? "";
if (requestPath.endsWith("/describe/")) {
describeRequests.push(requestPath);
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
name: "items",
version: 1,
schema: {
fields: [
{
name: "id",
type: { type: "int64" },
nullable: true,
},
],
},
}),
);
return;
}
if (requestPath.endsWith("/insert/")) {
insertRequests++;
req.resume();
req.on("end", () => {
res
.writeHead(200, { "Content-Type": "application/json" })
.end(JSON.stringify({ version: insertRequests + 1 }));
});
return;
}
res.writeHead(404).end();
},
async (db) => {
const table = await db.openTable("items");
await expect(
table.add([{ id: 1n }], {
allowExternalBlobOutsideBases: true,
}),
).rejects.toThrow("only supported on local tables");
expect(insertRequests).toBe(0);
await expect(table.add([{ id: 2n }])).resolves.toMatchObject({
version: 2,
});
await expect(
table.add([{ id: 3n }], {
allowExternalBlobOutsideBases: false,
}),
).resolves.toMatchObject({ version: 3 });
},
);
expect(describeRequests.length).toBeGreaterThan(0);
expect(insertRequests).toBe(2);
});
it("lists materialized views through the namespace route", async () => {
await withMockDatabase(
(req, res) => {
nodejs/__test__/table.test.ts
+147
View File
@@ -2,6 +2,7 @@
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import * as fs from "fs";
import { pathToFileURL } from "node:url";
import * as path from "path";
import * as tmp from "tmp";
@@ -416,6 +417,26 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
expect(await table.countRows("id == 10")).toBe(1);
});
it("should let me update float values to NaN and Infinity", async () => {
const db = await connect(tmpDir.name);
const floats = await db.createTable("floats", [
{ id: 1, price: 1.5 },
{ id: 2, price: 2.5 },
{ id: 3, price: 3.5 },
{ id: 4, price: 4.5 },
]);
await floats.update({ values: { price: NaN }, where: "id = 1" });
await floats.update({ values: { price: Infinity }, where: "id = 2" });
await floats.update({ values: { price: -Infinity }, where: "id = 3" });
const rows = await floats.query().toArray();
const price = (id: number) => rows.find((row) => row.id === id).price;
expect(price(1)).toBeNaN();
expect(price(2)).toBe(Infinity);
expect(price(3)).toBe(-Infinity);
expect(price(4)).toBe(4.5);
});
// https://github.com/lancedb/lancedb/issues/1293
test.each([new arrow.Float16(), new arrow.Float32(), new arrow.Float64()])(
"can create empty table with non default float type: %s",
@@ -2429,6 +2450,118 @@ describe("when dealing with blob columns", () => {
await expect(table.countRows()).resolves.toBe(0);
});
it("requires an explicit opt-in for blob URIs outside registered bases", async () => {
const { table, uri } = await openExternalBlobTable();
await expect(table.add([{ id: 1n, payload: uri }])).rejects.toThrow(
"allow_external_blob_outside_bases",
);
await expect(table.countRows()).resolves.toBe(0);
await expect(
table.add([{ id: 2n, payload: uri }], {
allowExternalBlobOutsideBases: false,
}),
).rejects.toThrow("allow_external_blob_outside_bases");
await expect(table.countRows()).resolves.toBe(0);
});
it.each([
["URI string", (uri: string) => uri],
["URI struct", (uri: string) => ({ uri })],
])(
"round-trips an external blob from a %s after reopening",
async (_label, blobValue) => {
const { db, table, uri, payload } = await openExternalBlobTable();
const result = await table.add([{ id: 1n, payload: blobValue(uri) }], {
allowExternalBlobOutsideBases: true,
});
expect(result.version).toBeGreaterThan(0);
await expect(table.countRows()).resolves.toBe(1);
table.close();
const reopened = await db.openTable("external_blobs");
const rows = await reopened.query().withRowId().toArray();
expect(rows).toHaveLength(1);
expect(rows[0].payload.blob_uri).toBe(uri);
const bytes = await reopened.fetchBlobs("payload", [
rows[0]._rowid as bigint,
]);
expect(bytes[0]).toEqual(payload);
},
);
it("applies the external blob opt-in to Arrow table input", async () => {
const { table, schema, uri, payload } = await openExternalBlobTable();
const input = makeArrowTable([{ id: 1n, payload: uri }], { schema });
const result = await table.add(input, {
allowExternalBlobOutsideBases: true,
});
expect(result.version).toBeGreaterThan(0);
const rows = await table.query().withRowId().toArray();
expect(rows[0].payload.blob_uri).toBe(uri);
const bytes = await table.fetchBlobs("payload", [rows[0]._rowid as bigint]);
expect(bytes[0]).toEqual(payload);
});
it("applies the external blob opt-in only to the current add", async () => {
const { table, uri } = await openExternalBlobTable();
await table.add([{ id: 1n, payload: uri }], {
allowExternalBlobOutsideBases: true,
});
await expect(table.countRows()).resolves.toBe(1);
await expect(table.add([{ id: 2n, payload: uri }])).rejects.toThrow(
"allow_external_blob_outside_bases",
);
await expect(table.countRows()).resolves.toBe(1);
});
it("supports progress and overwrite with the external blob opt-in", async () => {
const { table, uri, payload } = await openExternalBlobTable();
await table.add([{ id: 1n, payload: Buffer.from("old") }]);
let resolveFinal!: (progress: import("../lancedb").WriteProgress) => void;
let rejectFinal!: (error: Error) => void;
const finalProgress = new Promise<import("../lancedb").WriteProgress>(
(resolve, reject) => {
resolveFinal = resolve;
rejectFinal = reject;
},
);
const timeout = setTimeout(
() => rejectFinal(new Error("timed out waiting for final progress")),
5_000,
);
try {
const result = await table.add([{ id: 2n, payload: uri }], {
mode: "overwrite",
allowExternalBlobOutsideBases: true,
progress: (progress) => {
if (progress.done) {
resolveFinal(progress);
}
},
});
expect(result.version).toBeGreaterThan(0);
const progress = await finalProgress;
expect(progress.outputRows).toBe(1);
expect(progress.totalRows).toBe(1);
} finally {
clearTimeout(timeout);
}
const rows = await table.query().withRowId().toArray();
expect(rows).toHaveLength(1);
expect(rows[0].id).toBe(2n);
const bytes = await table.fetchBlobs("payload", [rows[0]._rowid as bigint]);
expect(bytes[0]).toEqual(payload);
});
it("discovers blob columns", async () => {
const { table } = await openBlobTable();
expect(await table.blobColumns()).toEqual(["image"]);
@@ -2688,6 +2821,20 @@ describe("when dealing with blob columns", () => {
const rowIds = [1, 2, 3].map((id) => rowIdById.get(id)!);
return { table, rowIds, alpha, beta };
}
async function openExternalBlobTable() {
const db = await connect(path.join(tmpDir.name, "db"));
const payload = Buffer.from("external-payload");
const payloadPath = path.join(tmpDir.name, "payload.bin");
fs.writeFileSync(payloadPath, payload);
const uri = pathToFileURL(payloadPath).href;
const schema = new Schema([
new Field("id", new Int64(), true),
blob("payload"),
]);
const table = await db.createEmptyTable("external_blobs", schema);
return { db, table, schema, uri, payload };
}
});
describe("when dealing with tags", () => {
nodejs/__test__/util.test.ts
+3
View File
@@ -14,6 +14,9 @@ test.each([
[new ArrayBuffer(8), "X'0000000000000000'"],
[Buffer.from("hello"), "X'68656c6c6f'"],
["Hello 'world'", "'Hello ''world'''"],
[NaN, "CAST('NaN' AS DOUBLE)"],
[Infinity, "CAST('Infinity' AS DOUBLE)"],
[-Infinity, "CAST('-Infinity' AS DOUBLE)"],
])("toSQL(%p) === %p", (value, expected) => {
expect(toSQL(value)).toBe(expected);
});
nodejs/lancedb/table.ts
+21 -1
View File
@@ -124,6 +124,17 @@ export interface AddDataOptions {
* ```
*/
progress: (progress: WriteProgress) => void;
/**
* Whether blob URIs outside registered external bases may be written.
*
* Defaults to `false`. This option is supported only for local/native
* tables; remote tables return an error when it is enabled. An enabled write
* stores the absolute URI reference without registering a base or copying
* the external object into the database. The object must remain accessible
* when the blob is read later.
*/
allowExternalBlobOutsideBases?: boolean;
}
export interface UpdateOptions {
@@ -526,6 +537,10 @@ export abstract class Table {
* Reads the table's current checkout. IDs from another version can fail after
* compaction unless stable row ids are enabled. Results keep input order and
* duplicates. Null blobs are `null`. Empty blobs are empty buffers.
* Remote servers limit each request to 1024 row IDs and 64 MiB of blob bytes.
* The client splits requests automatically and reads an individual larger
* blob through the Range route. This method still materializes all bytes in
* memory; use {@link Table.fetchBlobFiles} for large values.
*/
abstract fetchBlobs(
column: string,
@@ -1091,7 +1106,12 @@ export class LocalTable extends Table {
}
}
: undefined;
return await this.inner.add(buffer, mode, progress);
return await this.inner.add(
buffer,
mode,
progress,
options?.allowExternalBlobOutsideBases,
);
}
async update(
nodejs/lancedb/util.ts
+8
View File
@@ -15,6 +15,14 @@ export function toSQL(value: IntoSql): string {
if (typeof value === "string") {
return `'${value.replace(/'/g, "''")}'`;
} else if (typeof value === "number") {
// toString() gives "NaN" / "Infinity", which SQL reads as column names.
if (Number.isNaN(value)) {
return "CAST('NaN' AS DOUBLE)";
} else if (!Number.isFinite(value)) {
return value > 0
? "CAST('Infinity' AS DOUBLE)"
: "CAST('-Infinity' AS DOUBLE)";
}
return value.toString();
} else if (typeof value === "boolean") {
return value ? "TRUE" : "FALSE";
nodejs/src/query.rs
+1
View File
@@ -658,6 +658,7 @@ impl JsFullTextQuery {
FtsQuery::Phrase(_) => "phrase".to_string(),
FtsQuery::Boost(_) => "boost".to_string(),
FtsQuery::MultiMatch(_) => "multi_match".to_string(),
FtsQuery::CombinedFields(_) => "combined_fields".to_string(),
FtsQuery::Boolean(_) => "boolean".to_string(),
}
}
nodejs/src/table.rs
+5 -1
View File
@@ -74,13 +74,14 @@ impl Table {
#[napi(
catch_unwind,
ts_args_type = "buf: Buffer, mode: string, progressCallback?: (progress: WriteProgressInfo) => void"
ts_args_type = "buf: Buffer, mode: string, progressCallback?: (progress: WriteProgressInfo) => void, allowExternalBlobOutsideBases?: boolean"
)]
pub async fn add(
&self,
buf: Buffer,
mode: String,
progress_callback: Option<ProgressFn>,
allow_external_blob_outside_bases: Option<bool>,
) -> napi::Result<AddResult> {
let batches = ipc_file_to_batches(buf.to_vec())
.map_err(|e| napi::Error::from_reason(format!("Failed to read IPC file: {}", e)))?;
@@ -105,6 +106,9 @@ impl Table {
return Err(napi::Error::from_reason(format!("Invalid mode: {}", mode)));
};
op = op
.allow_external_blob_outside_bases(allow_external_blob_outside_bases.unwrap_or(false));
if let Some(tsfn) = progress_callback {
op = op.progress(move |p| {
// NonBlocking: dispatch onto the JS event loop without
python/python/lancedb/query.py
+8 -8
View File
@@ -2296,7 +2296,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
norm=self._norm,
fts_query=self._fts_query._query,
reranker=self._reranker,
limit=self._limit,
limit=self._limit or DEFAULT_HYBRID_LIMIT,
with_row_ids=True,
offset=self._offset,
)
@@ -2748,13 +2748,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
)
# Apply common configurations
if self._limit:
# The final offset/limit window is sliced out of the combined,
# reranked results, so each sub-query must fetch enough rows to
# cover the skipped prefix as well as the window itself.
sub_query_limit = self._limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
# The final offset/limit window is sliced out of the combined,
# reranked results, so each sub-query must fetch enough rows to
# cover the skipped prefix as well as the window itself.
limit = self._limit or DEFAULT_HYBRID_LIMIT
sub_query_limit = limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
# WAL-PK-FUSION: without the fallback, select `self._columns` as is.
self._pk_fusion = None
columns = self._columns
python/python/lancedb/rerankers/rrf.py
+20 -4
View File
@@ -107,8 +107,24 @@ class RRFReranker(Reranker):
`search().with_row_id(True)`"
)
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
empty_table = pa.Table.from_arrays([], names=[])
reranked = self.rerank_hybrid(query, combined, empty_table)
# RRF ranks each result list on its own, so a document's rank is its
# position within the list it came from, not within the concatenation.
rrf_score_map = defaultdict(float)
for result in vector_results:
for i, result_id in enumerate(result["_rowid"].to_pylist(), 1):
rrf_score_map[result_id] += 1 / (i + self.K)
return reranked
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
combined = self._deduplicate(combined)
relevance_scores = [
rrf_score_map[row_id] for row_id in combined["_rowid"].to_pylist()
]
combined = combined.append_column(
"_relevance_score", pa.array(relevance_scores, type=pa.float32())
)
combined = combined.sort_by([("_relevance_score", "descending")])
if self.score == "relevance":
combined = self._keep_relevance_score(combined)
return combined
python/python/lancedb/table.py
+168 -21
View File
@@ -134,6 +134,34 @@ _MODEL_BACKED_TOKENIZER_ERRORS = (
)
def _optimize_cleanup_since_ms(
cleanup_older_than: Optional[timedelta], retrain: bool
) -> Optional[int]:
# Called directly by both the sync and async optimize so stacklevel=3
# names the user's call site rather than the background event loop.
cleanup_since_ms: Optional[int] = None
if cleanup_older_than is not None:
cleanup_since_ms = round(cleanup_older_than.total_seconds() * 1000)
if cleanup_since_ms <= 0:
warnings.warn(
"optimize(cleanup_older_than=0) removes every version except "
"the latest. Any concurrent reader or writer still using an "
"older version will fail. Use a longer cleanup_older_than "
"unless no other process is working on this table.",
UserWarning,
stacklevel=3,
)
if retrain:
warnings.warn(
"The 'retrain' parameter is deprecated and will be removed in a "
"future version.",
DeprecationWarning,
stacklevel=3,
)
return cleanup_since_ms
def _add_unique_note(exception: BaseException, note: str) -> None:
existing_notes = getattr(exception, "__notes__", ()) or ()
message = (
@@ -233,6 +261,101 @@ IndexConfigType = Union[
KNOWN_METRICS = {"l2", "cosine", "dot", "hamming"}
def _blob_value_to_storage(value: Any) -> Optional[dict]:
"""Keep a Python blob's inline data or external URI before Arrow infers it."""
if value is None:
return None
if isinstance(value, (bytes, bytearray, memoryview)):
return {"data": bytes(value)}
if isinstance(value, str):
if not value:
raise ValueError("Blob uri cannot be empty")
return {"uri": value}
if isinstance(value, dict):
unknown = value.keys() - {"data", "uri", "position", "size"}
if unknown:
raise ValueError(f"Unknown blob fields: {sorted(unknown)}")
return value
try:
from lance.blob import Blob
except ModuleNotFoundError as err:
if err.name not in ("lance", "lance.blob"):
raise
else:
if isinstance(value, Blob):
return {
"data": value.data,
"uri": value.uri,
"position": value.position,
"size": value.size,
}
raise TypeError(f"Unsupported blob value: {type(value).__name__}")
def _blob_input_to_arrow(data: Any, schema: Optional[pa.Schema]) -> Optional[pa.Table]:
"""Convert Python rows with blob fields before Arrow loses their value types."""
if schema is None:
return None
if isinstance(data, list) and data and isinstance(data[0], dict):
names = {
field.name
for field in schema
if is_blob_v2_field(field)
and any(isinstance(row, dict) and field.name in row for row in data)
}
if not names:
return None
values = {name: [row.get(name) for row in data] for name in names}
rows = [{**row, **{name: None for name in names}} for row in data]
table = pa.Table.from_pylist(rows)
elif _check_for_pandas(data) and isinstance(data, pd.DataFrame):
names = {
field.name
for field in schema
if is_blob_v2_field(field) and field.name in data.columns
}
if not names:
return None
values = {
name: [
None
if value is pd.NA or isinstance(value, float) and np.isnan(value)
else value
for value in data[name]
]
for name in names
}
table = pa.Table.from_pandas(
data.assign(**{name: None for name in names}), preserve_index=False
).replace_schema_metadata(None)
else:
return None
for field in schema:
if field.name not in names:
continue
storage_type = (
field.type.storage_type
if isinstance(field.type, pa.ExtensionType)
else field.type
)
storage = pa.array(
[_blob_value_to_storage(value) for value in values[field.name]],
type=storage_type,
)
column = (
pa.ExtensionArray.from_storage(field.type, storage)
if isinstance(field.type, pa.ExtensionType)
else storage
)
table = table.set_column(
table.schema.get_field_index(field.name), field, column
)
return table
def _into_pyarrow_reader(
data, schema: Optional[pa.Schema] = None
) -> pa.RecordBatchReader:
@@ -275,9 +398,14 @@ def _into_pyarrow_reader(
return pa.Table.from_batches(data).to_reader()
else:
data = _serialize_json_values(data, schema)
return pa.Table.from_pylist(data).to_reader()
table = _blob_input_to_arrow(data, schema)
return (
table if table is not None else pa.Table.from_pylist(data)
).to_reader()
elif _check_for_pandas(data) and isinstance(data, pd.DataFrame):
table = pa.Table.from_pandas(data, preserve_index=False)
table = _blob_input_to_arrow(data, schema)
if table is None:
table = pa.Table.from_pandas(data, preserve_index=False)
# Do not serialize Pandas metadata
meta = table.schema.metadata if table.schema.metadata is not None else {}
meta = {k: v for k, v in meta.items() if k != b"pandas"}
@@ -1935,8 +2063,10 @@ class Table(ABC):
``_rowid`` values stay valid after compaction when the table has stable
row ids.
Convenience for small payloads. For large values use
:meth:`fetch_blob_files`.
Remote servers limit each request to 1024 row IDs and 64 MiB of blob
bytes. The client splits requests automatically and reads an individual
larger blob through the Range route. This method still materializes all
bytes in memory; for large values use :meth:`fetch_blob_files`.
"""
@abstractmethod
@@ -2217,6 +2347,13 @@ class Table(ABC):
All files belonging to versions older than this will be removed. Set
to 0 days to remove all versions except the latest. The latest version
is never removed.
.. warning::
Setting this to 0 deletes the data files of every older
version, so any other reader or writer still using an older
version of the table will fail. Only set it to 0 if no other
process is working on this dataset.
delete_unverified: bool, default False
Files leftover from a failed transaction may appear to be part of an
in-progress operation (e.g. appending new data) and these files will not
@@ -4423,6 +4560,13 @@ class LanceTable(Table):
All files belonging to versions older than this will be removed. Set
to 0 days to remove all versions except the latest. The latest version
is never removed.
.. warning::
Setting this to 0 deletes the data files of every older
version, so any other reader or writer still using an older
version of the table will fail. Only set it to 0 if no other
process is working on this dataset.
delete_unverified: bool, default False
Files leftover from a failed transaction may appear to be part of an
in-progress operation (e.g. appending new data) and these files will not
@@ -4447,10 +4591,9 @@ class LanceTable(Table):
modification operations.
"""
LOOP.run(
self._table.optimize(
cleanup_older_than=cleanup_older_than,
delete_unverified=delete_unverified,
retrain=retrain,
self._table._do_optimize(
_optimize_cleanup_since_ms(cleanup_older_than, retrain),
delete_unverified,
)
)
@@ -5848,6 +5991,9 @@ class AsyncTable:
"""
schema = await self.schema()
data = _serialize_json_values(data, schema)
blob_table = _blob_input_to_arrow(data, schema)
if blob_table is not None:
data = blob_table
if on_bad_vectors is None:
on_bad_vectors = "error"
if fill_value is None:
@@ -6958,6 +7104,13 @@ class AsyncTable:
All files belonging to versions older than this will be removed. Set
to 0 days to remove all versions except the latest. The latest version
is never removed.
.. warning::
Setting this to 0 deletes the data files of every older
version, so any other reader or writer still using an older
version of the table will fail. Only set it to 0 if no other
process is working on this dataset.
delete_unverified: bool, default False
Files leftover from a failed transaction may appear to be part of an
in-progress operation (e.g. appending new data) and these files will not
@@ -6981,20 +7134,14 @@ class AsyncTable:
you have added or modified 100,000 or more records or run more than 20 data
modification operations.
"""
cleanup_since_ms: Optional[int] = None
if cleanup_older_than is not None:
cleanup_since_ms = round(cleanup_older_than.total_seconds() * 1000)
if retrain:
import warnings
warnings.warn(
"The 'retrain' parameter is deprecated and will be removed in a "
"future version.",
DeprecationWarning,
stacklevel=2,
)
return await self._do_optimize(
_optimize_cleanup_since_ms(cleanup_older_than, retrain),
delete_unverified,
)
async def _do_optimize(
self, cleanup_since_ms: Optional[int], delete_unverified: bool
) -> OptimizeStats:
return await self._inner.optimize(
cleanup_since_ms=cleanup_since_ms,
delete_unverified=delete_unverified,
python/python/lancedb/util.py
+8
View File
@@ -5,6 +5,7 @@
import binascii
import functools
import importlib
import math
import os
import pathlib
import warnings
@@ -351,6 +352,13 @@ def _(value: int):
@value_to_sql.register(float)
def _(value: float):
# str() gives "nan" / "inf", which SQL would read as column names.
if math.isnan(value):
return "CAST('NaN' AS DOUBLE)"
if math.isinf(value):
return (
"CAST('Infinity' AS DOUBLE)" if value > 0 else "CAST('-Infinity' AS DOUBLE)"
)
return str(value)
python/python/tests/test_blob.py
+128 -1
View File
@@ -7,10 +7,11 @@ import sys
import textwrap
import lance
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import BlobType as LanceBlobType
from lance.blob import Blob, BlobType as LanceBlobType
import lancedb
from lancedb._blob import (
@@ -1432,3 +1433,129 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
@pytest.mark.parametrize("as_pandas", [False, True])
def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pandas):
payload = b"external-payload"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
rows = [
{"id": 1, "image": b"inline"},
{"id": 2, "image": blob_path.as_uri()},
]
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("mixed", schema=schema)
table.add(
pd.DataFrame(rows) if as_pandas else rows,
allow_external_blob_outside_bases=True,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
b"inline",
payload,
]
@pytest.mark.parametrize("as_pandas", [False, True])
def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_pandas):
payload = b"external-payload"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
uri = blob_path.as_uri()
rows = [
{"id": 1, "image": b"inline"},
{"id": 2, "image": uri},
{"id": 3, "image": {"data": b"from-dict"}},
{"id": 4, "image": Blob.from_bytes(b"from-blob")},
{"id": 5, "image": Blob.from_uri(uri)},
{"id": 6, "image": {"uri": uri}},
{"id": 7, "image": None},
]
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("mixed", schema=schema)
table.add(
pd.DataFrame(rows) if as_pandas else rows,
allow_external_blob_outside_bases=True,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[i] for i in range(1, 8)]).to_pylist() == [
b"inline",
payload,
b"from-dict",
b"from-blob",
payload,
payload,
None,
]
@pytest.mark.parametrize("as_pandas", [False, True])
def test_merge_insert_mixed_python_blobs_validates_external_uri(tmp_path, as_pandas):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"external-payload")
table = _blob_table("merge_mixed_uri", [{"id": 1, "image": b"before"}])
rows = [
{"id": 1, "image": b"updated"},
{"id": 2, "image": blob_path.as_uri()},
]
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
(
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(pd.DataFrame(rows) if as_pandas else rows)
)
assert table.count_rows() == 1
@pytest.mark.parametrize("as_pandas", [False, True])
def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
table = _blob_table("merge_python_blobs", [{"id": 1, "image": b"before"}])
rows = [
{"id": 1, "image": b"updated"},
{"id": 2, "image": {"data": b"from-dict"}},
{"id": 3, "image": Blob.from_bytes(b"from-blob")},
]
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(pd.DataFrame(rows) if as_pandas else rows)
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 2
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[i] for i in range(1, 4)]).to_pylist() == [
b"updated",
b"from-dict",
b"from-blob",
]
def test_create_table_with_python_blob_dict_and_object():
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table(
"create_python_blobs",
data=[
{"id": 1, "image": {"data": b"from-dict"}},
{"id": 2, "image": Blob.from_bytes(b"from-blob")},
],
schema=schema,
)
by_id = _row_ids_by_id(table)
assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
b"from-dict",
b"from-blob",
]
python/python/tests/test_hybrid_query.py
+27
View File
@@ -315,6 +315,33 @@ def test_hybrid_query_offset(sync_table: Table):
assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:]
def test_hybrid_query_default_limit(sync_table: Table):
# The vector and FTS legs match disjoint rows, so fusing their results
# yields more rows than the default limit.
new_rows = []
for i in range(20):
new_rows.append({"text": "close_vec", "vector": [0.1, 0.1]})
new_rows.append({"text": "dog", "vector": [50.0 + i, 50.0 + i]})
sync_table.add(new_rows)
def query():
return (
sync_table.search(query_type="hybrid")
.vector([0.1, 0.1])
.text("dog")
.with_row_id(True)
)
# Like the async hybrid query, the default limit is 10.
result = query().to_arrow()
assert len(result) == 10
# The offset window is taken after the default limit is applied.
full = query().limit(15).to_arrow()
offset_result = query().offset(5).to_arrow()
assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[5:]
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
# minimum_nprobes(0) must raise the same validation error a plain vector
# query raises, not silently no-op because 0 is falsy.
python/python/tests/test_remote_db.py
+16 -4
View File
@@ -2403,11 +2403,16 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
elif request.path.startswith("/v1/table/test/blob/image/"):
path = request.path.partition("?")[0]
row_id = int(path.split("/")[-2])
payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id]
payload = {10: b"alpha", 20: None, 30: b"gamma", 40: b""}[row_id]
if payload is None:
request.send_response(204)
request.end_headers()
return
if not payload:
request.send_response(416)
request.send_header("Content-Range", "bytes */0")
request.end_headers()
return
byte_range = request.headers["Range"].removeprefix("bytes=")
start_text, end_text = byte_range.split("-", maxsplit=1)
start = int(start_text)
@@ -2416,6 +2421,9 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
request.send_response(206)
request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
request.send_header("Content-Length", str(len(chunk)))
request.send_header(
"x-lancedb-version", str(BLOB_DESCRIBE_RESPONSE["version"])
)
request.end_headers()
request.wfile.write(chunk)
elif request.path == "/v1/table/test/query/":
@@ -2459,17 +2467,21 @@ def test_remote_blob_columns_and_fetch():
def test_remote_blob_files_are_lazy_seekable_handles():
with blob_remote_table() as table:
files = table.fetch_blob_files("image", [10, 20, 30])
files = table.fetch_blob_files("image", [10, 20, 30, 40])
assert len(files) == 3
alpha, null_row, gamma = files
assert len(files) == 4
alpha, null_row, gamma, empty = files
assert null_row is None
assert alpha is not None
assert gamma is not None
assert empty is not None
assert alpha.size() == 5
assert alpha.read_range(1, 3) == b"lph"
gamma.seek(2)
assert gamma.read() == b"mma"
assert empty.size() == 0
assert empty.read() == b""
assert empty.read_range(0, 0) == b""
alpha.close()
assert alpha.closed
with pytest.raises(RuntimeError, match="already closed"):
python/python/tests/test_rerankers.py
+32
View File
@@ -395,6 +395,38 @@ def test_mrr_multivector_rewards_consensus():
assert result["_rowid"].to_pylist()[0] == 2
def test_rrf_multivector_ranks_each_result_list():
# RRF scores a document by its rank *within each* result list. The top hit
# of the second list must score the same as the top hit of the first list,
# not as if it were ranked after every row of the first list.
reranker = RRFReranker(K=60)
def ranking(row_ids):
return pa.table({"_rowid": pa.array(row_ids, type=pa.int64())})
rs1 = ranking([1, 2, 3])
rs2 = ranking([4, 5, 6])
result = reranker.rerank_multivector([rs1, rs2])
scores = dict(
zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
)
assert scores[1] == pytest.approx(1 / 61)
assert scores[4] == pytest.approx(1 / 61)
assert scores[5] == pytest.approx(1 / 62)
# the second-ranked hit of the second list beats the third hit of the first
assert scores[5] > scores[3]
# A document found by both lists sums its reciprocal ranks from each list.
result = reranker.rerank_multivector([ranking([1, 2]), ranking([2, 3])])
scores = dict(
zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
)
assert scores[2] == pytest.approx(1 / 62 + 1 / 61)
assert result["_rowid"].to_pylist() == [2, 1, 3]
def test_rrf_reranker_distance():
data = pa.table(
{
python/python/tests/test_table.py
+28 -4
View File
@@ -4307,7 +4307,8 @@ async def test_optimize(mem_db_async: AsyncConnection):
assert stats.prune.bytes_removed == 0
assert stats.prune.old_versions_removed == 0
stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
with pytest.warns(UserWarning, match="concurrent"):
stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
assert stats.prune.bytes_removed > 0
assert stats.prune.old_versions_removed == 3
@@ -4335,12 +4336,35 @@ async def test_optimize_delete_unverified(tmp_db_async: AsyncConnection, tmp_pat
stats = await table.optimize(delete_unverified=False)
assert stats.prune.old_versions_removed == 0
stats = await table.optimize(
cleanup_older_than=timedelta(seconds=0), delete_unverified=True
)
with pytest.warns(UserWarning, match="concurrent"):
stats = await table.optimize(
cleanup_older_than=timedelta(seconds=0), delete_unverified=True
)
assert stats.prune.old_versions_removed == 2
@pytest.mark.asyncio
async def test_optimize_warns_on_zero_cleanup(mem_db_async: AsyncConnection):
table = await mem_db_async.create_table("test", data=[{"x": [1]}])
with pytest.warns(UserWarning, match="concurrent"):
await table.optimize(cleanup_older_than=timedelta(0))
with warnings.catch_warnings():
warnings.simplefilter("error")
await table.optimize(cleanup_older_than=timedelta(days=1))
await table.optimize()
def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection):
table = mem_db.create_table("test", data=[{"x": [1]}])
with warnings.catch_warnings(record=True) as seen:
warnings.simplefilter("default")
table.optimize(cleanup_older_than=timedelta(0))
table.optimize(cleanup_older_than=timedelta(0))
assert [w.filename for w in seen] == [__file__, __file__]
assert all("concurrent" in str(w.message) for w in seen)
def test_replace_field_metadata(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("my_table", data=[{"x": 0}])
python/python/tests/test_util.py
+29
View File
@@ -213,6 +213,35 @@ def test_value_to_sql_numpy_scalars():
assert value_to_sql(np.bool_(False)) == "FALSE"
def test_value_to_sql_nan_and_infinity(tmp_path):
# str(float("nan")) is "nan", which SQL reads as a column named nan, so
# table.update(values={"x": float("nan")}) failed with "No field named nan".
import math
import numpy as np
db = lancedb.connect(tmp_path)
table = db.create_table(
"test",
pa.table(
{
"id": [1, 2, 3, 4],
"f64": [1.0, 2.0, 3.0, 4.0],
"f32": pa.array([1.0, 2.0, 3.0, 4.0], pa.float32()),
}
),
)
table.update(where="id = 1", values={"f64": float("nan"), "f32": np.nan})
table.update(where="id = 2", values={"f64": math.inf, "f32": np.float32("inf")})
table.update(where="id = 3", values={"f64": -math.inf, "f32": -math.inf})
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert math.isnan(rows[1]["f64"]) and math.isnan(rows[1]["f32"])
assert rows[2]["f64"] == math.inf and rows[2]["f32"] == math.inf
assert rows[3]["f64"] == -math.inf and rows[3]["f32"] == -math.inf
assert rows[4]["f64"] == 4.0 and rows[4]["f32"] == 4.0
def test_append_vector_columns():
registry = EmbeddingFunctionRegistry.get_instance()
registry.register("test")(MockTextEmbeddingFunction)
python/src/query.rs
+3
View File
@@ -180,6 +180,9 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
.expect("Failed to import namespace");
match self.0 {
FtsQuery::CombinedFields(_) => Err(PyValueError::new_err(
"Combined-fields FTS queries are not yet supported in Python",
)),
FtsQuery::Match(query) => {
let kwargs = PyDict::new(py);
kwargs.set_item("boost", query.boost)?;
python/vendor/lance-arrow/Cargo.toml
Vendored
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lance-arrow"
version = "13.0.0-beta.15"
version = "13.0.0-beta.18"
edition = "2024"
authors = ["Lance Devs <dev@lance.org>"]
license = "Apache-2.0"
python/vendor/lance-arrow/README.lancedb.md
Vendored
+2 -2
View File
@@ -1,5 +1,5 @@
This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.15`
(commit `db53dcf367088b32ae3a7461407d43593759ae3e`), licensed under Apache-2.0.
This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.18`
(commit `b646eda6a891f481b6f6d6baef9cd17656aef65b`), licensed under Apache-2.0.
LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps
until a Lance release includes the fix. The patch is selected in the root
rust/lancedb/src/materialized_view.rs
+1 -1
View File
@@ -13,7 +13,7 @@ mod duplicate_pairs;
mod grouped;
mod grouped_units;
mod partitioned;
mod vector_dedup;
pub(crate) mod vector_dedup;
pub use duplicate_pairs::{VectorSource, VectorSourceKind};
pub use grouped::IVF_PARTITION;
/// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index
rust/lancedb/src/materialized_view/vector_dedup.rs
+1 -1
View File
@@ -156,7 +156,7 @@ fn edge_schema() -> SchemaRef {
/// Sorting canonical edges makes the first unremoved endpoint the smallest
/// unassigned representative. The sort spills; only vertex masks stay in RAM.
async fn select_representatives(
pub async fn select_representatives(
stream: SendableRecordBatchStream,
memory: usize,
) -> Result<(RoaringTreemap, RoaringTreemap)> {
rust/lancedb/src/remote/table.rs
+236
View File
@@ -96,6 +96,8 @@ const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching sch
fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool {
match query {
// Combined-fields queries do not expose a document granularity option.
FtsQuery::CombinedFields(_) => false,
FtsQuery::Match(query) => query
.document_granularity
.is_some_and(|granularity| granularity.is_list_element()),
@@ -5187,6 +5189,193 @@ mod tests {
assert_eq!(blobs.value(2), b"gamma");
}
#[tokio::test]
async fn test_fetch_blobs_splits_row_ids_at_one_version_and_preserves_order() {
let request_sizes = Arc::new(std::sync::Mutex::new(Vec::new()));
let seen = request_sizes.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
if request.url().path() == "/v1/table/my_table/describe/" {
return http::Response::builder()
.status(200)
.body(br#"{"version":7,"schema":{"fields":[]}}"#.to_vec())
.unwrap();
}
assert_eq!(request.url().path(), "/v1/table/my_table/fetch_blobs/");
let body = request_body_json(&request);
assert_eq!(body["version"], 7);
let ids = body["row_ids"].as_array().unwrap();
seen.lock().unwrap().push(ids.len());
if ids.len() > 1024 {
return http::Response::builder()
.status(400)
.body(b"fetch_blobs accepts at most 1024 row IDs".to_vec())
.unwrap();
}
let mut builder = LargeBinaryBuilder::new();
for id in ids {
let id = id.as_u64().unwrap();
if id == 1023 {
builder.append_null();
} else {
builder.append_value(id.to_string().as_bytes());
}
}
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(builder.finish())],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(write_ipc_stream_uncompressed(&batch))
.unwrap()
},
);
let ids: Vec<u64> = (0..1024).chain([42]).collect();
let blobs = table.fetch_blobs("image", &ids).await.unwrap();
assert_eq!(blobs.len(), 1025);
assert_eq!(blobs.value(0), b"0");
assert_eq!(blobs.value(1022), b"1022");
assert!(blobs.is_null(1023));
assert_eq!(blobs.value(1024), b"42");
assert_eq!(request_sizes.lock().unwrap().as_slice(), &[1024, 1]);
}
#[tokio::test]
async fn test_fetch_blobs_splits_byte_limited_requests_and_reads_large_blob_by_range() {
// Simulate a lower byte cap so this test exercises the same 400 response
// without allocating 64 MiB of blob data.
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let seen = requests.clone();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |request| {
let path = request.url().path();
if path == "/v1/table/my_table/describe/" {
return http::Response::builder()
.status(200)
.body(br#"{"version":42,"schema":{"fields":[]}}"#.to_vec())
.unwrap();
}
if path == "/v1/table/my_table/fetch_blobs/" {
let body = request_body_json(&request);
let ids = body["row_ids"].as_array().unwrap();
if body["version"].is_null() {
// Only the initial failed request may read live latest.
assert_eq!(ids.len(), 5);
} else {
assert_eq!(body["version"], 42);
}
seen.lock().unwrap().push(format!("POST {}", ids.len()));
let mut builder = LargeBinaryBuilder::new();
let mut total_bytes = 0;
for id in ids {
let value: Option<&[u8]> = match id.as_u64().unwrap() {
10 => Some(b"aaaa"),
20 => Some(b"bbb"),
30 => None,
40 => Some(b"0123456789"),
id => panic!("unexpected row id {id}"),
};
if let Some(value) = value {
total_bytes += value.len();
builder.append_value(value);
} else {
builder.append_null();
}
}
if total_bytes > 6 {
return http::Response::builder()
.status(400)
.body(br#"{"error":"Bad request: fetch_blobs accepts at most 67108864 total blob bytes"}"#.to_vec())
.unwrap();
}
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(builder.finish())],
)
.unwrap();
return http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(write_ipc_stream_uncompressed(&batch))
.unwrap();
}
assert_eq!(path, "/v1/table/my_table/blob/image/40/bytes");
assert!(request.url().query().unwrap().contains("version=42"));
let range = request
.headers()
.get(reqwest::header::RANGE)
.unwrap()
.to_str()
.unwrap();
seen.lock().unwrap().push(format!("GET {range}"));
match range {
"bytes=0-0" => http::Response::builder()
.status(206)
.header(reqwest::header::CONTENT_RANGE, "bytes 0-0/10")
.header(VERSION_HEADER, "42")
.body(b"0".to_vec())
.unwrap(),
"bytes=0-" => http::Response::builder()
.status(206)
.header(reqwest::header::CONTENT_RANGE, "bytes 0-9/10")
.body(b"0123456789".to_vec())
.unwrap(),
_ => panic!("unexpected range: {range}"),
}
},
);
let blobs = table
.fetch_blobs("image", &[10, 20, 30, 40, 20])
.await
.unwrap();
assert_eq!(blobs.len(), 5);
assert_eq!(blobs.value(0), b"aaaa");
assert_eq!(blobs.value(1), b"bbb");
assert!(blobs.is_null(2));
assert_eq!(blobs.value(3), b"0123456789");
assert_eq!(blobs.value(4), b"bbb");
let requests = requests.lock().unwrap();
assert!(requests.contains(&"GET bytes=0-0".to_string()));
assert!(requests.contains(&"GET bytes=0-".to_string()));
}
#[tokio::test]
async fn test_fetch_blobs_does_not_split_unrelated_bad_requests() {
let requests = Arc::new(AtomicUsize::new(0));
let seen = requests.clone();
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 5, 0), move |_| {
seen.fetch_add(1, Ordering::SeqCst);
http::Response::builder()
.status(400)
.body(b"unknown blob column".to_vec())
.unwrap()
});
assert_fetch_blobs_http_error(
table.fetch_blobs("missing", &[10, 20]).await.unwrap_err(),
"unknown blob column",
);
assert_eq!(requests.load(Ordering::SeqCst), 1);
}
fn table_with_fetch_blobs_response(body: Vec<u8>) -> Table {
table_with_fetch_blobs_content_type(Some(ARROW_STREAM_CONTENT_TYPE), body)
}
@@ -6152,6 +6341,53 @@ mod tests {
.unwrap();
}
#[tokio::test]
async fn test_query_combined_fields_uses_structured_fts() {
use lance_index::scalar::inverted::query::CombinedFieldsQuery;
let table =
Table::new_with_handler_version("my_table", semver::Version::new(0, 3, 0), |request| {
let body = request.body().unwrap().as_bytes().unwrap();
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
body["full_text_query"]["query"],
serde_json::json!({
"combined_fields": {
"query": "hello world",
"columns": ["title", "text"],
"boost": [1.0, 1.0],
"operator": "Or"
}
})
);
let data = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&data))
.unwrap()
});
table
.query()
.full_text_search(FullTextSearchQuery::new_query(
CombinedFieldsQuery::try_new(
"hello world".into(),
vec!["title".into(), "text".into()],
)
.unwrap()
.into(),
))
.execute()
.await
.unwrap();
}
#[tokio::test]
async fn test_query_row_document_granularity_uses_structured_fts() {
let table =
rust/lancedb/src/remote/table/blobs.rs
+216 -13
View File
@@ -21,7 +21,23 @@ use crate::error::Result;
use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient};
use crate::table::BaseTable;
use super::{FreshnessHeaders, FreshnessState, RemoteTable, freshness_headers_snapshot};
use super::{
FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, VERSION_HEADER,
freshness_headers_snapshot,
};
// The Cloud route rejects larger row-id lists. Its separate 64 MiB byte limit
// is handled by splitting a rejected request below.
const MAX_FETCH_BLOBS_ROW_IDS: usize = 1024;
fn is_fetch_blobs_byte_limit_error(error: &Error) -> bool {
matches!(error, Error::Http {
source,
status_code: Some(StatusCode::BAD_REQUEST),
..
} if source.to_string().contains("fetch_blobs accepts at most")
&& source.to_string().contains("total blob bytes"))
}
#[derive(Debug, Clone, Copy)]
enum RangeRequestMode {
@@ -35,6 +51,7 @@ trait BlobRangeRequester: Send + Sync + std::fmt::Debug {
&self,
range_header: &str,
mode: RangeRequestMode,
version: Option<u64>,
) -> Result<(String, Response)>;
}
@@ -56,13 +73,14 @@ impl<S: HttpSend> BlobRangeRequester for TableBlobRangeRequester<S> {
&self,
range_header: &str,
mode: RangeRequestMode,
version: Option<u64>,
) -> Result<(String, Response)> {
let freshness_request =
freshness_headers_snapshot(&self.freshness, self.read_consistency_interval);
let mut request = freshness_request
.apply(self.client.get(&self.path))
.header(header::RANGE, range_header);
if let Some(version) = self.version {
if let Some(version) = version.or(self.version) {
request = request.query(&[("version", version)]);
}
if let Some(branch) = &self.branch {
@@ -97,21 +115,25 @@ struct RemoteBlobState {
}
/// Seekable Cloud blob handle over HTTP Range.
///
/// Nonempty handles read the table version returned by their size probe.
#[derive(Debug)]
pub struct RemoteBlobFile {
requester: Arc<dyn BlobRangeRequester>,
state: Mutex<RemoteBlobState>,
closed: AtomicBool,
size: u64,
version: Option<u64>,
}
impl RemoteBlobFile {
fn new(requester: Arc<dyn BlobRangeRequester>, size: u64) -> Self {
fn new(requester: Arc<dyn BlobRangeRequester>, size: u64, version: Option<u64>) -> Self {
Self {
requester,
state: Mutex::new(RemoteBlobState::default()),
closed: AtomicBool::new(false),
size,
version,
}
}
@@ -160,7 +182,7 @@ impl RemoteBlobFile {
let range_header = format!("bytes={}-{}", range.start, range.end - 1);
let (request_id, response) = self
.requester
.request_range(&range_header, RangeRequestMode::DataRead)
.request_range(&range_header, RangeRequestMode::DataRead, self.version)
.await
.map_err(remote_blob_error)?;
self.ensure_open()?;
@@ -236,7 +258,7 @@ impl RemoteBlobFile {
let range_header = format!("bytes={cursor}-");
let (request_id, response) = self
.requester
.request_range(&range_header, RangeRequestMode::DataRead)
.request_range(&range_header, RangeRequestMode::DataRead, self.version)
.await
.map_err(remote_blob_error)?;
self.ensure_open()?;
@@ -369,7 +391,62 @@ impl<S: HttpSend> RemoteTable<S> {
message: "fetch_blobs is not supported on this LanceDB Cloud server".into(),
});
}
let read_snapshot = self.snapshot_read_state().await;
let mut read_snapshot = self.snapshot_read_state().await;
// A selection spanning requests must use one exact dataset version.
// Resolve latest before the first chunk; a checked-out version is exact already.
if row_ids.len() > MAX_FETCH_BLOBS_ROW_IDS && read_snapshot.version.is_none() {
read_snapshot.version = Some(self.describe_read_snapshot(read_snapshot).await?.version);
}
let mut pending: Vec<&[u64]> = row_ids.chunks(MAX_FETCH_BLOBS_ROW_IDS).rev().collect();
let mut chunks = Vec::new();
while let Some(ids) = pending.pop() {
match self.fetch_blobs_chunk(column, ids, read_snapshot).await {
Ok(blobs) => chunks.push(blobs),
Err(error) if is_fetch_blobs_byte_limit_error(&error) => {
// A single-request call may turn into several requests after a
// byte-cap error. No bytes from the failed request were used.
if read_snapshot.version.is_none() {
read_snapshot.version =
Some(self.describe_read_snapshot(read_snapshot).await?.version);
}
if ids.len() == 1 {
// The whole-byte route cannot serve this blob. The Range route
// has no aggregate response limit and preserves null alignment.
let mut files = self
.fetch_blob_files_with_snapshot(column, ids, read_snapshot)
.await?;
let blob = match files.pop().unwrap() {
Some(file) => Some(file.read().await?),
None => None,
};
chunks.push(LargeBinaryArray::from(vec![blob.as_deref()]));
} else {
let mid = ids.len() / 2;
pending.push(&ids[mid..]);
pending.push(&ids[..mid]);
}
}
Err(error) => return Err(error),
}
}
if chunks.len() == 1 {
return Ok(chunks.pop().unwrap());
}
let chunk_refs: Vec<&dyn Array> = chunks.iter().map(|chunk| chunk as &dyn Array).collect();
Ok(arrow::compute::concat(&chunk_refs)?
.as_any()
.downcast_ref::<LargeBinaryArray>()
.expect("concatenating LargeBinary arrays returns LargeBinary")
.clone())
}
async fn fetch_blobs_chunk(
&self,
column: &str,
row_ids: &[u64],
read_snapshot: ReadSnapshot,
) -> Result<LargeBinaryArray> {
let mut body = serde_json::json!({
"version": read_snapshot.version,
"column": column,
@@ -460,6 +537,16 @@ impl<S: HttpSend> RemoteTable<S> {
}
let read_snapshot = self.snapshot_read_state().await;
self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot)
.await
}
async fn fetch_blob_files_with_snapshot(
&self,
column: &str,
row_ids: &[u64],
read_snapshot: ReadSnapshot,
) -> Result<Vec<Option<BlobFile>>> {
let encoded_column = urlencoding::encode(column);
let requesters = row_ids
.iter()
@@ -499,12 +586,25 @@ async fn probe_blob_files(
const BLOB_REQUEST_CONCURRENCY: usize = 8;
fn probe_blob_version(response: &Response, request_id: &str) -> Result<u64> {
response
.headers()
.get(&VERSION_HEADER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse().ok())
.ok_or_else(|| Error::Http {
source: "blob size probe returned a missing or invalid x-lancedb-version header".into(),
request_id: request_id.to_string(),
status_code: Some(response.status()),
})
}
/// Probe one blob's size.
///
/// `204` represents null. `416` with `bytes */0` represents an empty blob.
async fn probe_blob_file(requester: Arc<dyn BlobRangeRequester>) -> Result<Option<BlobFile>> {
let (request_id, response) = requester
.request_range("bytes=0-0", RangeRequestMode::SizeProbe)
.request_range("bytes=0-0", RangeRequestMode::SizeProbe, None)
.await?;
match response.status() {
StatusCode::NO_CONTENT => {
@@ -537,7 +637,8 @@ async fn probe_blob_file(requester: Arc<dyn BlobRangeRequester>) -> Result<Optio
});
}
}
Ok(Some(RemoteBlobFile::new(requester, 0).into()))
// An empty handle never makes a data range request, so no version is needed.
Ok(Some(RemoteBlobFile::new(requester, 0, None).into()))
}
StatusCode::PARTIAL_CONTENT => {
let size = response
@@ -553,6 +654,7 @@ async fn probe_blob_file(requester: Arc<dyn BlobRangeRequester>) -> Result<Optio
request_id: request_id.clone(),
status_code: Some(StatusCode::PARTIAL_CONTENT),
})?;
let version = probe_blob_version(&response, &request_id)?;
let probe_body = response.bytes().await.err_to_http(request_id.clone())?;
if probe_body.len() != 1 {
return Err(Error::Http {
@@ -565,7 +667,9 @@ async fn probe_blob_file(requester: Arc<dyn BlobRangeRequester>) -> Result<Optio
status_code: Some(StatusCode::PARTIAL_CONTENT),
});
}
Ok(Some(RemoteBlobFile::new(requester, size).into()))
Ok(Some(
RemoteBlobFile::new(requester, size, Some(version)).into(),
))
}
status => Err(Error::Http {
source: format!("blob size probe expected HTTP 206 Partial Content, got {status}")
@@ -578,7 +682,7 @@ async fn probe_blob_file(requester: Arc<dyn BlobRangeRequester>) -> Result<Optio
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use std::time::Duration;
@@ -628,6 +732,7 @@ mod tests {
header::CONTENT_RANGE,
format!("bytes {start}-{end}/{}", payload.len()),
)
.header(VERSION_HEADER, "5")
.body(payload[start..=end].to_vec())
.unwrap()
}
@@ -698,6 +803,68 @@ mod tests {
assert!(requests.lock().unwrap().contains(&"bytes=5-11".to_string()));
}
#[tokio::test]
async fn remote_blob_file_reads_the_probed_version_after_table_changes() {
let changed = Arc::new(AtomicBool::new(false));
let requests = Arc::new(StdMutex::new(Vec::new()));
let table = RemoteTable::new_mock(
"my_table".to_string(),
{
let changed = changed.clone();
let requests = requests.clone();
move |request| {
let range = request
.headers()
.get(header::RANGE)
.unwrap()
.to_str()
.unwrap();
let version = request
.url()
.query_pairs()
.find(|(name, _)| name == "version")
.map(|(_, value)| value.into_owned());
requests
.lock()
.unwrap()
.push((range.to_string(), version.clone()));
if changed.load(Ordering::SeqCst) && version.as_deref() != Some("5") {
return http::Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(b"row id is absent from the latest version".to_vec())
.unwrap();
}
range_response(&request, PAYLOAD)
}
},
Some(Version::new(0, 5, 0)),
);
let file = table
.fetch_blob_files_impl("image", &[10])
.await
.unwrap()
.pop()
.flatten()
.unwrap();
changed.store(true, Ordering::SeqCst);
assert_eq!(file.read_range(5..9).await.unwrap(), &PAYLOAD[5..9]);
assert_eq!(file.read_up_to(4).await.unwrap(), &PAYLOAD[..4]);
file.seek(10).await.unwrap();
assert_eq!(file.read().await.unwrap(), &PAYLOAD[10..]);
assert_eq!(
requests.lock().unwrap().as_slice(),
[
("bytes=0-0".to_string(), None),
("bytes=5-8".to_string(), Some("5".to_string())),
("bytes=0-".to_string(), Some("5".to_string())),
("bytes=10-".to_string(), Some("5".to_string())),
]
);
}
#[tokio::test]
async fn remote_blob_file_keeps_the_open_timeline_after_parent_checkout() {
let range_requests = Arc::new(StdMutex::new(Vec::new()));
@@ -762,7 +929,7 @@ mod tests {
}
#[tokio::test]
async fn remote_blob_files_return_a_handle_for_an_empty_blob() {
async fn remote_blob_files_allow_an_empty_blob_without_a_version_header() {
let requests = Arc::new(StdMutex::new(Vec::new()));
let table = mock_remote_blob_table(requests.clone());
@@ -775,8 +942,10 @@ mod tests {
assert!(files[1].is_none());
let empty = files[2].take().unwrap();
assert_eq!(empty.size(), 0);
let probe_requests = requests.lock().unwrap().len();
assert!(empty.read_range(0..0).await.unwrap().is_empty());
assert!(empty.read().await.unwrap().is_empty());
assert_eq!(requests.lock().unwrap().len(), probe_requests);
let nonempty = files[0].take().unwrap();
assert_eq!(nonempty.size(), PAYLOAD.len() as u64);
}
@@ -834,6 +1003,30 @@ mod tests {
);
}
#[tokio::test]
async fn remote_blob_files_reject_a_probe_without_a_version() {
let table = RemoteTable::new_mock(
"my_table".to_string(),
|request| {
let mut response = range_response(&request, PAYLOAD);
response.headers_mut().remove(VERSION_HEADER);
response
},
Some(Version::new(0, 5, 0)),
);
let error = table
.fetch_blob_files_impl("image", &[10])
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("missing or invalid x-lancedb-version"),
"got: {error}"
);
}
#[tokio::test]
async fn remote_blob_files_reject_unsupported_server_version() {
let table = RemoteTable::new_mock(
@@ -1025,6 +1218,7 @@ mod tests {
&self,
range_header: &str,
_mode: RangeRequestMode,
_version: Option<u64>,
) -> Result<(String, Response)> {
assert_eq!(range_header, "bytes=0-0");
let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
@@ -1039,6 +1233,7 @@ mod tests {
header::CONTENT_RANGE,
format!("bytes 0-0/{}", 100 + self.index),
)
.header(VERSION_HEADER, "5")
.body(vec![0u8])
.unwrap();
Ok((format!("probe-{}", self.index), Response::from(response)))
@@ -1157,6 +1352,7 @@ mod tests {
header::CONTENT_RANGE,
format!("bytes 0-0/{}", PAYLOAD.len()),
)
.header(VERSION_HEADER, "5")
.body(vec![PAYLOAD[0]])
.unwrap();
}
@@ -1203,6 +1399,7 @@ mod tests {
&self,
range_header: &str,
_mode: RangeRequestMode,
_version: Option<u64>,
) -> Result<(String, Response)> {
let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
self.max_in_flight.fetch_max(now, Ordering::SeqCst);
@@ -1235,7 +1432,7 @@ mod tests {
in_flight,
max_in_flight: max_in_flight.clone(),
});
let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64);
let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64, Some(5));
let ranges: Vec<_> = (0..16u64).map(|start| start..start + 2).collect();
let output = file.read_ranges(&ranges).await.unwrap();
@@ -1292,6 +1489,7 @@ mod tests {
header::CONTENT_RANGE,
format!("bytes 0-0/{}", PAYLOAD.len()),
)
.header(VERSION_HEADER, "5")
.body(vec![PAYLOAD[0]])
.unwrap();
}
@@ -1322,6 +1520,7 @@ mod tests {
&self,
range_header: &str,
_mode: RangeRequestMode,
_version: Option<u64>,
) -> Result<(String, Response)> {
if range_header.ends_with('-') {
self.started.notify_one();
@@ -1347,7 +1546,11 @@ mod tests {
release: release.clone(),
started: started.clone(),
});
let file = Arc::new(RemoteBlobFile::new(requester, PAYLOAD.len() as u64));
let file = Arc::new(RemoteBlobFile::new(
requester,
PAYLOAD.len() as u64,
Some(5),
));
let reader = {
let file = file.clone();
rust/lancedb/src/table.rs
+18 -2
View File
@@ -1251,7 +1251,10 @@ impl Table {
/// Materialize blob bytes for the given row ids.
///
/// Output matches `row_ids` in length and order. Null blobs are null;
/// valid empty blobs contain empty byte strings. Prefer
/// valid empty blobs contain empty byte strings. Cloud limits individual
/// requests to 1024 row ids and 64 MiB of blob bytes; the remote client
/// splits requests and reads a single larger blob through the Range route.
/// This method still materializes all bytes in memory, so prefer
/// [`Self::fetch_blob_files`] for large selections.
///
/// `_rowid` values stay valid after compaction when the table has stable
@@ -3788,7 +3791,20 @@ impl BaseTable for NativeTable {
let Some(segment) = segments.first() else {
continue;
};
let params = load_segment_params(&dataset, segment).await?;
// The listing itself only needs the manifest. Missing index files must
// not hide every other index, or callers cannot find the one to repair.
let params = match load_segment_params(&dataset, segment).await {
Ok(params) => params,
Err(err) => {
log::warn!(
"Failed to read full text search configuration for index '{}': {}",
index.name,
err
);
index.index_details = None;
continue;
}
};
let details = serde_json::to_string(&params).map_err(|source| Error::Other {
message: format!(
"Failed to serialize full text search configuration for index '{}'",
rust/lancedb/src/table/create_index.rs
+44
View File
@@ -2078,6 +2078,50 @@ mod tests {
assert_eq!(index.columns, vec!["tags".to_string()]);
}
#[tokio::test]
async fn test_list_indices_with_missing_fts_index_files() {
let tmp_dir = tempdir().unwrap();
let uri = tmp_dir.path().to_str().unwrap();
let conn = connect(uri).execute().await.unwrap();
let batch = record_batch!(
("id", Int32, [1, 2, 3]),
("text", Utf8, ["alpha", "beta", "gamma"])
)
.unwrap();
let table = conn.create_table("t", batch).execute().await.unwrap();
table
.create_index(&["text"], Index::FTS(FtsIndexBuilder::default()))
.execute()
.await
.unwrap();
table
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
.execute()
.await
.unwrap();
let fts_uuid = table
.list_indices()
.await
.unwrap()
.into_iter()
.find(|index| index.name == "text_idx")
.and_then(|index| index.index_uuid)
.unwrap();
std::fs::remove_dir_all(tmp_dir.path().join("t.lance/_indices").join(fts_uuid)).unwrap();
// A new connection, so the index cache cannot serve the deleted files.
let conn = connect(uri).execute().await.unwrap();
let table = conn.open_table("t").execute().await.unwrap();
let mut indices = table.list_indices().await.unwrap();
indices.sort_by(|a, b| a.name.cmp(&b.name));
let names = indices.iter().map(|i| i.name.as_str()).collect::<Vec<_>>();
assert_eq!(names, vec!["id_idx", "text_idx"]);
assert_eq!(indices[1].index_type, crate::index::IndexType::FTS);
assert_eq!(indices[1].index_details, None);
assert!(indices[0].index_details.is_some());
}
#[tokio::test]
async fn test_create_inverted_index() {
let conn = connect("memory://").execute().await.unwrap();
rust/lancedb/src/table/datafusion/udtf.rs
+2
View File
@@ -6,3 +6,5 @@
pub mod fts;
pub mod duplicate_pairs;
pub mod vector_query;
rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs
+118 -6
View File
@@ -330,6 +330,96 @@ pub async fn plan_duplicate_pairs(
Ok(tasks)
}
/// Optional query columns. The default preserves the primitive's three-column
/// schema; SQL convenience calls include physical scope identity and may map
/// endpoints to an application key without fetching source vectors.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct DuplicatePairsOutput {
pub scope: bool,
pub id_column: Option<String>,
}
impl DuplicatePairsOutput {
pub fn schema(&self, dataset: &Dataset) -> Result<SchemaRef> {
let mut fields = duplicate_pair_schema().fields().to_vec();
if self.scope {
fields.extend([
Arc::new(Field::new("segment_id", DataType::Utf8, false)),
Arc::new(Field::new("partition_id", DataType::UInt64, false)),
Arc::new(Field::new("dataset_version", DataType::UInt64, false)),
]);
}
if let Some(column) = &self.id_column {
let schema = Schema::from(dataset.schema());
let field = schema.field_with_name(column)?;
if column.contains('.')
|| matches!(
field.data_type(),
DataType::FixedSizeList(..)
| DataType::List(..)
| DataType::LargeList(..)
| DataType::Struct(..)
| DataType::Map(..)
)
{
return plan_err!("id_column must be a top-level scalar column");
}
fields.push(Arc::new(Field::new(
"id_a",
field.data_type().clone(),
field.is_nullable(),
)));
fields.push(Arc::new(Field::new(
"id_b",
field.data_type().clone(),
field.is_nullable(),
)));
}
Ok(Arc::new(Schema::new(fields)))
}
async fn enrich(
&self,
batch: arrow_array::RecordBatch,
dataset: &Dataset,
task: &DuplicatePairTask,
) -> Result<arrow_array::RecordBatch> {
use arrow_array::{StringArray, UInt64Array};
let mut columns = batch.columns().to_vec();
let n = batch.num_rows();
if self.scope {
columns.extend([
Arc::new(StringArray::from(vec![task.segment_id.to_string(); n]))
as arrow_array::ArrayRef,
Arc::new(UInt64Array::from(vec![task.partition_id as u64; n])),
Arc::new(UInt64Array::from(vec![dataset.version().version; n])),
]);
}
if let Some(column) = &self.id_column {
for endpoint in 0..2 {
let ids = batch
.column(endpoint)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row IDs");
let projection = lance::dataset::ProjectionRequest::from_columns(
[column.as_str()],
dataset.schema(),
);
let values = dataset.take_rows(ids.values(), projection).await?;
if values.num_rows() != n {
return plan_err!("pair row ID is missing from the pinned snapshot");
}
columns.push(values.column(0).clone());
}
}
Ok(arrow_array::RecordBatch::try_new(
self.schema(dataset)?,
columns,
)?)
}
}
/// Streaming native operator over partition descriptors, never source-vector
/// rows. Assigned descriptors are consumed sequentially; distributed hosts send
/// one descriptor per worker request. No Python or per-vector ANN calls.
@@ -341,6 +431,7 @@ pub struct DuplicatePairsExec {
tasks: Vec<DuplicatePairTask>,
properties: Arc<PlanProperties>,
metrics: ExecutionPlanMetricsSet,
output: DuplicatePairsOutput,
}
impl DuplicatePairsExec {
@@ -374,8 +465,23 @@ impl DuplicatePairsExec {
tasks,
properties,
metrics: ExecutionPlanMetricsSet::new(),
output: DuplicatePairsOutput::default(),
})
}
/// Configure optional output fields before dispatching this plan.
pub fn with_output(mut self, output: DuplicatePairsOutput) -> Result<Self> {
self.properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(output.schema(&self.dataset)?),
Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
));
self.output = output;
Ok(self)
}
pub fn output(&self) -> &DuplicatePairsOutput {
&self.output
}
pub fn dataset(&self) -> &Arc<Dataset> {
&self.dataset
}
@@ -430,6 +536,7 @@ impl ExecutionPlan for DuplicatePairsExec {
if partition != 0 {
return plan_err!("duplicate-pair output partition out of bounds");
}
let output = self.output.clone();
let tasks = self.tasks.clone();
let dataset = self.dataset.clone();
let config = self.config.clone();
@@ -467,6 +574,7 @@ impl ExecutionPlan for DuplicatePairsExec {
// into one-task PE requests before execution.
let stream = stream::iter(tasks)
.then(move |task| {
let output = output.clone();
let dataset = dataset.clone();
let config = config.clone();
let reservation = reservation.clone();
@@ -474,7 +582,7 @@ impl ExecutionPlan for DuplicatePairsExec {
let completed = completed.clone();
async move {
let reader = lance::index::vector::dedup::find_duplicate_pairs_in_partition_with_options(
dataset,
dataset.clone(),
&config.column,
task.segment_id,
task.partition_id,
@@ -506,15 +614,19 @@ impl ExecutionPlan for DuplicatePairsExec {
}
},
);
Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
measured,
)) as SendableRecordBatchStream)
let schema = output.schema(&dataset)?;
let measured = measured.and_then(move |batch| {
let output = output.clone();
let dataset = dataset.clone();
let task = task.clone();
async move { output.enrich(batch, &dataset, &task).await }
});
Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(schema, measured)) as SendableRecordBatchStream)
}
})
.try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
self.schema(),
stream,
)))
}
rust/lancedb/src/table/datafusion/udtf/vector_query.rs
+352
View File
@@ -0,0 +1,352 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Composable index-backed pair and dedup relations. Hosts bind a snapshot once
//! during planning, distribute the pair child, and use their ordinary SQL sink.
use std::sync::Arc;
use arrow_schema::SchemaRef;
use async_trait::async_trait;
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::{Result, plan_err};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::{Expr, TableType};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
coalesce_partitions::CoalescePartitionsExec,
empty::EmptyExec,
execution_plan::{Boundedness, EmissionType},
projection::ProjectionExec,
stream::RecordBatchStreamAdapter,
};
use futures::{TryStreamExt, stream};
use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap};
use lance::{Dataset, index::DatasetIndexInternalExt};
use lance_core::datatypes::BlobHandling;
use lance_index::metrics::NoOpMetricsCollector;
use super::duplicate_pairs::{
DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput,
plan_duplicate_pairs,
};
use crate::table::{NativeTableExt, datafusion::BaseTableAdapter};
const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MATERIALIZE_BATCH_ROWS: usize = 1024;
/// Select nonempty physical IVF scopes before dispatch, never by filtering the
/// resulting pairs. A count spans all segments; partition IDs are segment-local.
#[derive(Debug, Clone, Default)]
pub enum PartitionSelection {
#[default]
All,
First {
count: usize,
},
Random {
count: usize,
seed: u64,
},
}
/// Query-only options. Snapshot binding is separate so multiple calls in one
/// statement can share the same immutable dataset.
#[derive(Debug, Clone)]
pub struct VectorQueryOptions {
pub column: String,
pub distance_threshold: f32,
pub selection: PartitionSelection,
pub output: DuplicatePairsOutput,
pub dedup: bool,
}
impl BaseTableAdapter {
/// Open a retained version, or refresh and capture the latest manifest once.
/// No worker is allowed to independently resolve "latest".
pub async fn vector_query_snapshot(&self, version: Option<u64>) -> Result<Arc<Dataset>> {
let native = self.table.as_native().ok_or_else(|| {
datafusion_common::DataFusionError::Plan(
"vector queries require a native indexed table".into(),
)
})?;
if version == Some(0) {
return plan_err!("dataset_version must be positive");
}
let current = native
.dataset
.get()
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
Ok(match version {
Some(version) => Arc::new(current.checkout_version(version).await?),
None => {
// Do not change a shared table handle's pinned version.
let mut latest = (*current).clone();
latest.checkout_latest().await?;
Arc::new(latest)
}
})
}
}
/// A table provider bound to one source snapshot. The host retains the user's
/// unbound SQL for subsequent refreshes; this provider lives only for this run.
#[derive(Debug)]
pub struct VectorQueryTable {
dataset: Arc<Dataset>,
config: DuplicatePairsConfig,
tasks: Vec<DuplicatePairTask>,
options: VectorQueryOptions,
schema: SchemaRef,
}
impl VectorQueryTable {
pub async fn try_new(dataset: Arc<Dataset>, options: VectorQueryOptions) -> Result<Self> {
if options.dedup
&& (!matches!(options.selection, PartitionSelection::All)
|| options.output.id_column.is_some())
{
return plan_err!(
"vector_dedup requires all partitions and returns original rows; sampling and id_column are pair-query options"
);
}
let config = DuplicatePairsConfig {
dataset_version: dataset.version().version,
column: options.column.clone(),
distance_threshold: options.distance_threshold,
};
let mut tasks = plan_duplicate_pairs(dataset.clone(), &config).await?;
// Canonical scope order is independent of metadata enumeration order.
tasks.sort_by_key(|task| (task.segment_id, task.partition_id));
let mut nonempty = Vec::new();
for group in tasks.chunk_by(|a, b| a.segment_id == b.segment_id) {
let index = dataset
.open_vector_index(&config.column, &group[0].segment_id, &NoOpMetricsCollector)
.await?;
nonempty.extend(
group
.iter()
.filter(|task| index.partition_size(task.partition_id) > 0)
.cloned(),
);
}
match options.selection {
PartitionSelection::All => {}
PartitionSelection::First { count } | PartitionSelection::Random { count, .. }
if count == 0 =>
{
return plan_err!("partition count must be positive");
}
PartitionSelection::First { count } => nonempty.truncate(count),
PartitionSelection::Random { count, seed } => {
// A single segment matches ivf_partition's SplitMix64 sampling.
// With multiple segments the canonical scope ordinal extends the
// same policy without conflating equal partition IDs.
let segments = tasks
.iter()
.map(|t| t.segment_id)
.collect::<std::collections::BTreeSet<_>>()
.len();
let mut ranked = nonempty
.into_iter()
.enumerate()
.map(|(ordinal, task)| {
let key = if segments == 1 {
task.partition_id as u64
} else {
ordinal as u64
};
(sample_rank(seed, key), task)
})
.collect::<Vec<_>>();
ranked.sort_by_key(|(rank, task)| (*rank, task.segment_id, task.partition_id));
ranked.truncate(count);
nonempty = ranked.into_iter().map(|(_, task)| task).collect();
nonempty.sort_by_key(|task| (task.segment_id, task.partition_id));
}
}
let schema = if options.dedup {
let mut scanner = dataset.scan();
scanner.blob_handling(BlobHandling::AllBinary);
scanner.schema().await?
} else {
options.output.schema(&dataset)?
};
Ok(Self {
dataset,
config,
tasks: nonempty,
options,
schema,
})
}
}
fn sample_rank(seed: u64, ordinal: u64) -> u64 {
let mut value = seed.wrapping_add(ordinal).wrapping_add(0x9E3779B97F4A7C15);
value = (value ^ (value >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
value = (value ^ (value >> 27)).wrapping_mul(0x94D049BB133111EB);
value ^ (value >> 31)
}
#[async_trait]
impl TableProvider for VectorQueryTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn table_type(&self) -> TableType {
TableType::View
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
_filters: &[Expr],
_limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
let output = if self.options.dedup {
DuplicatePairsOutput::default()
} else {
self.options.output.clone()
};
let pairs: Arc<dyn ExecutionPlan> = if self.tasks.is_empty() {
Arc::new(EmptyExec::new(output.schema(&self.dataset)?))
} else {
Arc::new(
DuplicatePairsExec::try_new(
self.dataset.clone(),
self.config.clone(),
self.tasks.clone(),
)?
.with_output(output)?,
)
};
let plan: Arc<dyn ExecutionPlan> = if self.options.dedup {
Arc::new(VectorDedupExec::new(
self.dataset.clone(),
pairs,
self.schema.clone(),
))
} else {
pairs
};
if let Some(indices) = projection {
let expressions: Vec<_> = indices
.iter()
.map(|&i| {
(
Arc::new(datafusion_physical_expr::expressions::Column::new(
self.schema.field(i).name(),
i,
))
as Arc<dyn datafusion_physical_expr::PhysicalExpr>,
self.schema.field(i).name().clone(),
)
})
.collect();
Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?))
} else {
Ok(plan)
}
}
}
/// Greedy direct representatives over the complete pair relation. Sorting spills
/// to disk; only vertex masks are retained. This has no durable checkpoint and
/// publishes nothing: failure is handled by the enclosing query/MV transaction.
#[derive(Debug)]
pub struct VectorDedupExec {
dataset: Arc<Dataset>,
pairs: Arc<dyn ExecutionPlan>,
properties: Arc<PlanProperties>,
}
impl VectorDedupExec {
fn new(dataset: Arc<Dataset>, pairs: Arc<dyn ExecutionPlan>, schema: SchemaRef) -> Self {
Self {
dataset,
pairs,
properties: Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema),
Partitioning::UnknownPartitioning(1),
EmissionType::Final,
Boundedness::Bounded,
)),
}
}
}
impl DisplayAs for VectorDedupExec {
fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"VectorDedupExec: source_version={}, policy=direct_representative",
self.dataset.version().version
)
}
}
impl ExecutionPlan for VectorDedupExec {
fn name(&self) -> &str {
"VectorDedupExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![&self.pairs]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
if children.len() != 1 {
return plan_err!("VectorDedupExec requires one pair input");
}
Ok(Arc::new(Self::new(
self.dataset.clone(),
children[0].clone(),
self.schema(),
)))
}
fn execute(
&self,
partition: usize,
context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
if partition != 0 {
return plan_err!("invalid dedup output partition");
}
let dataset = self.dataset.clone();
let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?;
let stream = stream::once(async move {
let (removed, _) = crate::materialized_view::vector_dedup::select_representatives(
pairs,
REPRESENTATIVE_SORT_MEMORY_BYTES,
)
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
let mut scanner = dataset.scan();
scanner
.with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter(
removed,
)))
.blob_handling(BlobHandling::AllBinary)
.batch_size(MATERIALIZE_BATCH_ROWS)
.strict_batch_size(true);
Ok::<_, datafusion_common::DataFusionError>(
scanner
.try_into_stream()
.await?
.map_err(datafusion_common::DataFusionError::from),
)
})
.try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.schema(),
stream,
)))
}
}
rust/lancedb/tests/duplicate_pairs.rs
+190 -1
View File
@@ -76,7 +76,7 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
.to_str()
.unwrap()
.to_string();
let n = 1025;
let n = 8193;
let vectors = (0..n * 2)
.flat_map(|i| {
let i = i % n;
@@ -196,6 +196,89 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
);
let actual = ctx.sql(&sql).await?.collect().await?;
assert_eq!(pairs(&actual), pairs(&expected));
use lancedb::table::datafusion::udtf::{
duplicate_pairs::DuplicatePairsOutput,
vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable},
};
let query_options = VectorQueryOptions {
column: "vector".into(),
distance_threshold: threshold,
selection: PartitionSelection::All,
output: DuplicatePairsOutput {
scope: true,
id_column: Some("id".into()),
},
dedup: false,
};
let provider = VectorQueryTable::try_new(ds.clone(), query_options.clone()).await?;
let extended = ctx.read_table(Arc::new(provider))?.collect().await?;
assert_eq!(pairs(&extended), pairs(&expected));
for batch in &extended {
let rowids = batch
.column_by_name("row_id_a")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let ids = batch
.column_by_name("id_a")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
for i in 0..batch.num_rows() {
let expected_id = (rowids.value(i) >> 32) as i64 * n as i64
+ (rowids.value(i) & 0xffffffff) as i64;
assert_eq!(ids.value(i), expected_id);
}
}
let mut sampled_options = query_options.clone();
sampled_options.selection = PartitionSelection::Random { count: 1, seed: 42 };
let sample = ctx
.read_table(Arc::new(
VectorQueryTable::try_new(ds.clone(), sampled_options.clone()).await?,
))?
.collect()
.await?;
let replay = ctx
.read_table(Arc::new(
VectorQueryTable::try_new(ds.clone(), sampled_options).await?,
))?
.collect()
.await?;
assert_eq!(sample, replay);
let scopes = sample
.iter()
.flat_map(|batch| {
let segments = batch
.column_by_name("segment_id")
.unwrap()
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap();
let partitions = batch
.column_by_name("partition_id")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
(0..batch.num_rows())
.map(|i| (segments.value(i).to_string(), partitions.value(i)))
.collect::<Vec<_>>()
})
.collect::<std::collections::BTreeSet<_>>();
assert!(scopes.len() <= 1);
let mut dedup_options = query_options;
dedup_options.dedup = true;
dedup_options.output = DuplicatePairsOutput::default();
let clean_query = ctx
.read_table(Arc::new(
VectorQueryTable::try_new(ds.clone(), dedup_options).await?,
))?
.select(vec![datafusion::prelude::col("id")])?
.collect()
.await?;
assert_eq!(
pairs(&actual).len(),
match threshold {
@@ -388,6 +471,21 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
.collect::<Vec<_>>();
actual_ids.sort_unstable();
expected_ids.sort_unstable();
let mut query_ids = clean_query
.iter()
.flat_map(|batch| {
batch
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec()
})
.collect::<Vec<_>>();
query_ids.sort_unstable();
assert_eq!(query_ids, expected_ids);
assert_eq!(actual_ids, expected_ids);
if threshold == 1.0 {
assert!(actual_ids.contains(&0) && actual_ids.contains(&2));
@@ -510,6 +608,90 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
)
.await?;
let version = ds.version().version;
use lancedb::table::datafusion::udtf::{
duplicate_pairs::DuplicatePairsOutput,
vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable},
};
let options = VectorQueryOptions {
column: "vector".into(),
distance_threshold: 1.0,
selection: PartitionSelection::All,
output: DuplicatePairsOutput::default(),
dedup: true,
};
let mapped = VectorQueryTable::try_new(
Arc::new(ds.clone()),
VectorQueryOptions {
dedup: false,
output: DuplicatePairsOutput {
scope: true,
id_column: Some("id".into()),
},
..options.clone()
},
)
.await?;
let mapped = SessionContext::new()
.read_table(Arc::new(mapped))?
.collect()
.await?;
let mut mapped_edges = Vec::new();
for batch in mapped {
let a = batch
.column_by_name("id_a")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let b = batch
.column_by_name("id_b")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
for i in 0..batch.num_rows() {
mapped_edges.push((a.value(i).min(b.value(i)), a.value(i).max(b.value(i))));
}
}
mapped_edges.sort_unstable();
assert_eq!(mapped_edges, [(0, 1), (1, 2)]);
let provider = VectorQueryTable::try_new(Arc::new(ds.clone()), options).await?;
let query_rows = SessionContext::new()
.read_table(Arc::new(provider))?
.collect()
.await?;
let mut query_ids = query_rows
.iter()
.flat_map(|batch| {
batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec()
})
.collect::<Vec<_>>();
query_ids.sort_unstable();
assert_eq!(query_ids, [0, 2, 3, 4]);
for batch in &query_rows {
let ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let images = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
for i in 0..batch.num_rows() {
assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]);
}
}
let definition = MaterializedViewDefinition::from_sql(&format!(
"SELECT * FROM vector_dedup('source', {version}, 'vector', 1)"
))?;
@@ -519,6 +701,13 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
.await?;
// A later source delete does not change the declared result snapshot.
table.delete("id = 0").await?;
let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?;
let explicit = adapter.vector_query_snapshot(Some(version)).await?;
let latest = adapter.vector_query_snapshot(None).await?;
assert_eq!(explicit.version().version, version);
assert!(latest.version().version > version);
assert_eq!(explicit.count_rows(None).await?, 5);
assert_eq!(latest.count_rows(None).await?, 4);
clean.refresh().execute().await?;
assert_eq!(clean.table().count_rows(None).await?, 4);
let output = Dataset::open(&clean.table().uri().await?).await?;