Compare commits

..

3 Commits

Author SHA1 Message Date
Lu Qiu 86d3f23f53 chore: keep remote write types crate-visible 2026-07-31 17:45:27 -04:00
Lu Qiu 1c130f61da Merge remote-tracking branch 'upstream/main' into codex/update-lance-10-1-0-beta-1
# Conflicts:
#	Cargo.lock
#	Cargo.toml
#	java/pom.xml
2026-07-31 17:22:33 -04:00
Lu Qiu b36cdc72d9 chore: update lance dependency to v10.1.0-beta.1 2026-07-31 16:47:25 -04:00
35 changed files with 236 additions and 2579 deletions
Generated
+70 -70
View File
@@ -601,7 +601,7 @@ dependencies = [
"bytes",
"fastrand",
"hex",
"http 1.5.0",
"http 1.4.2",
"sha1 0.10.6",
"time",
"tokio",
@@ -664,7 +664,7 @@ dependencies = [
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"percent-encoding",
@@ -694,7 +694,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body-util",
"regex-lite",
"tracing",
@@ -719,7 +719,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"regex-lite",
"tracing",
]
@@ -743,7 +743,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"regex-lite",
"tracing",
]
@@ -773,7 +773,7 @@ dependencies = [
"hex",
"hmac 0.13.0",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"lru",
"percent-encoding",
@@ -802,7 +802,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"regex-lite",
"tracing",
]
@@ -826,7 +826,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"regex-lite",
"tracing",
]
@@ -851,7 +851,7 @@ dependencies = [
"aws-types",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"regex-lite",
"tracing",
]
@@ -873,7 +873,7 @@ dependencies = [
"hex",
"hmac 0.13.0",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"p256",
"percent-encoding",
"ring",
@@ -906,7 +906,7 @@ dependencies = [
"bytes",
"crc-fast",
"hex",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"md-5 0.11.0",
@@ -940,7 +940,7 @@ dependencies = [
"bytes-utils",
"futures-core",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"percent-encoding",
@@ -961,7 +961,7 @@ dependencies = [
"h2 0.3.27",
"h2 0.4.14",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body 0.4.6",
"hyper 0.14.32",
"hyper 1.9.0",
@@ -1023,7 +1023,7 @@ dependencies = [
"bytes",
"fastrand",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"http-body-util",
@@ -1044,7 +1044,7 @@ dependencies = [
"aws-smithy-types",
"bytes",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"pin-project-lite",
"tokio",
"tracing",
@@ -1070,7 +1070,7 @@ checksum = "7442cb268338f0eb8278140a107c046756aa01093d8ef5e99628d34ae09c94f5"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-types",
"http 1.5.0",
"http 1.4.2",
]
[[package]]
@@ -1084,7 +1084,7 @@ dependencies = [
"bytes-utils",
"futures-core",
"http 0.2.12",
"http 1.5.0",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"http-body-util",
@@ -1132,7 +1132,7 @@ dependencies = [
"axum-core",
"bytes",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.9.0",
@@ -1165,7 +1165,7 @@ dependencies = [
"async-trait",
"bytes",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"mime",
@@ -3842,7 +3842,7 @@ dependencies = [
"fnv",
"futures-core",
"futures-sink",
"http 1.5.0",
"http 1.4.2",
"indexmap 2.14.0",
"slab",
"tokio",
@@ -3953,7 +3953,7 @@ checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
dependencies = [
"dirs",
"futures",
"http 1.5.0",
"http 1.4.2",
"indicatif",
"libc",
"log",
@@ -3976,7 +3976,7 @@ checksum = "430b33fa84f92796d4d263070b6c0d3ca219df7b9a0e1853ee431029b1612bcd"
dependencies = [
"async-trait",
"bytes",
"http 1.5.0",
"http 1.4.2",
"more-asserts",
"serde",
"thiserror 2.0.18",
@@ -4041,9 +4041,9 @@ dependencies = [
[[package]]
name = "http"
version = "1.5.0"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0"
checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425"
dependencies = [
"bytes",
"itoa",
@@ -4067,7 +4067,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c"
dependencies = [
"bytes",
"http 1.5.0",
"http 1.4.2",
]
[[package]]
@@ -4078,7 +4078,7 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a"
dependencies = [
"bytes",
"futures-core",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"pin-project-lite",
]
@@ -4145,7 +4145,7 @@ dependencies = [
"futures-channel",
"futures-core",
"h2 0.4.14",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"httparse",
"httpdate",
@@ -4177,7 +4177,7 @@ version = "0.27.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f"
dependencies = [
"http 1.5.0",
"http 1.4.2",
"hyper 1.9.0",
"hyper-util",
"rustls 0.23.40",
@@ -4211,7 +4211,7 @@ dependencies = [
"bytes",
"futures-channel",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"hyper 1.9.0",
"ipnet",
@@ -5182,7 +5182,7 @@ dependencies = [
"chrono",
"futures",
"goosefs-sdk",
"http 1.5.0",
"http 1.4.2",
"io-uring",
"lance-arrow",
"lance-core",
@@ -5414,7 +5414,7 @@ dependencies = [
"goosefs-sdk",
"half",
"hf-hub",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"lance",
"lance-arrow",
@@ -6080,15 +6080,15 @@ dependencies = [
[[package]]
name = "napi-build"
version = "2.4.0"
version = "2.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5282704fbe8d49b0cf8b08e3f33233416a528658f205c7e5ace63b582de0b11c"
checksum = "c9c366d2c8c60b86fa632df75f745509b52f9128f91a6bad4c796e44abb505e1"
[[package]]
name = "napi-derive"
version = "3.6.1"
version = "3.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4d5c9c02556ea6dc99dffd36c1ce60141411657438501a125b675776d011ce92"
checksum = "a49c513341a61a16a10af6efcce46b30d0822ba2d4fb197d24d33dfc199c78d5"
dependencies = [
"convert_case",
"ctor 1.0.5",
@@ -6100,9 +6100,9 @@ dependencies = [
[[package]]
name = "napi-derive-backend"
version = "6.1.1"
version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d60b5d773ad46c698c8cc2cd9fde0b283d39cbb7f71c04bee633c7bdba4423bd"
checksum = "4747005fa3e2c9989ac45a723a514c5db2411238b72981a3cda4c701a9dfea17"
dependencies = [
"convert_case",
"proc-macro2",
@@ -6356,7 +6356,7 @@ dependencies = [
"futures-channel",
"futures-core",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body-util",
"httparse",
"humantime",
@@ -6477,7 +6477,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"futures",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"jiff",
"log",
@@ -6502,7 +6502,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0d6f81ba6960e3fae1882f253b114b21d7e444e1534f209c7737a79f6243eb6f"
dependencies = [
"futures",
"http 1.5.0",
"http 1.4.2",
"mea",
"opendal-core",
]
@@ -6546,7 +6546,7 @@ checksum = "0030644366ef5d8cbe3a4a5822bf99a4aafddc1666e9d24b44d158d9062fc76a"
dependencies = [
"base64 0.22.1",
"bytes",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"opendal-service-azure-common",
@@ -6567,7 +6567,7 @@ checksum = "6dea4908d490143a9b0b7f7a790e139ff829b06a023f670455ed3d44f664b361"
dependencies = [
"base64 0.22.1",
"bytes",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"opendal-service-azure-common",
@@ -6585,7 +6585,7 @@ version = "0.57.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b489f13c42e69d69bdd72952b634356ec43a7881a20259b38b540fcecdf4051"
dependencies = [
"http 1.5.0",
"http 1.4.2",
"opendal-core",
]
@@ -6596,7 +6596,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aa8cafe9729213375c7331019b0cb756ad3e1aff7f45cd32c45eae91ebde8901"
dependencies = [
"bytes",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"quick-xml 0.39.4",
@@ -6614,7 +6614,7 @@ checksum = "48de101aac565ed06af4b47903c24eafd249075553ec1fb18256751c45148d47"
dependencies = [
"async-trait",
"bytes",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"percent-encoding",
@@ -6649,7 +6649,7 @@ checksum = "c4922661976a1d40794a2adfbdb888cc3c23097690f825a92f773af38908a848"
dependencies = [
"bytes",
"hf-xet",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"percent-encoding",
@@ -6665,7 +6665,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "328fa55e8888cbdfe00826bfea2a79042422b720e8369e9e021e46121dea5ace"
dependencies = [
"bytes",
"http 1.5.0",
"http 1.4.2",
"log",
"opendal-core",
"quick-xml 0.39.4",
@@ -6684,7 +6684,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"crc32c",
"http 1.5.0",
"http 1.4.2",
"log",
"md-5 0.11.0",
"opendal-core",
@@ -7583,7 +7583,7 @@ version = "0.14.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
dependencies = [
"heck 0.4.1",
"heck 0.5.0",
"itertools 0.14.0",
"log",
"multimap",
@@ -8226,7 +8226,7 @@ checksum = "57ac2757f3140aa2e213b554148ae0b52733e624fc6723f0cc6bb3d440176c95"
dependencies = [
"anyhow",
"form_urlencoded",
"http 1.5.0",
"http 1.4.2",
"log",
"percent-encoding",
"reqsign-core",
@@ -8244,7 +8244,7 @@ dependencies = [
"anyhow",
"bytes",
"form_urlencoded",
"http 1.5.0",
"http 1.4.2",
"log",
"percent-encoding",
"quick-xml 0.39.4",
@@ -8266,7 +8266,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"form_urlencoded",
"http 1.5.0",
"http 1.4.2",
"jsonwebtoken",
"log",
"pem",
@@ -8291,7 +8291,7 @@ dependencies = [
"futures",
"hex",
"hmac 0.12.1",
"http 1.5.0",
"http 1.4.2",
"jiff",
"log",
"percent-encoding",
@@ -8318,7 +8318,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "35cc609b49c69e76ecaceb775a03f792d1ed3e7755ab3548d4534fd801e3242e"
dependencies = [
"form_urlencoded",
"http 1.5.0",
"http 1.4.2",
"jsonwebtoken",
"log",
"percent-encoding",
@@ -8338,7 +8338,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e128f19525861dbded59e1e7c17653a8ed63d573ca04aed708d552dbef5bb32a"
dependencies = [
"anyhow",
"http 1.5.0",
"http 1.4.2",
"log",
"percent-encoding",
"reqsign-core",
@@ -8358,7 +8358,7 @@ dependencies = [
"futures-core",
"futures-util",
"h2 0.4.14",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.9.0",
@@ -8402,7 +8402,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.9.0",
@@ -8456,7 +8456,7 @@ checksum = "199dda04a536b532d0cc04d7979e39b1c763ea749bf91507017069c00b96056f"
dependencies = [
"anyhow",
"async-trait",
"http 1.5.0",
"http 1.4.2",
"reqwest 0.13.3",
"thiserror 2.0.18",
"tower-service",
@@ -8498,9 +8498,9 @@ dependencies = [
[[package]]
name = "rkyv"
version = "0.8.17"
version = "0.8.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "815cc8a37159a463064825246cadb07961e25cd9885908606f6d08a98d8f8874"
checksum = "73389e0c99e664f919275ab5b5b0471391fe9a8de61e1dff9b1eaf56a90f16e3"
dependencies = [
"bytecheck",
"bytes",
@@ -8517,9 +8517,9 @@ dependencies = [
[[package]]
name = "rkyv_derive"
version = "0.8.17"
version = "0.8.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c0ed1a78a1b19d184b0daa629dd9a024573173ec7d485b287cb369fb3607cc1c"
checksum = "5d2ed0b54125315fb36bd021e82d314d1c126548f871634b483f46b31d13cac6"
dependencies = [
"proc-macro2",
"quote",
@@ -9277,7 +9277,7 @@ version = "0.8.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451"
dependencies = [
"heck 0.4.1",
"heck 0.5.0",
"proc-macro2",
"quote",
"syn 2.0.117",
@@ -9289,7 +9289,7 @@ version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40"
dependencies = [
"heck 0.4.1",
"heck 0.5.0",
"proc-macro2",
"quote",
"syn 2.0.117",
@@ -9730,7 +9730,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.3.4",
"getrandom 0.4.2",
"once_cell",
"rustix",
"windows-sys 0.61.2",
@@ -10058,7 +10058,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"h2 0.4.14",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.9.0",
@@ -10114,7 +10114,7 @@ checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5"
dependencies = [
"bitflags 2.11.1",
"bytes",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"pin-project-lite",
@@ -10134,7 +10134,7 @@ dependencies = [
"bytes",
"futures-core",
"futures-util",
"http 1.5.0",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"pin-project-lite",
@@ -11155,7 +11155,7 @@ dependencies = [
"clap",
"crc32fast",
"futures",
"http 1.5.0",
"http 1.4.2",
"hyper 1.9.0",
"lazy_static",
"more-asserts",
@@ -11229,7 +11229,7 @@ dependencies = [
"chrono",
"clap",
"gearhash",
"http 1.5.0",
"http 1.4.2",
"itertools 0.14.0",
"lazy_static",
"more-asserts",
@@ -10,24 +10,6 @@ Reranks the results using the Reciprocal Rank Fusion (RRF) algorithm.
## Methods
### outputSchema()
```ts
outputSchema(inputSchema): Promise<Schema<any>>
```
Declare the RRF output schema for vector-only query execution.
#### Parameters
* **inputSchema**: `Schema`&lt;`any`&gt;
#### Returns
`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
***
### rerankHybrid()
```ts
@@ -8,27 +8,6 @@
## Methods
### outputSchema()?
```ts
optional outputSchema(inputSchema): Promise<Schema<any>>
```
Declare the schema returned when reranking a vector-only query.
This is required for vector-only reranking so query schema introspection
and execution agree. Hybrid-only rerankers may omit it.
#### Parameters
* **inputSchema**: `Schema`&lt;`any`&gt;
#### Returns
`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
***
### rerankHybrid()
```ts
-16
View File
@@ -79,22 +79,6 @@ describe("rerankers", function () {
expect(result).toHaveLength(2);
});
it("returns relevance scores when reranking a vector search", async function () {
const query = table
.vectorSearch([0.1, 0.1])
.limit(2)
.rerank(await RRFReranker.create());
const schema = await query.outputSchema();
const result = await query.toArray();
expect(schema.fields.map((field) => field.name)).toContain(
"_relevance_score",
);
expect(result).toHaveLength(2);
expect(result[0]._relevance_score).toBeCloseTo(1 / 60);
expect(result[1]._relevance_score).toBeCloseTo(1 / 61);
});
it("does not keep process alive after rerank query", async function () {
const script = `
import * as lancedb from "./dist/index.js";
+13 -25
View File
@@ -5,11 +5,9 @@ import {
Table as ArrowTable,
type IntoVector,
RecordBatch,
createEmptyTable,
extractVectorBuffer,
fromBufferToRecordBatch,
fromRecordBatchToBuffer,
fromTableToBuffer,
tableFromIPC,
} from "./arrow";
import { type IvfPqOptions } from "./indices";
@@ -746,30 +744,20 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
}
rerank(reranker: Reranker): VectorQuery {
super.doCall((inner) => {
const outputSchema = reranker.outputSchema?.bind(reranker);
inner.rerank(
async (args) => {
const vecResults = await fromBufferToRecordBatch(args.vecResults);
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
const result = await reranker.rerankHybrid(
args.query,
vecResults as RecordBatch,
ftsResults as RecordBatch,
);
super.doCall((inner) =>
inner.rerank(async (args) => {
const vecResults = await fromBufferToRecordBatch(args.vecResults);
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
const result = await reranker.rerankHybrid(
args.query,
vecResults as RecordBatch,
ftsResults as RecordBatch,
);
const buffer = fromRecordBatchToBuffer(result);
return buffer;
},
outputSchema
? async (args) => {
const inputSchema = tableFromIPC(args.inputSchema).schema;
const result = await outputSchema(inputSchema);
return fromTableToBuffer(createEmptyTable(result));
}
: undefined,
);
});
const buffer = fromRecordBatchToBuffer(result);
return buffer;
}),
);
return this;
}
+4 -12
View File
@@ -1,22 +1,14 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import { RecordBatch, Schema } from "apache-arrow";
import { RecordBatch } from "apache-arrow";
export * from "./rrf";
// Interface for a reranker. A reranker is used to rerank vector and hybrid
// search results. For vector-only searches, query is empty and ftsResults is an
// empty batch with the same schema as vecResults.
// Interface for a reranker. A reranker is used to rerank the results from a
// vector and FTS search. This is useful for combining the results from both
// search methods.
export interface Reranker {
/**
* Declare the schema returned when reranking a vector-only query.
*
* This is required for vector-only reranking so query schema introspection
* and execution agree. Hybrid-only rerankers may omit it.
*/
outputSchema?(inputSchema: Schema): Promise<Schema>;
rerankHybrid(
query: string,
vecResults: RecordBatch,
+1 -12
View File
@@ -1,7 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import { Field, Float32, RecordBatch, Schema } from "apache-arrow";
import { RecordBatch } from "apache-arrow";
import { fromBufferToRecordBatch, fromRecordBatchToBuffer } from "../arrow";
import { RrfReranker as NativeRRFReranker } from "../native";
@@ -24,17 +24,6 @@ export class RRFReranker {
);
}
/** Declare the RRF output schema for vector-only query execution. */
async outputSchema(inputSchema: Schema): Promise<Schema> {
return new Schema(
[
...inputSchema.fields,
new Field("_relevance_score", new Float32(), false),
],
inputSchema.metadata,
);
}
async rerankHybrid(
query: string,
vecResults: RecordBatch,
+2 -3
View File
@@ -6,8 +6,8 @@ use std::sync::Arc;
use crate::error::NapiErrorExt;
use crate::error::convert_error;
use crate::iterator::RecordBatchIterator;
use crate::rerankers::RerankHybridCallbackArgs;
use crate::rerankers::Reranker;
use crate::rerankers::{RerankHybridCallbackArgs, RerankOutputSchemaCallbackArgs};
use crate::util::{parse_distance_type, schema_to_buffer};
use arrow_array::{
Array, Float16Array as ArrowFloat16Array, Float32Array as ArrowFloat32Array,
@@ -388,9 +388,8 @@ impl VectorQuery {
pub fn rerank(
&mut self,
rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
) -> napi::Result<()> {
let reranker = Reranker::new(rerank_hybrid, output_schema)?;
let reranker = Reranker::new(rerank_hybrid)?;
self.inner = self.inner.clone().rerank(Arc::new(reranker));
Ok(())
}
+2 -51
View File
@@ -6,7 +6,7 @@ use async_trait::async_trait;
use napi::{bindgen_prelude::*, threadsafe_function::ThreadsafeFunction};
use napi_derive::napi;
use lancedb::ipc::{batches_to_ipc_file, ipc_file_to_schema, schema_to_ipc_file};
use lancedb::ipc::batches_to_ipc_file;
use lancedb::rerankers::Reranker as LanceDBReranker;
use lancedb::{error::Error, ipc::ipc_file_to_batches};
@@ -21,72 +21,28 @@ type RerankHybridFn = ThreadsafeFunction<
true,
>;
type RerankOutputSchemaFn = ThreadsafeFunction<
RerankOutputSchemaCallbackArgs,
Promise<Buffer>,
RerankOutputSchemaCallbackArgs,
Status,
false,
true,
>;
/// Reranker implementation that "wraps" a NodeJS Reranker implementation.
/// This contains references to the callbacks that can be used to invoke the
/// reranking methods on the NodeJS implementation and handles serializing the
/// record batches to Arrow IPC buffers.
pub struct Reranker {
rerank_hybrid: RerankHybridFn,
output_schema: Option<RerankOutputSchemaFn>,
}
impl Reranker {
pub fn new(
rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
) -> napi::Result<Self> {
let rerank_hybrid = rerank_hybrid
.build_threadsafe_function()
.weak::<true>()
.build()?;
let output_schema = output_schema
.map(|output_schema| {
output_schema
.build_threadsafe_function()
.weak::<true>()
.build()
})
.transpose()?;
Ok(Self {
rerank_hybrid,
output_schema,
})
Ok(Self { rerank_hybrid })
}
}
#[async_trait]
impl lancedb::rerankers::Reranker for Reranker {
async fn output_schema(
&self,
input: &arrow_schema::SchemaRef,
) -> lancedb::error::Result<arrow_schema::SchemaRef> {
let output_schema = self.output_schema.as_ref().ok_or(Error::NotSupported {
message: "vector rerankers must declare their output schema".to_string(),
})?;
let callback_args = RerankOutputSchemaCallbackArgs {
input_schema: Buffer::from(schema_to_ipc_file(input.as_ref())?),
};
let promised_buffer: Promise<Buffer> = output_schema
.call_async(callback_args)
.await
.map_err(|e| Error::Runtime {
message: format!("napi error status={}, reason={}", e.status, e.reason),
})?;
let buffer = promised_buffer.await.map_err(|e| Error::Runtime {
message: format!("napi error status={}, reason={}", e.status, e.reason),
})?;
ipc_file_to_schema(buffer.to_vec())
}
async fn rerank_hybrid(
&self,
query: &str,
@@ -130,11 +86,6 @@ pub struct RerankHybridCallbackArgs {
pub fts_results: Buffer,
}
#[napi(object)]
pub struct RerankOutputSchemaCallbackArgs {
pub input_schema: Buffer,
}
fn buffer_to_record_batch(buffer: Buffer) -> Result<RecordBatch> {
let mut reader = ipc_file_to_batches(buffer.to_vec()).default_error()?;
reader
+2 -6
View File
@@ -339,9 +339,7 @@ impl Table {
let transforms = NewColumnTransform::SqlExpressions(transforms);
let res = self
.inner_ref()?
.add_columns()
.transform(transforms)
.execute()
.add_columns(transforms, None)
.await
.default_error()?;
Ok(res.into())
@@ -358,9 +356,7 @@ impl Table {
let transforms = NewColumnTransform::AllNulls(schema);
let res = self
.inner_ref()?
.add_columns()
.transform(transforms)
.execute()
.add_columns(transforms, None)
.await
.default_error()?;
Ok(res.into())
+3 -17
View File
@@ -707,9 +707,6 @@ class LanceDBConnection(DBConnection):
self._namespace_client_properties = namespace_client_properties
if _inner is not None:
self._conn = _inner
# Native-derived wrappers resolve this in their async reconstruction
# path so construction never synchronously re-enters LOOP.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client = None
return
@@ -759,14 +756,11 @@ class LanceDBConnection(DBConnection):
# storage_options. Also, this class really shouldn't be holding any state
# beyond _conn.
self._conn = AsyncConnection(LOOP.run(do_connect()))
# Keep property access synchronous so debugger introspection cannot wait on
# the background loop while that thread is suspended at a breakpoint.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client: Optional[LanceNamespace] = None
@property
def read_consistency_interval(self) -> Optional[timedelta]:
return self._read_consistency_interval
return LOOP.run(self._conn.get_read_consistency_interval())
@property
def session(self) -> Optional[Session]:
@@ -777,16 +771,8 @@ class LanceDBConnection(DBConnection):
return self._conn.uri
@classmethod
def from_inner(
cls,
inner: LanceDbConnection,
read_consistency_interval: Optional[timedelta],
):
return cls(
None,
read_consistency_interval=read_consistency_interval,
_inner=inner,
)
def from_inner(cls, inner: LanceDbConnection):
return cls(None, _inner=inner)
def __repr__(self) -> str:
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
+1 -3
View File
@@ -92,10 +92,8 @@ class LanceMergeInsertBuilder(object):
self._when_not_matched_by_source_delete = True
if isinstance(condition, Expr):
self._when_not_matched_by_source_condition_expr = condition._inner
self._when_not_matched_by_source_condition = None
else:
elif condition is not None:
self._when_not_matched_by_source_condition = condition
self._when_not_matched_by_source_condition_expr = None
return self
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
+1 -1
View File
@@ -226,7 +226,7 @@ class PermutationBuilder:
async def do_execute():
inner_tbl = await self._async.execute()
return await LanceTable.from_inner(inner_tbl)
return LanceTable.from_inner(inner_tbl)
return LOOP.run(do_execute())
+1 -1
View File
@@ -2697,7 +2697,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._fts_query.phrase_query(True)
if self._distance_type:
self._vector_query.metric(self._distance_type)
if self._minimum_nprobes is not None:
if self._minimum_nprobes:
self._vector_query.minimum_nprobes(self._minimum_nprobes)
if self._maximum_nprobes is not None:
self._vector_query.maximum_nprobes(self._maximum_nprobes)
+5 -9
View File
@@ -1595,8 +1595,8 @@ class Table(ABC):
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
newer.
rows are ``None``. Unsupported on LanceDB Cloud, where
:meth:`fetch_blobs` returns full bytes instead.
"""
@abstractmethod
@@ -2182,15 +2182,11 @@ class LanceTable(Table):
return self.name
@classmethod
async def from_inner(cls, tbl: LanceDBTable):
from .db import AsyncConnection, LanceDBConnection
def from_inner(cls, tbl: LanceDBTable):
from .db import LanceDBConnection
async_tbl = AsyncTable(tbl)
inner_conn = tbl.database()
read_consistency_interval = await AsyncConnection(
inner_conn
).get_read_consistency_interval()
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
conn = LanceDBConnection.from_inner(tbl.database())
return cls(
conn,
async_tbl.name,
-17
View File
@@ -77,23 +77,6 @@ def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
from lancedb.background_loop import LOOP
from lancedb.db import LanceDBConnection
consistency_interval = timedelta(seconds=5)
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
def fail_run(*args, **kwargs):
raise AssertionError("properties should not use the Python background loop")
monkeypatch.setattr(LOOP, "run", fail_run)
assert db.read_consistency_interval == consistency_interval
assert db_from_inner.read_consistency_interval == consistency_interval
def test_ingest_pd(tmp_path):
db = lancedb.connect(tmp_path)
-13
View File
@@ -123,19 +123,6 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
assert texts.count("a") == 1
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.
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
(
sync_table.search(query_type="hybrid")
.vector([0.0, 0.4])
.text("dog")
.minimum_nprobes(0)
.to_arrow()
)
def test_hybrid_query_distance_range(sync_table: Table):
reranker = RRFReranker(return_score="all")
result = (
-20
View File
@@ -6,7 +6,6 @@ import math
import pytest
from lancedb import DBConnection, Table, connect
from lancedb.background_loop import LOOP
from lancedb.permutation import Permutation, Permutations, permutation_builder
@@ -32,25 +31,6 @@ def test_split_random_ratios(mem_db):
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
import threading
db = connect(tmp_path)
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
original_run = LOOP.run
def fail_on_reentry(future):
assert threading.current_thread() is not LOOP.thread
return original_run(future)
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
permutation_tbl = permutation_builder(tbl).execute()
assert permutation_tbl.count_rows() == 10
assert permutation_tbl._conn.read_consistency_interval is None
def test_split_random_counts(mem_db):
"""Test random splitting with absolute counts."""
tbl = mem_db.create_table(
+2 -33
View File
@@ -2055,24 +2055,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
request.send_header("phalanx-version", str(server_version))
request.end_headers()
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
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]
if payload is None:
request.send_response(204)
request.end_headers()
return
byte_range = request.headers["Range"].removeprefix("bytes=")
start_text, end_text = byte_range.split("-", maxsplit=1)
start = int(start_text)
end = int(end_text) if end_text else len(payload) - 1
chunk = payload[start : end + 1]
request.send_response(206)
request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
request.send_header("Content-Length", str(len(chunk)))
request.end_headers()
request.wfile.write(chunk)
elif request.path == "/v1/table/test/query/":
content_len = int(request.headers.get("Content-Length", 0))
body = json.loads(request.rfile.read(content_len))
@@ -2110,21 +2092,8 @@ def test_remote_blob_columns_and_fetch():
assert table.blob_columns() == ["image"]
blobs = table.fetch_blobs("image", [10, 20, 30])
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
def test_remote_blob_files_are_lazy_seekable_handles():
with blob_remote_table() as table:
files = table.fetch_blob_files("image", [10, 20, 30])
assert len(files) == 3
alpha, null_row, gamma = files
assert null_row is None
assert alpha is not None
assert gamma is not None
assert alpha.size() == 5
assert alpha.read_range(1, 3) == b"lph"
gamma.seek(2)
assert gamma.read() == b"mma"
with pytest.raises(NotImplementedError, match="Use fetch_blobs for full bytes"):
table.fetch_blob_files("image", [10, 20, 30])
def test_remote_blob_fetch_accepts_query_table():
-45
View File
@@ -6,7 +6,6 @@ import os
import sys
import threading
import warnings
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from time import sleep
from typing import List
@@ -2125,27 +2124,6 @@ def test_delete(mem_db: DBConnection):
assert table.to_arrow()["id"].to_pylist() == [1]
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
num_workers = 8
table = mem_db.create_table(
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
)
barrier = threading.Barrier(num_workers)
def delete(row_id: int):
barrier.wait()
return table.delete(f"id = {row_id}")
with ThreadPoolExecutor(max_workers=num_workers) as pool:
results = list(pool.map(delete, range(num_workers)))
assert all(result.num_deleted_rows == 1 for result in results)
assert sorted(result.version for result in results) == list(
range(2, num_workers + 2)
)
assert table.count_rows() == 0
def test_delete_expr(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2386,29 +2364,6 @@ def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
assert table.to_arrow().sort_by("a") == expected
def test_merge_insert_by_source_delete_reconfigure(mem_db: DBConnection):
# Calling when_not_matched_by_source_delete() again with no condition must
# widen the delete to unconditional, not keep the earlier condition around.
table = mem_db.create_table(
"my_table",
data=pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]}),
)
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
merge_insert_res = (
table.merge_insert("a")
.when_matched_update_all()
.when_not_matched_insert_all()
.when_not_matched_by_source_delete("a > 2")
.when_not_matched_by_source_delete()
.execute(new_data)
)
assert merge_insert_res.num_deleted_rows == 2
expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
assert table.to_arrow().sort_by("a") == expected
@pytest.mark.asyncio
async def test_merge_insert_by_source_delete_expr_async(
mem_db_async: AsyncConnection,
+13 -40
View File
@@ -426,11 +426,9 @@ pub struct PyBlobFile {
impl PyBlobFile {
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
let inner = self_.inner.clone();
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read().await }))
let bytes = block_on(async move { inner.read().await })
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
}
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
@@ -446,32 +444,24 @@ impl PyBlobFile {
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
let inner = self_.inner.clone();
self_
.py()
.detach(move || block_on(async move { inner.close().await }))
block_on(async move { inner.close().await })
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
}
fn is_closed(self_: PyRef<'_, Self>) -> bool {
let inner = self_.inner.clone();
self_
.py()
.detach(move || block_on(async move { inner.is_closed().await }))
block_on(async move { inner.is_closed().await })
}
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
let inner = self_.inner.clone();
self_
.py()
.detach(move || block_on(async move { inner.seek(position).await }))
block_on(async move { inner.seek(position).await })
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
}
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
let inner = self_.inner.clone();
self_
.py()
.detach(move || block_on(async move { inner.tell().await }))
block_on(async move { inner.tell().await })
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
}
@@ -485,20 +475,16 @@ impl PyBlobFile {
.checked_add(length as u64)
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
let inner = self_.inner.clone();
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
let bytes = block_on(async move { inner.read_range(offset..end).await })
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
}
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
let inner = self_.inner.clone();
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read_up_to(length).await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
let bytes = block_on(async move { inner.read_up_to(length).await })
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
}
}
@@ -745,9 +731,6 @@ impl Table {
#[allow(private_interfaces)]
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
// Do not hold the Python borrow across the await. The cloned Rust table
// handle is thread-safe and allows deletes on the same Python table to
// run concurrently without PyO3 reporting "Already borrowed".
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = match &condition {
@@ -1378,12 +1361,7 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner
.add_columns()
.transform(definitions)
.execute()
.await
.infer_error()?;
let result = inner.add_columns(definitions, None).await.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
@@ -1397,12 +1375,7 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner
.add_columns()
.transform(transform)
.execute()
.await
.infer_error()?;
let result = inner.add_columns(transform, None).await.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
+1 -1
View File
@@ -1,2 +1,2 @@
[toolchain]
channel = "1.97.0"
channel = "1.95.0"
+2 -198
View File
@@ -9,7 +9,6 @@
//!
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
use std::ops::Range;
use std::sync::Arc;
use arrow_array::LargeBinaryArray;
@@ -18,202 +17,10 @@ use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
use lance_arrow::FieldExt;
use lance_encoding::version::LanceFileVersion;
use lance_io::object_store::ObjectStore;
use object_store::path::Path;
use crate::error::{Error, Result};
/// Seekable handle for one blob value, backed by local storage or a remote
/// HTTP byte-range endpoint.
#[derive(Debug)]
pub struct BlobFile {
inner: BlobFileInner,
}
#[derive(Debug)]
enum BlobFileInner {
Native(lance::dataset::BlobFile),
#[cfg(feature = "remote")]
Remote(Box<crate::remote::table::blobs::RemoteBlobFile>),
}
impl From<lance::dataset::BlobFile> for BlobFile {
fn from(value: lance::dataset::BlobFile) -> Self {
Self {
inner: BlobFileInner::Native(value),
}
}
}
#[cfg(feature = "remote")]
impl From<crate::remote::table::blobs::RemoteBlobFile> for BlobFile {
fn from(value: crate::remote::table::blobs::RemoteBlobFile) -> Self {
Self {
inner: BlobFileInner::Remote(Box::new(value)),
}
}
}
impl BlobFile {
/// Inline reader over a data-file slice.
pub fn new_inline(
object_store: Arc<ObjectStore>,
path: Path,
position: u64,
size: u64,
) -> Self {
lance::dataset::BlobFile::new_inline(object_store, path, position, size).into()
}
/// Dedicated sidecar-file reader.
pub fn new_dedicated(object_store: Arc<ObjectStore>, path: Path, size: u64) -> Self {
lance::dataset::BlobFile::new_dedicated(object_store, path, size).into()
}
/// Packed reader for a slice in a shared sidecar.
pub fn new_packed(
object_store: Arc<ObjectStore>,
path: Path,
position: u64,
size: u64,
) -> Self {
lance::dataset::BlobFile::new_packed(object_store, path, position, size).into()
}
/// External reader at a resolved object location.
pub fn new_external(
object_store: Arc<ObjectStore>,
path: Path,
uri: String,
position: u64,
size: u64,
) -> Self {
lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into()
}
/// Close the handle.
pub async fn close(&self) -> lance_core::Result<()> {
match &self.inner {
BlobFileInner::Native(file) => file.close().await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.close().await,
}
}
/// Whether the handle is closed.
pub async fn is_closed(&self) -> bool {
match &self.inner {
BlobFileInner::Native(file) => file.is_closed().await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.is_closed(),
}
}
/// Read a range without moving the cursor.
pub async fn read_range(&self, range: Range<u64>) -> lance_core::Result<bytes::Bytes> {
match &self.inner {
BlobFileInner::Native(file) => file.read_range(range).await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.read_range(range).await,
}
}
/// Read ranges without moving the cursor.
pub async fn read_ranges(
&self,
ranges: &[Range<u64>],
) -> lance_core::Result<Vec<bytes::Bytes>> {
match &self.inner {
BlobFileInner::Native(file) => file.read_ranges(ranges).await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.read_ranges(ranges).await,
}
}
/// Read from the cursor to the end.
pub async fn read(&self) -> lance_core::Result<bytes::Bytes> {
match &self.inner {
BlobFileInner::Native(file) => file.read().await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.read().await,
}
}
/// Read up to `len` bytes and advance the cursor.
pub async fn read_up_to(&self, len: usize) -> lance_core::Result<bytes::Bytes> {
match &self.inner {
BlobFileInner::Native(file) => file.read_up_to(len).await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.read_up_to(len).await,
}
}
/// Move the cursor to `new_cursor`.
pub async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> {
match &self.inner {
BlobFileInner::Native(file) => file.seek(new_cursor).await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.seek(new_cursor).await,
}
}
/// Current cursor position.
pub async fn tell(&self) -> lance_core::Result<u64> {
match &self.inner {
BlobFileInner::Native(file) => file.tell().await,
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.tell().await,
}
}
/// Blob length in bytes.
pub fn size(&self) -> u64 {
match &self.inner {
BlobFileInner::Native(file) => file.size(),
#[cfg(feature = "remote")]
BlobFileInner::Remote(file) => file.size(),
}
}
/// Physical byte offset in the data file. `None` on remote handles. The
/// Cloud byte-range route does not expose storage layout.
pub fn position(&self) -> Option<u64> {
match &self.inner {
BlobFileInner::Native(file) => Some(file.position()),
#[cfg(feature = "remote")]
BlobFileInner::Remote(_) => None,
}
}
/// Path of the data file holding the blob. `None` on remote handles. The
/// Cloud byte-range route does not expose storage layout.
pub fn data_path(&self) -> Option<&Path> {
match &self.inner {
BlobFileInner::Native(file) => Some(file.data_path()),
#[cfg(feature = "remote")]
BlobFileInner::Remote(_) => None,
}
}
/// Native storage layout. `None` on remote handles. The Cloud byte-range
/// route does not expose layout.
pub fn kind(&self) -> Option<lance_core::datatypes::BlobKind> {
match &self.inner {
BlobFileInner::Native(file) => Some(file.kind()),
#[cfg(feature = "remote")]
BlobFileInner::Remote(_) => None,
}
}
/// External URI for native handles. Remote handles do not expose storage URIs.
pub fn uri(&self) -> Option<&str> {
match &self.inner {
BlobFileInner::Native(file) => file.uri(),
#[cfg(feature = "remote")]
BlobFileInner::Remote(_) => None,
}
}
}
pub use lance::dataset::BlobFile;
/// One row-specific blob range read request.
///
@@ -457,10 +264,7 @@ pub(crate) async fn take_blob_files_aligned(
let handles = dataset.take_blobs(row_ids, column).await?;
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
Ok(handles
.into_iter()
.map(|handle| handle.map(Into::into))
.collect())
Ok(handles)
}
#[cfg(test)]
+3 -3
View File
@@ -264,7 +264,7 @@ pub fn compute_output_schema(
let field_name = ed
.dest_column
.clone()
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
sb.push(Field::new(
field_name,
@@ -291,7 +291,7 @@ pub fn compute_embeddings_for_batch(
let dst_field_name = fld
.dest_column
.clone()
.unwrap_or_else(|| format!("{}_embedding", fld.source_column));
.unwrap_or_else(|| format!("{}_embedding", &fld.source_column));
let dst_field = Field::new(
dst_field_name,
@@ -315,7 +315,7 @@ impl<R: RecordBatchReader> WithEmbeddings<R> {
let field_name = ed
.dest_column
.clone()
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
Ok(Field::new(
field_name,
func.dest_type()?.into_owned(),
+2 -219
View File
@@ -511,9 +511,7 @@ pub trait QueryBase {
/// Rerank the results using the specified reranker.
///
/// For vector-only searches, the reranker receives the vector results and an
/// empty full-text result set and must declare its output schema. Reranking
/// multiple query vectors in one query is not supported.
/// This is currently only supported for Hybrid Search.
fn rerank(self, reranker: Arc<dyn Reranker>) -> Self;
/// The method to normalize the scores. Can be "rank" or "Score". If "Rank",
@@ -1140,44 +1138,6 @@ pub struct VectorQuery {
}
impl VectorQuery {
fn check_vector_rerank_supported(&self) -> Result<()> {
if self.request.query_vector.len() > 1 {
return Err(Error::NotSupported {
message: "reranking multiple query vectors is not supported; execute one query per vector"
.to_string(),
});
}
Ok(())
}
async fn vector_rerank_output_schema(&self) -> Result<SchemaRef> {
self.check_vector_rerank_supported()?;
// Rerankers receive row IDs internally. Apply their schema transform to
// that exact input and then hide the row ID from the declared public
// schema unless it was explicitly requested.
let vector_query = self.clone().with_row_id();
let plan = vector_query
.create_plan(QueryExecutionOptions::default())
.await?;
let reranker = self
.request
.base
.reranker
.as_ref()
.expect("vector_rerank_output_schema requires a reranker");
let input_schema = plan.schema();
let output_schema = reranker.output_schema(&input_schema).await?;
if self.request.base.with_row_id {
Ok(output_schema)
} else {
Ok(RecordBatch::new_empty(output_schema)
.drop_column(ROW_ID)?
.schema())
}
}
fn new(base: Query) -> Self {
Self {
parent: base.parent,
@@ -1483,61 +1443,6 @@ impl VectorQuery {
Ok(single_batch_stream(results, max_batch_length))
}
async fn execute_vector_rerank(
&self,
options: QueryExecutionOptions,
) -> Result<SendableRecordBatchStream> {
self.check_vector_rerank_supported()?;
let max_batch_length = options.max_batch_length as usize;
let internal_options = options.without_output_batch_length_limit();
// RRF needs row IDs to assign and preserve scores. Keep them internal unless
// the caller explicitly requested them.
let vector_query = self.clone().with_row_id();
let vector_results = vector_query
.inner_execute_with_options(internal_options)
.await?;
let schema = vector_results.schema();
let vector_results = vector_results.try_collect::<Vec<_>>().await?;
let vector_results = concat_batches(&schema, vector_results.iter())?;
let vector_schema = vector_results.schema();
let fts_results = RecordBatch::new_empty(vector_schema.clone());
let reranker = self
.request
.base
.reranker
.as_ref()
.expect("execute_vector_rerank requires a reranker");
let expected_schema = reranker.output_schema(&vector_schema).await?;
let mut results = reranker
.rerank_hybrid("", vector_results, fts_results)
.await?;
check_reranker_result(&results)?;
if results.schema() != expected_schema {
return Err(Error::Schema {
message: format!(
"reranker returned schema {:?}, but declared {:?}",
results.schema(),
expected_schema
),
});
}
let limit = self.request.base.limit.unwrap_or(DEFAULT_TOP_K);
if results.num_rows() > limit {
results = results.slice(0, limit);
}
if !self.request.base.with_row_id {
results = results.drop_column(ROW_ID)?;
}
Ok(single_batch_stream(results, max_batch_length))
}
async fn inner_execute_with_options(
&self,
options: QueryExecutionOptions,
@@ -1590,23 +1495,6 @@ impl ExecutableQuery for VectorQuery {
return Ok(hybrid_result);
}
if self.request.base.reranker.is_some() {
let timeout = options.timeout;
let mut rerank_options = options;
// A single outer deadline covers planning, candidate collection,
// schema declaration, and the complete reranker callback.
rerank_options.timeout = None;
let execution = self.execute_vector_rerank(rerank_options);
return match timeout {
Some(timeout) => tokio::time::timeout(timeout, execution)
.await
.map_err(|_| Error::Timeout {
message: format!("Query timeout after {} ms", timeout.as_millis()),
})?,
None => execution.await,
};
}
self.inner_execute_with_options(options).await
}
@@ -1619,15 +1507,6 @@ impl ExecutableQuery for VectorQuery {
let query = AnyQuery::VectorQuery(self.request.clone());
self.parent.analyze_plan(&query, options).await
}
async fn output_schema(&self) -> Result<SchemaRef> {
if self.request.base.full_text_search.is_none() && self.request.base.reranker.is_some() {
self.vector_rerank_output_schema().await
} else {
let plan = self.create_plan(QueryExecutionOptions::default()).await?;
Ok(plan.schema())
}
}
}
impl HasQuery for VectorQuery {
@@ -1764,13 +1643,7 @@ impl ExecutableQuery for TakeQuery {
#[cfg(test)]
mod tests {
use std::{
collections::HashSet,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
};
use std::{collections::HashSet, sync::Arc};
use super::*;
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
@@ -1786,31 +1659,6 @@ mod tests {
use crate::{Table, connect, database::CreateTableMode, index::Index};
#[derive(Debug)]
struct SlowReranker {
invoked: Arc<AtomicBool>,
}
#[async_trait::async_trait]
impl Reranker for SlowReranker {
async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
RRFReranker::default().output_schema(input).await
}
async fn rerank_hybrid(
&self,
query: &str,
vector_results: RecordBatch,
fts_results: RecordBatch,
) -> Result<RecordBatch> {
self.invoked.store(true, Ordering::SeqCst);
tokio::time::sleep(Duration::from_secs(2)).await;
RRFReranker::default()
.rerank_hybrid(query, vector_results, fts_results)
.await
}
}
#[tokio::test]
async fn test_setters_getters() {
// TODO: Switch back to memory://foo after https://github.com/lancedb/lancedb/issues/1051
@@ -2522,71 +2370,6 @@ mod tests {
// We don't guarantee order.
assert!(query_index.values().contains(&0));
assert!(query_index.values().contains(&1));
let reranked = query.rerank(Arc::new(RRFReranker::default()));
let Err(execute_error) = reranked.execute().await else {
panic!("multi-vector reranking should be rejected");
};
assert!(
execute_error
.to_string()
.contains("reranking multiple query vectors is not supported")
);
let schema_error = reranked.output_schema().await.unwrap_err();
assert!(
schema_error
.to_string()
.contains("reranking multiple query vectors is not supported")
);
}
#[tokio::test]
async fn test_vector_rerank_timeout_covers_reranker() {
let tmp_dir = tempdir().unwrap();
let table = make_test_table(&tmp_dir).await;
let invoked = Arc::new(AtomicBool::new(false));
let reranker = SlowReranker {
invoked: invoked.clone(),
};
let result = table
.vector_search(&[0.1, 0.2, 0.3, 0.4])
.unwrap()
.limit(1)
.rerank(Arc::new(reranker))
.execute_with_options(QueryExecutionOptions {
timeout: Some(Duration::from_secs(1)),
..Default::default()
})
.await;
assert!(invoked.load(Ordering::SeqCst));
assert!(matches!(result, Err(Error::Timeout { .. })));
}
#[tokio::test]
async fn test_vector_rerank_output_schema_matches_execution() {
let tmp_dir = tempdir().unwrap();
let table = make_test_table(&tmp_dir).await;
let query = table
.vector_search(&[0.1, 0.2, 0.3, 0.4])
.unwrap()
.limit(1)
.rerank(Arc::new(RRFReranker::default()));
let promised = query.output_schema().await.unwrap();
let actual = query
.execute()
.await
.unwrap()
.next()
.await
.unwrap()
.unwrap()
.schema();
assert_eq!(promised, actual);
assert!(promised.column_with_name("_relevance_score").is_some());
}
#[tokio::test]
+2 -2
View File
@@ -706,7 +706,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
.err_to_http(request_id.clone())?;
debug!(
"Received response for request_id={}: {:?}",
request_id, response
request_id, &response
);
Ok((request_id, response))
}
@@ -768,7 +768,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
Ok((status, response)) if status.is_success() => {
debug!(
"Received response for request_id={}: {:?}",
retry_counter.request_id, response
retry_counter.request_id, &response
);
return Ok((retry_counter.request_id, response));
}
+44 -26
View File
@@ -1,7 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
pub mod blobs;
mod blobs;
pub mod insert;
use self::insert::{RemoteWriteExec, WriteOp};
@@ -3089,12 +3089,10 @@ mod tests {
Box::pin(table.delete("false").map_ok(|_| ())),
Box::pin(
table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"x".into(),
"y".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("x".into(), "y".into())]),
None,
)
.map_ok(|_| ()),
),
Box::pin(async {
@@ -4302,9 +4300,32 @@ mod tests {
"fetch_blobs",
);
let message = table
.fetch_blob_files("image", &[1])
.await
.unwrap_err()
.to_string();
assert!(
message.contains("fetch_blob_files is not supported on LanceDB Cloud"),
"got: {message}"
);
assert!(
!message.contains("Use fetch_blobs"),
"old server must not be told to use fetch_blobs, got: {message}"
);
}
#[tokio::test]
async fn test_blob_files_point_at_fetch_blobs_on_a_blob_capable_server() {
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
|_| -> http::Response<String> { panic!("fetch_blob_files must not reach the server") },
);
assert_not_supported_error(
table.fetch_blob_files("image", &[1]).await.unwrap_err(),
"requires LanceDB Cloud server 0.5.0 or newer",
"Use fetch_blobs for full bytes",
);
}
@@ -6390,12 +6411,13 @@ mod tests {
});
let result = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![
("b".into(), "a + 1".into()),
("x".into(), "cast(NULL as int32)".into()),
]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![
("b".into(), "a + 1".into()),
("x".into(), "cast(NULL as int32)".into()),
]),
None,
)
.await
.unwrap();
@@ -7120,12 +7142,10 @@ mod tests {
}
"add_columns" => {
let _ = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"c".into(),
"a + 1".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + 1".into())]),
None,
)
.await;
}
"drop_columns" => {
@@ -9883,12 +9903,10 @@ mod tests {
.await
.unwrap();
branch
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"b".into(),
"a + 1".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]),
None,
)
.await
.unwrap();
branch
File diff suppressed because it is too large Load Diff
+5 -18
View File
@@ -8,7 +8,6 @@ use arrow::{
compute::{concat_batches, filter_record_batch},
};
use arrow_array::{BooleanArray, RecordBatch, UInt64Array};
use arrow_schema::SchemaRef;
use async_trait::async_trait;
use lance::dataset::ROW_ID;
@@ -48,28 +47,16 @@ impl std::fmt::Display for NormalizeMethod {
}
}
/// Interface for a reranker. A reranker is used to rerank vector and hybrid
/// search results. This is useful for combining results from multiple search
/// methods or assigning a relevance score to vector search results.
/// Interface for a reranker. A reranker is used to rerank the results from a
/// vector and FTS search. This is useful for combining the results from both
/// search methods.
#[async_trait]
pub trait Reranker: std::fmt::Debug + Sync + Send {
/// Declare the schema returned by [`Self::rerank_hybrid`] for a vector-only
/// query.
///
/// Vector reranking validates the returned batch against this schema so
/// [`crate::query::ExecutableQuery::output_schema`] and execution cannot
/// disagree. Rerankers that only support hybrid search do not need to
/// implement this method.
async fn output_schema(&self, _input: &SchemaRef) -> Result<SchemaRef> {
Err(Error::NotSupported {
message: "vector rerankers must declare their output schema".to_string(),
})
}
// TODO support vector reranking and FTS reranking. Currently only hybrid reranking is supported.
/// Rerank function receives the individual results from the vector and FTS search
/// results. You can choose to use any of the results to generate the final results,
/// allowing maximum flexibility. For a vector-only search, `query` is empty and
/// `fts_results` is an empty batch with the same schema as `vector_results`.
/// allowing maximum flexibility.
async fn rerank_hybrid(
&self,
query: &str,
+9 -16
View File
@@ -9,7 +9,7 @@ use arrow::{
compute::{sort_to_indices, take},
};
use arrow_array::{Float32Array, RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions};
use arrow_schema::{DataType, Field, Schema, SortOptions};
use async_trait::async_trait;
use lance::dataset::ROW_ID;
@@ -44,19 +44,6 @@ impl Default for RRFReranker {
#[async_trait]
impl Reranker for RRFReranker {
async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
let mut fields = input.fields().to_vec();
fields.push(Arc::new(Field::new(
RELEVANCE_SCORE,
DataType::Float32,
false,
)));
Ok(Arc::new(Schema::new_with_metadata(
fields,
input.metadata().clone(),
)))
}
async fn rerank_hybrid(
&self,
_query: &str,
@@ -148,9 +135,15 @@ impl Reranker for RRFReranker {
.collect();
// add relevance score to schema
let schema = self.output_schema(&combined_results.schema()).await?;
let mut fields = combined_results.schema().fields().to_vec();
fields.push(Arc::new(Field::new(
RELEVANCE_SCORE,
DataType::Float32,
false,
)));
let schema = Schema::new(fields);
let combined_results = RecordBatch::try_new(schema, columns)?;
let combined_results = RecordBatch::try_new(Arc::new(schema), columns)?;
Ok(combined_results)
}
+7 -5
View File
@@ -3,7 +3,6 @@
//! LanceDB Table APIs
use crate::blob::BlobFile;
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
use arrow_schema::{Schema, SchemaRef};
use async_trait::async_trait;
@@ -13,6 +12,7 @@ use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::display::DisplayableExecutionPlan;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use lance::dataset::BlobFile;
pub use lance::dataset::ColumnAlteration;
pub use lance::dataset::NewColumnTransform;
pub use lance::dataset::ReadParams;
@@ -65,7 +65,6 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
use self::dataset::DatasetConsistencyWrapper;
use self::merge::MergeInsertBuilder;
pub mod add_columns;
mod add_data;
pub mod branch_merge;
mod create_index;
@@ -80,7 +79,6 @@ pub mod schema_evolution;
pub mod update;
pub mod write_progress;
use crate::index::waiter::wait_for_index;
pub use add_columns::AddColumnsBuilder;
#[cfg(feature = "remote")]
pub(crate) use add_data::PreprocessingOutput;
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
@@ -1622,8 +1620,12 @@ impl Table {
}
/// Add new columns to the table, providing values to fill in.
pub fn add_columns(&self) -> AddColumnsBuilder {
AddColumnsBuilder::new(self.inner.clone())
pub async fn add_columns(
&self,
transforms: NewColumnTransform,
read_columns: Option<Vec<String>>,
) -> Result<AddColumnsResult> {
self.inner.add_columns(transforms, read_columns).await
}
/// Change a column's name or nullability.
-161
View File
@@ -1,161 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Builder for adding columns to a table.
use std::sync::Arc;
use lance::dataset::NewColumnTransform;
use super::BaseTable;
use super::schema_evolution::AddColumnsResult;
use crate::{Error, Result};
/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns).
pub struct AddColumnsBuilder {
parent: Arc<dyn BaseTable>,
transform: Option<NewColumnTransform>,
read_columns: Option<Vec<String>>,
}
impl std::fmt::Debug for AddColumnsBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AddColumnsBuilder")
.field("parent", &self.parent)
.field("has_transform", &self.transform.is_some())
.field("read_columns", &self.read_columns)
.finish()
}
}
impl AddColumnsBuilder {
pub(crate) fn new(parent: Arc<dyn BaseTable>) -> Self {
Self {
parent,
transform: None,
read_columns: None,
}
}
/// Set how the new columns' values are produced. Required.
pub fn transform(mut self, transform: NewColumnTransform) -> Self {
self.transform = Some(transform);
self
}
/// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper
/// receives. Every other transform determines what it reads, so setting
/// this alongside one is an error rather than a silent no-op.
pub fn read_columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.read_columns = Some(columns.into_iter().map(Into::into).collect());
self
}
/// Add the columns.
pub async fn execute(self) -> Result<AddColumnsResult> {
let Self {
parent,
transform,
read_columns,
} = self;
let Some(transform) = transform else {
return Err(Error::InvalidInput {
message: "add_columns requires a transform".into(),
});
};
if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) {
return Err(Error::InvalidInput {
message: "read_columns applies only to a BatchUDF transform; \
every other transform determines what it reads"
.into(),
});
}
parent.add_columns(transform, read_columns).await
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{Int32Array, RecordBatch, record_batch};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{BatchUDF, NewColumnTransform};
use crate::Table;
use crate::connect;
async fn table_with_two_columns(name: &str) -> Table {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap();
conn.create_table(name, batch).execute().await.unwrap()
}
#[tokio::test]
async fn test_requires_a_transform() {
let table = table_with_two_columns("no_transform").await;
let err = table.add_columns().execute().await.unwrap_err();
assert!(
err.to_string().contains("requires a transform"),
"got: {err}"
);
}
#[tokio::test]
async fn test_read_columns_with_sql_expressions_is_rejected() {
let table = table_with_two_columns("read_cols_sql").await;
let err = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"doubled".into(),
"x * 2".into(),
)]))
.read_columns(["x"])
.execute()
.await
.unwrap_err();
assert!(err.to_string().contains("BatchUDF"), "got: {err}");
let schema = table.schema().await.unwrap();
assert!(
schema.field_with_name("doubled").is_err(),
"a rejected call must not commit"
);
}
#[tokio::test]
async fn test_read_columns_limits_what_a_batch_udf_sees() {
let table = table_with_two_columns("read_cols_udf").await;
let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)]));
let mapper_schema = output_schema.clone();
let udf = BatchUDF {
mapper: Box::new(move |batch: &RecordBatch| {
assert!(batch.column_by_name("x").is_some());
assert!(batch.column_by_name("y").is_none(), "y was not requested");
let x = batch["x"].as_any().downcast_ref::<Int32Array>().unwrap();
let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect();
Ok(RecordBatch::try_new(
mapper_schema.clone(),
vec![Arc::new(doubled)],
)?)
}),
output_schema,
result_checkpoint: None,
};
table
.add_columns()
.transform(NewColumnTransform::BatchUDF(udf))
.read_columns(["x"])
.execute()
.await
.unwrap();
let schema = table.schema().await.unwrap();
assert!(schema.field_with_name("sum").is_ok());
}
}
+5 -9
View File
@@ -576,12 +576,10 @@ mod tests {
// Add a new physical column AFTER the embedding column.
table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"score".into(),
"42.0".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
None,
)
.await
.unwrap();
@@ -685,9 +683,7 @@ mod tests {
true,
)]));
table
.add_columns()
.transform(NewColumnTransform::AllNulls(nested_schema))
.execute()
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
.await
.unwrap();
+19 -24
View File
@@ -193,12 +193,10 @@ mod tests {
// Add a computed column
let result = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"doubled".into(),
"id * 2".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
None,
)
.await
.unwrap();
@@ -253,12 +251,13 @@ mod tests {
// Add multiple columns at once
table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![
("y".into(), "x + 1".into()),
("z".into(), "x * x".into()),
]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![
("y".into(), "x + 1".into()),
("z".into(), "x * x".into()),
]),
None,
)
.await
.unwrap();
@@ -284,12 +283,10 @@ mod tests {
// Add a column with a constant value
table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"constant".into(),
"42".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
None,
)
.await
.unwrap();
@@ -662,12 +659,10 @@ mod tests {
// Add column increments version
let add_result = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"c".into(),
"a + b".into(),
)]))
.execute()
.add_columns(
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
None,
)
.await
.unwrap();
assert!(add_result.version > v1);
+2 -254
View File
@@ -9,17 +9,14 @@ use arrow_array::{
};
use arrow_schema::{DataType, Field, Fields, Schema};
use futures::TryStreamExt;
use lance::Dataset;
use lance_encoding::version::LanceFileVersion;
use lancedb::{
Connection, Error, Result, Table,
blob::{BlobRangeRequest, blob},
connect, connect_namespace,
database::listing::{
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
},
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
query::{ExecutableQuery, QueryBase},
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
table::{AddDataMode, CompactionOptions, OptimizeAction},
};
use tempfile::tempdir;
@@ -1078,252 +1075,3 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
}
Ok(())
}
/// Rows exercising the null/empty interleavings from
/// <https://github.com/lancedb/lancedb/issues/3744>: a payload, a null, a valid
/// empty value, then payloads whose descriptors a fragment rewrite used to zero.
fn null_empty_input_batch() -> RecordBatch {
let owned = [
Some(dedicated_blob_bytes(1)),
None,
Some(Vec::new()),
Some(dedicated_blob_bytes(4)),
Some(dedicated_blob_bytes(5)),
Some(dedicated_blob_bytes(6)),
];
let payloads: Vec<Option<&[u8]>> = owned.iter().map(|payload| payload.as_deref()).collect();
binary_input_batch(&[1, 2, 3, 4, 5, 6], &payloads)
}
/// One `(id, Some((payload length, first byte)))` per live row, or `(id, None)`
/// for a null blob. Comparing lengths and first bytes keeps failure output
/// readable where comparing whole payloads would not.
type BlobSummary = Vec<(i64, Option<(usize, Option<u8>)>)>;
/// The rows [`null_empty_input_batch`] leaves behind after `id IN (1, 4)` is
/// deleted: a null, a valid empty value, and the two payloads that follow them.
fn expected_null_empty_survivors() -> BlobSummary {
vec![
(2, None),
(3, Some((0, None))),
(5, Some((DEDICATED_BLOB_LEN, Some(5)))),
(6, Some((DEDICATED_BLOB_LEN, Some(6)))),
]
}
/// `optimize()` only rewrites a fragment when lance's compaction planner selects
/// it — here because the delete pushes the fragment past
/// `materialize_deletions_threshold` (0.1 by default; these tests delete 2 of 6
/// rows). Without this check, a planner or threshold change upstream would leave
/// both regression tests green while no rewrite happened at all.
fn assert_compacted(stats: &OptimizeStats) {
let metrics = stats
.compaction
.as_ref()
.expect("OptimizeAction::All runs compaction");
assert!(
metrics.fragments_removed >= 1,
"optimize() rewrote no fragment, so this test proves nothing: {metrics:?}"
);
}
fn summarize(rows: &[(i64, Option<Vec<u8>>)]) -> BlobSummary {
rows.iter()
.map(|(id, payload)| {
(
*id,
payload
.as_ref()
.map(|bytes| (bytes.len(), bytes.first().copied())),
)
})
.collect()
}
async fn sorted_id_rowid(table: &Table) -> Result<Vec<(i64, u64)>> {
let mut pairs = collect_id_rowid(table).await?;
pairs.sort_by_key(|(id, _)| *id);
Ok(pairs)
}
/// `{position, size}` descriptors of a legacy v1 blob column, keyed by `id`.
async fn v1_blob_descriptors(table: &Table) -> Result<Vec<(i64, Option<(u64, u64)>)>> {
let batches = table
.query()
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
let ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let descriptors = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<StructArray>()
.expect("v1 blob column reads back as a descriptor struct");
let position = descriptors
.column_by_name("position")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let size = descriptors
.column_by_name("size")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let mut rows: Vec<(i64, Option<(u64, u64)>)> = (0..batch.num_rows())
.map(|row| {
let descriptor =
(!descriptors.is_null(row)).then(|| (position.value(row), size.value(row)));
(ids.value(row), descriptor)
})
.collect();
rows.sort_by_key(|(id, _)| *id);
Ok(rows)
}
/// Payload bytes of every live row of a legacy v1 blob column, keyed by `id`.
/// [`Table::fetch_blobs`] rejects v1 columns, so read them through lance.
async fn v1_blob_payloads(dataset_uri: &str, table: &Table) -> Result<Vec<(i64, Option<Vec<u8>>)>> {
let pairs = sorted_id_rowid(table).await?;
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
let dataset = Arc::new(Dataset::open(dataset_uri).await?);
let files = dataset.take_blobs(&row_ids, "image").await?;
assert_eq!(
files.len(),
pairs.len(),
"take_blobs returned {} handles for {} live rows",
files.len(),
pairs.len()
);
let mut rows = Vec::with_capacity(pairs.len());
for ((id, _), file) in pairs.iter().zip(files) {
let payload = match file {
Some(file) => Some(file.read().await?.to_vec()),
None => None,
};
rows.push((*id, payload));
}
Ok(rows)
}
/// Length and first byte of every live blob v2 value, keyed by `id`.
async fn blob_v2_values(table: &Table) -> Result<BlobSummary> {
let pairs = sorted_id_rowid(table).await?;
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
let bytes = table.fetch_blobs("image", &row_ids).await?;
Ok(pairs
.iter()
.enumerate()
.map(|(slot, (id, _))| {
let value = (!bytes.is_null(slot))
.then(|| (bytes.value(slot).len(), bytes.value(slot).first().copied()));
(*id, value)
})
.collect())
}
/// Regression test for [#3744]: on storage 2.0 (legacy v1 descriptors),
/// compaction rewrote every payload following a null or empty value in the same
/// fragment as `{position: 0, size: 0}`, so the payload bytes read back as `b""`
/// and the new fragment no longer referenced them at all.
///
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
#[tokio::test]
async fn optimize_preserves_v1_blob_payloads_with_null_and_empty() -> Result<()> {
let tmp = tempdir().unwrap();
let db_uri = tmp.path().to_str().unwrap().to_string();
let db = connect(&db_uri)
.database_options(&ListingDatabaseOptions {
new_table_config: NewTableConfig {
data_storage_version: Some(LanceFileVersion::V2_0),
..Default::default()
},
..Default::default()
})
.execute()
.await?;
let legacy = Field::new("image", DataType::LargeBinary, true).with_metadata(
std::collections::HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]),
);
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
legacy,
]));
let table = db.create_empty_table("t", schema).execute().await?;
table.add(null_empty_input_batch()).execute().await?;
assert_eq!(
storage_format_version(&table).await,
LanceFileVersion::V2_0.resolve(),
"v1 blob descriptors only exist below storage 2.2"
);
let dataset_uri = table.uri().await?;
// Any rewrite triggers it; deleting rows is the shape from the issue.
table.delete("id IN (1, 4)").await?;
let descriptors_before = v1_blob_descriptors(&table).await?;
let before = v1_blob_payloads(&dataset_uri, &table).await?;
assert_eq!(
summarize(&before),
expected_null_empty_survivors(),
"test setup no longer produces the null/empty/payload mix"
);
let stats = table.optimize(OptimizeAction::All).await?;
assert_compacted(&stats);
let descriptors_after = v1_blob_descriptors(&table).await?;
let after = v1_blob_payloads(&dataset_uri, &table).await?;
assert_eq!(
summarize(&after),
summarize(&before),
"optimize() lost blob payloads; descriptors before={descriptors_before:?} after={descriptors_after:?}"
);
assert!(after == before, "optimize() changed blob payload bytes");
Ok(())
}
/// Regression test for the blob v2 half of [#3744]: compaction rewrote a valid
/// empty value as null, destroying the null-vs-empty distinction.
///
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
#[tokio::test]
async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let table = db
.create_empty_table("t", blob_table_schema())
.execute()
.await?;
table.add(null_empty_input_batch()).execute().await?;
assert!(
storage_format_version(&table).await >= LanceFileVersion::V2_2,
"blob v2 columns require storage >= 2.2"
);
table.delete("id IN (1, 4)").await?;
let before = blob_v2_values(&table).await?;
assert_eq!(
before,
expected_null_empty_survivors(),
"test setup no longer produces the null/empty/payload mix"
);
let stats = table.optimize(OptimizeAction::All).await?;
assert_compacted(&stats);
assert_eq!(
blob_v2_values(&table).await?,
before,
"optimize() changed blob v2 values"
);
Ok(())
}