fix: lower ST_DWithin in native query planning

This commit is contained in:
Xuanwo
2026-10-27 08:37:35 +08:00
parent b0620b22f6
commit c70b9f16c3
4 changed files with 34 additions and 14 deletions
+8 -5
View File
@@ -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();
+12 -3
View File
@@ -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
+7 -1
View File
@@ -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
{
+7 -5
View File
@@ -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()),