From c70b9f16c3cca5c3457b3d4f8937a0939ca4d748 Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Tue, 27 Oct 2026 08:37:35 +0800
Subject: [PATCH] fix: lower ST_DWithin in native query planning

---
 rust/lancedb/src/query.rs           | 13 ++++++++-----
 rust/lancedb/src/query/sql.rs       | 15 ++++++++++++---
 rust/lancedb/src/table/query.rs     |  8 +++++++-
 rust/lancedb/src/table/query/lsm.rs | 12 +++++++-----
 4 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs
index 09485e414..9db74ff00 100644
--- a/rust/lancedb/src/query.rs
+++ b/rust/lancedb/src/query.rs
@@ -48,8 +48,8 @@ use crate::{
 };
 
 mod hybrid;
-pub(crate) mod wal_fusion; // WAL-PK-FUSION: delete with the module.
 pub(crate) mod sql;
+pub(crate) mod wal_fusion; // WAL-PK-FUSION: delete with the module.
 
 pub(crate) const DEFAULT_TOP_K: usize = 10;
 
@@ -582,8 +582,8 @@ impl<T: HasQuery> QueryBase for T {
     }
 
     fn only_if(mut self, filter: impl AsRef<str>) -> Self {
-        let filter = sql::rewrite_st_dwithin(filter.as_ref());
-        self.mut_query().add_filter(QueryFilter::Sql(filter));
+        self.mut_query()
+            .add_filter(QueryFilter::Sql(filter.as_ref().to_string()));
         self
     }
 
@@ -2403,9 +2403,11 @@ mod tests {
         compute::concat_batches,
         datatypes::{Int32Type, UInt8Type},
     };
+    #[cfg(feature = "geo")]
+    use arrow_array::Float64Array;
     use arrow_array::{
-        FixedSizeListArray, Float32Array, Float64Array, Int32Array, RecordBatch,
-        RecordBatchIterator, 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 datafusion_physical_plan::display::DisplayableExecutionPlan;
@@ -3208,6 +3210,7 @@ mod tests {
         assert!(result.is_err());
     }
 
+    #[cfg(feature = "geo")]
     #[tokio::test]
     async fn test_st_dwithin_filter() {
         let tmp_dir = tempdir().unwrap();
diff --git a/rust/lancedb/src/query/sql.rs b/rust/lancedb/src/query/sql.rs
index 80f9ceefe..1940908df 100644
--- a/rust/lancedb/src/query/sql.rs
+++ b/rust/lancedb/src/query/sql.rs
@@ -41,7 +41,7 @@ impl Dialect for LanceDialect {
 /// 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.
-pub(super) fn rewrite_st_dwithin(filter: &str) -> String {
+pub(crate) fn rewrite_st_dwithin(filter: &str) -> String {
     if !filter.to_ascii_lowercase().contains("st_dwithin") {
         return filter.to_string();
     }
@@ -111,11 +111,20 @@ 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 statement = Parser::new(&dialect)
+    let mut statements = Parser::new(&dialect)
         .with_tokens(tokens)
-        .parse_statement()
+        .parse_statements()
         .ok()?;
+    if statements.len() != 1 {
+        return None;
+    }
+    let statement = statements.pop()?;
     if let Statement::Query(query) = statement
         && let SetExpr::Select(select) = *query.body
         && let Some(expr) = select.selection
diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs
index dbb3dbc68..0212fdbbe 100644
--- a/rust/lancedb/src/table/query.rs
+++ b/rust/lancedb/src/table/query.rs
@@ -162,7 +162,13 @@ pub async fn create_plan(
     query: &AnyQuery,
     options: QueryExecutionOptions,
 ) -> Result<Arc<dyn ExecutionPlan>> {
-    let query = query.canonicalized()?;
+    let mut query = query.canonicalized()?;
+    // Remote and namespace queries retain their original SQL. Native execution
+    // lowers ST_DWithin after canonicalization, before take/RTE routing clones
+    // or consumes the query.
+    if let Some(QueryFilter::Sql(predicate)) = &mut query.base_mut().filter {
+        *predicate = crate::query::sql::rewrite_st_dwithin(predicate);
+    }
     if let AnyQuery::Query(request) = &query
         && let Some(offsets) = &request.take_offsets
     {
diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs
index 7af52e572..8c0f081fc 100644
--- a/rust/lancedb/src/table/query/lsm.rs
+++ b/rust/lancedb/src/table/query/lsm.rs
@@ -407,11 +407,13 @@ fn base_scanner(
                 let df_schema = schema.clone().to_dfschema().map_err(|error| {
                     enrich_filter_error(error, &schema, "Failed to create DFSchema")
                 })?;
-                let expr = SessionContext::new()
-                    .parse_sql_expr(sql, &df_schema)
-                    .map_err(|error| {
-                        enrich_filter_error(error, &schema, "Failed to parse filter expression")
-                    })?;
+                let ctx = SessionContext::new();
+                // Match the native scanner's SQL environment so the same
+                // lowered filter functions resolve on the RTE read path.
+                lance_datafusion::udf::register_functions(&ctx);
+                let expr = ctx.parse_sql_expr(sql, &df_schema).map_err(|error| {
+                    enrich_filter_error(error, &schema, "Failed to parse filter expression")
+                })?;
                 scanner.filter_expr(expr)
             }
             QueryFilter::Datafusion(expr) => scanner.filter_expr(expr.clone()),