diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index 6ebaebd75..7fc7ee668 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -16,18 +16,6 @@ A builder for LanceDB queries. - `StandardQueryBase`<`NativeQuery`> -## Properties - -### inner - -```ts -protected inner: Query | Promise; -``` - -#### Inherited from - -`StandardQueryBase.inner` - ## Methods ### analyzePlan() diff --git a/docs/src/js/classes/QueryBase.md b/docs/src/js/classes/QueryBase.md index 35c071525..31b154525 100644 --- a/docs/src/js/classes/QueryBase.md +++ b/docs/src/js/classes/QueryBase.md @@ -25,14 +25,6 @@ Common methods supported by all query types - `AsyncIterable`<`RecordBatch`> -## Properties - -### inner - -```ts -protected inner: NativeQueryType | Promise; -``` - ## Methods ### analyzePlan() diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 6ae2f9c8c..c00a68844 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -12,18 +12,6 @@ A query that returns a subset of the rows in the table. - [`QueryBase`](QueryBase.md)<`NativeTakeQuery`> -## Properties - -### inner - -```ts -protected inner: TakeQuery | Promise; -``` - -#### Inherited from - -[`QueryBase`](QueryBase.md).[`inner`](QueryBase.md#inner) - ## Methods ### analyzePlan() diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index f9412c76d..1e81c1716 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -18,18 +18,6 @@ This builder can be reused to execute the query many times. - `StandardQueryBase`<`NativeVectorQuery`> -## Properties - -### inner - -```ts -protected inner: VectorQuery | Promise; -``` - -#### Inherited from - -`StandardQueryBase.inner` - ## Methods ### addQueryVector() diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index d6064a02f..78e8e8016 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -47,6 +47,7 @@ import { BooleanQuery, Occur, Operator, + VectorQuery, instanceOfFullTextQuery, } from "../lancedb/query"; @@ -2337,9 +2338,23 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( expect(results2[0].text).toBe(data[1].text); }); - test("auto search follows embedding metadata across executions", async () => { + test("auto search stays consistent with the active revision", async () => { + let initCalls = 0; + let queryCalls = 0; + let markStarted!: () => void; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + let releaseEmbedding!: () => void; + const embeddingReleased = new Promise((resolve) => { + releaseEmbedding = resolve; + }); + @register("refresh-test") class TestEmbedding extends EmbeddingFunction { + async init() { + initCalls += 1; + } ndims() { return 1; } @@ -2347,6 +2362,11 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( 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[]) { @@ -2362,7 +2382,10 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( readConsistencyInterval: 0, }); const tracked = await reader.openTable("test"); - const autoQuery = tracked.search("greetings").select(["text"]).limit(1); + const autoQuery = (tracked.search("greetings") as VectorQuery) + .nprobes(1) + .select(["text"]) + .limit(1); const func = new TestEmbedding(); const schema = LanceSchema({ @@ -2371,18 +2394,52 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( }); 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); - const ftsData = [{ text: "greetings from full text", vector: [0.0] }]; + const repeatedResults = await autoQuery.toArray(); + expect(repeatedResults[0].text).toBe(data[0].text); + expect(initCalls).toBe(baselineInitCalls + 1); + expect(queryCalls).toBe(1); + + const multiVectorResults = await ( + tracked.search("greetings") as VectorQuery + ) + .addQueryVector(Promise.resolve([0.2])) + .select(["text"]) + .limit(1) + .toArray(); + expect(multiVectorResults).toHaveLength(2); + expect(multiVectorResults.map((row) => row.text).sort()).toEqual( + data.map((row) => row.text).sort(), + ); + + 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(ftsData[1].text); expect( (await tracked.schema()).metadata.get("embedding_functions"), diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 3003fc2e1..63095ab1e 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -103,7 +103,7 @@ export interface FullTextSearchOptions { type NativeQueryLike = NativeQuery | NativeVectorQuery | NativeTakeQuery; class DeferredNativeQuery { - private readonly calls: Array<(inner: NativeQueryType) => void> = []; + protected readonly calls: Array<(inner: NativeQueryType) => void> = []; constructor(private readonly factory: () => Promise) {} @@ -131,6 +131,18 @@ function nearestToNative( return inner.nearestTo(Float32Array.from(vector as number[])); } +function addQueryVectorToNative( + inner: NativeVectorQuery, + vector: Awaited, +) { + 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} @@ -142,13 +154,24 @@ export class QueryBase< NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery, > implements AsyncIterable { + /** + * @hidden + */ + protected inner: + | NativeQueryType + | Promise + | DeferredNativeQuery; + /** * @hidden */ protected constructor( - protected inner: NativeQueryType | Promise, + inner: + | NativeQueryType + | Promise + | DeferredNativeQuery, ) { - // intentionally empty + this.inner = inner; } // call a function on the inner (either a promise or the actual object) @@ -156,10 +179,8 @@ export class QueryBase< * @hidden */ protected doCall(fn: (inner: NativeQueryType) => void) { - if ((this.inner as unknown) instanceof DeferredNativeQuery) { - const deferred = this - .inner as unknown as DeferredNativeQuery; - deferred.doCall(fn); + if (this.inner instanceof DeferredNativeQuery) { + this.inner.doCall(fn); } else if (this.inner instanceof Promise) { this.inner = this.inner.then((inner) => { fn(inner); @@ -174,10 +195,8 @@ export class QueryBase< * @hidden */ protected resolveInner(): NativeQueryType | Promise { - if ((this.inner as unknown) instanceof DeferredNativeQuery) { - const deferred = this - .inner as unknown as DeferredNativeQuery; - return deferred.resolve(); + if (this.inner instanceof DeferredNativeQuery) { + return this.inner.resolve(); } return this.inner; } @@ -400,7 +419,12 @@ export class StandardQueryBase< extends QueryBase implements ExecutableQuery { - constructor(inner: NativeQueryType | Promise) { + constructor( + inner: + | NativeQueryType + | Promise + | DeferredNativeQuery, + ) { super(inner); } @@ -550,10 +574,22 @@ export class VectorQuery extends StandardQueryBase { /** * @hidden */ - constructor(inner: NativeVectorQuery | Promise) { + constructor( + inner: + | NativeVectorQuery + | Promise + | DeferredNativeQuery, + ) { super(inner); } + /** + * @hidden + */ + protected doVectorCall(fn: (inner: NativeVectorQuery) => void) { + super.doCall(fn); + } + /** * Set the number of partitions to search (probe) * @@ -581,7 +617,7 @@ export class VectorQuery extends StandardQueryBase { * 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; } @@ -595,7 +631,7 @@ export class VectorQuery extends StandardQueryBase { * but will also increase latency. */ minimumNprobes(minimumNprobes: number): VectorQuery { - super.doCall((inner) => inner.minimumNprobes(minimumNprobes)); + this.doVectorCall((inner) => inner.minimumNprobes(minimumNprobes)); return this; } @@ -609,7 +645,7 @@ export class VectorQuery extends StandardQueryBase { * potential false negatives. */ maximumNprobes(maximumNprobes: number): VectorQuery { - super.doCall((inner) => inner.maximumNprobes(maximumNprobes)); + this.doVectorCall((inner) => inner.maximumNprobes(maximumNprobes)); return this; } @@ -622,7 +658,7 @@ export class VectorQuery extends StandardQueryBase { * `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; } @@ -636,7 +672,7 @@ export class VectorQuery extends StandardQueryBase { * 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; } @@ -650,7 +686,7 @@ export class VectorQuery extends StandardQueryBase { * 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; } @@ -671,7 +707,7 @@ export class VectorQuery extends StandardQueryBase { distanceType( distanceType: Required["distanceType"], ): VectorQuery { - super.doCall((inner) => inner.distanceType(distanceType)); + this.doVectorCall((inner) => inner.distanceType(distanceType)); return this; } @@ -705,7 +741,7 @@ export class VectorQuery extends StandardQueryBase { * 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; } @@ -730,7 +766,7 @@ export class VectorQuery extends StandardQueryBase { * 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; } @@ -744,7 +780,7 @@ export class VectorQuery extends StandardQueryBase { * calculate your recall to select an appropriate value for nprobes. */ bypassVectorIndex(): VectorQuery { - super.doCall((inner) => inner.bypassVectorIndex()); + this.doVectorCall((inner) => inner.bypassVectorIndex()); return this; } @@ -761,34 +797,19 @@ export class VectorQuery extends StandardQueryBase { addQueryVector(vector: IntoVector): VectorQuery { 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.addQueryVector(v); - const inner = value.inner as - | NativeVectorQuery - | Promise; - return inner; - } catch (e) { - return Promise.reject(e); - } + const inner = (await this.resolveInner()) as NativeVectorQuery; + addQueryVectorToNative(inner, await vector); + 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); @@ -807,11 +828,54 @@ export class VectorQuery extends StandardQueryBase { } } +type AutoQueryResolution = { + inner: NativeQuery | NativeVectorQuery; + route: "fts" | "vector"; +}; + +class DeferredAutoNativeQuery extends DeferredNativeQuery { + private readonly vectorCalls: Array< + (inner: NativeVectorQuery) => void | Promise + > = []; + + constructor( + private readonly autoFactory: () => Promise, + ) { + super(async () => (await autoFactory()).inner as NativeVectorQuery); + } + + doVectorCall(fn: (inner: NativeVectorQuery) => void | Promise) { + this.vectorCalls.push(fn); + } + + async resolve(): Promise { + const resolution = await this.autoFactory(); + for (const call of this.calls) { + call(resolution.inner as NativeVectorQuery); + } + if (resolution.route === "vector") { + for (const call of this.vectorCalls) { + await call(resolution.inner as NativeVectorQuery); + } + } + return resolution.inner as NativeVectorQuery; + } +} + class DeferredAutoQuery extends VectorQuery { - constructor(factory: () => Promise) { - super( - new DeferredNativeQuery(factory) as unknown as Promise, - ); + constructor(private readonly deferred: DeferredAutoNativeQuery) { + super(deferred); + } + + protected doVectorCall(fn: (inner: NativeVectorQuery) => void) { + this.deferred.doVectorCall(fn); + } + + addQueryVector(vector: IntoVector): VectorQuery { + this.deferred.doVectorCall(async (inner) => { + addQueryVectorToNative(inner, await vector); + }); + return this; } } @@ -825,17 +889,74 @@ export function createAutoQuery( table: NativeTable, query: string, columns: string[] | null, - getVector: () => Promise | undefined>, + getVector: (metadata: string) => Promise>, ): VectorQuery { - return new DeferredAutoQuery(async () => { - const vector = await getVector(); - const inner = table.query(); - if (vector === undefined) { - inner.fullTextSearch({ query, columns }); - return inner; + type RouteSnapshot = { + table: NativeTable; + embeddingMetadata: string | undefined; + }; + type CachedPreparation = { + metadata: string; + vector: Promise>; + }; + + let cachedPreparation: CachedPreparation | undefined; + + const snapshotRoute = async (): Promise => { + const snapshot = await table.querySnapshot(); + const schema = tableFromIPC(await snapshot.schema()).schema; + return { + table: snapshot, + embeddingMetadata: schema.metadata.get("embedding_functions"), + }; + }; + + const deferred = new DeferredAutoNativeQuery(async () => { + while (true) { + const initial = await snapshotRoute(); + if (initial.embeddingMetadata === undefined) { + cachedPreparation = undefined; + const inner = initial.table.query(); + inner.fullTextSearch({ query, columns }); + return { inner, route: "fts" }; + } + + const metadata = initial.embeddingMetadata; + if (cachedPreparation?.metadata !== metadata) { + cachedPreparation = { + metadata, + vector: getVector(metadata), + }; + } + + const preparation = cachedPreparation; + let vector: Awaited; + try { + vector = await preparation.vector; + } catch (error) { + if (cachedPreparation === preparation) { + cachedPreparation = undefined; + } + throw error; + } + + // Provider preparation can perform arbitrary asynchronous work. Take a + // fresh pinned snapshot afterwards and only use the prepared vector if + // that snapshot has the same embedding configuration. + const current = await snapshotRoute(); + if (current.embeddingMetadata !== metadata) { + cachedPreparation = undefined; + continue; + } + + return { + inner: nearestToNative(current.table.query(), vector), + route: "vector", + }; } - return nearestToNative(inner, vector); }); + + return new DeferredAutoQuery(deferred); } /** diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 80671ffe2..c148c63d4 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -1044,15 +1044,18 @@ export class LocalTable extends Table { const columns = typeof ftsColumns === "string" ? [ftsColumns] : (ftsColumns ?? null); - return createAutoQuery(this.inner, query, columns, async () => { - const functions = await this.getEmbeddingFunctions(); + 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) { - 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); }); } diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 2ac2fecb2..5d871eebf 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -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 { + let snapshot = self.inner_ref()?.query_snapshot().await.default_error()?; + Ok(Table::new(snapshot)) + } + #[napi(catch_unwind)] pub fn take_offsets(&self, offsets: Vec) -> napi::Result { Ok(TakeQuery::new( diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 5fefabeb2..b1b41cae7 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1678,6 +1678,14 @@ impl BaseTable for RemoteTable { fn id(&self) -> &str { &self.identifier } + async fn query_snapshot(&self) -> Result> { + let description = self.describe().await?; + let schema: arrow_schema::Schema = description.schema.try_into()?; + let snapshot = self.with_branch(self.branch.clone()); + *snapshot.version.write().await = Some(description.version); + snapshot.schema_cache.seed(Arc::new(schema)); + Ok(Arc::new(snapshot)) + } async fn version(&self) -> Result { self.describe().await.map(|desc| desc.version) } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 0d8a8e8b9..a084374e4 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -581,6 +581,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; + /// 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>; /// Count the number of rows in this table. async fn count_rows(&self, filter: Option) -> Result; /// Create a physical plan for the query. @@ -1068,6 +1075,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 { + 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 @@ -2835,6 +2852,15 @@ impl BaseTable for NativeTable { &self.id } + async fn query_snapshot(&self) -> Result> { + let dataset = self.dataset.get().await?; + let snapshot = dataset::DatasetConsistencyWrapper::new_time_travel( + (*dataset).clone(), + self.read_consistency_interval, + ); + Ok(Arc::new(self.with_dataset(snapshot))) + } + async fn version(&self) -> Result { Ok(self.dataset.get().await?.version().version) }