mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-28 14:36:13 +00:00
fix: lower ST_DWithin in native query planning
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
{
|
||||
|
||||
@@ -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()),
|
||||
|
||||
Reference in New Issue
Block a user