fix: keep ST_DWithin SQL intact across query routes

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", () => {

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")

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::<Float32Type, _, _>(
+            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::<Int32Type>();
 
         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::<Vec<_>>()
+            .await
+            .unwrap();
+        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+        let ids = batch["id"].as_primitive::<Int32Type>();
+
+        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::<Vec<_>>()
+            .await
+            .unwrap();
+        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+        let ids = batch["id"].as_primitive::<Int32Type>();
+
+        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::<Vec<_>>()
+            .await
+            .unwrap();
+        let plain = concat_batches(&plain[0].schema(), &plain).unwrap();
+        let mut plain_ids = plain["id"].as_primitive::<Int32Type>().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::<Vec<_>>()
+            .await
+            .unwrap();
+        let vector = concat_batches(&vector[0].schema(), &vector).unwrap();
+        let mut vector_ids = vector["id"].as_primitive::<Int32Type>().values().to_vec();
+        vector_ids.sort_unstable();
+        assert_eq!(vector_ids, vec![1, 4]);
     }
 
     fn make_non_empty_batches() -> Box<dyn arrow_array::RecordBatchReader + Send> {

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<Expr> {
-    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<Expr> {
         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);
+        }
     }
 }

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::<Vec<_>>()
+            .await
+            .unwrap();
+        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+        assert_eq!(batch["id"].as_primitive::<Int32Type>().values(), &[1, 2]);
+
+        let batches = table
+            .query()
+            .nearest_to([0.0, 0.0])
+            .unwrap()
+            .only_if(FILTER)
+            .limit(2)
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+        assert_eq!(batch["id"].as_primitive::<Int32Type>().values(), &[1, 2]);
+    }
+
     fn schema_describe_response(schema: &Schema) -> http::Response<String> {
         let json_schema = JsonSchema::try_from(schema).unwrap();
         http::Response::builder()

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<Vec<Option<String>>>,
+        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<bytes::Bytes> {
+            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::<Vec<_>>()
+                .await
+                .unwrap();
+            let batch = arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap();
+            assert_eq!(
+                batch["id"]
+                    .as_primitive::<arrow::datatypes::Int32Type>()
+                    .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());