From 8b7e13b0c610c586029b763337d5baf2c3edab16 Mon Sep 17 00:00:00 2001 From: Dan Tasse <105866+dantasse@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:19:44 -0400 Subject: [PATCH 1/3] docs: add comments about metadata conventions (#4054) In LanceDB Enterprise, we've adopted these conventions to give some "canonical" metadata paths. This lets us display them in a certain way in the UI or let agents standardize on them, to assume they'll find info in a certain place. This PR (only comments/docs) just documents those choices. --- docs/src/js/classes/Table.md | 12 ++++++++++++ docs/src/js/interfaces/FieldMetadataUpdate.md | 3 ++- nodejs/lancedb/table.ts | 15 ++++++++++++++- python/python/lancedb/table.py | 13 +++++++++++++ rust/lancedb/src/table.rs | 18 +++++++++++++++++- rust/lancedb/src/table/schema_evolution.rs | 4 +++- 6 files changed, 61 insertions(+), 4 deletions(-) diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 9a85d0d96..159348450 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -1292,6 +1292,18 @@ abstract updateFieldMetadata(updates): Promise Update per-field (column) metadata. +The following keys are treated specially, by convention, and should be +used when appropriate: + +- `lancedb:description`: for a human-readable description of a field. +- `lancedb:tag:`: for a user-defined key-value tag, where the suffix + names the tag category; e.g. `lancedb:tag:model: "clip"`. +- `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and + `feature_v2` might be in the same logical column. +- `lancedb:status`: for status options (`production`, `candidate`, + `deprecated`, `archived`) to designate the current life cycle state of + this column. + #### Parameters * **updates**: [`FieldMetadataUpdate`](../interfaces/FieldMetadataUpdate.md)[] diff --git a/docs/src/js/interfaces/FieldMetadataUpdate.md b/docs/src/js/interfaces/FieldMetadataUpdate.md index 38c675630..a85e3e7a0 100644 --- a/docs/src/js/interfaces/FieldMetadataUpdate.md +++ b/docs/src/js/interfaces/FieldMetadataUpdate.md @@ -17,7 +17,8 @@ metadata: Record; ``` Metadata key/value pairs. Merged into the field's existing metadata by -default; a value of `null` deletes that key. +default; a value of `null` deletes that key. See +[Table.updateFieldMetadata](../classes/Table.md#updatefieldmetadata) for the conventional `lancedb:*` keys. *** diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 82f23e2f2..dc062e337 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -630,6 +630,18 @@ export abstract class Table { /** * Update per-field (column) metadata. + * + * The following keys are treated specially, by convention, and should be + * used when appropriate: + * + * - `lancedb:description`: for a human-readable description of a field. + * - `lancedb:tag:`: for a user-defined key-value tag, where the suffix + * names the tag category; e.g. `lancedb:tag:model: "clip"`. + * - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and + * `feature_v2` might be in the same logical column. + * - `lancedb:status`: for status options (`production`, `candidate`, + * `deprecated`, `archived`) to designate the current life cycle state of + * this column. * @param {FieldMetadataUpdate[]} updates One or more per-field updates. Each * update's metadata is merged into the field's existing metadata by default; * a value of `null` deletes that key, and `replace: true` swaps the whole map. @@ -1555,7 +1567,8 @@ export interface FieldMetadataUpdate { path: string; /** * Metadata key/value pairs. Merged into the field's existing metadata by - * default; a value of `null` deletes that key. + * default; a value of `null` deletes that key. See + * {@link Table.updateFieldMetadata} for the conventional `lancedb:*` keys. */ metadata: Record; /** If true, replace the field's entire metadata map instead of merging. */ diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 76d5fc825..c354a944e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2127,12 +2127,25 @@ class Table(ABC): ---------- updates : dict One or more dicts, each with: + - "path": str — dot-path to the field (e.g. "embedding" or "a.b.c"). - "metadata": dict[str, str | None] — keys to set; a value of ``None`` deletes that key. - "replace": bool, optional — replace the field's whole metadata map instead of merging (default False). + The following keys are treated specially, by convention, and should + be used when appropriate: + + - "lancedb:description": for a human-readable description of a field. + - ``"lancedb:tag:"`` for a user-defined key-value tag, where the + suffix names the tag category; e.g. "lancedb:tag:model": "clip". + - "lancedb:logical-column" for a column grouping; e.g. "feature_v1" + and "feature_v2" might be in the same logical column. + - "lancedb:status" for status options ("production", "candidate", + "deprecated", "archived") to designate the current life cycle + state of this column. + Returns ------- UpdateFieldMetadataResult diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 182e3e1a7..af8bcb5e2 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1777,7 +1777,23 @@ impl Table { self.inner.alter_columns(alterations).await } - /// Update per-field metadata (merges by default). + /// Update per-field (column) metadata. + /// + /// Each [`FieldMetadataUpdate`] is merged into the field's existing metadata + /// by default; use [`FieldMetadataUpdate::remove`] to delete a key, or + /// [`FieldMetadataUpdate::replace`] to swap the field's entire metadata map. + /// + /// The following keys are treated specially, by convention, and should be + /// used when appropriate: + /// + /// - `lancedb:description`: for a human-readable description of a field. + /// - `lancedb:tag:`: for a user-defined key-value tag, where the suffix + /// names the tag category; e.g. `lancedb:tag:model: "clip"`. + /// - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and + /// `feature_v2` might be in the same logical column. + /// - `lancedb:status`: for status options (`production`, `candidate`, + /// `deprecated`, `archived`) to designate the current life cycle state of + /// this column. pub async fn update_field_metadata( &self, updates: &[FieldMetadataUpdate], diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index d10a45eea..4f8dc811a 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -55,7 +55,9 @@ pub struct DropColumnsResult { pub struct FieldMetadataUpdate { /// Dot-separated path to the field (e.g. `"embedding"` or `"address.zip"`). pub path: String, - /// Keys to set (`Some`) or delete (`None`). + /// Keys to set (`Some`) or delete (`None`). See + /// [`Table::update_field_metadata`](crate::Table::update_field_metadata) for + /// the conventional `lancedb:*` keys. pub metadata: HashMap>, /// If `true`, replace the field's entire metadata map instead of merging. pub replace: bool, From ae81d735638d7efdea39c1f29055134383b46abb Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 04:23:28 +0800 Subject: [PATCH 2/3] fix: share scans across batched vector queries (#3805) ## Summary - use the Lance native batch KNN path so fixed-size batch vector searches share one flat table scan - validate consistent query-vector dimensions and retain the per-vector plan when offsets require its existing semantics - add Rust and Python regressions and update Rust, Python, and TypeScript API documentation ## Root cause LanceDB expanded every vector in a batch into a separate scan plan and joined the plans with `UnionExec`. For unindexed tables on S3, a batch of ten vectors therefore ran ten concurrent full scans, amplifying CPU and retained data enough to produce the reported memory spike. The native Lance batch KNN path performs bounded-memory selection for all query vectors over one flat scan. LanceDB now supplies the vectors as a batch and avoids applying a global scanner limit to the combined per-query results. Batch queries with a nonzero offset keep the previous plan because the native batch API does not support per-query offsets. ## Validation - targeted Rust batch-query plan and execution tests - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples` - `cargo fmt --all -- --check` - targeted Python batch-vector regression after rebuilding the extension - Ruff formatting/checks for the touched Python files - Node.js build, lint, docs generation, and targeted batch-vector Jest test - `git diff --check` Fixes #2468 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- nodejs/lancedb/query.ts | 10 +-- python/python/lancedb/query.py | 11 +-- python/python/tests/test_query.py | 17 +++++ rust/lancedb/src/query.rs | 113 ++++++++++++++++++++++++++++-- rust/lancedb/src/table/query.rs | 72 ++++++++++++++----- 5 files changed, 187 insertions(+), 36 deletions(-) diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index f1d31eae1..75a787fcd 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -727,11 +727,11 @@ export class VectorQuery extends StandardQueryBase { * Add a query vector to the search * * This method can be called multiple times to add multiple query vectors - * to the search. If multiple query vectors are added, then they will be searched - * in parallel, and the results will be concatenated. A column called `query_index` - * will be added to indicate the index of the query vector that produced the result. - * - * Performance wise, this is equivalent to running multiple queries concurrently. + * to the search. A column called `query_index` will be added to indicate the index + * of the query vector that produced the result. Flat searches share one table scan + * across the query vectors, avoiding the scan and memory amplification of running + * multiple queries concurrently. Indexed searches may still perform per-vector + * index work. */ addQueryVector(vector: IntoVector): VectorQuery { if (vector instanceof Promise) { diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 5dff2537e..9301d7df8 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -3401,9 +3401,10 @@ class AsyncQuery(AsyncStandardQuery): pass in multiple vectors. When multiple vectors are passed in, if the vector column is with multivector type, then the vectors will be treated as a single query. Or the vectors will be treated as multiple queries, this can be useful - if you want to find the nearest vectors to multiple query vectors. - This is not expected to be faster than making multiple queries concurrently; - it is just a convenience method. If multiple vectors are passed in then + if you want to find the nearest vectors to multiple query vectors. Flat + searches share one table scan across the query vectors, avoiding the scan + and memory amplification of making multiple queries concurrently. If + multiple vectors are passed in then an additional column `query_index` will be added to the results. This column will contain the index of the query vector that the result is nearest to. """ @@ -3532,8 +3533,8 @@ class AsyncFTSQuery(AsyncStandardQuery): Typically, a single vector is passed in as the query. However, you can also pass in multiple vectors. This can be useful if you want to find the nearest - vectors to multiple query vectors. This is not expected to be faster than - making multiple queries concurrently; it is just a convenience method. + vectors to multiple query vectors. Flat searches share one table scan across + the query vectors instead of issuing concurrent full scans. If multiple vectors are passed in then an additional column `query_index` will be added to the results. This column will contain the index of the query vector that the result is nearest to. diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index d2629d1a8..4758f0e2d 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -897,6 +897,23 @@ def test_query_builder_batches(table): assert rs_list["id"][1] == 2 +def test_batch_vector_query_shares_filtered_flat_scan(table): + query = ( + table.search([[1.0, 2.0], [3.0, 4.0]]) + .where("id > 0", prefilter=True) + .limit(1) + .select(["id"]) + ) + + plan = query.explain_plan(verbose=True) + assert "KNNVectorDistance: queries=2" in plan + assert "UnionExec" not in plan + + results = query.to_arrow() + assert len(results) == 2 + assert results["query_index"].to_pylist() == [0, 1] + + def test_dynamic_projection(table): rs = ( LanceVectorQueryBuilder(table, [0, 0], "vector") diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 654777adb..2a1283f22 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1174,12 +1174,12 @@ impl VectorQuery { /// Add another query vector to the search. /// - /// Multiple searches will be dispatched as part of the query. - /// This is a convenience method for adding multiple query vectors - /// to the search. It is not expected to be faster than issuing - /// multiple queries concurrently. + /// Multiple searches will be dispatched as a batch. Flat searches share + /// one table scan across the query vectors, avoiding the scan and memory + /// amplification of issuing the searches concurrently. Indexed searches + /// may still perform per-vector index work. /// - /// The output data will contain an additional columns `query_index` which + /// The output data will contain an additional column `query_index` which /// will contain the index of the query vector that was used to generate the /// result. pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result { @@ -1646,7 +1646,11 @@ mod tests { use std::{collections::HashSet, sync::Arc}; use super::*; - use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type}; + use arrow::{ + array::downcast_array, + compute::concat_batches, + datatypes::{Int32Type, UInt8Type}, + }; use arrow_array::{ FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray, types::Float32Type, @@ -2334,7 +2338,8 @@ mod tests { .limit(1); let plan = query.explain_plan(true).await.unwrap(); - assert!(plan.contains("UnionExec")); + assert!(plan.contains("KNNVectorDistance: queries=2")); + assert!(!plan.contains("UnionExec")); let results = query .execute() @@ -2349,6 +2354,100 @@ mod tests { // We don't guarantee order. assert!(query_index.values().contains(&0)); assert!(query_index.values().contains(&1)); + + // Batch KNN does not support a per-query offset, so offset queries keep + // the legacy per-vector plan to preserve their result semantics. + let offset_query = table + .query() + .nearest_to(&[0.1, 0.2, 0.3, 0.4]) + .unwrap() + .add_query_vector(&[0.5, 0.6, 0.7, 0.8]) + .unwrap() + .limit(1) + .offset(1); + assert!( + offset_query + .explain_plan(true) + .await + .unwrap() + .contains("UnionExec") + ); + let offset_results = offset_query + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + offset_results + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + } + + #[tokio::test] + async fn test_multiple_binary_query_vectors() { + let vectors = FixedSizeListArray::from_iter_primitive::( + vec![ + Some(vec![Some(0), Some(0)]), + Some(vec![Some(255), Some(255)]), + ], + 2, + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)], + ) + .unwrap(); + + let conn = connect("memory://").execute().await.unwrap(); + let table = conn + .create_table("binary_batch", batch) + .execute() + .await + .unwrap(); + let query = table + .query() + .nearest_to(&[0.0, 0.0]) + .unwrap() + .add_query_vector(&[255.0, 255.0]) + .unwrap() + .distance_type(DistanceType::Hamming) + .limit(1); + + // Binary queries retain the per-vector plan because Lance's binary + // nearest path requires primitive UInt8 query arrays. + assert!( + query + .explain_plan(true) + .await + .unwrap() + .contains("UnionExec") + ); + + let results = query + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let results = concat_batches(&results[0].schema(), &results).unwrap(); + assert_eq!(results.num_rows(), 2); + + let ids = results["id"].as_primitive::(); + assert!(ids.values().contains(&0)); + assert!(ids.values().contains(&1)); + let query_index = results["query_index"].as_primitive::(); + assert!(query_index.values().contains(&0)); + assert!(query_index.values().contains(&1)); } #[tokio::test] diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 629cb4e6f..2684ac5e2 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -21,7 +21,6 @@ use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::union::UnionExec; -use futures::future::try_join_all; use lance::dataset::mem_wal::DatasetMemWalExt; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::scanner::Scanner; @@ -170,6 +169,7 @@ pub async fn create_plan( let mut column = query.column.clone(); let mut query_vector = query.query_vector.first().cloned(); + let mut is_batch_query = false; if query.query_vector.len() > 1 { if column.is_none() { // Infer a vector column with the same dimension of the query vector. @@ -180,16 +180,37 @@ pub async fn create_plan( )?); } let vector_field = schema.field(column.as_ref().unwrap()).unwrap(); - if let DataType::List(_) = vector_field.data_type() { - // Multivector handling: concatenate into FixedSizeList> + let (_, element_type) = + lance::index::vector::utils::get_vector_type(schema, column.as_ref().unwrap())?; + let is_binary = matches!(element_type, DataType::UInt8); + if matches!(vector_field.data_type(), DataType::List(_)) + || (query.base.offset.unwrap_or(0) == 0 && !is_binary) + { + // Lance distinguishes these cases from the vector column type: a + // list-like query against a List column is one multivector query, + // while the same query against a FixedSizeList column is a batch of + // independent queries. The batch path shares a single flat scan and + // bounds retained candidate data instead of running one scan per + // query vector. let vectors = query .query_vector .iter() .map(|arr| arr.as_ref()) .collect::>(); let dim = vectors[0].len(); + if let Some((query_index, actual_dim)) = vectors + .iter() + .enumerate() + .find_map(|(index, vector)| (vector.len() != dim).then_some((index, vector.len()))) + { + return Err(Error::InvalidInput { + message: format!( + "query vector at index {query_index} has dimension {actual_dim}, expected {dim}" + ), + }); + } let mut fsl_builder = FixedSizeListBuilder::with_capacity( - Float32Builder::with_capacity(dim), + Float32Builder::with_capacity(dim * vectors.len()), dim as i32, vectors.len(), ); @@ -200,8 +221,12 @@ pub async fn create_plan( fsl_builder.append(true); } query_vector = Some(Arc::new(fsl_builder.finish())); + is_batch_query = !matches!(vector_field.data_type(), DataType::List(_)); } else { - // Multiple query vectors: create a plan for each and union them + // Lance's batch path has no per-query offset, and its binary path + // requires primitive UInt8 queries rather than a fixed-size list. + // Keep the prior plan shape for these cases so offsets are applied + // per query and binary query vectors retain their primitive shape. let query_vecs = query.query_vector.clone(); let plan_futures = query_vecs .into_iter() @@ -214,7 +239,7 @@ pub async fn create_plan( } }) .collect::>(); - let plans = try_join_all(plan_futures).await?; + let plans = futures::future::try_join_all(plan_futures).await?; return create_multi_vector_plan(plans); } } @@ -251,10 +276,14 @@ pub async fn create_plan( } } - scanner.limit( - query.base.limit.map(|limit| limit as i64), - query.base.offset.map(|offset| offset as i64), - )?; + // For a batch query, `nearest` already applies k to each query vector. + // Adding Scanner's global limit would truncate the combined result to k rows. + if !is_batch_query { + scanner.limit( + query.base.limit.map(|limit| limit as i64), + query.base.offset.map(|offset| offset as i64), + )?; + } if let Some(ef) = query.ef { scanner.ef(ef); @@ -1088,7 +1117,7 @@ mod tests { } #[tokio::test] - async fn test_create_plan_multivector_structure() { + async fn test_create_plan_batch_vector_uses_shared_scan() { use arrow_array::{Float32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use datafusion_physical_plan::display::DisplayableExecutionPlan; @@ -1115,11 +1144,18 @@ mod tests { .unwrap(); let native_table = table.as_native().unwrap(); - // This triggers the "create_multi_vector_plan" logic branch + // A batch of vectors against a fixed-size vector column should use + // Lance's native batch KNN path instead of independent scan plans. let q1 = Arc::new(Float32Array::from(vec![1.0, 2.0])); let q2 = Arc::new(Float32Array::from(vec![3.0, 4.0])); let req = VectorQueryRequest { + base: QueryRequest { + filter: Some(QueryFilter::Sql("id >= 0".to_string())), + limit: Some(1), + select: Select::Columns(vec!["id".to_string()]), + ..Default::default() + }, column: Some("vector".to_string()), query_vector: vec![q1, q2], ..Default::default() @@ -1136,19 +1172,17 @@ mod tests { .indent(true) .to_string(); - // We expect a RepartitionExec wrapping a UnionExec assert!( - display.contains("RepartitionExec"), - "Plan should include Repartitioning" + display.contains("KNNVectorDistance: queries=2"), + "plan should use native batch KNN, got:\n{display}" ); assert!( - display.contains("UnionExec"), - "Plan should include a Union of multiple searches" + !display.contains("UnionExec"), + "flat batch KNN should share one scan, got:\n{display}" ); - // We expect the projection to add the 'query_index' column (logic inside multi_vector_plan) assert!( display.contains("query_index"), - "Plan should add query_index column" + "plan should add query_index column, got:\n{display}" ); } From 79f626b09edc71b41fb285ee695e6714e14eb63c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 06:17:28 +0800 Subject: [PATCH 3/3] fix: support double-quoted filter identifiers (#3825) ## Summary - tokenize predicates with the same GenericDialect lexical rules Lance delegates to - rewrite only SQL-standard double-quoted identifier tokens to Lance backticks - apply one predicate contract to query, count, update, delete, and both merge conditions - cover mixed-case identifiers, ordinary literals, comments, and every filter-bearing table operation ## Root cause Lance plans double-quoted tokens as string literals for compatibility. As a result, `"PartyAbbrev" = 'D'` compared two literals and silently evaluated to false instead of filtering the mixed-case column. ## Validation - `cargo fmt --all -- --check` - `cargo test --locked --quiet --features remote -p lancedb expr::sql::tests` - `cargo test --locked --quiet --features remote -p lancedb test_double_quoted_predicates_across_table_operations` - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples` Fixes #2057 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- nodejs/examples/basic.test.ts | 2 +- python/python/tests/docs/test_basic.py | 4 +- python/python/tests/docs/test_guide_tables.py | 4 +- rust/lancedb/src/expr.rs | 1 + rust/lancedb/src/expr/sql.rs | 113 ++++++++++- rust/lancedb/src/materialized_view.rs | 13 +- rust/lancedb/src/materialized_view/refresh.rs | 184 +++++++++++++++--- rust/lancedb/src/query.rs | 169 +++++++++++++++- rust/lancedb/src/remote/table.rs | 60 +++--- rust/lancedb/src/table.rs | 18 +- rust/lancedb/src/table/delete.rs | 3 +- rust/lancedb/src/table/merge.rs | 28 ++- rust/lancedb/src/table/query.rs | 28 ++- rust/lancedb/src/table/update.rs | 15 +- 14 files changed, 570 insertions(+), 72 deletions(-) diff --git a/nodejs/examples/basic.test.ts b/nodejs/examples/basic.test.ts index b56bb2f95..e45fd622a 100644 --- a/nodejs/examples/basic.test.ts +++ b/nodejs/examples/basic.test.ts @@ -170,7 +170,7 @@ test("basic table examples", async () => { // --8<-- [end:create_index] // --8<-- [start:delete_rows] - await tbl.delete('item = "fizz"'); + await tbl.delete("item = 'fizz'"); // --8<-- [end:delete_rows] // --8<-- [start:drop_table] diff --git a/python/python/tests/docs/test_basic.py b/python/python/tests/docs/test_basic.py index 2a824371f..35d7aac10 100644 --- a/python/python/tests/docs/test_basic.py +++ b/python/python/tests/docs/test_basic.py @@ -105,7 +105,7 @@ def test_quickstart(tmp_path): tbl.create_index(num_sub_vectors=1) # --8<-- [end:create_index] # --8<-- [start:delete_rows] - tbl.delete('item = "fizz"') + tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows] # --8<-- [start:drop_table] db.drop_table("my_table") @@ -201,7 +201,7 @@ async def test_quickstart_async(tmp_path): await tbl.create_index("vector") # --8<-- [end:create_index_async] # --8<-- [start:delete_rows_async] - await tbl.delete('item = "fizz"') + await tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows_async] # --8<-- [start:drop_table_async] await db.drop_table("my_table_async") diff --git a/python/python/tests/docs/test_guide_tables.py b/python/python/tests/docs/test_guide_tables.py index 9ae86d167..dab8d43e9 100644 --- a/python/python/tests/docs/test_guide_tables.py +++ b/python/python/tests/docs/test_guide_tables.py @@ -266,7 +266,7 @@ def test_table(): tbl.add(pydantic_model_items) # --8<-- [end:add_table_from_pydantic] # --8<-- [start:delete_row] - tbl.delete('item = "fizz"') + tbl.delete("item = 'fizz'") # --8<-- [end:delete_row] # --8<-- [start:delete_specific_row] data = [ @@ -538,7 +538,7 @@ async def test_table_async(): await async_tbl.add(pydantic_model_items) # --8<-- [end:add_table_async_from_pydantic] # --8<-- [start:delete_row_async] - await async_tbl.delete('item = "fizz"') + await async_tbl.delete("item = 'fizz'") # --8<-- [end:delete_row_async] # --8<-- [start:delete_specific_row_async] data = [ diff --git a/rust/lancedb/src/expr.rs b/rust/lancedb/src/expr.rs index 75cce443d..da69914e3 100644 --- a/rust/lancedb/src/expr.rs +++ b/rust/lancedb/src/expr.rs @@ -19,6 +19,7 @@ mod sql; +pub(crate) use sql::canonicalize_sql_predicate; pub use sql::expr_to_sql_string; use std::sync::Arc; diff --git a/rust/lancedb/src/expr/sql.rs b/rust/lancedb/src/expr/sql.rs index 23b89821a..24a676485 100644 --- a/rust/lancedb/src/expr/sql.rs +++ b/rust/lancedb/src/expr/sql.rs @@ -1,10 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::any::TypeId; + use datafusion_common::ScalarValue; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_expr::Expr; -use datafusion_sql::unparser::{self, dialect::Dialect}; +use datafusion_sql::sqlparser::{ + dialect::{Dialect as SqlParserDialect, GenericDialect}, + tokenizer::{Token, Tokenizer}, +}; +use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect}; /// Unparser dialect that matches the quoting style expected by the Lance SQL /// parser. Lance uses backtick (`` ` ``) as the only delimited-identifier @@ -19,7 +25,7 @@ use datafusion_sql::unparser::{self, dialect::Dialect}; /// lower-case by the SQL parser, which would break case-sensitive schemas). struct LanceSqlDialect; -impl Dialect for LanceSqlDialect { +impl UnparserDialect for LanceSqlDialect { fn identifier_quote_style(&self, identifier: &str) -> Option { let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase()) || !identifier @@ -30,6 +36,61 @@ impl Dialect for LanceSqlDialect { } } +/// Lance's tokenizer dialect with SQL-standard double-quoted identifiers added. +/// +/// Keep this deliberately small: Lance's parser wraps `GenericDialect` and +/// delegates only identifier recognition, leaving every other dialect option at +/// its default. In particular, `/*! ... */` remains an ordinary block comment. +#[derive(Debug, Default)] +struct PredicateDialect(GenericDialect); + +impl SqlParserDialect for PredicateDialect { + fn dialect(&self) -> TypeId { + self.0.dialect() + } + + fn is_identifier_start(&self, ch: char) -> bool { + self.0.is_identifier_start(ch) + } + + fn is_identifier_part(&self, ch: char) -> bool { + self.0.is_identifier_part(ch) + } + + fn is_delimited_identifier_start(&self, ch: char) -> bool { + ch == '"' || ch == '`' + } +} + +/// Canonicalize a raw SQL predicate for Lance's parser. +/// +/// Lance wraps [`GenericDialect`] for identifier recognition while retaining the +/// default dialect behavior for every other lexical option. [`PredicateDialect`] +/// mirrors that contract and additionally recognizes `"` as an identifier +/// delimiter, allowing this function to rewrite only those identifier tokens. +pub fn canonicalize_sql_predicate(predicate: &str) -> crate::Result { + let dialect = PredicateDialect::default(); + let tokens = Tokenizer::new(&dialect, predicate) + .with_unescape(false) + .tokenize() + .map_err(|err| crate::Error::InvalidInput { + message: format!("invalid SQL predicate: {err}"), + })?; + + Ok(tokens + .into_iter() + .map(|token| match token { + Token::Word(word) if word.quote_style == Some('"') => { + // with_unescape(false) retains doubled double quotes. Decode + // those before escaping any backticks for Lance's delimiter. + let identifier = word.value.replace("\"\"", "\"").replace('`', "``"); + format!("`{identifier}`") + } + other => other.to_string(), + }) + .collect()) +} + /// Prefix for placeholder strings inserted in place of binary literals. Chosen /// to be extremely unlikely to occur in user data. const BINARY_PLACEHOLDER_PREFIX: &str = "__lancedb_binary_placeholder_"; @@ -113,3 +174,51 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { } Ok(sql) } + +#[cfg(test)] +mod tests { + use super::canonicalize_sql_predicate; + + #[test] + fn normalizes_double_quoted_identifiers() { + assert_eq!( + canonicalize_sql_predicate(r#""PartyAbbrev" = 'D'"#).unwrap(), + "`PartyAbbrev` = 'D'" + ); + assert_eq!( + canonicalize_sql_predicate(r#""MetaData"."userId" = 5"#).unwrap(), + "`MetaData`.`userId` = 5" + ); + assert_eq!( + canonicalize_sql_predicate(r#""a""b" = 1"#).unwrap(), + "`a\"b` = 1" + ); + } + + #[test] + fn preserves_quotes_inside_literals_and_backticks() { + let filter = r#"name = 'Alice "Ace"' AND `quoted"field` = 1"#; + assert_eq!(canonicalize_sql_predicate(filter).unwrap(), filter); + } + + #[test] + fn preserves_literals_and_comments_using_lance_dialect_rules() { + let predicate = r#"path = '\' AND "PartyAbbrev" = 'D' -- unmatched " in comment"#; + assert_eq!( + canonicalize_sql_predicate(predicate).unwrap(), + r#"path = '\' AND `PartyAbbrev` = 'D' -- unmatched " in comment"# + ); + + let predicate = r#"id = 1 /* unmatched " in block comment */"#; + assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate); + + let predicate = r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#; + assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate); + } + + #[test] + fn rejects_unterminated_double_quoted_identifier() { + let error = canonicalize_sql_predicate(r#""PartyAbbrev = 'D'"#).unwrap_err(); + assert!(matches!(error, crate::Error::InvalidInput { .. })); + } +} diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index b28d52931..08d6c921e 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -170,6 +170,15 @@ pub(crate) fn plan( filter: Option<&str>, limit: Option, ) -> Result<(MaterializedViewDefinition, Vec, Lineage)> { + let filter = filter + .map(crate::expr::canonicalize_sql_predicate) + .transpose() + .map_err(|err| match err { + Error::InvalidInput { message } => Error::InvalidInput { + message: format!("invalid view filter: {message}"), + }, + err => err, + })?; let projections: Vec<(String, String)> = if projections.is_empty() { source_schema .fields() @@ -274,7 +283,7 @@ pub(crate) fn plan( declared.push(output); } - if let Some(filter) = filter { + if let Some(filter) = filter.as_deref() { let expr = planner .parse_filter(filter) .map_err(|e| Error::InvalidInput { @@ -314,7 +323,7 @@ pub(crate) fn plan( .into_iter() .map(|(output, expression)| ViewProjection { output, expression }) .collect(), - filter: filter.map(String::from), + filter, limit, inputs, }; diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 735751c27..b967e81f8 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -46,8 +46,9 @@ use lance_table::format::Fragment; use serde::{Deserialize, Serialize}; use super::{ - INCARNATION_META_KEY, MaterializedViewDefinition, REFRESHED_AT_MS_META_KEY, - SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY, + DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition, + REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY, + definition_to_metadata, }; use crate::database::OpenTableRequest; use crate::table::{NativeTable, NativeTableExt, Table}; @@ -197,8 +198,28 @@ pub(crate) async fn execute_refresh( ), }); } + let definition_changed = + definition.filter != replanned.filter || definition.inputs != replanned.inputs; let definition = &replanned; + // A watermark written for a legacy raw filter certifies the rows that + // filter produced, not the canonical predicate above. Rebuild instead of + // accepting or advancing it, and persist the migrated definition in the + // same metadata commit that certifies the replacement rows. + if definition_changed { + return rebuild( + view_native, + &view_ds, + &source_ds, + source_version, + source_ts, + definition, + true, + expected_incarnation, + ) + .await; + } + let metadata = &view_ds.schema().metadata; let watermark: Option = metadata .get(SOURCE_VERSION_META_KEY) @@ -257,6 +278,7 @@ pub(crate) async fn execute_refresh( source_version, source_ts, definition, + false, expected_incarnation, ) .await @@ -271,6 +293,7 @@ pub(crate) async fn execute_refresh( source_version, source_ts, definition, + false, expected_incarnation, ) .await @@ -683,6 +706,7 @@ async fn incremental( view_ds.clone(), source_version, source_ts, + None, expected_incarnation, ) .await?; @@ -704,6 +728,7 @@ async fn incremental( published, source_version, source_ts, + None, expected_incarnation, ) .await?; @@ -775,6 +800,7 @@ async fn incremental( published, source_version, source_ts, + None, expected_incarnation, ) .await?; @@ -824,12 +850,14 @@ async fn incremental( appended, source_version, source_ts, + None, expected_incarnation, ) .await?; Ok(Some(result)) } +#[allow(clippy::too_many_arguments)] async fn rebuild( view_native: &NativeTable, view_ds: &Dataset, @@ -837,6 +865,7 @@ async fn rebuild( source_version: u64, source_ts: u128, definition: &MaterializedViewDefinition, + persist_definition: bool, expected_incarnation: Option<&str>, ) -> Result { let rows_written = Arc::new(AtomicU64::new(0)); @@ -867,6 +896,7 @@ async fn rebuild( replaced, source_version, source_ts, + persist_definition.then_some(definition), expected_incarnation, ) .await?; @@ -981,6 +1011,7 @@ async fn stamp_watermark( mut dataset: Dataset, source_version: u64, source_ts: u128, + definition: Option<&MaterializedViewDefinition>, expected_incarnation: Option<&str>, ) -> Result { ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?; @@ -993,27 +1024,32 @@ async fn stamp_watermark( .get(INCARNATION_META_KEY) .cloned() .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - dataset - .update_schema_metadata([ - (INCARNATION_META_KEY.to_string(), Some(incarnation)), - ( - SOURCE_VERSION_META_KEY.to_string(), - Some(source_version.to_string()), - ), - ( - SOURCE_VERSION_TS_META_KEY.to_string(), - Some(source_ts.to_string()), - ), - ( - REFRESHED_AT_MS_META_KEY.to_string(), - Some(now_ms().to_string()), - ), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]) - .await?; + let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))]; + if let Some(definition) = definition { + metadata.push(( + DEFINITION_META_KEY.to_string(), + Some(definition_to_metadata(definition)?), + )); + } + metadata.extend([ + ( + SOURCE_VERSION_META_KEY.to_string(), + Some(source_version.to_string()), + ), + ( + SOURCE_VERSION_TS_META_KEY.to_string(), + Some(source_ts.to_string()), + ), + ( + REFRESHED_AT_MS_META_KEY.to_string(), + Some(now_ms().to_string()), + ), + ( + VIEW_VERSION_META_KEY.to_string(), + Some(predicted.to_string()), + ), + ]); + dataset.update_schema_metadata(metadata).await?; let actual = dataset.version().version; if actual != predicted { return Err(Error::Runtime { @@ -1585,6 +1621,106 @@ mod tests { assert_eq!(read(view.table(), "x").await, vec![20, 40]); } + #[tokio::test] + async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!( + ("id", Int32, [1, 2, 3]), + ("PartyAbbrev", Utf8, ["D", "R", "D"]) + ) + .unwrap(); + conn.create_table("src", batch) + .write_options(crate::materialized_view::tests::stable_row_ids()) + .execute() + .await + .unwrap(); + conn.create_materialized_view("democrats", "src") + .select([("id", "id")]) + .only_if(r#""PartyAbbrev" = 'D'"#) + .execute() + .await + .unwrap(); + + // Reopen from schema metadata so these assertions cover the stored + // predicate and lineage, not only the declaration-time handle. + let view = conn.open_materialized_view("democrats").await.unwrap(); + assert_eq!( + view.definition().filter.as_deref(), + Some("`PartyAbbrev` = 'D'") + ); + assert_eq!(view.definition().inputs, ["PartyAbbrev", "id"]); + + let result = view.refresh().execute().await.unwrap(); + assert_eq!(result.rows_written, 2); + assert_eq!(read(view.table(), "id").await, vec![1, 3]); + } + + #[tokio::test] + async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!( + ("id", Int32, [1, 2, 3]), + ("PartyAbbrev", Utf8, ["D", "R", "D"]) + ) + .unwrap(); + conn.create_table("legacy_src", batch) + .write_options(crate::materialized_view::tests::stable_row_ids()) + .execute() + .await + .unwrap(); + let view = conn + .create_materialized_view("legacy_view", "legacy_src") + .select([("id", "id")]) + .only_if(r#""PartyAbbrev" = 'X'"#) + .execute() + .await + .unwrap(); + assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0); + + // Model a definition and up-to-date watermark written before filter + // canonicalization was applied to materialized views. + let mut legacy = view.definition().clone(); + legacy.filter = Some(r#""PartyAbbrev" = 'D'"#.into()); + legacy.inputs = vec!["id".into()]; + let native = view.table().as_native().unwrap(); + let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); + let predicted = dataset.version().version + 1; + dataset + .update_schema_metadata([ + ( + DEFINITION_META_KEY.to_string(), + Some(definition_to_metadata(&legacy).unwrap()), + ), + ( + VIEW_VERSION_META_KEY.to_string(), + Some(predicted.to_string()), + ), + ]) + .await + .unwrap(); + native.dataset.update(dataset); + + let reopened = conn.open_materialized_view("legacy_view").await.unwrap(); + let result = reopened.refresh().execute().await.unwrap(); + assert_eq!(result.mode, RefreshMode::Rebuild); + assert_eq!(result.rows_written, 2); + assert_eq!(read(reopened.table(), "id").await, vec![1, 3]); + + // A fresh handle proves the migration was stored alongside the new + // watermark and therefore happens only once. + let migrated = conn.open_materialized_view("legacy_view").await.unwrap(); + assert_eq!( + migrated.definition().filter.as_deref(), + Some("`PartyAbbrev` = 'D'") + ); + assert_eq!(migrated.definition().inputs, ["PartyAbbrev", "id"]); + assert_eq!( + migrated.refresh().execute().await.unwrap().mode, + RefreshMode::NoOp + ); + assert_eq!(read(migrated.table(), "id").await, vec![1, 3]); + } + #[tokio::test] async fn test_append_refreshes_incrementally() { let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; @@ -2767,7 +2903,7 @@ mod tests { let stale = view_native.dataset.get().await.unwrap().as_ref().clone(); view.table().delete("x = 1").await.unwrap(); - let err = stamp_watermark(view_native, stale, 99, 99, None).await; + let err = stamp_watermark(view_native, stale, 99, 99, None, None).await; assert!(err.is_err()); let result = view.refresh().execute().await.unwrap(); diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 2a1283f22..cd346f42e 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -399,6 +399,9 @@ pub trait QueryBase { /// x > 5 OR y = 'test' /// ``` /// + /// Identifiers may be delimited with SQL-standard double quotes or + /// backticks. String literals must use single quotes. + /// /// Filtering performance can often be improved by creating a scalar index /// on the filter column(s). /// @@ -913,6 +916,17 @@ impl QueryRequest { /// use different representations) the error is recorded and surfaced later /// by [`Self::check_filter`]. pub(crate) fn add_filter(&mut self, new: QueryFilter) { + let new = match new { + QueryFilter::Sql(filter) => match crate::expr::canonicalize_sql_predicate(&filter) { + Ok(filter) => QueryFilter::Sql(filter), + Err(err) => { + self.filter_error = Some(err.to_string()); + return; + } + }, + other => other, + }; + self.filter = Some(match self.filter.take() { None => new, Some(existing) => match and_filters(existing, new) { @@ -1652,8 +1666,8 @@ mod tests { datatypes::{Int32Type, UInt8Type}, }; use arrow_array::{ - FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray, - types::Float32Type, + FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchIterator, + StringArray, cast::AsArray, types::Float32Type, }; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use futures::{StreamExt, TryStreamExt}; @@ -1882,6 +1896,157 @@ mod tests { query.execute().await.unwrap(); } + #[tokio::test] + async fn test_double_quoted_predicates_across_table_operations() { + let tmp_dir = tempdir().unwrap(); + let dataset_path = tmp_dir.path().join("test.lance"); + let uri = dataset_path.to_str().unwrap(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("PartyAbbrev", DataType::Utf8, false), + ArrowField::new("path", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["D", "R", "R", "D"])), + Arc::new(StringArray::from(vec!["\\", "\\", "x", "x"])), + ], + ) + .unwrap(); + + let conn = connect(uri).execute().await.unwrap(); + let table = conn.create_table("parties", batch).execute().await.unwrap(); + let batches = table + .query() + .only_if(r#""PartyAbbrev" = 'D'"#) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!( + table + .count_rows(Some(r#""PartyAbbrev" = 'D'"#.to_string())) + .await + .unwrap(), + 2 + ); + + // Public BaseTable dispatch cannot bypass canonicalization. + let query = AnyQuery::Query(QueryRequest { + filter: Some(QueryFilter::Sql(r#""PartyAbbrev" = 'D'"#.to_string())), + ..Default::default() + }); + let batches = table + .base_table() + .query(&query, Default::default()) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!( + table + .base_table() + .count_rows(Some(crate::table::Filter::Sql( + r#""PartyAbbrev" = 'D'"#.to_string(), + ))) + .await + .unwrap(), + 2 + ); + + for predicate in [ + r#"id = 1 -- unmatched " in a valid SQL comment"#, + r#"id = 1 /* unmatched " in a valid SQL comment */"#, + r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#, + r#"path = '\' AND "PartyAbbrev" = 'D'"#, + ] { + let batches = table + .query() + .only_if(predicate) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + + // The same canonical predicate contract applies to both merge filters. + let source = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["D", "R", "R"])), + Arc::new(StringArray::from(vec!["\\", "\\", "x"])), + ], + ) + .unwrap(); + let mut merge = table.merge_insert(&["id"]); + merge.when_not_matched_by_source_delete(Some(r#""PartyAbbrev" = 'D'"#.to_string())); + let result = table + .base_table() + .merge_insert( + merge, + Box::new(RecordBatchIterator::new(vec![Ok(source)], schema.clone())), + ) + .await + .unwrap(); + assert_eq!(result.num_deleted_rows, 1); + + let source = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["U", "U", "U"])), + Arc::new(StringArray::from(vec!["\\", "\\", "x"])), + ], + ) + .unwrap(); + let mut merge = table.merge_insert(&["id"]); + merge.when_matched_update_all(Some(r#"target."PartyAbbrev" = 'D'"#.to_string())); + merge + .execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema))) + .await + .unwrap(); + assert_eq!( + table + .count_rows(Some(r#""PartyAbbrev" = 'U'"#.to_string())) + .await + .unwrap(), + 1 + ); + + let update = table + .update() + .only_if(r#""PartyAbbrev" = 'R'"#) + .column("PartyAbbrev", "'X'"); + table.base_table().update(update).await.unwrap(); + assert_eq!( + table + .count_rows(Some(r#""PartyAbbrev" = 'X'"#.to_string())) + .await + .unwrap(), + 2 + ); + + let result = table + .base_table() + .delete(crate::table::Predicate::String(r#""PartyAbbrev" = 'X'"#)) + .await + .unwrap(); + assert_eq!(result.num_deleted_rows, 2); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + } + #[tokio::test] async fn test_select_with_transform() { let batches = make_non_empty_batches(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index fad04a098..1afc2615a 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1379,10 +1379,11 @@ impl RemoteTable { query: &AnyQuery, version: Option, ) -> Result> { + let query = query.canonicalized()?; let mut base_body = serde_json::json!({ "version": version }); self.apply_branch_body(&mut base_body); - match query { + match &query { AnyQuery::Query(query) => { let mut body = base_body.clone(); self.apply_query_params(&mut body, query)?; @@ -2491,7 +2492,7 @@ impl BaseTable for RemoteTable { let mut body = if let Some(filter) = filter { let filter_sql = match filter { - Filter::Sql(sql) => sql.clone(), + Filter::Sql(sql) => crate::expr::canonicalize_sql_predicate(&sql)?, Filter::Datafusion(expr) => expr_to_sql_string(&expr)?, }; serde_json::json!({ "predicate": filter_sql, "version": read_snapshot.version }) @@ -2747,7 +2748,8 @@ impl BaseTable for RemoteTable { Ok(final_analyze) } - async fn update(&self, update: UpdateBuilder) -> Result { + async fn update(&self, mut update: UpdateBuilder) -> Result { + update.canonicalize_filter()?; self.check_mutable().await?; let request = self .client @@ -2794,7 +2796,7 @@ impl BaseTable for RemoteTable { async fn delete(&self, predicate: Predicate<'_>) -> Result { self.check_mutable().await?; let predicate_sql = match predicate { - Predicate::String(s) => s.to_string(), + Predicate::String(s) => crate::expr::canonicalize_sql_predicate(s)?, Predicate::Expr(expr) => expr_to_sql_string(expr)?, }; let mut body = serde_json::json!({ "predicate": predicate_sql }); @@ -2851,9 +2853,10 @@ impl BaseTable for RemoteTable { async fn merge_insert( &self, - params: MergeInsertBuilder, + mut params: MergeInsertBuilder, new_data: Box, ) -> Result { + params.canonicalize_filters()?; self.check_mutable().await?; let timeout = params.timeout; @@ -3864,13 +3867,17 @@ mod tests { ); assert_eq!( request.body().unwrap().as_bytes().unwrap(), - br#"{"predicate":"a > 10","version":null}"# + br#"{"predicate":"`A` > 10","version":null}"# ); http::Response::builder().status(200).body("42").unwrap() }); - let count = table.count_rows(Some("a > 10".into())).await.unwrap(); + let count = table + .base_table() + .count_rows(Some(Filter::Sql(r#""A" > 10"#.into()))) + .await + .unwrap(); assert_eq!(count, 42); } @@ -4353,7 +4360,7 @@ mod tests { assert_eq!(expression, "b - 1"); let only_if = value.get("predicate").unwrap().as_str().unwrap(); - assert_eq!(only_if, "b > 10"); + assert_eq!(only_if, "`B` > 10"); } if old_server { @@ -4369,14 +4376,12 @@ mod tests { } }); - let result = table + let update = table .update() .column("a", "a + 1") .column("b", "b - 1") - .only_if("b > 10") - .execute() - .await - .unwrap(); + .only_if(r#""B" > 10"#); + let result = table.base_table().update(update).await.unwrap(); assert_eq!(result.version, if old_server { 0 } else { 43 }); assert_eq!(result.rows_updated, if old_server { 0 } else { 5 }); @@ -4463,10 +4468,10 @@ mod tests { let params = request.url().query_pairs().collect::>(); assert_eq!(params["on"], "some_col"); - assert_eq!(params["when_matched_update_all"], "false"); + assert_eq!(params["when_matched_update_all"], "true"); assert_eq!(params["when_not_matched_insert_all"], "false"); assert_eq!(params["when_not_matched_by_source_delete"], "false"); - assert!(!params.contains_key("when_matched_update_all_filt")); + assert_eq!(params["when_matched_update_all_filt"], "target.`A` > 0"); assert!(!params.contains_key("when_not_matched_by_source_delete_filt")); assert!(!params.contains_key("use_index")); @@ -4483,11 +4488,9 @@ mod tests { } }); - let result = table - .merge_insert(&["some_col"]) - .execute(data) - .await - .unwrap(); + let mut merge = table.merge_insert(&["some_col"]); + merge.when_matched_update_all(Some(r#"target."A" > 0"#.into())); + let result = table.base_table().merge_insert(merge, data).await.unwrap(); assert_eq!(result.version, if old_server { 0 } else { 43 }); if !old_server { @@ -4549,7 +4552,7 @@ mod tests { let body = request.body().unwrap().as_bytes().unwrap(); let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let predicate = body.get("predicate").unwrap().as_str().unwrap(); - assert_eq!(predicate, "id in (1, 2, 3)"); + assert_eq!(predicate, "`ID` in (1, 2, 3)"); if old_server { http::Response::builder() @@ -4567,7 +4570,11 @@ mod tests { } }); - let result = table.delete("id in (1, 2, 3)").await.unwrap(); + let result = table + .base_table() + .delete(Predicate::String(r#""ID" in (1, 2, 3)"#)) + .await + .unwrap(); assert_eq!(result.version, if old_server { 0 } else { 43 }); } @@ -4659,6 +4666,7 @@ mod tests { let body = request.body().unwrap().as_bytes().unwrap(); let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let expected_body = serde_json::json!({ + "filter": "`A` > 0", "k": isize::MAX as usize, "prefilter": true, "vector": [], // Empty vector means no vector query. @@ -4674,9 +4682,13 @@ mod tests { .unwrap() }); + let query = AnyQuery::Query(QueryRequest { + filter: Some(QueryFilter::Sql(r#""A" > 0"#.into())), + ..Default::default() + }); let data = table - .query() - .execute() + .base_table() + .query(&query, Default::default()) .await .unwrap() .collect::>() diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index af8bcb5e2..8436657ca 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1164,7 +1164,10 @@ impl Table { /// /// * `filter` if present, only count rows matching the filter pub async fn count_rows(&self, filter: Option) -> Result { - self.inner.count_rows(filter.map(Filter::Sql)).await + let filter = filter + .map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate).map(Filter::Sql)) + .transpose()?; + self.inner.count_rows(filter).await } /// Names of the blob v2 columns in this table, in declaration order. @@ -1364,7 +1367,13 @@ impl Table { /// # }); /// ``` pub async fn delete(&self, predicate: impl Into>) -> Result { - self.inner.delete(predicate.into()).await + match predicate.into() { + Predicate::String(predicate) => { + let predicate = crate::expr::canonicalize_sql_predicate(predicate)?; + self.inner.delete(Predicate::String(&predicate)).await + } + predicate @ Predicate::Expr(_) => self.inner.delete(predicate).await, + } } /// Create an index on the provided column(s). @@ -3239,7 +3248,10 @@ impl BaseTable for NativeTable { let dataset = self.dataset.get().await?; match filter { None => Ok(dataset.count_rows(None).await?), - Some(Filter::Sql(sql)) => Ok(dataset.count_rows(Some(sql)).await?), + Some(Filter::Sql(sql)) => { + let sql = crate::expr::canonicalize_sql_predicate(&sql)?; + Ok(dataset.count_rows(Some(sql)).await?) + } Some(Filter::Datafusion(_)) => Err(Error::NotSupported { message: "Datafusion filters are not yet supported".to_string(), }), diff --git a/rust/lancedb/src/table/delete.rs b/rust/lancedb/src/table/delete.rs index 8f11ee019..cb1da03ae 100644 --- a/rust/lancedb/src/table/delete.rs +++ b/rust/lancedb/src/table/delete.rs @@ -31,8 +31,9 @@ pub(crate) async fn execute_delete( table.dataset.ensure_mutable()?; match predicate { Predicate::String(s) => { + let predicate = crate::expr::canonicalize_sql_predicate(s)?; let mut dataset = (*table.dataset.get().await?).clone(); - let delete_result = dataset.delete(s).boxed().await?; + let delete_result = dataset.delete(&predicate).boxed().await?; let num_deleted_rows = delete_result.num_deleted_rows; let version = dataset.version().version; table.dataset.update(dataset); diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 3227e3edf..ef2af8fe0 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -220,9 +220,32 @@ impl MergeInsertBuilder { /// /// Returns version and statistics about the merge operation including the number of rows /// inserted, updated, and deleted. - pub async fn execute(self, new_data: Box) -> Result { + pub async fn execute( + mut self, + new_data: Box, + ) -> Result { + self.canonicalize_filters()?; self.table.clone().merge_insert(self, new_data).await } + + pub(crate) fn canonicalize_filters(&mut self) -> Result<()> { + self.when_matched_update_all_filt = + canonicalize_merge_filter(self.when_matched_update_all_filt.take())?; + self.when_not_matched_by_source_delete_filt = + canonicalize_merge_filter(self.when_not_matched_by_source_delete_filt.take())?; + Ok(()) + } +} + +fn canonicalize_merge_filter(filter: Option) -> Result> { + filter + .map(|filter| match filter { + MergeFilter::Sql(predicate) => { + crate::expr::canonicalize_sql_predicate(&predicate).map(MergeFilter::Sql) + } + filter @ MergeFilter::Expr(_) => Ok(filter), + }) + .transpose() } /// Internal implementation of the merge insert logic @@ -230,9 +253,10 @@ impl MergeInsertBuilder { /// This logic was moved from NativeTable::merge_insert to keep table.rs clean. pub(crate) async fn execute_merge_insert( table: &NativeTable, - params: MergeInsertBuilder, + mut params: MergeInsertBuilder, new_data: Box, ) -> Result { + params.canonicalize_filters()?; super::computed_columns::ensure_no_function_bindings_for_mutation( table.schema().await?.as_ref(), "merge_insert", diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 2684ac5e2..b413b2e17 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -44,6 +44,22 @@ impl AnyQuery { Self::VectorQuery(query) => &query.base, } } + + fn base_mut(&mut self) -> &mut QueryRequest { + match self { + Self::Query(query) => query, + Self::VectorQuery(query) => &mut query.base, + } + } + + /// Canonicalize any raw SQL filter immediately before backend dispatch. + pub(crate) fn canonicalized(&self) -> Result { + let mut query = self.clone(); + if let Some(QueryFilter::Sql(predicate)) = &mut query.base_mut().filter { + *predicate = crate::expr::canonicalize_sql_predicate(predicate)?; + } + Ok(query) + } } //Decide between namespace or local @@ -52,15 +68,16 @@ pub async fn execute_query( query: &AnyQuery, options: QueryExecutionOptions, ) -> Result { + let query = query.canonicalized()?; // QueryTable pushdown runs the query server-side, but only on the main // branch: the namespace request carries no branch yet, so a branch handle // must fall through to local execution. - if can_execute_namespace_query(table, query).await? + if can_execute_namespace_query(table, &query).await? && let Some(ref namespace_client) = table.namespace_client { - return execute_namespace_query(table, namespace_client.clone(), query, options).await; + return execute_namespace_query(table, namespace_client.clone(), &query, options).await; } - execute_generic_query(table, query, options).await + execute_generic_query(table, &query, options).await } async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> Result { @@ -135,9 +152,10 @@ pub async fn create_plan( query: &AnyQuery, options: QueryExecutionOptions, ) -> Result> { + let query = query.canonicalized()?; let query = match query { - AnyQuery::VectorQuery(query) => query.clone(), - AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query.clone()), + AnyQuery::VectorQuery(query) => query, + AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query), }; query.base.check_filter()?; diff --git a/rust/lancedb/src/table/update.rs b/rust/lancedb/src/table/update.rs index 98050dfe8..f10594f23 100644 --- a/rust/lancedb/src/table/update.rs +++ b/rust/lancedb/src/table/update.rs @@ -62,22 +62,33 @@ impl UpdateBuilder { } /// Executes the update operation. - pub async fn execute(self) -> Result { + pub async fn execute(mut self) -> Result { if self.columns.is_empty() { Err(Error::InvalidInput { message: "at least one column must be specified in an update operation".to_string(), }) } else { + self.canonicalize_filter()?; self.parent.clone().update(self).await } } + + pub(crate) fn canonicalize_filter(&mut self) -> Result<()> { + self.filter = self + .filter + .take() + .map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate)) + .transpose()?; + Ok(()) + } } /// Internal implementation of the update logic pub(crate) async fn execute_update( table: &NativeTable, - update: UpdateBuilder, + mut update: UpdateBuilder, ) -> Result { + update.canonicalize_filter()?; table.dataset.ensure_mutable()?; // 1. Snapshot the current dataset