diff --git a/docs/src/js/classes/AutoQuery.md b/docs/src/js/classes/AutoQuery.md index f62c5e84d..1d7ea6952 100644 --- a/docs/src/js/classes/AutoQuery.md +++ b/docs/src/js/classes/AutoQuery.md @@ -9,8 +9,8 @@ A builder for automatic string searches. Automatic search determines whether to use full-text or vector search from -the table revision selected for execution. This builder exposes the common -operations supported by both query families. +the table revision selected for each execution. This builder exposes the +common operations supported by both query families. ## Extends diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 7ace628ee..aaa31f1db 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -1889,6 +1889,37 @@ describe("automatic search schema consistency", () => { 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 () { diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 55a18ff38..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); } @@ -792,17 +781,44 @@ 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 execution. This builder exposes the common - * operations supported by both query families. + * 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(inner: Promise) { - super(inner); + 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; } } @@ -822,14 +838,15 @@ export class Query extends StandardQueryBase { /** @hidden */ static autoSearch( - tbl: Promise, + tbl: () => Promise, query: string, vector: (tbl: NativeTable) => Promise | undefined>, columns?: string[], ): AutoQuery { - const nativeQuery = Promise.resolve(tbl).then(async (tbl) => { - const resolved = await vector(tbl); - const inner = tbl.query(); + const nativeQuery = async () => { + const snapshot = await Promise.resolve(tbl()); + const resolved = await vector(snapshot); + const inner = snapshot.query(); if (resolved === undefined) { inner.fullTextSearch({ query, @@ -845,7 +862,7 @@ export class Query extends StandardQueryBase { return inner.nearestToRaw(raw.data, raw.dtype); } return inner.nearestTo(Float32Array.from(resolved as number[])); - }); + }; return new AutoQuery(nativeQuery); } diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 09f07bca6..233dadc64 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -1058,7 +1058,7 @@ export class LocalTable extends Table { const columns = typeof ftsColumns === "string" ? [ftsColumns] : ftsColumns; return Query.autoSearch( - this.inner.checkoutCurrent(), + () => this.inner.checkoutCurrent(), query, vector, columns, diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index edc42b412..095a1c3ea 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1683,9 +1683,17 @@ impl BaseTable for RemoteTable { } async fn checkout_current(&self) -> Result> { - let version = self.version().await?; + 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.checkout(version).await?; + *snapshot.version.write().await = Some(version); + *snapshot.location.write().await = location; + snapshot.schema_cache.seed(schema); Ok(Arc::new(snapshot)) } @@ -7234,6 +7242,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/query.rs b/rust/lancedb/src/table/query.rs index 9feb9d5ab..2619735da 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -70,6 +70,9 @@ async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> R .contains(&NamespaceClientPushdownOperation::QueryTable) && table.namespace_client.is_some() && table.dataset.current_branch().is_none() + // QueryTableRequest cannot carry a pinned dataset version. Pushing a + // time-travel query down would silently execute against latest. + && table.dataset.time_travel_version().is_none() && !requires_local_namespace_execution(query)) { return Ok(false); @@ -694,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() @@ -886,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;