fix(node): stabilize deferred auto search

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<void>((resolve) => {
+        markStarted = resolve;
+      });
+      let releaseEmbedding!: () => void;
+      const embeddingReleased = new Promise<void>((resolve) => {
+        releaseEmbedding = resolve;
+      });
+
       @register("refresh-test")
       class TestEmbedding extends EmbeddingFunction<string> {
+        async init() {
+          initCalls += 1;
+        }
         ndims() {
           return 1;
         }
@@ -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"),

nodejs/lancedb/query.ts

@@ -103,7 +103,7 @@ export interface FullTextSearchOptions {
 type NativeQueryLike = NativeQuery | NativeVectorQuery | NativeTakeQuery;
 
 class DeferredNativeQuery<NativeQueryType extends NativeQueryLike> {
-  private readonly calls: Array<(inner: NativeQueryType) => void> = [];
+  protected readonly calls: Array<(inner: NativeQueryType) => void> = [];
 
   constructor(private readonly factory: () => Promise<NativeQueryType>) {}
 
@@ -131,6 +131,18 @@ function nearestToNative(
   return inner.nearestTo(Float32Array.from(vector as number[]));
 }
 
+function addQueryVectorToNative(
+  inner: NativeVectorQuery,
+  vector: Awaited<IntoVector>,
+) {
+  const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
+  if (raw) {
+    inner.addQueryVectorRaw(raw.data, raw.dtype);
+  } else {
+    inner.addQueryVector(Float32Array.from(vector as number[]));
+  }
+}
+
 /** Common methods supported by all query types
  *
  * @see {@link Query}
@@ -142,13 +154,24 @@ export class QueryBase<
   NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery,
 > implements AsyncIterable<RecordBatch>
 {
+  /**
+   * @hidden
+   */
+  protected inner:
+    | NativeQueryType
+    | Promise<NativeQueryType>
+    | DeferredNativeQuery<NativeQueryType>;
+
   /**
    * @hidden
    */
   protected constructor(
-    protected inner: NativeQueryType | Promise<NativeQueryType>,
+    inner:
+      | NativeQueryType
+      | Promise<NativeQueryType>
+      | DeferredNativeQuery<NativeQueryType>,
   ) {
-    // 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<NativeQueryType>;
-      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<NativeQueryType> {
-    if ((this.inner as unknown) instanceof DeferredNativeQuery) {
-      const deferred = this
-        .inner as unknown as DeferredNativeQuery<NativeQueryType>;
-      return deferred.resolve();
+    if (this.inner instanceof DeferredNativeQuery) {
+      return this.inner.resolve();
     }
     return this.inner;
   }
@@ -400,7 +419,12 @@ export class StandardQueryBase<
   extends QueryBase<NativeQueryType>
   implements ExecutableQuery
 {
-  constructor(inner: NativeQueryType | Promise<NativeQueryType>) {
+  constructor(
+    inner:
+      | NativeQueryType
+      | Promise<NativeQueryType>
+      | DeferredNativeQuery<NativeQueryType>,
+  ) {
     super(inner);
   }
 
@@ -550,10 +574,22 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
   /**
    * @hidden
    */
-  constructor(inner: NativeVectorQuery | Promise<NativeVectorQuery>) {
+  constructor(
+    inner:
+      | NativeVectorQuery
+      | Promise<NativeVectorQuery>
+      | DeferredNativeQuery<NativeVectorQuery>,
+  ) {
     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<NativeVectorQuery> {
    * the minimum and maximum to the same value.
    */
   nprobes(nprobes: number): VectorQuery {
-    super.doCall((inner) => inner.nprobes(nprobes));
+    this.doVectorCall((inner) => inner.nprobes(nprobes));
 
     return this;
   }
@@ -595,7 +631,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * but will also increase latency.
    */
   minimumNprobes(minimumNprobes: number): VectorQuery {
-    super.doCall((inner) => inner.minimumNprobes(minimumNprobes));
+    this.doVectorCall((inner) => inner.minimumNprobes(minimumNprobes));
     return this;
   }
 
@@ -609,7 +645,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * potential false negatives.
    */
   maximumNprobes(maximumNprobes: number): VectorQuery {
-    super.doCall((inner) => inner.maximumNprobes(maximumNprobes));
+    this.doVectorCall((inner) => inner.maximumNprobes(maximumNprobes));
     return this;
   }
 
@@ -622,7 +658,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * `undefined` means no lower or upper bound.
    */
   distanceRange(lowerBound?: number, upperBound?: number): VectorQuery {
-    super.doCall((inner) => inner.distanceRange(lowerBound, upperBound));
+    this.doVectorCall((inner) => inner.distanceRange(lowerBound, upperBound));
     return this;
   }
 
@@ -636,7 +672,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * also increase the latency of your query. The default value is 1.5*limit.
    */
   ef(ef: number): VectorQuery {
-    super.doCall((inner) => inner.ef(ef));
+    this.doVectorCall((inner) => inner.ef(ef));
     return this;
   }
 
@@ -650,7 +686,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * whose data type is a fixed-size-list of floats.
    */
   column(column: string): VectorQuery {
-    super.doCall((inner) => inner.column(column));
+    this.doVectorCall((inner) => inner.column(column));
     return this;
   }
 
@@ -671,7 +707,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
   distanceType(
     distanceType: Required<IvfPqOptions>["distanceType"],
   ): VectorQuery {
-    super.doCall((inner) => inner.distanceType(distanceType));
+    this.doVectorCall((inner) => inner.distanceType(distanceType));
     return this;
   }
 
@@ -705,7 +741,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * distance between the query vector and the actual uncompressed vector.
    */
   refineFactor(refineFactor: number): VectorQuery {
-    super.doCall((inner) => inner.refineFactor(refineFactor));
+    this.doVectorCall((inner) => inner.refineFactor(refineFactor));
     return this;
   }
 
@@ -730,7 +766,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * factor can often help restore some of the results lost by post filtering.
    */
   postfilter(): VectorQuery {
-    super.doCall((inner) => inner.postfilter());
+    this.doVectorCall((inner) => inner.postfilter());
     return this;
   }
 
@@ -744,7 +780,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * calculate your recall to select an appropriate value for nprobes.
    */
   bypassVectorIndex(): VectorQuery {
-    super.doCall((inner) => inner.bypassVectorIndex());
+    this.doVectorCall((inner) => inner.bypassVectorIndex());
     return this;
   }
 
@@ -761,34 +797,19 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
   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<NativeVectorQuery>;
-          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<NativeVectorQuery> {
   }
 }
 
+type AutoQueryResolution = {
+  inner: NativeQuery | NativeVectorQuery;
+  route: "fts" | "vector";
+};
+
+class DeferredAutoNativeQuery extends DeferredNativeQuery<NativeVectorQuery> {
+  private readonly vectorCalls: Array<
+    (inner: NativeVectorQuery) => void | Promise<void>
+  > = [];
+
+  constructor(
+    private readonly autoFactory: () => Promise<AutoQueryResolution>,
+  ) {
+    super(async () => (await autoFactory()).inner as NativeVectorQuery);
+  }
+
+  doVectorCall(fn: (inner: NativeVectorQuery) => void | Promise<void>) {
+    this.vectorCalls.push(fn);
+  }
+
+  async resolve(): Promise<NativeVectorQuery> {
+    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<NativeQuery | NativeVectorQuery>) {
-    super(
-      new DeferredNativeQuery(factory) as unknown as Promise<NativeVectorQuery>,
-    );
+  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<Awaited<IntoVector> | undefined>,
+  getVector: (metadata: string) => Promise<Awaited<IntoVector>>,
 ): 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<Awaited<IntoVector>>;
+  };
+
+  let cachedPreparation: CachedPreparation | undefined;
+
+  const snapshotRoute = async (): Promise<RouteSnapshot> => {
+    const snapshot = await table.querySnapshot();
+    const schema = tableFromIPC(await snapshot.schema()).schema;
+    return {
+      table: snapshot,
+      embeddingMetadata: schema.metadata.get("embedding_functions"),
+    };
+  };
+
+  const 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<IntoVector>;
+      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);
 }
 
 /**

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);
       });
     }

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<Table> {
+        let snapshot = self.inner_ref()?.query_snapshot().await.default_error()?;
+        Ok(Table::new(snapshot))
+    }
+
     #[napi(catch_unwind)]
     pub fn take_offsets(&self, offsets: Vec<i64>) -> napi::Result<TakeQuery> {
         Ok(TakeQuery::new(