mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 07:58:31 +00:00
Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2fbf6d6211 | |||
| 391cac9034 | |||
| 21530432a0 | |||
| 9b825c5f29 | |||
| 8083232dd5 | |||
| 302b21aa94 |
Generated
+42
-42
@@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4888,8 +4888,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4911,7 +4911,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4925,7 +4925,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -4934,8 +4934,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4945,8 +4945,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4983,8 +4983,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5013,8 +5013,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5031,8 +5031,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5041,8 +5041,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5075,8 +5075,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5107,8 +5107,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5172,8 +5172,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5195,8 +5195,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5236,8 +5236,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5251,8 +5251,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5264,8 +5264,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5318,8 +5318,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5333,8 +5333,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5374,8 +5374,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5388,8 +5388,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "11.0.0-beta.22"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"frostem",
|
||||
"icu_segmenter",
|
||||
|
||||
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lancedb = { path = "rust/lancedb", default-features = false }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
|
||||
+1
-1
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>11.0.0-beta.22</lance-core.version>
|
||||
<lance-core.version>12.0.0-beta.2</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>
|
||||
|
||||
@@ -685,6 +685,56 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
},
|
||||
);
|
||||
|
||||
// https://github.com/lancedb/lancedb/issues/1963
|
||||
it("should query documents with LangChain PDF metadata", async () => {
|
||||
const tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
try {
|
||||
const db = await connect(tmpDir.name);
|
||||
const documents = [
|
||||
{
|
||||
text: "first page",
|
||||
vector: [1, 0],
|
||||
source: "first.pdf",
|
||||
loc: { pageNumber: 1, lines: { from: 1, to: 12 } },
|
||||
pdf: {
|
||||
version: "1.10.100",
|
||||
info: {
|
||||
format: "PDF 1.7",
|
||||
producer: "pdf.js",
|
||||
creator: "Writer",
|
||||
},
|
||||
totalPages: 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
text: "second page",
|
||||
vector: [0, 1],
|
||||
source: "second.pdf",
|
||||
loc: { pageNumber: 2, lines: { from: 13, to: 24 } },
|
||||
pdf: {
|
||||
version: "1.10.100",
|
||||
info: {
|
||||
format: "PDF 1.7",
|
||||
producer: "pdf.js",
|
||||
creator: "Writer",
|
||||
},
|
||||
totalPages: 2,
|
||||
},
|
||||
},
|
||||
];
|
||||
const documentsTable = await db.createTable("documents", documents);
|
||||
|
||||
const results = await documentsTable.query().toArray();
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].source).toBe("first.pdf");
|
||||
expect(results[0].pdf.info.producer).toBe("pdf.js");
|
||||
expect(results[1].loc.pageNumber).toBe(2);
|
||||
} finally {
|
||||
tmpDir.removeCallback();
|
||||
}
|
||||
});
|
||||
|
||||
describe("merge insert", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
@@ -2535,7 +2585,24 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
);
|
||||
});
|
||||
|
||||
test("full text search if no embedding function provided", async () => {
|
||||
test("full text search if only an unrelated embedding function is registered", async () => {
|
||||
register("unused")(
|
||||
class extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
{ text: "hello world", vector: [0.1, 0.2, 0.3] },
|
||||
@@ -2557,6 +2624,306 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
expect(results2[0].text).toBe(data[1].text);
|
||||
});
|
||||
|
||||
test("auto search stays consistent with the active revision", async () => {
|
||||
let initCalls = 0;
|
||||
let queryCalls = 0;
|
||||
let markStarted!: () => void;
|
||||
const started = new Promise<void>((resolve) => {
|
||||
markStarted = resolve;
|
||||
});
|
||||
let releaseEmbedding!: () => void;
|
||||
const embeddingReleased = new Promise<void>((resolve) => {
|
||||
releaseEmbedding = resolve;
|
||||
});
|
||||
|
||||
@register("refresh-test")
|
||||
class TestEmbedding extends EmbeddingFunction<string> {
|
||||
async init() {
|
||||
initCalls += 1;
|
||||
}
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(value: string) {
|
||||
queryCalls += 1;
|
||||
if (value === "blocked") {
|
||||
markStarted();
|
||||
await embeddingReleased;
|
||||
}
|
||||
return value === "greetings" ? [0.1] : [0.2];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map((value) =>
|
||||
value === "hello world" ? [0.1] : [0.2],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
await writer.createTable("test", [{ text: "plain", vector: [0.0] }]);
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("test");
|
||||
type SnapshotCountingNative = {
|
||||
querySnapshot: () => Promise<unknown>;
|
||||
};
|
||||
const native = (tracked as unknown as { inner: SnapshotCountingNative })
|
||||
.inner;
|
||||
const querySnapshot = native.querySnapshot.bind(native);
|
||||
let snapshotCalls = 0;
|
||||
native.querySnapshot = async () => {
|
||||
snapshotCalls += 1;
|
||||
return await querySnapshot();
|
||||
};
|
||||
const autoQuery = tracked.search("greetings").select(["text"]).limit(1);
|
||||
|
||||
const func = new TestEmbedding();
|
||||
const schema = LanceSchema({
|
||||
text: func.sourceField(new arrow.Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
const data = [{ text: "hello world" }, { text: "goodbye world" }];
|
||||
await writer.createTable("test", data, { mode: "overwrite", schema });
|
||||
const baselineInitCalls = initCalls;
|
||||
|
||||
expect(
|
||||
(await tracked.schema()).metadata.get("embedding_functions"),
|
||||
).toBeDefined();
|
||||
const results = await autoQuery.toArray();
|
||||
expect(results[0].text).toBe(data[0].text);
|
||||
expect(initCalls).toBe(baselineInitCalls + 1);
|
||||
expect(queryCalls).toBe(1);
|
||||
expect(snapshotCalls).toBe(1);
|
||||
|
||||
const repeatedResults = await autoQuery.toArray();
|
||||
expect(repeatedResults[0].text).toBe(data[0].text);
|
||||
expect(initCalls).toBe(baselineInitCalls + 1);
|
||||
expect(queryCalls).toBe(1);
|
||||
expect(snapshotCalls).toBe(2);
|
||||
|
||||
const pending = tracked
|
||||
.search("blocked")
|
||||
.select(["text"])
|
||||
.limit(1)
|
||||
.toArray();
|
||||
await started;
|
||||
|
||||
const ftsData = [
|
||||
{ text: "greetings from full text", vector: [0.0] },
|
||||
{ text: "blocked from full text", vector: [0.0] },
|
||||
];
|
||||
const ftsTable = await writer.createTable("test", ftsData, {
|
||||
mode: "overwrite",
|
||||
});
|
||||
await ftsTable.createIndex("text", { config: Index.fts() });
|
||||
releaseEmbedding();
|
||||
|
||||
const pendingResults = await pending;
|
||||
expect(pendingResults[0].text).toBe(data[1].text);
|
||||
|
||||
expect(
|
||||
(await tracked.schema()).metadata.get("embedding_functions"),
|
||||
).toBeUndefined();
|
||||
const ftsResults = await autoQuery.toArray();
|
||||
expect(ftsResults[0].text).toBe(ftsData[0].text);
|
||||
});
|
||||
|
||||
test("auto search keeps newer preparation during a revision race", async () => {
|
||||
let aCalls = 0;
|
||||
let bCalls = 0;
|
||||
let markAStarted!: () => void;
|
||||
const aStarted = new Promise<void>((resolve) => {
|
||||
markAStarted = resolve;
|
||||
});
|
||||
let releaseA!: () => void;
|
||||
const aReleased = new Promise<void>((resolve) => {
|
||||
releaseA = resolve;
|
||||
});
|
||||
let markBStarted!: () => void;
|
||||
const bStarted = new Promise<void>((resolve) => {
|
||||
markBStarted = resolve;
|
||||
});
|
||||
let releaseB!: () => void;
|
||||
const bReleased = new Promise<void>((resolve) => {
|
||||
releaseB = resolve;
|
||||
});
|
||||
|
||||
@register("race-a")
|
||||
class EmbeddingA extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
aCalls += 1;
|
||||
markAStarted();
|
||||
await aReleased;
|
||||
return [0.1];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.1]);
|
||||
}
|
||||
}
|
||||
|
||||
@register("race-b")
|
||||
class EmbeddingB extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
bCalls += 1;
|
||||
markBStarted();
|
||||
await bReleased;
|
||||
return [0.2];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.2]);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
const embeddingA = new EmbeddingA();
|
||||
const schemaA = LanceSchema({
|
||||
text: embeddingA.sourceField(new arrow.Utf8()),
|
||||
vector: embeddingA.vectorField(),
|
||||
});
|
||||
await writer.createTable("race", [{ text: "revision a" }], {
|
||||
schema: schemaA,
|
||||
});
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("race");
|
||||
const query = tracked.search("query");
|
||||
|
||||
const first = query.toArray();
|
||||
await aStarted;
|
||||
|
||||
const embeddingB = new EmbeddingB();
|
||||
const schemaB = LanceSchema({
|
||||
text: embeddingB.sourceField(new arrow.Utf8()),
|
||||
vector: embeddingB.vectorField(),
|
||||
});
|
||||
await writer.createTable("race", [{ text: "revision b" }], {
|
||||
mode: "overwrite",
|
||||
schema: schemaB,
|
||||
});
|
||||
const second = query.toArray();
|
||||
await bStarted;
|
||||
|
||||
releaseA();
|
||||
releaseB();
|
||||
await Promise.all([first, second]);
|
||||
expect(aCalls).toBe(1);
|
||||
expect(bCalls).toBe(1);
|
||||
});
|
||||
|
||||
test("stale FTS routing keeps newer vector preparation", async () => {
|
||||
let vectorCalls = 0;
|
||||
let markVectorStarted!: () => void;
|
||||
const vectorStarted = new Promise<void>((resolve) => {
|
||||
markVectorStarted = resolve;
|
||||
});
|
||||
let releaseVector!: () => void;
|
||||
const vectorReleased = new Promise<void>((resolve) => {
|
||||
releaseVector = resolve;
|
||||
});
|
||||
|
||||
@register("stale-fts-race")
|
||||
class RaceEmbedding extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
vectorCalls += 1;
|
||||
markVectorStarted();
|
||||
await vectorReleased;
|
||||
return [0.1];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.1]);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
const ftsTable = await writer.createTable("stale_fts", [
|
||||
{ text: "hello", vector: [0.0] },
|
||||
]);
|
||||
await ftsTable.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("stale_fts");
|
||||
type Snapshot = {
|
||||
schema: () => Promise<Buffer>;
|
||||
};
|
||||
type NativeWithSnapshot = {
|
||||
querySnapshot: () => Promise<Snapshot>;
|
||||
};
|
||||
const native = (tracked as unknown as { inner: NativeWithSnapshot })
|
||||
.inner;
|
||||
const querySnapshot = native.querySnapshot.bind(native);
|
||||
let snapshotCalls = 0;
|
||||
let markStaleSchemaStarted!: () => void;
|
||||
const staleSchemaStarted = new Promise<void>((resolve) => {
|
||||
markStaleSchemaStarted = resolve;
|
||||
});
|
||||
let releaseStaleSchema!: () => void;
|
||||
const staleSchemaReleased = new Promise<void>((resolve) => {
|
||||
releaseStaleSchema = resolve;
|
||||
});
|
||||
native.querySnapshot = async () => {
|
||||
const snapshot = await querySnapshot();
|
||||
snapshotCalls += 1;
|
||||
if (snapshotCalls === 1) {
|
||||
const schema = snapshot.schema.bind(snapshot);
|
||||
snapshot.schema = async () => {
|
||||
markStaleSchemaStarted();
|
||||
await staleSchemaReleased;
|
||||
return await schema();
|
||||
};
|
||||
}
|
||||
return snapshot;
|
||||
};
|
||||
|
||||
const query = tracked.search("hello");
|
||||
const staleFtsExecution = query.toArray();
|
||||
await staleSchemaStarted;
|
||||
|
||||
const embedding = new RaceEmbedding();
|
||||
const vectorSchema = LanceSchema({
|
||||
text: embedding.sourceField(new arrow.Utf8()),
|
||||
vector: embedding.vectorField(),
|
||||
});
|
||||
await writer.createTable("stale_fts", [{ text: "hello" }], {
|
||||
mode: "overwrite",
|
||||
schema: vectorSchema,
|
||||
});
|
||||
|
||||
const vectorExecution = query.toArray();
|
||||
await vectorStarted;
|
||||
releaseStaleSchema();
|
||||
await staleFtsExecution;
|
||||
releaseVector();
|
||||
await vectorExecution;
|
||||
|
||||
await query.toArray();
|
||||
expect(vectorCalls).toBe(1);
|
||||
});
|
||||
|
||||
test("tokenizes FTS queries by column or index name", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
@@ -3107,6 +3474,30 @@ describe("column name options", () => {
|
||||
expect(results[1].query_index).toBe(1);
|
||||
});
|
||||
|
||||
test("observes promised additional vectors while the query is pending", async () => {
|
||||
const initialVector = new Promise<number[]>(() => undefined);
|
||||
const query = table.query().nearestTo(initialVector);
|
||||
const unhandled: unknown[] = [];
|
||||
const onUnhandled = (reason: unknown) => unhandled.push(reason);
|
||||
process.on("unhandledRejection", onUnhandled);
|
||||
|
||||
try {
|
||||
query.addQueryVector(Promise.reject(new Error("extra vector failed")));
|
||||
await new Promise<void>((resolve) => setImmediate(resolve));
|
||||
expect(unhandled).toEqual([]);
|
||||
|
||||
const rejectedQuery = table
|
||||
.query()
|
||||
.nearestTo([0.1, 0.2])
|
||||
.addQueryVector(Promise.reject(new Error("consumed vector failed")));
|
||||
await expect(rejectedQuery.toArray()).rejects.toThrow(
|
||||
"consumed vector failed",
|
||||
);
|
||||
} finally {
|
||||
process.off("unhandledRejection", onUnhandled);
|
||||
}
|
||||
});
|
||||
|
||||
test("index and search multivectors", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [];
|
||||
|
||||
+129
-95
@@ -100,6 +100,29 @@ export interface FullTextSearchOptions {
|
||||
columns?: string | string[];
|
||||
}
|
||||
|
||||
function nearestToNative(
|
||||
inner: NativeQuery,
|
||||
vector: Awaited<IntoVector>,
|
||||
): NativeVectorQuery {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
return inner.nearestToRaw(raw.data, raw.dtype);
|
||||
}
|
||||
return inner.nearestTo(Float32Array.from(vector as number[]));
|
||||
}
|
||||
|
||||
function addQueryVectorToNative(
|
||||
inner: NativeVectorQuery,
|
||||
vector: Awaited<IntoVector>,
|
||||
) {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
inner.addQueryVectorRaw(raw.data, raw.dtype);
|
||||
} else {
|
||||
inner.addQueryVector(Float32Array.from(vector as number[]));
|
||||
}
|
||||
}
|
||||
|
||||
/** Common methods supported by all query types
|
||||
*
|
||||
* @see {@link Query}
|
||||
@@ -499,6 +522,13 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
super(inner);
|
||||
}
|
||||
|
||||
/**
|
||||
* @hidden
|
||||
*/
|
||||
protected doVectorCall(fn: (inner: NativeVectorQuery) => void) {
|
||||
super.doCall(fn);
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the number of partitions to search (probe)
|
||||
*
|
||||
@@ -526,7 +556,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* the minimum and maximum to the same value.
|
||||
*/
|
||||
nprobes(nprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.nprobes(nprobes));
|
||||
this.doVectorCall((inner) => inner.nprobes(nprobes));
|
||||
|
||||
return this;
|
||||
}
|
||||
@@ -540,7 +570,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* but will also increase latency.
|
||||
*/
|
||||
minimumNprobes(minimumNprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.minimumNprobes(minimumNprobes));
|
||||
this.doVectorCall((inner) => inner.minimumNprobes(minimumNprobes));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -554,7 +584,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* potential false negatives.
|
||||
*/
|
||||
maximumNprobes(maximumNprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.maximumNprobes(maximumNprobes));
|
||||
this.doVectorCall((inner) => inner.maximumNprobes(maximumNprobes));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -567,7 +597,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* `undefined` means no lower or upper bound.
|
||||
*/
|
||||
distanceRange(lowerBound?: number, upperBound?: number): VectorQuery {
|
||||
super.doCall((inner) => inner.distanceRange(lowerBound, upperBound));
|
||||
this.doVectorCall((inner) => inner.distanceRange(lowerBound, upperBound));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -581,7 +611,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* also increase the latency of your query. The default value is 1.5*limit.
|
||||
*/
|
||||
ef(ef: number): VectorQuery {
|
||||
super.doCall((inner) => inner.ef(ef));
|
||||
this.doVectorCall((inner) => inner.ef(ef));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -595,7 +625,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* whose data type is a fixed-size-list of floats.
|
||||
*/
|
||||
column(column: string): VectorQuery {
|
||||
super.doCall((inner) => inner.column(column));
|
||||
this.doVectorCall((inner) => inner.column(column));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -616,7 +646,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
distanceType(
|
||||
distanceType: Required<IvfPqOptions>["distanceType"],
|
||||
): VectorQuery {
|
||||
super.doCall((inner) => inner.distanceType(distanceType));
|
||||
this.doVectorCall((inner) => inner.distanceType(distanceType));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -650,7 +680,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* distance between the query vector and the actual uncompressed vector.
|
||||
*/
|
||||
refineFactor(refineFactor: number): VectorQuery {
|
||||
super.doCall((inner) => inner.refineFactor(refineFactor));
|
||||
this.doVectorCall((inner) => inner.refineFactor(refineFactor));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -675,7 +705,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* factor can often help restore some of the results lost by post filtering.
|
||||
*/
|
||||
postfilter(): VectorQuery {
|
||||
super.doCall((inner) => inner.postfilter());
|
||||
this.doVectorCall((inner) => inner.postfilter());
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -689,7 +719,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* calculate your recall to select an appropriate value for nprobes.
|
||||
*/
|
||||
bypassVectorIndex(): VectorQuery {
|
||||
super.doCall((inner) => inner.bypassVectorIndex());
|
||||
this.doVectorCall((inner) => inner.bypassVectorIndex());
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -705,35 +735,31 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
*/
|
||||
addQueryVector(vector: IntoVector): VectorQuery {
|
||||
if (vector instanceof Promise) {
|
||||
// Observe the promise as soon as it is accepted. The existing native
|
||||
// query may still be pending, and delaying observation until it resolves
|
||||
// can otherwise surface a fast rejection as unhandled.
|
||||
const settledVector = vector.then(
|
||||
(value) => ({ status: "fulfilled" as const, value }),
|
||||
(reason) => ({ status: "rejected" as const, reason }),
|
||||
);
|
||||
const res = (async () => {
|
||||
try {
|
||||
const v = await vector;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
|
||||
const value: any = this.addQueryVector(v);
|
||||
const inner = value.inner as
|
||||
| NativeVectorQuery
|
||||
| Promise<NativeVectorQuery>;
|
||||
return inner;
|
||||
} catch (e) {
|
||||
return Promise.reject(e);
|
||||
const inner = await this.getInner();
|
||||
const outcome = await settledVector;
|
||||
if (outcome.status === "rejected") {
|
||||
throw outcome.reason;
|
||||
}
|
||||
addQueryVectorToNative(inner, outcome.value);
|
||||
return inner;
|
||||
})();
|
||||
return new VectorQuery(res);
|
||||
} else {
|
||||
super.doCall((inner) => {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
inner.addQueryVectorRaw(raw.data, raw.dtype);
|
||||
} else {
|
||||
inner.addQueryVector(Float32Array.from(vector as number[]));
|
||||
}
|
||||
});
|
||||
this.doVectorCall((inner) => addQueryVectorToNative(inner, vector));
|
||||
return this;
|
||||
}
|
||||
}
|
||||
|
||||
rerank(reranker: Reranker): VectorQuery {
|
||||
super.doCall((inner) =>
|
||||
this.doVectorCall((inner) =>
|
||||
inner.rerank(async (args) => {
|
||||
const vecResults = await fromBufferToRecordBatch(args.vecResults);
|
||||
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
|
||||
@@ -752,6 +778,71 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a string query whose vector/FTS routing is resolved against the active
|
||||
* table schema when the query executes.
|
||||
*
|
||||
* @hidden
|
||||
*/
|
||||
export function createAutoQuery(
|
||||
table: NativeTable,
|
||||
query: string,
|
||||
columns: string[] | null,
|
||||
getVector: (metadata: string) => Promise<Awaited<IntoVector>>,
|
||||
): AutoQuery {
|
||||
type RouteSnapshot = {
|
||||
table: NativeTable;
|
||||
embeddingMetadata: string | undefined;
|
||||
};
|
||||
type CachedPreparation = {
|
||||
metadata: string;
|
||||
vector: Promise<Awaited<IntoVector>>;
|
||||
};
|
||||
|
||||
let cachedPreparation: CachedPreparation | undefined;
|
||||
|
||||
const snapshotRoute = async (): Promise<RouteSnapshot> => {
|
||||
const snapshot = await table.querySnapshot();
|
||||
const schema = tableFromIPC(await snapshot.schema()).schema;
|
||||
return {
|
||||
table: snapshot,
|
||||
embeddingMetadata: schema.metadata.get("embedding_functions"),
|
||||
};
|
||||
};
|
||||
|
||||
const createInner = async (): Promise<NativeQuery | NativeVectorQuery> => {
|
||||
const route = await snapshotRoute();
|
||||
if (route.embeddingMetadata === undefined) {
|
||||
const inner = route.table.query();
|
||||
inner.fullTextSearch({ query, columns });
|
||||
return inner;
|
||||
}
|
||||
|
||||
const metadata = route.embeddingMetadata;
|
||||
if (cachedPreparation?.metadata !== metadata) {
|
||||
cachedPreparation = {
|
||||
metadata,
|
||||
vector: Promise.resolve().then(() => getVector(metadata)),
|
||||
};
|
||||
}
|
||||
|
||||
const preparation = cachedPreparation;
|
||||
let vector: Awaited<IntoVector>;
|
||||
try {
|
||||
vector = await preparation.vector;
|
||||
} catch (error) {
|
||||
if (cachedPreparation === preparation) {
|
||||
cachedPreparation = undefined;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
|
||||
return nearestToNative(route.table.query(), vector);
|
||||
};
|
||||
|
||||
return new AutoQuery(createInner);
|
||||
}
|
||||
|
||||
/**
|
||||
* A query that returns a subset of the rows in the table.
|
||||
*
|
||||
@@ -836,37 +927,6 @@ export class Query extends StandardQueryBase<NativeQuery> {
|
||||
super(tbl.query());
|
||||
}
|
||||
|
||||
/** @hidden */
|
||||
static autoSearch(
|
||||
tbl: () => Promise<NativeTable>,
|
||||
query: string,
|
||||
vector: (tbl: NativeTable) => Promise<Awaited<IntoVector> | undefined>,
|
||||
columns?: string[],
|
||||
): AutoQuery {
|
||||
const nativeQuery = async () => {
|
||||
const snapshot = await Promise.resolve(tbl());
|
||||
const resolved = await vector(snapshot);
|
||||
const inner = snapshot.query();
|
||||
if (resolved === undefined) {
|
||||
inner.fullTextSearch({
|
||||
query,
|
||||
columns: columns ?? null,
|
||||
});
|
||||
return inner;
|
||||
}
|
||||
|
||||
const raw = Array.isArray(resolved)
|
||||
? null
|
||||
: extractVectorBuffer(resolved);
|
||||
if (raw) {
|
||||
return inner.nearestToRaw(raw.data, raw.dtype);
|
||||
}
|
||||
return inner.nearestTo(Float32Array.from(resolved as number[]));
|
||||
};
|
||||
|
||||
return new AutoQuery(nativeQuery);
|
||||
}
|
||||
|
||||
/**
|
||||
* Find the nearest vectors to the given query vector.
|
||||
*
|
||||
@@ -905,45 +965,19 @@ export class Query extends StandardQueryBase<NativeQuery> {
|
||||
* a default `limit` of 10 will be used. @see {@link Query#limit}
|
||||
*/
|
||||
nearestTo(vector: IntoVector): VectorQuery {
|
||||
const callNearestTo = (
|
||||
inner: NativeQuery,
|
||||
resolved: Float32Array | Float64Array | Uint8Array | number[],
|
||||
): NativeVectorQuery => {
|
||||
const raw = Array.isArray(resolved)
|
||||
? null
|
||||
: extractVectorBuffer(resolved);
|
||||
if (raw) {
|
||||
return inner.nearestToRaw(raw.data, raw.dtype);
|
||||
}
|
||||
return inner.nearestTo(Float32Array.from(resolved as number[]));
|
||||
};
|
||||
|
||||
if (this.inner instanceof Promise) {
|
||||
const nativeQuery = this.inner.then(async (inner) => {
|
||||
const resolved = vector instanceof Promise ? await vector : vector;
|
||||
return callNearestTo(inner, resolved);
|
||||
});
|
||||
const inner = this.inner;
|
||||
if (inner instanceof Promise) {
|
||||
const nativeQuery = inner.then(async (resolvedInner) =>
|
||||
nearestToNative(resolvedInner, await vector),
|
||||
);
|
||||
return new VectorQuery(nativeQuery);
|
||||
}
|
||||
if (vector instanceof Promise) {
|
||||
const res = (async () => {
|
||||
try {
|
||||
const v = await vector;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
|
||||
const value: any = this.nearestTo(v);
|
||||
const inner = value.inner as
|
||||
| NativeVectorQuery
|
||||
| Promise<NativeVectorQuery>;
|
||||
return inner;
|
||||
} catch (e) {
|
||||
return Promise.reject(e);
|
||||
}
|
||||
})();
|
||||
return new VectorQuery(res);
|
||||
} else {
|
||||
const vectorQuery = callNearestTo(this.inner, vector);
|
||||
return new VectorQuery(vectorQuery);
|
||||
return new VectorQuery(
|
||||
vector.then((resolvedVector) => nearestToNative(inner, resolvedVector)),
|
||||
);
|
||||
}
|
||||
return new VectorQuery(nearestToNative(inner, vector));
|
||||
}
|
||||
|
||||
nearestToText(query: string | FullTextQuery, columns?: string[]): Query {
|
||||
|
||||
+18
-21
@@ -48,6 +48,7 @@ import {
|
||||
Query,
|
||||
TakeQuery,
|
||||
VectorQuery,
|
||||
createAutoQuery,
|
||||
instanceOfFullTextQuery,
|
||||
} from "./query";
|
||||
import { sanitizeType } from "./sanitize";
|
||||
@@ -1177,33 +1178,29 @@ export class LocalTable extends Table {
|
||||
});
|
||||
}
|
||||
|
||||
if (queryType === "auto" && typeof query !== "string") {
|
||||
return this.query().fullTextSearch(query, {
|
||||
columns: ftsColumns,
|
||||
});
|
||||
}
|
||||
if (queryType === "auto") {
|
||||
if (instanceOfFullTextQuery(query)) {
|
||||
return this.query().fullTextSearch(query, {
|
||||
columns: ftsColumns,
|
||||
});
|
||||
}
|
||||
|
||||
if (queryType === "auto" && typeof query === "string") {
|
||||
const vector = async (snapshot: _NativeTable) => {
|
||||
const functions = await this.getEmbeddingFunctions(snapshot);
|
||||
const columns =
|
||||
typeof ftsColumns === "string" ? [ftsColumns] : (ftsColumns ?? null);
|
||||
return createAutoQuery(this.inner, query, columns, async (metadata) => {
|
||||
const functions = await getRegistry().parseFunctions(
|
||||
new Map([["embedding_functions", metadata]]),
|
||||
);
|
||||
// TODO: Support multiple embedding functions
|
||||
const embeddingFunc: EmbeddingFunctionConfig | undefined = functions
|
||||
.values()
|
||||
.next().value;
|
||||
if (embeddingFunc === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
// The route only calls this callback when embedding metadata exists.
|
||||
// parseFunctions either yields a provider or reports malformed metadata.
|
||||
if (!embeddingFunc)
|
||||
throw new Error("Invalid embedding function metadata");
|
||||
return await embeddingFunc.function.computeQueryEmbeddings(query);
|
||||
};
|
||||
|
||||
const columns =
|
||||
typeof ftsColumns === "string" ? [ftsColumns] : ftsColumns;
|
||||
return Query.autoSearch(
|
||||
() => this.inner.checkoutCurrent(),
|
||||
query,
|
||||
vector,
|
||||
columns,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
const queryPromise = this.getEmbeddingFunctions().then(
|
||||
|
||||
@@ -278,6 +278,13 @@ impl Table {
|
||||
Ok(Query::new(self.inner_ref()?.query()))
|
||||
}
|
||||
|
||||
/// Return a read-only table handle pinned to the current query revision.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn query_snapshot(&self) -> napi::Result<Self> {
|
||||
let snapshot = self.inner_ref()?.query_snapshot().await.default_error()?;
|
||||
Ok(Self::new(snapshot))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn take_offsets(&self, offsets: Vec<i64>) -> napi::Result<TakeQuery> {
|
||||
Ok(TakeQuery::new(
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
import threading
|
||||
|
||||
@@ -86,6 +87,25 @@ def test_s3_lifecycle(s3_bucket: str):
|
||||
asyncio.run(test())
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_concurrent_open_table(s3_bucket: str):
|
||||
uri = f"s3://{s3_bucket}/test_concurrent_open_table"
|
||||
db = lancedb.connect(uri, storage_options=copy.copy(CONFIG))
|
||||
db.create_table("test", pa.table({"x": [1, 2, 3]}))
|
||||
|
||||
num_workers = 32
|
||||
barrier = threading.Barrier(num_workers)
|
||||
|
||||
def open_and_count(_):
|
||||
barrier.wait()
|
||||
return db.open_table("test").count_rows()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as pool:
|
||||
row_counts = list(pool.map(open_and_count, range(num_workers)))
|
||||
|
||||
assert row_counts == [3] * num_workers
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def kms_key():
|
||||
kms = get_boto3_client("kms", endpoint_url=CONFIG["aws_endpoint"])
|
||||
|
||||
@@ -1476,7 +1476,7 @@ mod tests {
|
||||
use crate::table::{AnyQuery, WriteOptions};
|
||||
use arrow_array::{Int32Array, RecordBatch, StringArray};
|
||||
use arrow_schema::{DataType, Field, Schema, SchemaRef};
|
||||
use futures::{TryStreamExt, stream::once};
|
||||
use futures::{TryStreamExt, future::try_join_all, stream::once};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
@@ -1614,6 +1614,59 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_concurrent_open_table_reuses_connection_object_store() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let uri = tempdir.path().to_str().unwrap();
|
||||
let session = Arc::new(lance::session::Session::default());
|
||||
let request = ConnectRequest {
|
||||
uri: uri.to_string(),
|
||||
#[cfg(feature = "remote")]
|
||||
client_config: Default::default(),
|
||||
options: Default::default(),
|
||||
namespace_client_properties: Default::default(),
|
||||
manifest_enabled: false,
|
||||
read_consistency_interval: None,
|
||||
session: Some(session.clone()),
|
||||
};
|
||||
let db = ListingDatabase::connect_with_options(&request)
|
||||
.await
|
||||
.unwrap();
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
|
||||
db.create_table(CreateTableRequest {
|
||||
name: "test".to_string(),
|
||||
namespace_path: vec![],
|
||||
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
|
||||
mode: CreateTableMode::Create,
|
||||
write_options: Default::default(),
|
||||
location: None,
|
||||
namespace_client: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let before = session.store_registry().stats();
|
||||
let opened_tables = try_join_all((0..32).map(|_| {
|
||||
db.open_table(OpenTableRequest {
|
||||
name: "test".to_string(),
|
||||
namespace_path: vec![],
|
||||
index_cache_size: None,
|
||||
lance_read_params: None,
|
||||
location: None,
|
||||
namespace_client: None,
|
||||
managed_versioning: None,
|
||||
})
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let after = session.store_registry().stats();
|
||||
|
||||
assert_eq!(opened_tables.len(), 32);
|
||||
assert_eq!(after.misses, before.misses);
|
||||
assert_eq!(after.active_stores, before.active_stores);
|
||||
assert!(after.hits >= before.hits + 32);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_listing_database_root_ops_do_not_create_manifest() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
|
||||
@@ -10,7 +10,7 @@ use lance::io::WrappingObjectStore;
|
||||
use object_store::{
|
||||
CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
|
||||
ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result,
|
||||
UploadPart, path::Path,
|
||||
UploadPart, list::PaginatedListStore, path::Path,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -187,6 +187,14 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper {
|
||||
secondary: self.secondary.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn PaginatedListStore>,
|
||||
) -> Option<Arc<dyn PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
// windows pathing can't be simply concatenated
|
||||
|
||||
@@ -12,7 +12,7 @@ use lance::io::WrappingObjectStore;
|
||||
use object_store::{
|
||||
CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
|
||||
PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult,
|
||||
UploadPart, path::Path,
|
||||
UploadPart, list::PaginatedListStore, path::Path,
|
||||
};
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -57,6 +57,14 @@ impl WrappingObjectStore for IoStatsHolder {
|
||||
stats: self.0.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn PaginatedListStore>,
|
||||
) -> Option<Arc<dyn PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
impl IoTrackingStore {
|
||||
|
||||
@@ -47,6 +47,10 @@ impl TerminalResult {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn value(&self) -> Option<&Value> {
|
||||
self.value.as_ref()
|
||||
}
|
||||
|
||||
fn decode<T: DeserializeOwned>(self) -> Result<T> {
|
||||
let value = self.value.ok_or_else(|| match &self.request_id {
|
||||
Some(request_id) => Error::Http {
|
||||
|
||||
+1089
-279
File diff suppressed because it is too large
Load Diff
@@ -6,6 +6,7 @@
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::{Array, LargeBinaryArray};
|
||||
use arrow_schema::DataType;
|
||||
@@ -20,7 +21,7 @@ use crate::error::Result;
|
||||
use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient};
|
||||
use crate::table::BaseTable;
|
||||
|
||||
use super::{FreshnessHeaders, RemoteTable};
|
||||
use super::{FreshnessHeaders, FreshnessState, RemoteTable, freshness_headers_snapshot};
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum RangeRequestMode {
|
||||
@@ -43,7 +44,10 @@ struct TableBlobRangeRequester<S: HttpSend> {
|
||||
path: String,
|
||||
version: Option<u64>,
|
||||
branch: Option<String>,
|
||||
freshness: FreshnessHeaders,
|
||||
freshness: Arc<std::sync::Mutex<FreshnessState>>,
|
||||
parent_freshness: Arc<std::sync::Mutex<FreshnessState>>,
|
||||
parent_freshness_request: FreshnessHeaders,
|
||||
read_consistency_interval: Option<Duration>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
@@ -53,8 +57,9 @@ impl<S: HttpSend> BlobRangeRequester for TableBlobRangeRequester<S> {
|
||||
range_header: &str,
|
||||
mode: RangeRequestMode,
|
||||
) -> Result<(String, Response)> {
|
||||
let mut request = self
|
||||
.freshness
|
||||
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 {
|
||||
@@ -71,6 +76,9 @@ impl<S: HttpSend> BlobRangeRequester for TableBlobRangeRequester<S> {
|
||||
return Ok((request_id, response));
|
||||
}
|
||||
let response = self.client.check_response(&request_id, response).await?;
|
||||
freshness_request.observe_headers(&self.freshness, response.headers());
|
||||
self.parent_freshness_request
|
||||
.observe_headers(&self.parent_freshness, response.headers());
|
||||
Ok((request_id, response))
|
||||
}
|
||||
}
|
||||
@@ -361,18 +369,21 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
message: "fetch_blobs is not supported on this LanceDB Cloud server".into(),
|
||||
});
|
||||
}
|
||||
let version = self.current_version().await;
|
||||
let read_snapshot = self.snapshot_read_state().await;
|
||||
let mut body = serde_json::json!({
|
||||
"version": version,
|
||||
"version": read_snapshot.version,
|
||||
"column": column,
|
||||
"row_ids": row_ids,
|
||||
});
|
||||
self.apply_branch_body(&mut body);
|
||||
|
||||
let request = self
|
||||
.post_read(&format!("/v1/table/{}/fetch_blobs/", self.identifier))
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/fetch_blobs/", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let (request_id, response) = self
|
||||
.send_with_freshness(request, true, read_snapshot.freshness)
|
||||
.await?;
|
||||
let mut stream = self.read_arrow_response(&request_id, response).await?;
|
||||
|
||||
let mut blob_chunks: Vec<Arc<dyn Array>> = Vec::new();
|
||||
@@ -448,8 +459,7 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
});
|
||||
}
|
||||
|
||||
let version = self.current_version().await;
|
||||
let freshness = self.snapshot_freshness_headers();
|
||||
let read_snapshot = self.snapshot_read_state().await;
|
||||
let encoded_column = urlencoding::encode(column);
|
||||
let requesters = row_ids
|
||||
.iter()
|
||||
@@ -461,9 +471,12 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
let requester: Arc<dyn BlobRangeRequester> = Arc::new(TableBlobRangeRequester {
|
||||
client: self.client.clone(),
|
||||
path,
|
||||
version,
|
||||
version: read_snapshot.version,
|
||||
branch: self.branch.clone(),
|
||||
freshness,
|
||||
freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)),
|
||||
parent_freshness: self.freshness.clone(),
|
||||
parent_freshness_request: read_snapshot.freshness,
|
||||
read_consistency_interval: self.client.read_consistency_interval,
|
||||
});
|
||||
requester
|
||||
})
|
||||
@@ -685,6 +698,46 @@ mod tests {
|
||||
assert!(requests.lock().unwrap().contains(&"bytes=5-11".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()));
|
||||
let captured = range_requests.clone();
|
||||
let table = RemoteTable::new_mock(
|
||||
"my_table".to_string(),
|
||||
move |request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version":5,"schema":{"fields":[]}}"#.as_bytes().to_vec())
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/blob/image/10/bytes" => {
|
||||
captured.lock().unwrap().push((
|
||||
request.url().query().unwrap_or_default().to_string(),
|
||||
request.headers().clone(),
|
||||
));
|
||||
range_response(&request, PAYLOAD)
|
||||
}
|
||||
path => panic!("unexpected path: {path}"),
|
||||
},
|
||||
Some(Version::new(0, 5, 0)),
|
||||
);
|
||||
|
||||
table.checkout(5).await.unwrap();
|
||||
let file = table
|
||||
.fetch_blob_files_impl("image", &[10])
|
||||
.await
|
||||
.unwrap()
|
||||
.pop()
|
||||
.flatten()
|
||||
.unwrap();
|
||||
table.checkout_latest().await.unwrap();
|
||||
file.read_range(5..12).await.unwrap();
|
||||
|
||||
let requests = range_requests.lock().unwrap();
|
||||
let (query, headers) = requests.last().unwrap();
|
||||
assert!(query.contains("version=5"));
|
||||
assert!(!headers.contains_key("x-lancedb-min-timestamp"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_blob_file_reuses_sequential_response_until_seek() {
|
||||
let requests = Arc::new(StdMutex::new(Vec::new()));
|
||||
|
||||
@@ -24,7 +24,10 @@ use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter;
|
||||
use crate::Error;
|
||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||
use crate::remote::client::{HttpSend, RestfulLanceDbClient, Sender};
|
||||
use crate::remote::table::{MergeInsertRequest, REQUEST_TIMEOUT_HEADER, RemoteTable};
|
||||
use crate::remote::table::{
|
||||
FreshnessHeaders, FreshnessState, MergeInsertRequest, REQUEST_TIMEOUT_HEADER, RemoteTable,
|
||||
freshness_headers_snapshot,
|
||||
};
|
||||
use crate::table::datafusion::insert::COUNT_SCHEMA;
|
||||
use crate::table::write_progress::WriteProgressTracker;
|
||||
use crate::table::{AddResult, MergeResult};
|
||||
@@ -54,6 +57,38 @@ pub enum WriteResult {
|
||||
Merge(MergeResult),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct WriteFreshness {
|
||||
state: Option<Arc<Mutex<FreshnessState>>>,
|
||||
read_consistency_interval: Option<Duration>,
|
||||
}
|
||||
|
||||
impl WriteFreshness {
|
||||
fn prepare(
|
||||
&self,
|
||||
request: reqwest::RequestBuilder,
|
||||
) -> (reqwest::RequestBuilder, Option<FreshnessHeaders>) {
|
||||
match &self.state {
|
||||
Some(state) => {
|
||||
let freshness_request =
|
||||
freshness_headers_snapshot(state, self.read_consistency_interval);
|
||||
(freshness_request.apply(request), Some(freshness_request))
|
||||
}
|
||||
None => (request, None),
|
||||
}
|
||||
}
|
||||
|
||||
fn observe(
|
||||
&self,
|
||||
freshness_request: Option<FreshnessHeaders>,
|
||||
headers: &reqwest::header::HeaderMap,
|
||||
) {
|
||||
if let (Some(state), Some(freshness_request)) = (&self.state, freshness_request) {
|
||||
freshness_request.observe_headers(state, headers);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ExecutionPlan for streaming a write (add or merge_insert) to a remote
|
||||
/// LanceDB table.
|
||||
///
|
||||
@@ -71,6 +106,7 @@ pub struct RemoteWriteExec<S: HttpSend = Sender> {
|
||||
table_name: String,
|
||||
identifier: String,
|
||||
client: RestfulLanceDbClient<S>,
|
||||
freshness: WriteFreshness,
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
op: WriteOp,
|
||||
properties: Arc<PlanProperties>,
|
||||
@@ -170,6 +206,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
table_name,
|
||||
identifier,
|
||||
client,
|
||||
freshness: WriteFreshness::default(),
|
||||
input,
|
||||
op,
|
||||
properties: Arc::new(properties),
|
||||
@@ -183,6 +220,18 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn with_freshness(
|
||||
mut self,
|
||||
state: Arc<Mutex<FreshnessState>>,
|
||||
read_consistency_interval: Option<Duration>,
|
||||
) -> Self {
|
||||
self.freshness = WriteFreshness {
|
||||
state: Some(state),
|
||||
read_consistency_interval,
|
||||
};
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the add result after execution, if this exec ran an insert.
|
||||
pub fn add_result(&self) -> Option<AddResult> {
|
||||
match self
|
||||
@@ -285,6 +334,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
/// each threading the same handful of arguments.
|
||||
struct PartRequestCtx<'a, S: HttpSend> {
|
||||
client: &'a RestfulLanceDbClient<S>,
|
||||
freshness: &'a WriteFreshness,
|
||||
identifier: &'a str,
|
||||
table_name: &'a str,
|
||||
upload_id: &'a str,
|
||||
@@ -352,7 +402,11 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
}
|
||||
|
||||
/// Build the `/insert` request for a single multipart part.
|
||||
fn build_part_request(&self, part_id: &str, body: reqwest::Body) -> reqwest::RequestBuilder {
|
||||
fn build_part_request(
|
||||
&self,
|
||||
part_id: &str,
|
||||
body: reqwest::Body,
|
||||
) -> (reqwest::RequestBuilder, Option<FreshnessHeaders>) {
|
||||
let mut request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/insert/", self.identifier))
|
||||
@@ -368,12 +422,16 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
if let Some(b) = self.branch {
|
||||
request = request.query(&[("branch", b)]);
|
||||
}
|
||||
request.body(body)
|
||||
self.freshness.prepare(request.body(body))
|
||||
}
|
||||
|
||||
/// Send a single part's request and drain the response, mapping HTTP and
|
||||
/// table-not-found errors into `DataFusionError`.
|
||||
async fn send_part_request(&self, request: reqwest::RequestBuilder) -> DataFusionResult<()> {
|
||||
async fn send_part_request(
|
||||
&self,
|
||||
request: reqwest::RequestBuilder,
|
||||
freshness_request: Option<FreshnessHeaders>,
|
||||
) -> DataFusionResult<()> {
|
||||
let (request_id, response) = self
|
||||
.client
|
||||
.send(request)
|
||||
@@ -388,6 +446,8 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
.check_response(&request_id, response)
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
self.freshness
|
||||
.observe(freshness_request, response.headers());
|
||||
response.bytes().await.map_err(|e| {
|
||||
DataFusionError::External(Box::new(Error::Http {
|
||||
source: Box::new(e),
|
||||
@@ -419,7 +479,7 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
let body = reqwest::Body::wrap_stream(chunk_rx);
|
||||
|
||||
let part_id = uuid::Uuid::new_v4().to_string();
|
||||
let request = self.build_part_request(&part_id, body);
|
||||
let (request, freshness_request) = self.build_part_request(&part_id, body);
|
||||
|
||||
// Measured from just before the request is sent, matching the window the
|
||||
// client read timeout applies to the upload.
|
||||
@@ -495,7 +555,7 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
Ok::<bool, DataFusionError>(input_ended)
|
||||
};
|
||||
|
||||
let send = self.send_part_request(request);
|
||||
let send = self.send_part_request(request, freshness_request);
|
||||
|
||||
// `join!` rather than `tokio::spawn`: the producer borrows `input` (and
|
||||
// `schema`), so it cannot satisfy the `'static` bound a spawned task
|
||||
@@ -569,7 +629,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
// Building a fresh exec (with a new, empty `result`) is what makes the
|
||||
// outer rescannable retry loop work: `reset_state()` clears the captured
|
||||
// result so a re-execution starts clean.
|
||||
Ok(Arc::new(Self::new_inner(
|
||||
let mut exec = Self::new_inner(
|
||||
self.table_name.clone(),
|
||||
self.identifier.clone(),
|
||||
self.client.clone(),
|
||||
@@ -580,7 +640,9 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
self.branch.clone(),
|
||||
self.max_bytes_per_request,
|
||||
self.max_request_duration,
|
||||
)))
|
||||
);
|
||||
exec.freshness = self.freshness.clone();
|
||||
Ok(Arc::new(exec))
|
||||
}
|
||||
|
||||
fn execute(
|
||||
@@ -613,6 +675,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
&self.metrics,
|
||||
));
|
||||
let client = self.client.clone();
|
||||
let freshness = self.freshness.clone();
|
||||
let identifier = self.identifier.clone();
|
||||
let op = self.op.clone();
|
||||
let result_slot = self.result.clone();
|
||||
@@ -634,6 +697,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
let overwrite = matches!(op, WriteOp::Insert { overwrite: true });
|
||||
let ctx = PartRequestCtx {
|
||||
client: &client,
|
||||
freshness: &freshness,
|
||||
identifier: &identifier,
|
||||
table_name: &table_name,
|
||||
upload_id,
|
||||
@@ -688,7 +752,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
|
||||
let (error_tx, mut error_rx) = tokio::sync::oneshot::channel();
|
||||
let body = Self::stream_as_http_body(input_stream, error_tx, tracker)?;
|
||||
let request = request.body(body);
|
||||
let (request, freshness_request) = freshness.prepare(request.body(body));
|
||||
|
||||
let result: DataFusionResult<(String, _)> = async {
|
||||
let (request_id, response) = client
|
||||
@@ -708,6 +772,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
.check_response(&request_id, response)
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
freshness.observe(freshness_request, response.headers());
|
||||
|
||||
Ok((request_id, response))
|
||||
}
|
||||
|
||||
@@ -560,6 +560,13 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
fn id(&self) -> &str;
|
||||
/// Get the arrow [Schema] of the table.
|
||||
async fn schema(&self) -> Result<SchemaRef>;
|
||||
/// Create a read-only handle pinned to the table's current active revision.
|
||||
///
|
||||
/// The returned handle is independent from later refreshes or checkouts on
|
||||
/// this handle. This is used by bindings that must prepare client-side
|
||||
/// query state from the same revision that the query will execute against.
|
||||
#[doc(hidden)]
|
||||
async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Count the number of rows in this table.
|
||||
async fn count_rows(&self, filter: Option<Filter>) -> Result<usize>;
|
||||
/// Create a physical plan for the query.
|
||||
@@ -1139,6 +1146,16 @@ impl Table {
|
||||
self.inner.schema().await
|
||||
}
|
||||
|
||||
/// Create a read-only handle pinned to the current active revision.
|
||||
#[doc(hidden)]
|
||||
pub async fn query_snapshot(&self) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: self.inner.query_snapshot().await?,
|
||||
database: self.database.clone(),
|
||||
embedding_registry: self.embedding_registry.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Count the number of rows in this dataset.
|
||||
///
|
||||
/// # Arguments
|
||||
@@ -3059,6 +3076,17 @@ impl BaseTable for NativeTable {
|
||||
&self.id
|
||||
}
|
||||
|
||||
async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>> {
|
||||
let snapshot = self.dataset.new_query_snapshot().await?;
|
||||
let mut table = self.with_dataset(snapshot);
|
||||
// QueryTable requests do not carry a revision. A pinned snapshot must
|
||||
// execute locally until the namespace API can accept that revision.
|
||||
table
|
||||
.pushdown_operations
|
||||
.remove(&NamespaceClientPushdownOperation::QueryTable);
|
||||
Ok(Arc::new(table))
|
||||
}
|
||||
|
||||
async fn version(&self) -> Result<u64> {
|
||||
Ok(self.dataset.get().await?.version().version)
|
||||
}
|
||||
@@ -4118,6 +4146,14 @@ mod tests {
|
||||
parent_list_calls: self.parent_list_calls.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
_original: Arc<dyn object_store::list::PaginatedListStore>,
|
||||
) -> Option<Arc<dyn object_store::list::PaginatedListStore>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -4221,6 +4257,14 @@ mod tests {
|
||||
self.called.store(true, Ordering::Relaxed);
|
||||
original
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn object_store::list::PaginatedListStore>,
|
||||
) -> Option<Arc<dyn object_store::list::PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -32,6 +32,10 @@ struct DatasetState {
|
||||
/// `Some(version)` = pinned to a specific version (time travel),
|
||||
/// `None` = tracking latest.
|
||||
pinned_version: Option<u64>,
|
||||
/// Whether the pin is an internal query snapshot rather than user-visible
|
||||
/// time travel. Query snapshots remain read-only but preserve MemWAL read
|
||||
/// semantics.
|
||||
query_snapshot: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -70,6 +74,7 @@ impl DatasetConsistencyWrapper {
|
||||
state: Arc::new(Mutex::new(DatasetState {
|
||||
dataset,
|
||||
pinned_version: None,
|
||||
query_snapshot: false,
|
||||
})),
|
||||
consistency,
|
||||
shard_writer: Arc::new(ShardWriterCache::default()),
|
||||
@@ -93,6 +98,36 @@ impl DatasetConsistencyWrapper {
|
||||
wrapper
|
||||
}
|
||||
|
||||
/// Create an independent read-only wrapper pinned to the current dataset
|
||||
/// while retaining this wrapper's live MemWAL read context.
|
||||
pub async fn new_query_snapshot(&self) -> Result<Self> {
|
||||
// Apply the configured consistency policy before taking the snapshot.
|
||||
// The returned dataset is intentionally discarded: a checkout may race
|
||||
// after this await, so the dataset and its pin provenance must instead
|
||||
// be cloned together from one authoritative state sample below.
|
||||
self.get().await?;
|
||||
|
||||
let (dataset, query_snapshot) = {
|
||||
let state = self.state.lock()?;
|
||||
// Preserve user time travel so the MemWAL safety guard still sees
|
||||
// it. Latest and already-internal snapshots remain internal pins.
|
||||
(
|
||||
state.dataset.clone(),
|
||||
state.query_snapshot || state.pinned_version.is_none(),
|
||||
)
|
||||
};
|
||||
let version = dataset.version().version;
|
||||
Ok(Self {
|
||||
state: Arc::new(Mutex::new(DatasetState {
|
||||
dataset,
|
||||
pinned_version: Some(version),
|
||||
query_snapshot,
|
||||
})),
|
||||
consistency: ConsistencyMode::Lazy,
|
||||
shard_writer: self.shard_writer.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The MemWAL `ShardWriter` cache co-located with this dataset.
|
||||
pub(crate) fn shard_writer(&self) -> &Arc<ShardWriterCache> {
|
||||
&self.shard_writer
|
||||
@@ -169,6 +204,7 @@ impl DatasetConsistencyWrapper {
|
||||
let mut state = self.state.lock()?;
|
||||
state.dataset = Arc::new(new_dataset);
|
||||
state.pinned_version = None;
|
||||
state.query_snapshot = false;
|
||||
drop(state);
|
||||
if let ConsistencyMode::Eventual(bg_cache) = &self.consistency {
|
||||
bg_cache.invalidate();
|
||||
@@ -202,10 +238,10 @@ impl DatasetConsistencyWrapper {
|
||||
|
||||
/// Returns the version, if in time travel mode, or None otherwise.
|
||||
pub fn time_travel_version(&self) -> Option<u64> {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.pinned_version
|
||||
let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
|
||||
(!state.query_snapshot)
|
||||
.then_some(state.pinned_version)
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Convert into a wrapper in latest version mode.
|
||||
@@ -225,6 +261,7 @@ impl DatasetConsistencyWrapper {
|
||||
if state.pinned_version.is_some() {
|
||||
state.dataset = Arc::new(new_dataset);
|
||||
state.pinned_version = None;
|
||||
state.query_snapshot = false;
|
||||
}
|
||||
drop(state);
|
||||
if let ConsistencyMode::Eventual(bg_cache) = &self.consistency {
|
||||
@@ -260,6 +297,7 @@ impl DatasetConsistencyWrapper {
|
||||
let mut state = self.state.lock()?;
|
||||
state.dataset = Arc::new(new_dataset);
|
||||
state.pinned_version = Some(version_value);
|
||||
state.query_snapshot = false;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -461,6 +499,29 @@ mod tests {
|
||||
assert_eq!(wrapper.time_travel_version(), Some(1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_snapshot_samples_dataset_and_pin_together() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap();
|
||||
let ds = create_test_dataset(uri).await;
|
||||
|
||||
let wrapper = DatasetConsistencyWrapper::new_latest(ds, None);
|
||||
wrapper.as_time_travel(1u64).await.unwrap();
|
||||
let stale_time_travel_dataset = wrapper.get().await.unwrap();
|
||||
|
||||
append_to_dataset(uri).await;
|
||||
wrapper.as_latest().await.unwrap();
|
||||
|
||||
let snapshot = wrapper.new_query_snapshot().await.unwrap();
|
||||
let snapshot_dataset = snapshot.get().await.unwrap();
|
||||
assert_eq!(snapshot_dataset.version().version, 2);
|
||||
assert_ne!(
|
||||
snapshot_dataset.version().version,
|
||||
stale_time_travel_dataset.version().version
|
||||
);
|
||||
assert_eq!(snapshot.time_travel_version(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_as_latest_from_time_travel() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -1056,6 +1056,44 @@ mod lsm_tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn query_snapshot_preserves_lsm_read_semantics() {
|
||||
let dir = tempdir().unwrap();
|
||||
let table = id_value_table(&dir).await;
|
||||
table
|
||||
.set_lsm_write_spec(LsmWriteSpec::unsharded())
|
||||
.await
|
||||
.unwrap();
|
||||
lsm_upsert(&table, vec![4, 5]).await;
|
||||
|
||||
let snapshot = table.query_snapshot().await.unwrap();
|
||||
let rows = collect_id_value(snapshot.query().execute().await.unwrap()).await;
|
||||
assert_eq!(
|
||||
rows.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
|
||||
vec![1, 2, 3, 4, 5]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn query_snapshot_preserves_time_travel_lsm_guard() {
|
||||
let dir = tempdir().unwrap();
|
||||
let table = id_value_table(&dir).await;
|
||||
table
|
||||
.set_lsm_write_spec(LsmWriteSpec::unsharded())
|
||||
.await
|
||||
.unwrap();
|
||||
lsm_upsert(&table, vec![4]).await;
|
||||
|
||||
let version = table.version().await.unwrap();
|
||||
table.checkout(version).await.unwrap();
|
||||
let direct_error = table.query().execute().await.err().unwrap();
|
||||
assert!(matches!(direct_error, Error::NotSupported { .. }));
|
||||
|
||||
let snapshot = table.query_snapshot().await.unwrap();
|
||||
let snapshot_error = snapshot.query().execute().await.err().unwrap();
|
||||
assert!(matches!(snapshot_error, Error::NotSupported { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lsm_read_dedup_newest_wins() {
|
||||
let dir = tempdir().unwrap();
|
||||
|
||||
@@ -1056,6 +1056,37 @@ mod tests {
|
||||
assert_eq!(namespace_client.query_table_calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_snapshot_disables_namespace_pushdown() {
|
||||
use crate::connect;
|
||||
use crate::table::BaseTable;
|
||||
use arrow_array::{Int32Array, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
|
||||
let batch =
|
||||
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap();
|
||||
let table = conn
|
||||
.create_table("test_snapshot_namespace_fallback", vec![batch])
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let mut native_table = table.as_native().unwrap().clone();
|
||||
native_table.namespace_client = Some(Arc::new(CountingNamespaceClient::default()));
|
||||
native_table
|
||||
.pushdown_operations
|
||||
.insert(NamespaceClientPushdownOperation::QueryTable);
|
||||
|
||||
let snapshot = BaseTable::query_snapshot(&native_table).await.unwrap();
|
||||
let snapshot = snapshot.as_any().downcast_ref::<NativeTable>().unwrap();
|
||||
assert!(
|
||||
!can_execute_namespace_query(snapshot, &AnyQuery::Query(QueryRequest::default()),)
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_plan_multivector_structure() {
|
||||
use arrow_array::{Float32Array, RecordBatch};
|
||||
|
||||
@@ -298,7 +298,7 @@ async fn build_read_context(
|
||||
for shard_id in shard_ids {
|
||||
let manifest_store =
|
||||
ShardManifestStore::new(store.clone(), &base_path, shard_id, scan_batch_size);
|
||||
if let Some(manifest) = manifest_store.read_latest().await? {
|
||||
if let Some(manifest) = manifest_store.latest().await? {
|
||||
snapshots.push(snapshot_from_manifest(shard_id, &manifest, &exclude));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user