From 3efc9187da78260923736bc3f64c1a549534cccf Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:33:42 +0000 Subject: [PATCH] fix(node): defer auto search routing --- nodejs/__test__/table.test.ts | 54 +++++++++++++ nodejs/lancedb/connection.ts | 8 +- nodejs/lancedb/permutation.ts | 2 +- nodejs/lancedb/query.ts | 139 ++++++++++++++++++++++++---------- nodejs/lancedb/table.ts | 56 ++++++-------- 5 files changed, 183 insertions(+), 76 deletions(-) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 1727f301e..d6064a02f 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2337,6 +2337,60 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( expect(results2[0].text).toBe(data[1].text); }); + test("auto search follows embedding metadata across executions", async () => { + @register("refresh-test") + class TestEmbedding extends EmbeddingFunction { + ndims() { + return 1; + } + embeddingDataType() { + return new arrow.Float32(); + } + async computeQueryEmbeddings(value: string) { + return value === "greetings" ? [0.1] : [0.2]; + } + async computeSourceEmbeddings(values: string[]) { + return values.map((value) => + value === "hello world" ? [0.1] : [0.2], + ); + } + } + + const writer = await connect(tmpDir.name); + await writer.createTable("test", [{ text: "plain", vector: [0.0] }]); + const reader = await connect(tmpDir.name, { + readConsistencyInterval: 0, + }); + const tracked = await reader.openTable("test"); + const autoQuery = tracked.search("greetings").select(["text"]).limit(1); + + const func = new TestEmbedding(); + const schema = LanceSchema({ + text: func.sourceField(new arrow.Utf8()), + vector: func.vectorField(), + }); + const data = [{ text: "hello world" }, { text: "goodbye world" }]; + await writer.createTable("test", data, { mode: "overwrite", schema }); + + expect( + (await tracked.schema()).metadata.get("embedding_functions"), + ).toBeDefined(); + const results = await autoQuery.toArray(); + expect(results[0].text).toBe(data[0].text); + + const ftsData = [{ text: "greetings from full text", vector: [0.0] }]; + const ftsTable = await writer.createTable("test", ftsData, { + mode: "overwrite", + }); + await ftsTable.createIndex("text", { config: Index.fts() }); + + expect( + (await tracked.schema()).metadata.get("embedding_functions"), + ).toBeUndefined(); + const ftsResults = await autoQuery.toArray(); + expect(ftsResults[0].text).toBe(ftsData[0].text); + }); + test("tokenizes FTS queries by column or index name", async () => { const db = await connect(tmpDir.name); const data = [ diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 52ff811f9..e63a7ae65 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -535,7 +535,7 @@ export class LocalConnection extends Connection { options?.indexCacheSize, ); - let table: Table = await LocalTable.create(innerTable); + let table: Table = new LocalTable(innerTable); // "main" is the default branch, so treat it as no branch. On a real branch, // scope and pin in one step (yielding "version V of branch B"); otherwise // pin the version, if any, against main. @@ -570,7 +570,7 @@ export class LocalConnection extends Connection { options?.isShallow ?? true, ); - return await LocalTable.create(innerTable); + return new LocalTable(innerTable); } private getStorageOptions( @@ -652,7 +652,7 @@ export class LocalConnection extends Connection { storageOptions, ); - return await LocalTable.create(innerTable); + return new LocalTable(innerTable); } async createEmptyTable( @@ -698,7 +698,7 @@ export class LocalConnection extends Connection { namespacePath ?? [], storageOptions, ); - return await LocalTable.create(innerTable); + return new LocalTable(innerTable); } async dropTable(name: string, namespacePath?: string[]): Promise { diff --git a/nodejs/lancedb/permutation.ts b/nodejs/lancedb/permutation.ts index 714caae94..e55b68553 100644 --- a/nodejs/lancedb/permutation.ts +++ b/nodejs/lancedb/permutation.ts @@ -172,7 +172,7 @@ export class PermutationBuilder { */ async execute(): Promise { const nativeTable: NativeTable = await this.inner.execute(); - return await LocalTable.create(nativeTable); + return new LocalTable(nativeTable); } } diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 843a1276f..3003fc2e1 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -100,6 +100,37 @@ export interface FullTextSearchOptions { columns?: string | string[]; } +type NativeQueryLike = NativeQuery | NativeVectorQuery | NativeTakeQuery; + +class DeferredNativeQuery { + private readonly calls: Array<(inner: NativeQueryType) => void> = []; + + constructor(private readonly factory: () => Promise) {} + + doCall(fn: (inner: NativeQueryType) => void) { + this.calls.push(fn); + } + + async resolve(): Promise { + const inner = await this.factory(); + for (const call of this.calls) { + call(inner); + } + return inner; + } +} + +function nearestToNative( + inner: NativeQuery, + vector: Awaited, +): NativeVectorQuery { + const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector); + if (raw) { + return inner.nearestToRaw(raw.data, raw.dtype); + } + return inner.nearestTo(Float32Array.from(vector as number[])); +} + /** Common methods supported by all query types * * @see {@link Query} @@ -125,7 +156,11 @@ export class QueryBase< * @hidden */ protected doCall(fn: (inner: NativeQueryType) => void) { - if (this.inner instanceof Promise) { + if ((this.inner as unknown) instanceof DeferredNativeQuery) { + const deferred = this + .inner as unknown as DeferredNativeQuery; + deferred.doCall(fn); + } else if (this.inner instanceof Promise) { this.inner = this.inner.then((inner) => { fn(inner); return inner; @@ -135,6 +170,18 @@ 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(); + } + return this.inner; + } + /** * Return only the specified columns. * @@ -210,12 +257,13 @@ export class QueryBase< protected nativeExecute( options?: Partial, ): Promise { - if (this.inner instanceof Promise) { - return this.inner.then((inner) => + const inner = this.resolveInner(); + if (inner instanceof Promise) { + return inner.then((inner) => inner.execute(options?.maxBatchLength, options?.timeoutMs), ); } else { - return this.inner.execute(options?.maxBatchLength, options?.timeoutMs); + return inner.execute(options?.maxBatchLength, options?.timeoutMs); } } @@ -245,12 +293,7 @@ export class QueryBase< /** Collect the results as an Arrow @see {@link ArrowTable}. */ async toArrow(options?: Partial): Promise { const batches = []; - let inner; - if (this.inner instanceof Promise) { - inner = await this.inner; - } else { - inner = this.inner; - } + const inner = await this.resolveInner(); for await (const batch of new RecordBatchIterable(inner, options)) { batches.push(batch); } @@ -279,10 +322,11 @@ export class QueryBase< * @returns A Promise that resolves to a string containing the query execution plan explanation. */ async explainPlan(verbose = false): Promise { - if (this.inner instanceof Promise) { - return this.inner.then((inner) => inner.explainPlan(verbose)); + const inner = this.resolveInner(); + if (inner instanceof Promise) { + return inner.then((inner) => inner.explainPlan(verbose)); } else { - return this.inner.explainPlan(verbose); + return inner.explainPlan(verbose); } } @@ -321,12 +365,11 @@ export class QueryBase< distributedMetrics?: AnalyzePlanDistributedMetrics, ): Promise { const distributedMetricsMode = distributedMetrics ?? "aggregate"; - if (this.inner instanceof Promise) { - return this.inner.then((inner) => - inner.analyzePlan(distributedMetricsMode), - ); + const inner = this.resolveInner(); + if (inner instanceof Promise) { + return inner.then((inner) => inner.analyzePlan(distributedMetricsMode)); } else { - return this.inner.analyzePlan(distributedMetricsMode); + return inner.analyzePlan(distributedMetricsMode); } } @@ -340,10 +383,11 @@ export class QueryBase< */ async outputSchema(): Promise { let schemaBuffer: Buffer; - if (this.inner instanceof Promise) { - schemaBuffer = await this.inner.then((inner) => inner.outputSchema()); + const inner = this.resolveInner(); + if (inner instanceof Promise) { + schemaBuffer = await inner.then((inner) => inner.outputSchema()); } else { - schemaBuffer = await this.inner.outputSchema(); + schemaBuffer = await inner.outputSchema(); } const schema = tableFromIPC(schemaBuffer).schema; return schema; @@ -763,6 +807,37 @@ export class VectorQuery extends StandardQueryBase { } } +class DeferredAutoQuery extends VectorQuery { + constructor(factory: () => Promise) { + super( + new DeferredNativeQuery(factory) as unknown as Promise, + ); + } +} + +/** + * Create a string query whose vector/FTS routing is resolved against the active + * table schema when the query executes. + * + * @hidden + */ +export function createAutoQuery( + table: NativeTable, + query: string, + columns: string[] | null, + getVector: () => Promise | undefined>, +): VectorQuery { + return new DeferredAutoQuery(async () => { + const vector = await getVector(); + const inner = table.query(); + if (vector === undefined) { + inner.fullTextSearch({ query, columns }); + return inner; + } + return nearestToNative(inner, vector); + }); +} + /** * A query that returns a subset of the rows in the table. * @@ -840,23 +915,11 @@ export class Query extends StandardQueryBase { * a default `limit` of 10 will be used. @see {@link Query#limit} */ nearestTo(vector: IntoVector): VectorQuery { - const callNearestTo = ( - inner: NativeQuery, - resolved: Float32Array | Float64Array | Uint8Array | number[], - ): NativeVectorQuery => { - const raw = Array.isArray(resolved) - ? null - : extractVectorBuffer(resolved); - if (raw) { - return inner.nearestToRaw(raw.data, raw.dtype); - } - return inner.nearestTo(Float32Array.from(resolved as number[])); - }; - - if (this.inner instanceof Promise) { - const nativeQuery = this.inner.then(async (inner) => { + const inner = this.resolveInner(); + if (inner instanceof Promise) { + const nativeQuery = inner.then(async (inner) => { const resolved = vector instanceof Promise ? await vector : vector; - return callNearestTo(inner, resolved); + return nearestToNative(inner, resolved); }); return new VectorQuery(nativeQuery); } @@ -876,7 +939,7 @@ export class Query extends StandardQueryBase { })(); return new VectorQuery(res); } else { - const vectorQuery = callNearestTo(this.inner, vector); + const vectorQuery = nearestToNative(inner, vector); return new VectorQuery(vectorQuery); } } diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index cde5344b0..80671ffe2 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -44,6 +44,7 @@ import { Query, TakeQuery, VectorQuery, + createAutoQuery, instanceOfFullTextQuery, } from "./query"; import { sanitizeType } from "./sanitize"; @@ -814,30 +815,10 @@ export abstract class Table { export class LocalTable extends Table { private readonly inner: _NativeTable; - private readonly hasEmbeddingFunctions: boolean; - private constructor(inner: _NativeTable, hasEmbeddingFunctions: boolean) { + constructor(inner: _NativeTable) { super(); this.inner = inner; - this.hasEmbeddingFunctions = hasEmbeddingFunctions; - } - - static async create(inner: _NativeTable): Promise { - const schemaBuf = await inner.schema(); - const schema = tableFromIPC(schemaBuf).schema; - const serializedFunctions = schema.metadata.get("embedding_functions"); - let hasEmbeddingFunctions = false; - if (serializedFunctions !== undefined) { - try { - const functions: unknown = JSON.parse(serializedFunctions); - hasEmbeddingFunctions = - !Array.isArray(functions) || functions.length > 0; - } catch { - // Let parseFunctions report malformed metadata when the query executes. - hasEmbeddingFunctions = true; - } - } - return new LocalTable(inner, hasEmbeddingFunctions); } get name(): string { return this.inner.name; @@ -1054,14 +1035,25 @@ export class LocalTable extends Table { }); } - // The query type is auto or vector - // fall back to full text search if no embedding functions are defined and the query is a string - if ( - queryType === "auto" && - (!this.hasEmbeddingFunctions || instanceOfFullTextQuery(query)) - ) { - return this.query().fullTextSearch(query, { - columns: ftsColumns, + if (queryType === "auto") { + if (instanceOfFullTextQuery(query)) { + return this.query().fullTextSearch(query, { + columns: ftsColumns, + }); + } + + const columns = + typeof ftsColumns === "string" ? [ftsColumns] : (ftsColumns ?? null); + return createAutoQuery(this.inner, query, columns, async () => { + const functions = await this.getEmbeddingFunctions(); + // TODO: Support multiple embedding functions + const embeddingFunc: EmbeddingFunctionConfig | undefined = functions + .values() + .next().value; + if (!embeddingFunc) { + return undefined; + } + return await embeddingFunc.function.computeQueryEmbeddings(query); }); } @@ -1481,9 +1473,7 @@ export class Branches { fromRef?: string, fromVersion?: number, ): Promise
{ - return await LocalTable.create( - await this.#inner.create(name, fromRef, fromVersion), - ); + return new LocalTable(await this.#inner.create(name, fromRef, fromVersion)); } /** @@ -1494,7 +1484,7 @@ export class Branches { * latest and stays writable. */ async checkout(name: string, version?: number): Promise
{ - return await LocalTable.create(await this.#inner.checkout(name, version)); + return new LocalTable(await this.#inner.checkout(name, version)); } /** Delete a branch. */