diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index d848676cd..fd904953a 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2006,6 +2006,27 @@ describe("When querying a table", () => { .toArrow({ timeoutMs: 0 }), ).rejects.toThrow("Query timeout"); }); + + it("should execute ST_DWithin filters for plain and vector queries", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("geo_dwithin", [ + { id: 1, x: 0.0, y: 0.0, vector: [0.0, 0.0] }, + { id: 2, x: 1.0, y: 1.0, vector: [1.0, 1.0] }, + { id: 3, x: 3.0, y: 4.0, vector: [3.0, 4.0] }, + ]); + const filter = "ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)"; + + const plain = await table.query().where(filter).toArray(); + expect(plain.map((row) => row.id).sort((a, b) => a - b)).toEqual([1, 2]); + + const vector = await table + .query() + .nearestTo([0.0, 0.0]) + .where(filter) + .limit(2) + .toArray(); + expect(vector.map((row) => row.id).sort((a, b) => a - b)).toEqual([1, 2]); + }); }); describe("Read consistency interval", () => { diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index cef985aa6..8185a79f2 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -300,6 +300,28 @@ async def test_query_to_pandas_kwargs(table, table_async): assert async_df["id"].tolist() == [1, 2] +def test_st_dwithin_plain_and_vector_queries(tmp_db): + data = pa.table( + { + "id": pa.array([1, 2, 3], pa.int32()), + "x": pa.array([0.0, 1.0, 3.0], pa.float64()), + "y": pa.array([0.0, 1.0, 4.0], pa.float64()), + "vector": pa.array( + [[0.0, 0.0], [1.0, 1.0], [3.0, 4.0]], + type=pa.list_(pa.float32(), list_size=2), + ), + } + ) + table = tmp_db.create_table("geo_dwithin", data) + filter = "ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)" + + plain = table.search().where(filter).to_arrow() + assert sorted(plain["id"].to_pylist()) == [1, 2] + + vector = table.search([0.0, 0.0]).where(filter).limit(2).to_arrow() + assert sorted(vector["id"].to_pylist()) == [1, 2] + + @pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"]) def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode): pytest.importorskip("lance") diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 9db74ff00..23310ec5a 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -3211,34 +3211,73 @@ mod tests { } #[cfg(feature = "geo")] - #[tokio::test] - async fn test_st_dwithin_filter() { - let tmp_dir = tempdir().unwrap(); - let uri = tmp_dir.path().to_str().unwrap(); - let batch = RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _), - ( - "x", - Arc::new(Float64Array::from_iter_values([0.0, 1.0, 3.0])) as _, - ), - ( - "y", - Arc::new(Float64Array::from_iter_values([0.0, 1.0, 4.0])) as _, - ), - ]) - .unwrap(); + const ST_DWITHIN_FILTER: &str = "ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)"; - let table = connect(uri) + #[cfg(feature = "geo")] + fn geo_query_batch(rows: &[(i32, f64, f64, [f32; 2])]) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("x", DataType::Float64, false), + ArrowField::new("y", DataType::Float64, false), + ArrowField::new( + "vector", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 2, + ), + false, + ), + ])); + let vectors = FixedSizeListArray::from_iter_primitive::( + rows.iter() + .map(|(_, _, _, vector)| Some(vector.iter().copied().map(Some).collect())), + 2, + ); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values( + rows.iter().map(|(id, _, _, _)| *id), + )), + Arc::new(Float64Array::from_iter_values( + rows.iter().map(|(_, x, _, _)| *x), + )), + Arc::new(Float64Array::from_iter_values( + rows.iter().map(|(_, _, y, _)| *y), + )), + Arc::new(vectors), + ], + ) + .unwrap() + } + + #[cfg(feature = "geo")] + async fn make_geo_query_table(tmp_dir: &tempfile::TempDir) -> Table { + connect(tmp_dir.path().to_str().unwrap()) .execute() .await .unwrap() - .create_table("points", batch) + .create_table( + "points", + geo_query_batch(&[ + (1, 0.0, 0.0, [0.0, 0.0]), + (2, 1.0, 1.0, [1.0, 1.0]), + (3, 3.0, 4.0, [3.0, 4.0]), + ]), + ) .execute() .await - .unwrap(); + .unwrap() + } + + #[cfg(feature = "geo")] + #[tokio::test] + async fn test_st_dwithin_filter() { + let tmp_dir = tempdir().unwrap(); + let table = make_geo_query_table(&tmp_dir).await; let batches = table .query() - .only_if("ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)") + .only_if(ST_DWITHIN_FILTER) .execute() .await .unwrap() @@ -3249,6 +3288,127 @@ mod tests { let ids = batch["id"].as_primitive::(); assert_eq!(ids.values(), &[1, 2]); + + let result = table + .query() + .only_if("ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0) LIMIT 0") + .execute() + .await; + assert!(result.is_err(), "trailing SQL must fail in the real parser"); + } + + #[cfg(feature = "geo")] + #[tokio::test] + async fn test_st_dwithin_vector_filter() { + let tmp_dir = tempdir().unwrap(); + let table = make_geo_query_table(&tmp_dir).await; + let batches = table + .query() + .nearest_to([0.0, 0.0]) + .unwrap() + .only_if(ST_DWITHIN_FILTER) + .limit(2) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + let ids = batch["id"].as_primitive::(); + + assert_eq!(ids.values(), &[1, 2]); + } + + #[cfg(feature = "geo")] + #[tokio::test] + async fn test_st_dwithin_explain_uses_lowered_expression() { + let tmp_dir = tempdir().unwrap(); + let table = make_geo_query_table(&tmp_dir).await; + let plan = table + .query() + .only_if(ST_DWITHIN_FILTER) + .explain_plan(true) + .await + .unwrap(); + let plan = plan.to_ascii_lowercase(); + + assert!(plan.contains("st_distance"), "{plan}"); + assert!(!plan.contains("st_dwithin"), "{plan}"); + } + + #[cfg(feature = "geo")] + #[tokio::test] + async fn test_st_dwithin_filter_survives_take_offsets_routing() { + let tmp_dir = tempdir().unwrap(); + let table = make_geo_query_table(&tmp_dir).await; + let batches = table + .take_offsets(vec![0, 1, 2]) + .only_if(ST_DWITHIN_FILTER) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + let ids = batch["id"].as_primitive::(); + + assert_eq!(ids.values(), &[1, 2]); + } + + #[cfg(feature = "geo")] + #[tokio::test] + async fn test_st_dwithin_filter_on_rte_base_and_active_memtable() { + let tmp_dir = tempdir().unwrap(); + let table = make_geo_query_table(&tmp_dir).await; + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .enable_rte(crate::table::RTEWriteSpec::unsharded()) + .await + .unwrap(); + + let source = geo_query_batch(&[(4, 0.25, 0.25, [0.25, 0.25])]); + let schema = source.schema(); + let mut merge = table.merge_insert(&["id"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + merge + .execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema))) + .await + .unwrap(); + + let plain = table + .query() + .only_if(ST_DWITHIN_FILTER) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let plain = concat_batches(&plain[0].schema(), &plain).unwrap(); + let mut plain_ids = plain["id"].as_primitive::().values().to_vec(); + plain_ids.sort_unstable(); + assert_eq!(plain_ids, vec![1, 4]); + + let vector = table + .query() + .nearest_to([0.0, 0.0]) + .unwrap() + .only_if(ST_DWITHIN_FILTER) + .limit(2) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let vector = concat_batches(&vector[0].schema(), &vector).unwrap(); + let mut vector_ids = vector["id"].as_primitive::().values().to_vec(); + vector_ids.sort_unstable(); + assert_eq!(vector_ids, vec![1, 4]); } fn make_non_empty_batches() -> Box { diff --git a/rust/lancedb/src/query/sql.rs b/rust/lancedb/src/query/sql.rs index 1940908df..2270bb35a 100644 --- a/rust/lancedb/src/query/sql.rs +++ b/rust/lancedb/src/query/sql.rs @@ -8,7 +8,7 @@ use std::{any::TypeId, ops::ControlFlow}; use datafusion_sql::sqlparser::{ ast::{ BinaryOperator, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, ObjectName, - SelectItem, SetExpr, Statement, visit_expressions_mut, + visit_expressions_mut, }, dialect::{Dialect, GenericDialect}, parser::Parser, @@ -39,8 +39,10 @@ impl Dialect for LanceDialect { /// Rewrite `ST_DWithin(a, b, distance)` as `ST_Distance(a, b) <= distance`. /// /// GeoDataFusion does not currently expose `ST_DWithin`, but its definition is -/// exactly this distance comparison. Invalid SQL and unsupported function -/// shapes are left untouched so the query planner can report the usual error. +/// exactly this distance comparison. Only unqualified, ordinary calls with +/// exactly three positional arguments are rewritten. Invalid SQL and other +/// function shapes are left untouched so the query planner can report the +/// usual error. pub(crate) fn rewrite_st_dwithin(filter: &str) -> String { if !filter.to_ascii_lowercase().contains("st_dwithin") { return filter.to_string(); @@ -54,12 +56,23 @@ pub(crate) fn rewrite_st_dwithin(filter: &str) -> String { let Expr::Function(function) = expr else { return ControlFlow::<()>::Continue(()); }; - if !function.name.to_string().eq_ignore_ascii_case("st_dwithin") { + if function.name.0.len() != 1 + || !function.name.to_string().eq_ignore_ascii_case("st_dwithin") + || function.uses_odbc_syntax + || !matches!(&function.parameters, FunctionArguments::None) + || function.filter.is_some() + || function.null_treatment.is_some() + || function.over.is_some() + || !function.within_group.is_empty() + { return ControlFlow::Continue(()); } let FunctionArguments::List(arguments) = &function.args else { return ControlFlow::Continue(()); }; + if arguments.duplicate_treatment.is_some() || !arguments.clauses.is_empty() { + return ControlFlow::Continue(()); + } let [ FunctionArg::Unnamed(FunctionArgExpr::Expr(first)), FunctionArg::Unnamed(FunctionArgExpr::Expr(second)), @@ -94,9 +107,8 @@ pub(crate) fn rewrite_st_dwithin(filter: &str) -> String { } fn parse_filter(filter: &str) -> Option { - let statement = format!("SELECT 1 FROM t WHERE {filter}"); let dialect = LanceDialect::default(); - let mut token_iter = Tokenizer::new(&dialect, &statement) + let mut token_iter = Tokenizer::new(&dialect, filter) .tokenize() .ok()? .into_iter(); @@ -111,29 +123,15 @@ fn parse_filter(filter: &str) -> Option { tokens.push(std::mem::replace(&mut previous, next)); } tokens.push(previous); - // A predicate must parse as one complete expression. Do not let a rewrite - // silently discard anything after a statement delimiter. - if tokens.iter().any(|token| matches!(token, Token::SemiColon)) { - return None; - } - let mut statements = Parser::new(&dialect) - .with_tokens(tokens) - .parse_statements() - .ok()?; - if statements.len() != 1 { + let mut parser = Parser::new(&dialect).with_tokens(tokens); + let expr = parser.parse_expr().ok()?; + // A filter is one expression. Reject trailing clauses or statements rather + // than rewriting a prefix and silently changing what the caller asked for. + if parser.peek_token().token != Token::EOF { return None; } - let statement = statements.pop()?; - if let Statement::Query(query) = statement - && let SetExpr::Select(select) = *query.body - && let Some(expr) = select.selection - && matches!(select.projection.as_slice(), [SelectItem::UnnamedExpr(_)]) - { - Some(expr) - } else { - None - } + Some(expr) } #[cfg(test)] @@ -149,11 +147,30 @@ mod tests { } #[test] - fn leaves_other_filters_unchanged() { - assert_eq!(rewrite_st_dwithin("id == 1"), "id == 1"); - assert_eq!( - rewrite_st_dwithin("ST_DWithin(point, origin)"), - "ST_DWithin(point, origin)" - ); + fn leaves_unsupported_function_shapes_unchanged() { + for filter in [ + "id == 1", + "'ST_DWithin(point, origin, 1.0)'", + "geo.ST_DWithin(point, origin, 1.0)", + "ST_DWithin(point, origin)", + "ST_DWithin(point, origin, 1.0, extra)", + "ST_DWithin(DISTINCT point, origin, 1.0)", + "ST_DWithin(point => point, origin => origin, distance => 1.0)", + "ST_DWithin(point, origin, 1.0) OVER ()", + ] { + assert_eq!(rewrite_st_dwithin(filter), filter); + } + } + + #[test] + fn leaves_malformed_or_multi_expression_filters_unchanged() { + for filter in [ + "ST_DWithin(point, origin, 1.0", + "ST_DWithin(point, origin, 1.0) LIMIT 0", + "ST_DWithin(point, origin, 1.0) ORDER BY id", + "ST_DWithin(point, origin, 1.0); SELECT 1", + ] { + assert_eq!(rewrite_st_dwithin(filter), filter); + } } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 181ce61f9..7c646d598 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -5669,6 +5669,54 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + #[tokio::test] + async fn test_query_preserves_st_dwithin_filter() { + const FILTER: &str = "ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)"; + let expected_data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .unwrap(); + let response_data = expected_data.clone(); + let table = Table::new_with_handler("my_table", move |request| { + let body = request_body_json(&request); + assert_eq!(body["filter"], FILTER); + + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&response_data)) + .unwrap() + }); + + let batches = table + .query() + .only_if(FILTER) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(batch["id"].as_primitive::().values(), &[1, 2]); + + let batches = table + .query() + .nearest_to([0.0, 0.0]) + .unwrap() + .only_if(FILTER) + .limit(2) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(batch["id"].as_primitive::().values(), &[1, 2]); + } + fn schema_describe_response(schema: &Schema) -> http::Response { let json_schema = JsonSchema::try_from(schema).unwrap(); http::Response::builder() diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 0212fdbbe..fc8e76e81 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -1153,7 +1153,7 @@ mod tests { use futures::TryStreamExt; use lance_arrow::FixedSizeListArrayExt; use std::sync::{ - Arc, + Arc, Mutex, atomic::{AtomicUsize, Ordering}, }; @@ -1368,6 +1368,123 @@ mod tests { assert!(ns_request.vector.single_vector.as_ref().unwrap().is_empty()); } + #[derive(Debug)] + struct CapturingNamespaceClient { + filters: Mutex>>, + response: bytes::Bytes, + } + + #[async_trait::async_trait] + impl LanceNamespace for CapturingNamespaceClient { + fn namespace_id(&self) -> String { + "capturing".to_string() + } + + async fn query_table(&self, request: NsQueryTableRequest) -> lance::Result { + self.filters.lock().unwrap().push(request.filter); + Ok(self.response.clone()) + } + } + + fn namespace_query_response() -> bytes::Bytes { + use arrow_ipc::writer::FileWriter; + use arrow_schema::{DataType, Field, Schema}; + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef], + ) + .unwrap(); + let mut response = Vec::new(); + { + let mut writer = FileWriter::try_new(&mut response, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + bytes::Bytes::from(response) + } + + #[tokio::test] + async fn test_namespace_query_pushdown_preserves_st_dwithin_filter() { + use crate::connect; + use arrow_schema::{DataType, Field, Schema}; + + const FILTER: &str = "ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)"; + + let vectors = fixed_size_list_array(vec![0.0, 0.0, 1.0, 1.0], 2); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("x", DataType::Float64, false), + Field::new("y", DataType::Float64, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(arrow_array::Float64Array::from(vec![0.0, 1.0])), + Arc::new(arrow_array::Float64Array::from(vec![0.0, 1.0])), + Arc::new(vectors), + ], + ) + .unwrap(); + let table = connect("memory://") + .execute() + .await + .unwrap() + .create_table("namespace_geo", batch) + .execute() + .await + .unwrap(); + let namespace_client = Arc::new(CapturingNamespaceClient { + filters: Mutex::new(Vec::new()), + response: namespace_query_response(), + }); + let mut native_table = table.as_native().unwrap().clone(); + native_table.namespace_client = Some(namespace_client.clone()); + native_table + .pushdown_operations + .insert(NamespaceClientPushdownOperation::QueryTable); + + for query in [ + AnyQuery::Query(QueryRequest { + filter: Some(QueryFilter::Sql(FILTER.to_string())), + ..Default::default() + }), + AnyQuery::VectorQuery(VectorQueryRequest { + base: QueryRequest { + filter: Some(QueryFilter::Sql(FILTER.to_string())), + limit: Some(2), + ..Default::default() + }, + column: Some("vector".to_string()), + query_vector: vec![Arc::new(Float32Array::from(vec![0.0, 0.0])) as ArrayRef], + ..Default::default() + }), + ] { + let batches = native_table + .query(&query, QueryExecutionOptions::default()) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!( + batch["id"] + .as_primitive::() + .values(), + &[1, 2] + ); + } + + assert_eq!( + namespace_client.filters.lock().unwrap().as_slice(), + &[Some(FILTER.to_string()), Some(FILTER.to_string())] + ); + } + #[test] fn test_convert_to_namespace_query_plain_query_without_limit() { let any_query = AnyQuery::Query(QueryRequest::default());