From 35b5d015ac8a38313d8322c26f4b2f86ef6f4e23 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 04:17:32 +0800 Subject: [PATCH] fix(node): preserve embedding registration in server bundles (#3806) ## Summary - lazily initialize built-in OpenAI and Hugging Face providers when consumers call the public embedding registry API - choose automatic vector versus FTS search from embedding metadata on a fresh pinned table revision for every execution - expose automatic string searches as an `AutoQuery` with only operations common to both native query families - keep the registry shared and built-in registration safe across duplicated module graphs ## Root cause Nitro treats dependency modules as side-effect-free and removes the bare OpenAI provider import from its generated route. Registration therefore never runs, so `getRegistry().get("openai")` remains undefined even when the registry itself is shared globally. Bundlers may also duplicate the provider and registry module graphs. The public embedding entry point now initializes built-in providers only when `getRegistry()` is explicitly called, keeping initialization on a live path that Nitro retains. Each terminal automatic-search execution pins the exact table revision visible at dispatch, reads embedding metadata and computes an embedding from that snapshot, replays the builder operations, and constructs and executes the selected native query against the same snapshot. Pinned native snapshots execute locally when namespace pushdown cannot carry their revision, while remote snapshots are seeded directly from one version-and-schema response. The public `AutoQuery` builder exposes only the operations shared by FTS and vector search, so runtime class narrowing cannot expose invalid vector-only methods. Repeated built-in registration replaces stale constructors from duplicated module graphs while public `register()` retains its duplicate-alias error. ## Validation - `cargo fmt --all` - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples` - `pnpm build` - `pnpm lint` - `pnpm run docs` - `pnpm test --runInBand` (783 passed, 5 skipped) - serial examples suite with a local OpenAI mock (11 passed), including `sentence-transformers.test.ts` - packaged Nitro 2.13.4 server route using the reported imports returned `{"registered":true}` - fresh-process FTS fixture initialized both public built-ins and confirmed automatic string search still returned the indexed row - schema-consistency regressions cover read-consistency refresh, checkout, checkoutLatest, restore, runtime class narrowing, concurrent overwrite during embedding computation, and reused automatic-search builders - focused regressions confirm pinned native snapshots bypass unversioned namespace pushdown and remote snapshots use one describe request Fixes #2429 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- docs/src/js/classes/AutoQuery.md | 518 ++++++++++++++++++ docs/src/js/classes/Table.md | 4 +- docs/src/js/globals.md | 1 + .../embedding/functions/getRegistry.md | 14 +- nodejs/__test__/embedding_registry.test.ts | 95 ++++ nodejs/__test__/fixtures/auto_fts_search.cjs | 33 ++ nodejs/__test__/table.test.ts | 191 +++++++ nodejs/lancedb/embedding/index.ts | 44 +- nodejs/lancedb/embedding/openai.ts | 5 +- nodejs/lancedb/embedding/registry.ts | 60 +- nodejs/lancedb/embedding/transformers.ts | 5 +- nodejs/lancedb/index.ts | 1 + nodejs/lancedb/query.ts | 139 +++-- nodejs/lancedb/table.ts | 44 +- nodejs/src/table.rs | 6 + rust/lancedb/src/remote/table.rs | 38 ++ rust/lancedb/src/table.rs | 32 ++ rust/lancedb/src/table/query.rs | 49 +- 18 files changed, 1211 insertions(+), 68 deletions(-) create mode 100644 docs/src/js/classes/AutoQuery.md create mode 100644 nodejs/__test__/embedding_registry.test.ts create mode 100644 nodejs/__test__/fixtures/auto_fts_search.cjs diff --git a/docs/src/js/classes/AutoQuery.md b/docs/src/js/classes/AutoQuery.md new file mode 100644 index 000000000..1d7ea6952 --- /dev/null +++ b/docs/src/js/classes/AutoQuery.md @@ -0,0 +1,518 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / AutoQuery + +# Class: AutoQuery + +A builder for automatic string searches. + +Automatic search determines whether to use full-text or vector search from +the table revision selected for each execution. This builder exposes the +common operations supported by both query families. + +## Extends + +- `StandardQueryBase`<`NativeQuery` \| `NativeVectorQuery`> + +## Properties + +### inner + +```ts +protected inner: Query | VectorQuery | Promise; +``` + +#### Inherited from + +`StandardQueryBase.inner` + +## Methods + +### analyzePlan() + +```ts +analyzePlan(distributedMetrics?): Promise +``` + +Executes the query and returns the physical query plan annotated with runtime metrics. + +This is useful for debugging and performance analysis, as it shows how the query was executed +and includes metrics such as elapsed time, rows processed, and I/O statistics. + +#### Parameters + +* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md) + How distributed worker metrics are displayed for remote query plans. + Defaults to `"aggregate"`. + +#### Returns + +`Promise`<`string`> + +A query execution plan with runtime metrics for each step. + +#### Example + +```ts +import * as lancedb from "@lancedb/lancedb" + +const db = await lancedb.connect("./.lancedb"); +const table = await db.createTable("my_table", [ + { vector: [1.1, 0.9], id: "1" }, +]); + +const plan = await table.query().nearestTo([0.5, 0.2]).analyzePlan(); + +Example output (with runtime metrics inlined): +AnalyzeExec verbose=true, metrics=[] + ProjectionExec: expr=[id@3 as id, vector@0 as vector, _distance@2 as _distance], metrics=[output_rows=1, elapsed_compute=3.292µs] + Take: columns="vector, _rowid, _distance, (id)", metrics=[output_rows=1, elapsed_compute=66.001µs, batches_processed=1, bytes_read=8, iops=1, requests=1] + CoalesceBatchesExec: target_batch_size=1024, metrics=[output_rows=1, elapsed_compute=3.333µs] + GlobalLimitExec: skip=0, fetch=10, metrics=[output_rows=1, elapsed_compute=167ns] + FilterExec: _distance@2 IS NOT NULL, metrics=[output_rows=1, elapsed_compute=8.542µs] + SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST], metrics=[output_rows=1, elapsed_compute=63.25µs, row_replacements=1] + KNNVectorDistance: metric=l2, metrics=[output_rows=1, elapsed_compute=114.333µs, output_batches=1] + LanceScan: uri=/path/to/data, projection=[vector], row_id=true, row_addr=false, ordered=false, metrics=[output_rows=1, elapsed_compute=103.626µs, bytes_read=549, iops=2, requests=2] +``` + +#### Inherited from + +`StandardQueryBase.analyzePlan` + +*** + +### execute() + +```ts +protected execute(options?): AsyncGenerator, void, unknown> +``` + +Execute the query and return the results as an + +#### Parameters + +* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)> + +#### Returns + +`AsyncGenerator`<`RecordBatch`<`any`>, `void`, `unknown`> + +#### See + + - AsyncIterator +of + - RecordBatch. + +By default, LanceDb will use many threads to calculate results and, when +the result set is large, multiple batches will be processed at one time. +This readahead is limited however and backpressure will be applied if this +stream is consumed slowly (this constrains the maximum memory used by a +single query) + +#### Inherited from + +`StandardQueryBase.execute` + +*** + +### explainPlan() + +```ts +explainPlan(verbose): Promise +``` + +Generates an explanation of the query execution plan. + +#### Parameters + +* **verbose**: `boolean` = `false` + If true, provides a more detailed explanation. Defaults to false. + +#### Returns + +`Promise`<`string`> + +A Promise that resolves to a string containing the query execution plan explanation. + +#### Example + +```ts +import * as lancedb from "@lancedb/lancedb" +const db = await lancedb.connect("./.lancedb"); +const table = await db.createTable("my_table", [ + { vector: [1.1, 0.9], id: "1" }, +]); +const plan = await table.query().nearestTo([0.5, 0.2]).explainPlan(); +``` + +#### Inherited from + +`StandardQueryBase.explainPlan` + +*** + +### fastSearch() + +```ts +fastSearch(): this +``` + +Skip searching un-indexed data. This can make search faster, but will miss +any data that is not yet indexed. + +Use [Table#optimize](Table.md#optimize) to index all un-indexed data. + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.fastSearch` + +*** + +### ~~filter()~~ + +```ts +filter(predicate): this +``` + +A filter statement to be applied to this query. + +#### Parameters + +* **predicate**: `string` + +#### Returns + +`this` + +#### See + +where + +#### Deprecated + +Use `where` instead + +#### Inherited from + +`StandardQueryBase.filter` + +*** + +### fullTextSearch() + +```ts +fullTextSearch(query, options?): this +``` + +#### Parameters + +* **query**: `string` \| [`FullTextQuery`](../interfaces/FullTextQuery.md) + +* **options?**: `Partial`<[`FullTextSearchOptions`](../interfaces/FullTextSearchOptions.md)> + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.fullTextSearch` + +*** + +### limit() + +```ts +limit(limit): this +``` + +Set the maximum number of results to return. + +By default, a plain search has no limit. If this method is not +called then every valid row from the table will be returned. + +#### Parameters + +* **limit**: `number` + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.limit` + +*** + +### offset() + +```ts +offset(offset): this +``` + +Set the number of rows to skip before returning results. + +This is useful for pagination. + +#### Parameters + +* **offset**: `number` + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.offset` + +*** + +### orderBy() + +```ts +orderBy(ordering): this +``` + +Sort the results by the specified column(s). + +#### Parameters + +* **ordering**: [`ColumnOrdering`](../interfaces/ColumnOrdering.md) \| [`ColumnOrdering`](../interfaces/ColumnOrdering.md)[] + +#### Returns + +`this` + +This query builder. + +#### Inherited from + +`StandardQueryBase.orderBy` + +*** + +### outputSchema() + +```ts +outputSchema(): Promise> +``` + +Returns the schema of the output that will be returned by this query. + +This can be used to inspect the types and names of the columns that will be +returned by the query before executing it. + +#### Returns + +`Promise`<`Schema`<`any`>> + +An Arrow Schema describing the output columns. + +#### Inherited from + +`StandardQueryBase.outputSchema` + +*** + +### select() + +```ts +select(columns): this +``` + +Return only the specified columns. + +By default a query will return all columns from the table. However, this can have +a very significant impact on latency. LanceDb stores data in a columnar fashion. This +means we can finely tune our I/O to select exactly the columns we need. + +As a best practice you should always limit queries to the columns that you need. If you +pass in an array of column names then only those columns will be returned. + +You can also use this method to create new "dynamic" columns based on your existing columns. +For example, you may not care about "a" or "b" but instead simply want "a + b". This is often +seen in the SELECT clause of an SQL query (e.g. `SELECT a+b FROM my_table`). + +To create dynamic columns you can pass in a Map. A column will be returned +for each entry in the map. The key provides the name of the column. The value is +an SQL string used to specify how the column is calculated. + +For example, an SQL query might state `SELECT a + b AS combined, c`. The equivalent +input to this method would be: + +#### Parameters + +* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> + +#### Returns + +`this` + +#### Example + +```ts +new Map([["combined", "a + b"], ["c", "c"]]) + +Columns will always be returned in the order given, even if that order is different than +the order used when adding the data. + +Note that you can pass in a `Record` (e.g. an object literal). This method +uses `Object.entries` which should preserve the insertion order of the object. However, +object insertion order is easy to get wrong and `Map` is more foolproof. +``` + +#### Inherited from + +`StandardQueryBase.select` + +*** + +### toArray() + +```ts +toArray(options?): Promise +``` + +Collect the results as an array of objects. + +#### Parameters + +* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)> + +#### Returns + +`Promise`<`any`[]> + +#### Inherited from + +`StandardQueryBase.toArray` + +*** + +### toArrow() + +```ts +toArrow(options?): Promise> +``` + +Collect the results as an Arrow + +#### Parameters + +* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)> + +#### Returns + +`Promise`<`Table`<`any`>> + +#### See + +ArrowTable. + +#### Inherited from + +`StandardQueryBase.toArrow` + +*** + +### useLsm() + +```ts +useLsm(enable): this +``` + +Control MemWAL read routing for this query. + +By default (unset), when the table carries a MemWAL write spec (see +[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so +they also return data written via the `mergeInsert` LSM path that has not yet +been compacted into the base table (the active/frozen in-memory memtables and +the flushed generations), deduplicated by primary key; a table without a spec +reads the base table. + +#### Parameters + +* **enable**: `boolean` + `true` forces the LSM scanner and errors if the table has no + MemWAL write spec. `false` bypasses the MemWAL and reads the base table only, + even when a spec is present. + Note: the LSM scanner does not support every query shape (e.g. reranking, + hybrid search, `orderBy`). On a MemWAL table those shapes error unless + `useLsm(false)` is set, because a base-only read would silently exclude + un-compacted MemWAL data. + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.useLsm` + +*** + +### where() + +```ts +where(predicate): this +``` + +A filter statement to be applied to this query. + +The filter should be supplied as an SQL query string. For example: + +#### Parameters + +* **predicate**: `string` + +#### Returns + +`this` + +#### Example + +```ts +x > 10 +y > 0 AND y < 100 +x > 5 OR y = 'test' + +Filtering performance can often be improved by creating a scalar index +on the filter column(s). + +Calling this multiple times combines the filters with a logical AND rather +than replacing the previous filter. +``` + +#### Inherited from + +`StandardQueryBase.where` + +*** + +### withRowId() + +```ts +withRowId(): this +``` + +Whether to return the row id in the results. + +This column can be used to match results between different queries. For +example, to match results from a full text search and a vector search in +order to perform hybrid search. + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.withRowId` diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 06dc8479e..9a85d0d96 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -942,7 +942,7 @@ Get the schema of the table. abstract search( query, queryType?, - ftsColumns?): Query | VectorQuery + ftsColumns?): Query | VectorQuery | AutoQuery ``` Create a search query to find the nearest neighbors @@ -964,7 +964,7 @@ of the given query #### Returns -[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md) +[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md) \| [`AutoQuery`](AutoQuery.md) *** diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index e0635ab65..beb9cbeff 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -18,6 +18,7 @@ ## Classes +- [AutoQuery](classes/AutoQuery.md) - [BooleanQuery](classes/BooleanQuery.md) - [BoostQuery](classes/BoostQuery.md) - [BranchContents](classes/BranchContents.md) diff --git a/docs/src/js/namespaces/embedding/functions/getRegistry.md b/docs/src/js/namespaces/embedding/functions/getRegistry.md index 331dbd60b..b149a4a88 100644 --- a/docs/src/js/namespaces/embedding/functions/getRegistry.md +++ b/docs/src/js/namespaces/embedding/functions/getRegistry.md @@ -10,16 +10,12 @@ function getRegistry(): EmbeddingFunctionRegistry ``` -Utility function to get the global instance of the registry +Get the global embedding function registry. + +LanceDB built-in providers are initialized when this public API is first +used, so importing the root package does not change automatic search +selection for tables without embedding metadata. ## Returns [`EmbeddingFunctionRegistry`](../classes/EmbeddingFunctionRegistry.md) - -`EmbeddingFunctionRegistry` The global instance of the registry - -## Example - -```ts -const registry = getRegistry(); -const openai = registry.get("openai").create(); diff --git a/nodejs/__test__/embedding_registry.test.ts b/nodejs/__test__/embedding_registry.test.ts new file mode 100644 index 000000000..83933399a --- /dev/null +++ b/nodejs/__test__/embedding_registry.test.ts @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import { execFileSync } from "node:child_process"; +import { resolve } from "node:path"; + +import type { OpenAIEmbeddingFunction } from "../lancedb/embedding/openai"; +import type { EmbeddingFunctionRegistry } from "../lancedb/embedding/registry"; + +type EmbeddingModule = typeof import("../lancedb/embedding"); +type OpenAIModule = typeof import("../lancedb/embedding/openai"); +type RegistryModule = typeof import("../lancedb/embedding/registry"); + +describe("embedding function registry", () => { + const registries: EmbeddingFunctionRegistry[] = []; + + afterEach(() => { + for (const registry of registries) { + registry.reset(); + } + registries.length = 0; + }); + + it("defers built-in providers until the public registry API is used", () => { + jest.isolateModules(() => { + const embedding = require("../lancedb/embedding") as EmbeddingModule; + const { getRegistry: getInternalRegistry } = + require("../lancedb/embedding/registry") as RegistryModule; + const registry = getInternalRegistry(); + registries.push(registry); + + expect(registry.length()).toBe(0); + expect(embedding.getRegistry()).toBe(registry); + expect(registry.get("openai")).toBeDefined(); + expect(registry.get("huggingface")).toBeDefined(); + }); + }); + + it("preserves automatic FTS search in a fresh process", () => { + execFileSync( + process.execPath, + [resolve(__dirname, "fixtures", "auto_fts_search.cjs")], + { stdio: "pipe" }, + ); + }); + + it("shares registrations across duplicated provider module graphs", () => { + let registeringRegistry: EmbeddingFunctionRegistry | undefined; + let latestOpenAIConstructor: typeof OpenAIEmbeddingFunction | undefined; + + jest.isolateModules(() => { + require("../lancedb/embedding/openai"); + const { getRegistry } = + require("../lancedb/embedding/registry") as RegistryModule; + registeringRegistry = getRegistry(); + registries.push(registeringRegistry); + expect(registeringRegistry.get("openai")).toBeDefined(); + }); + + expect(() => { + jest.isolateModules(() => { + const { OpenAIEmbeddingFunction } = + require("../lancedb/embedding/openai") as OpenAIModule; + latestOpenAIConstructor = OpenAIEmbeddingFunction; + const { getRegistry } = + require("../lancedb/embedding/registry") as RegistryModule; + registries.push(getRegistry()); + }); + }).not.toThrow(); + + const previousApiKey = process.env.OPENAI_API_KEY; + process.env.OPENAI_API_KEY = "test"; + try { + const latestOpenAI = registeringRegistry! + .get("openai")! + .create(); + expect(latestOpenAI).toBeInstanceOf(latestOpenAIConstructor!); + } finally { + if (previousApiKey === undefined) { + delete process.env.OPENAI_API_KEY; + } else { + process.env.OPENAI_API_KEY = previousApiKey; + } + } + + jest.isolateModules(() => { + const { getRegistry } = + require("../lancedb/embedding") as EmbeddingModule; + const publicRegistry = getRegistry(); + registries.push(publicRegistry); + expect(publicRegistry).toBe(registeringRegistry); + expect(publicRegistry.get("openai")).toBeDefined(); + }); + }); +}); diff --git a/nodejs/__test__/fixtures/auto_fts_search.cjs b/nodejs/__test__/fixtures/auto_fts_search.cjs new file mode 100644 index 000000000..b5ab060b3 --- /dev/null +++ b/nodejs/__test__/fixtures/auto_fts_search.cjs @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const assert = require("node:assert/strict"); +const tmp = require("tmp"); +const { connect, embedding, Index } = require("../../dist"); +const { getRegistry } = require("../../dist/embedding/registry"); + +async function main() { + assert.equal(typeof embedding.getRegistry, "function"); + assert.equal(getRegistry().length(), 0); + assert.equal(embedding.getRegistry(), getRegistry()); + assert.equal(getRegistry().length(), 2); + + const dir = tmp.dirSync({ unsafeCleanup: true }); + let db; + try { + db = await connect(dir.name); + const table = await db.createTable("docs", [{ text: "hello world" }]); + await table.createIndex("text", { config: Index.fts() }); + + const rows = await table.search("hello").toArray(); + assert.equal(rows[0].text, "hello world"); + } finally { + db?.close(); + dir.removeCallback(); + } +} + +main().catch((error) => { + console.error(error); + process.exitCode = 1; +}); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 0f6ed3615..0aee5caf0 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -11,10 +11,13 @@ import * as arrow17 from "apache-arrow-17"; import * as arrow18 from "apache-arrow-18"; import { + AutoQuery, Connection, MatchQuery, PhraseQuery, + Query, Table, + VectorQuery, connect, tokenize, } from "../lancedb"; @@ -1777,6 +1780,194 @@ describe("Read consistency interval", () => { }); }); +describe("automatic search schema consistency", () => { + let tmpDir: tmp.DirResult; + + class SchemaRefreshEmbedding extends EmbeddingFunction { + ndims() { + return 2; + } + + embeddingDataType() { + return new Float32(); + } + + async computeSourceEmbeddings(data: string[]) { + return data.map((value) => [value.length, 1]); + } + + async computeQueryEmbeddings(value: string) { + return [value.length, 1]; + } + } + + function embeddingSchema() { + const func = new SchemaRefreshEmbedding(); + return LanceSchema({ + text: func.sourceField(new Utf8()), + vector: func.vectorField(), + }); + } + + beforeEach(() => { + getRegistry().reset(); + register("schema-refresh")(SchemaRefreshEmbedding); + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + }); + + afterEach(() => { + getRegistry().reset(); + tmpDir.removeCallback(); + }); + + it("uses the schema refreshed from another connection", async () => { + const first = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + const second = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + + try { + const stale = await first.createTable("docs", [{ text: "before" }], { + schema: embeddingSchema(), + }); + const replacement = await second.createTable( + "docs", + [{ text: "after hello" }], + { mode: "overwrite" }, + ); + await replacement.createIndex("text", { config: Index.fts() }); + + const search = stale.search("hello"); + expect(search).toBeInstanceOf(AutoQuery); + expect(search).not.toBeInstanceOf(Query); + expect(search).not.toBeInstanceOf(VectorQuery); + expect("nprobes" in search).toBe(false); + + const rows = await search.toArray(); + expect(rows[0].text).toBe("after hello"); + expect((await stale.schema()).metadata.has("embedding_functions")).toBe( + false, + ); + } finally { + first.close(); + second.close(); + } + }); + + it("tracks embedding metadata across checkout and restore", async () => { + const first = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + const second = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + + try { + await first.createTable("docs", [{ text: "before" }], { + schema: embeddingSchema(), + }); + const table = await second.createTable( + "docs", + [{ text: "after hello" }], + { mode: "overwrite" }, + ); + await table.createIndex("text", { config: Index.fts() }); + + await table.checkout(1); + expect((await table.search("before").toArray())[0].text).toBe("before"); + + await table.checkoutLatest(); + expect((await table.search("hello").toArray())[0].text).toBe( + "after hello", + ); + + await table.checkout(1); + await table.restore(); + expect((await table.search("before").toArray())[0].text).toBe("before"); + } finally { + first.close(); + second.close(); + } + }); + + it("pins automatic search while computing an embedding", async () => { + let markStarted!: () => void; + let releaseEmbedding!: () => void; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const released = new Promise((resolve) => { + releaseEmbedding = resolve; + }); + + class BlockingEmbedding extends SchemaRefreshEmbedding { + async computeQueryEmbeddings(value: string) { + markStarted(); + await released; + return [value.length, 1]; + } + } + + register("schema-refresh-blocking")(BlockingEmbedding); + const func = new BlockingEmbedding(); + const schema = LanceSchema({ + text: func.sourceField(new Utf8()), + vector: func.vectorField(), + }); + const first = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + const second = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + + try { + const table = await first.createTable( + "docs", + [{ text: "hello before" }], + { schema }, + ); + const pending = table.search("hello").toArray(); + await started; + + const replacement = await second.createTable( + "docs", + [{ text: "hello after" }], + { mode: "overwrite" }, + ); + await replacement.createIndex("text", { config: Index.fts() }); + releaseEmbedding(); + + expect((await pending)[0].text).toBe("hello before"); + } finally { + releaseEmbedding(); + first.close(); + second.close(); + } + }); + + it("refreshes a reused automatic search for every execution", async () => { + const first = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + const second = await connect(tmpDir.name, { readConsistencyInterval: 0 }); + + try { + const table = await first.createTable("docs", [ + { text: "hello before", marker: "before" }, + ]); + await table.createIndex("text", { config: Index.fts() }); + const search = table.search("hello").select(["text"]); + + const before = (await search.toArray())[0]; + expect(before.text).toBe("hello before"); + expect(before.marker).toBeUndefined(); + + const replacement = await second.createTable( + "docs", + [{ text: "hello after", marker: "after" }], + { mode: "overwrite" }, + ); + await replacement.createIndex("text", { config: Index.fts() }); + + const after = (await search.toArray())[0]; + expect(after.text).toBe("hello after"); + expect(after.marker).toBeUndefined(); + } finally { + first.close(); + second.close(); + } + }); +}); + describe("schema evolution", function () { let tmpDir: tmp.DirResult; beforeEach(() => { diff --git a/nodejs/lancedb/embedding/index.ts b/nodejs/lancedb/embedding/index.ts index d0ffaec0d..d748e245a 100644 --- a/nodejs/lancedb/embedding/index.ts +++ b/nodejs/lancedb/embedding/index.ts @@ -4,7 +4,15 @@ import { Field, Schema } from "../arrow"; import { sanitizeType } from "../sanitize"; import { EmbeddingFunction } from "./embedding_function"; -import { EmbeddingFunctionConfig, getRegistry } from "./registry"; +import { + EmbeddingFunctionConfig, + EmbeddingFunctionRegistry, + getRegistry as getGlobalRegistry, + registerBuiltIn, +} from "./registry"; + +type OpenAIModule = typeof import("./openai"); +type TransformersModule = typeof import("./transformers"); export { FieldOptions, @@ -14,7 +22,39 @@ export { EmbeddingFunctionConstructor, } from "./embedding_function"; -export * from "./registry"; +export { + EmbeddingFunctionRegistry, + parseEmbeddingMetadata, + register, +} from "./registry"; +export type { + CreateReturnType, + EmbeddingFunctionConfig, + EmbeddingFunctionCreate, + EmbeddingMetadataEntry, + ResolvedEmbeddingFunctionConfig, +} from "./registry"; + +function initializeBuiltInProviders() { + const { OpenAIEmbeddingFunction } = require("./openai") as OpenAIModule; + const { TransformersEmbeddingFunction } = + require("./transformers") as TransformersModule; + + registerBuiltIn("openai", OpenAIEmbeddingFunction); + registerBuiltIn("huggingface", TransformersEmbeddingFunction); +} + +/** + * Get the global embedding function registry. + * + * LanceDB built-in providers are initialized when this public API is first + * used, so importing the root package does not change automatic search + * selection for tables without embedding metadata. + */ +export function getRegistry(): EmbeddingFunctionRegistry { + initializeBuiltInProviders(); + return getGlobalRegistry(); +} /** * Create a schema with embedding functions. diff --git a/nodejs/lancedb/embedding/openai.ts b/nodejs/lancedb/embedding/openai.ts index 5771cfeb5..2218d44bc 100644 --- a/nodejs/lancedb/embedding/openai.ts +++ b/nodejs/lancedb/embedding/openai.ts @@ -5,14 +5,13 @@ import type OpenAI from "openai"; import type { EmbeddingCreateParams } from "openai/resources/index"; import { Float, Float32 } from "../arrow"; import { EmbeddingFunction } from "./embedding_function"; -import { register } from "./registry"; +import { registerBuiltIn } from "./registry"; export type OpenAIOptions = { apiKey: string; model: EmbeddingCreateParams["model"]; }; -@register("openai") export class OpenAIEmbeddingFunction extends EmbeddingFunction< string, Partial @@ -100,3 +99,5 @@ export class OpenAIEmbeddingFunction extends EmbeddingFunction< return response.data[0].embedding; } } + +registerBuiltIn("openai", OpenAIEmbeddingFunction); diff --git a/nodejs/lancedb/embedding/registry.ts b/nodejs/lancedb/embedding/registry.ts index 5f32f683c..c9ee9135e 100644 --- a/nodejs/lancedb/embedding/registry.ts +++ b/nodejs/lancedb/embedding/registry.ts @@ -7,6 +7,10 @@ import { } from "./embedding_function"; import "reflect-metadata"; +const builtInFunctionsKey = Symbol.for( + "@lancedb/lancedb::embedding-built-in-functions::v1", +); + export type CreateReturnType = T extends { init: () => Promise } ? Promise : T; @@ -59,6 +63,15 @@ export class EmbeddingFunctionRegistry { }; } + /** @ignore */ + setBuiltIn< + T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor, + >(name: string, ctor: T): T { + this.#functions.set(name, ctor); + Reflect.defineMetadata("lancedb::embedding::name", name, ctor); + return ctor; + } + get>( name: string, ): EmbeddingFunctionCreate | undefined; @@ -96,6 +109,7 @@ export class EmbeddingFunctionRegistry { */ reset(this: EmbeddingFunctionRegistry) { this.#functions.clear(); + getBuiltInFunctions(this).clear(); } /** @@ -183,12 +197,56 @@ export class EmbeddingFunctionRegistry { } } -const _REGISTRY = new EmbeddingFunctionRegistry(); +function getBuiltInFunctions(registry: EmbeddingFunctionRegistry): Set { + const registryWithBuiltIns = registry as EmbeddingFunctionRegistry & { + [key: symbol]: Set | undefined; + }; + let builtInFunctions = registryWithBuiltIns[builtInFunctionsKey]; + if (builtInFunctions === undefined) { + builtInFunctions = new Set(); + registryWithBuiltIns[builtInFunctionsKey] = builtInFunctions; + } + return builtInFunctions; +} + +// Server bundlers can load the side-effect embedding entry points and the public +// embedding API from separate module graphs. Keep their registry shared. +const registryKey = Symbol.for( + "@lancedb/lancedb::embedding-function-registry::v1", +); +const registryGlobal = globalThis as typeof globalThis & { + [key: symbol]: EmbeddingFunctionRegistry | undefined; +}; + +function getGlobalRegistry(): EmbeddingFunctionRegistry { + const existingRegistry = registryGlobal[registryKey]; + if (existingRegistry !== undefined) { + return existingRegistry; + } + const registry = new EmbeddingFunctionRegistry(); + registryGlobal[registryKey] = registry; + return registry; +} + +const _REGISTRY = getGlobalRegistry(); export function register(name?: string) { return _REGISTRY.register(name); } +/** @ignore */ +export function registerBuiltIn< + T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor, +>(name: string, ctor: T): T { + const builtInFunctions = getBuiltInFunctions(_REGISTRY); + if (builtInFunctions.has(name)) { + return _REGISTRY.setBuiltIn(name, ctor); + } + _REGISTRY.register(name)(ctor); + builtInFunctions.add(name); + return ctor; +} + /** * Utility function to get the global instance of the registry * @returns `EmbeddingFunctionRegistry` The global instance of the registry diff --git a/nodejs/lancedb/embedding/transformers.ts b/nodejs/lancedb/embedding/transformers.ts index 161575285..06043ea7c 100644 --- a/nodejs/lancedb/embedding/transformers.ts +++ b/nodejs/lancedb/embedding/transformers.ts @@ -3,7 +3,7 @@ import { Float, Float32 } from "../arrow"; import { EmbeddingFunction } from "./embedding_function"; -import { register } from "./registry"; +import { registerBuiltIn } from "./registry"; export type XenovaTransformerOptions = { /** The wasm compatible model to use */ @@ -31,7 +31,6 @@ export type XenovaTransformerOptions = { }; }; -@register("huggingface") export class TransformersEmbeddingFunction extends EmbeddingFunction< string, Partial @@ -158,6 +157,8 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< } } +registerBuiltIn("huggingface", TransformersEmbeddingFunction); + const tensorDiv = ( src: import("@huggingface/transformers").Tensor, divBy: number, diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index ebc8cda8d..34d7ce4d9 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -103,6 +103,7 @@ export { } from "./native.js"; export { + AutoQuery, ExecutableQuery, Query, QueryBase, diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 843a1276f..3b9b286a0 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -111,13 +111,15 @@ export class QueryBase< NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery, > implements AsyncIterable { + protected inner!: NativeQueryType | Promise; + /** * @hidden */ - protected constructor( - protected inner: NativeQueryType | Promise, - ) { - // intentionally empty + protected constructor(inner?: NativeQueryType | Promise) { + if (inner !== undefined) { + this.inner = inner; + } } // call a function on the inner (either a promise or the actual object) @@ -135,6 +137,15 @@ export class QueryBase< } } + /** + * Return the native query used by the next terminal operation. + * + * @hidden + */ + protected async getInner(): Promise { + return this.inner; + } + /** * Return only the specified columns. * @@ -207,16 +218,11 @@ export class QueryBase< /** * @hidden */ - protected nativeExecute( + protected async nativeExecute( options?: Partial, ): Promise { - if (this.inner instanceof Promise) { - return this.inner.then((inner) => - inner.execute(options?.maxBatchLength, options?.timeoutMs), - ); - } else { - return this.inner.execute(options?.maxBatchLength, options?.timeoutMs); - } + const inner = await this.getInner(); + return inner.execute(options?.maxBatchLength, options?.timeoutMs); } /** @@ -245,12 +251,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.getInner(); for await (const batch of new RecordBatchIterable(inner, options)) { batches.push(batch); } @@ -279,11 +280,8 @@ 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)); - } else { - return this.inner.explainPlan(verbose); - } + const inner = await this.getInner(); + return inner.explainPlan(verbose); } /** @@ -321,13 +319,8 @@ export class QueryBase< distributedMetrics?: AnalyzePlanDistributedMetrics, ): Promise { const distributedMetricsMode = distributedMetrics ?? "aggregate"; - if (this.inner instanceof Promise) { - return this.inner.then((inner) => - inner.analyzePlan(distributedMetricsMode), - ); - } else { - return this.inner.analyzePlan(distributedMetricsMode); - } + const inner = await this.getInner(); + return inner.analyzePlan(distributedMetricsMode); } /** @@ -339,12 +332,8 @@ export class QueryBase< * @returns An Arrow Schema describing the output columns. */ async outputSchema(): Promise { - let schemaBuffer: Buffer; - if (this.inner instanceof Promise) { - schemaBuffer = await this.inner.then((inner) => inner.outputSchema()); - } else { - schemaBuffer = await this.inner.outputSchema(); - } + const inner = await this.getInner(); + const schemaBuffer = await inner.outputSchema(); const schema = tableFromIPC(schemaBuffer).schema; return schema; } @@ -356,7 +345,7 @@ export class StandardQueryBase< extends QueryBase implements ExecutableQuery { - constructor(inner: NativeQueryType | Promise) { + constructor(inner?: NativeQueryType | Promise) { super(inner); } @@ -788,6 +777,51 @@ export class TakeQuery extends QueryBase { } } +/** + * A builder for automatic string searches. + * + * Automatic search determines whether to use full-text or vector search from + * the table revision selected for each execution. This builder exposes the + * common operations supported by both query families. + * + * @hideconstructor + */ +export class AutoQuery extends StandardQueryBase< + NativeQuery | NativeVectorQuery +> { + private readonly calls: Array< + (inner: NativeQuery | NativeVectorQuery) => void + > = []; + + /** @hidden */ + constructor( + private readonly createInner: () => Promise< + NativeQuery | NativeVectorQuery + >, + ) { + super(); + } + + /** @hidden */ + protected override doCall( + fn: (inner: NativeQuery | NativeVectorQuery) => void, + ) { + this.calls.push(fn); + } + + /** @hidden */ + protected override async getInner(): Promise< + NativeQuery | NativeVectorQuery + > { + const calls = [...this.calls]; + const inner = await this.createInner(); + for (const call of calls) { + call(inner); + } + return inner; + } +} + /** A builder for LanceDB queries. * * @see {@link Table#query}, {@link Table#search} @@ -802,6 +836,37 @@ export class Query extends StandardQueryBase { super(tbl.query()); } + /** @hidden */ + static autoSearch( + tbl: () => Promise, + query: string, + vector: (tbl: NativeTable) => Promise | undefined>, + columns?: string[], + ): AutoQuery { + const nativeQuery = async () => { + const snapshot = await Promise.resolve(tbl()); + const resolved = await vector(snapshot); + const inner = snapshot.query(); + if (resolved === undefined) { + inner.fullTextSearch({ + query, + columns: columns ?? null, + }); + return inner; + } + + const raw = Array.isArray(resolved) + ? null + : extractVectorBuffer(resolved); + if (raw) { + return inner.nearestToRaw(raw.data, raw.dtype); + } + return inner.nearestTo(Float32Array.from(resolved as number[])); + }; + + return new AutoQuery(nativeQuery); + } + /** * Find the nearest vectors to the given query vector. * diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 28603cca9..a4fc76ef1 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -43,6 +43,7 @@ import { Table as _NativeTable, } from "./native"; import { + AutoQuery, FullTextQuery, Query, TakeQuery, @@ -523,7 +524,7 @@ export abstract class Table { query: string | IntoVector | MultiVector | FullTextQuery, queryType?: string, ftsColumns?: string | string[], - ): VectorQuery | Query; + ): VectorQuery | Query | AutoQuery; /** * Search the table with a given query vector. * @@ -975,10 +976,11 @@ export class LocalTable extends Table { return this.inner.display(); } - private async getEmbeddingFunctions(): Promise< - Map - > { - const schema = await this.schema(); + private async getEmbeddingFunctions( + inner: _NativeTable = this.inner, + ): Promise> { + const schemaBuf = await inner.schema(); + const schema = tableFromIPC(schemaBuf).schema; const registry = getRegistry(); return registry.parseFunctions(schema.metadata); } @@ -1160,7 +1162,7 @@ export class LocalTable extends Table { query: string | IntoVector | MultiVector | FullTextQuery, queryType: string = "auto", ftsColumns?: string | string[], - ): VectorQuery | Query { + ): VectorQuery | Query | AutoQuery { if (typeof query !== "string" && !instanceOfFullTextQuery(query)) { if (queryType === "fts") { throw new Error("Cannot perform full text search on a vector query"); @@ -1175,17 +1177,35 @@ 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" && - (getRegistry().length() === 0 || instanceOfFullTextQuery(query)) - ) { + if (queryType === "auto" && typeof query !== "string") { return this.query().fullTextSearch(query, { columns: ftsColumns, }); } + if (queryType === "auto" && typeof query === "string") { + const vector = async (snapshot: _NativeTable) => { + const functions = await this.getEmbeddingFunctions(snapshot); + // TODO: Support multiple embedding functions + const embeddingFunc: EmbeddingFunctionConfig | undefined = functions + .values() + .next().value; + if (embeddingFunc === undefined) { + return undefined; + } + return await embeddingFunc.function.computeQueryEmbeddings(query); + }; + + const columns = + typeof ftsColumns === "string" ? [ftsColumns] : ftsColumns; + return Query.autoSearch( + () => this.inner.checkoutCurrent(), + query, + vector, + columns, + ); + } + const queryPromise = this.getEmbeddingFunctions().then( async (functions) => { // TODO: Support multiple embedding functions diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 9d60b2056..ff16ac042 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -554,6 +554,12 @@ impl Table { .default_error() } + #[napi(catch_unwind)] + pub async fn checkout_current(&self) -> napi::Result { + let table = self.inner_ref()?.checkout_current().await.default_error()?; + Ok(Self::new(table)) + } + #[napi(catch_unwind)] pub async fn checkout(&self, version: i64) -> napi::Result<()> { self.inner_ref()? diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 44c92310f..77d5983eb 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1725,6 +1725,22 @@ impl BaseTable for RemoteTable { async fn version(&self) -> Result { self.describe().await.map(|desc| desc.version) } + + async fn checkout_current(&self) -> Result> { + let description = self.describe().await?; + let TableDescription { + version, + schema, + location, + } = description; + let schema = Arc::new(arrow_schema::Schema::try_from(schema)?); + let snapshot = self.with_branch(self.branch.clone()); + *snapshot.version.write().await = Some(version); + *snapshot.location.write().await = location; + snapshot.schema_cache.seed(schema); + Ok(Arc::new(snapshot)) + } + async fn checkout(&self, version: u64) -> Result<()> { // Validate the version exists. The describe is sent without freshness // headers so a stale `min_version` from a previous write doesn't ride @@ -8739,6 +8755,28 @@ mod tests { } } + /// A pinned snapshot should reuse the version and schema returned by its + /// initial describe instead of issuing two more describe requests. + #[tokio::test] + async fn test_checkout_current_seeds_schema_from_single_describe() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let calls = describe_calls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + calls.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version":42,"schema":{"fields":[{"name":"a","type":{"type":"int32"},"nullable":false}]}}"#, + ) + .unwrap() + }); + + let snapshot = table.checkout_current().await.unwrap(); + assert_eq!(snapshot.schema().await.unwrap().fields().len(), 1); + assert_eq!(describe_calls.load(Ordering::SeqCst), 1); + } + /// Test that schema cache is invalidated after checkout #[tokio::test] async fn test_schema_cache_invalidation_on_checkout() { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index efc36d260..5007a61d9 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -785,6 +785,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn drop_columns(&self, columns: &[&str]) -> Result; /// Get the version of the table. async fn version(&self) -> Result; + /// Return a new table handle pinned to the exact revision currently visible. + async fn checkout_current(&self) -> Result> { + Err(Error::NotSupported { + message: "checkout_current is not supported on this table type".into(), + }) + } /// Checkout a specific version of the table. async fn checkout(&self, version: u64) -> Result<()>; /// Checkout a table version referenced by a tag. @@ -1944,6 +1950,20 @@ impl Table { self.inner.version().await } + /// Return a new table handle pinned to the exact revision currently visible. + /// + /// This is used when asynchronous preparation must remain consistent with + /// the revision used for a later read. + #[doc(hidden)] + pub async fn checkout_current(&self) -> Result { + let inner = self.inner.checkout_current().await?; + Ok(Self { + inner, + database: self.database.clone(), + embedding_registry: self.embedding_registry.clone(), + }) + } + /// Checks out a specific version of the Table /// /// Any read operation on the table will now access the data at the checked out version. @@ -3043,6 +3063,18 @@ impl BaseTable for NativeTable { Ok(self.dataset.get().await?.version().version) } + async fn checkout_current(&self) -> Result> { + let current = self.dataset.get().await?; + let dataset = dataset::DatasetConsistencyWrapper::new_time_travel( + current.as_ref().clone(), + self.read_consistency_interval, + ); + Ok(Arc::new(Self { + dataset, + ..self.clone() + })) + } + async fn checkout(&self, version: u64) -> Result<()> { self.dataset.as_time_travel(version).await } diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 9b81786cd..d35286c84 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -697,6 +697,7 @@ mod tests { use super::*; use crate::query::{QueryExecutionOptions, QueryRequest}; + use crate::table::BaseTable; fn fixed_size_list_array(values: Vec, dimension: i32) -> FixedSizeListArray { FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap() @@ -889,10 +890,56 @@ mod tests { async fn query_table(&self, _request: NsQueryTableRequest) -> lance::Result { self.query_table_calls.fetch_add(1, Ordering::SeqCst); - panic!("approx_mode queries must not be pushed down to namespace query_table"); + panic!("query must not be pushed down to namespace query_table"); } } + #[tokio::test] + async fn test_execute_query_pinned_snapshot_with_namespace_pushdown_runs_locally() { + use crate::connect; + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + + let conn = connect("memory://").execute().await.unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + ) + .unwrap(); + let table = conn + .create_table("test_pinned_namespace_fallback", vec![batch]) + .execute() + .await + .unwrap(); + + let namespace_client = Arc::new(CountingNamespaceClient::default()); + let mut native_table = table.as_native().unwrap().clone(); + native_table.namespace_client = Some(namespace_client.clone()); + native_table + .pushdown_operations + .insert(NamespaceClientPushdownOperation::QueryTable); + + let snapshot = native_table.checkout_current().await.unwrap(); + let snapshot = snapshot.as_any().downcast_ref::().unwrap(); + assert!(snapshot.dataset.time_travel_version().is_some()); + + let query = AnyQuery::Query(QueryRequest { + filter: Some(QueryFilter::Sql("id > 3".to_string())), + ..Default::default() + }); + let stream = execute_query(snapshot, &query, QueryExecutionOptions::default()) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 2 + ); + assert_eq!(namespace_client.query_table_calls.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_execute_query_approx_mode_with_namespace_pushdown_runs_locally() { use crate::connect;