From aea22492d08214f61bb99925539ba8f42e6a0133 Mon Sep 17 00:00:00 2001 From: Kristin Cowalcijk Date: Fri, 13 Feb 2026 20:10:38 +0800 Subject: [PATCH 1/5] fix(rust/sedona-spatial-join): prevent filter pushdown past KNN joins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KNN joins have different semantics than regular spatial joins — pushing filters to the object (build) side changes which objects are the k nearest neighbors, producing incorrect results. Add KnnJoinEarlyRewrite optimizer rule that converts KNN joins to SpatialJoinPlanNode extension nodes before DataFusion's PushDownFilter rule runs, since extension nodes naturally block filter pushdown via prevent_predicate_push_down_columns(). Rule ordering: MergeSpatialProjectionIntoJoin → KnnJoinEarlyRewrite → PushDownFilter → ... → SpatialJoinLogicalRewrite (for non-KNN joins). Closes #605 --- .../src/planner/optimizer.rs | 254 +++++++++++++++--- .../tests/spatial_join_integration.rs | 105 ++++++++ 2 files changed, 317 insertions(+), 42 deletions(-) diff --git a/rust/sedona-spatial-join/src/planner/optimizer.rs b/rust/sedona-spatial-join/src/planner/optimizer.rs index 1edbc2d3da..cce68f5724 100644 --- a/rust/sedona-spatial-join/src/planner/optimizer.rs +++ b/rust/sedona-spatial-join/src/planner/optimizer.rs @@ -20,7 +20,7 @@ use crate::planner::logical_plan_node::SpatialJoinPlanNode; use crate::planner::spatial_expr_utils::collect_spatial_predicate_names; use crate::planner::spatial_expr_utils::is_spatial_predicate; use datafusion::execution::session_state::SessionStateBuilder; -use datafusion::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule}; +use datafusion::optimizer::{ApplyOrder, Optimizer, OptimizerConfig, OptimizerRule}; use datafusion_common::tree_node::Transformed; use datafusion_common::NullEquality; use datafusion_common::Result; @@ -29,20 +29,126 @@ use datafusion_expr::{BinaryExpr, Expr, Operator}; use datafusion_expr::{Filter, Join, JoinType, LogicalPlan}; use sedona_common::option::SedonaOptions; -/// Register only the logical spatial join optimizer rule. +/// Register the logical spatial join optimizer rules. /// -/// This enables building `Join(filter=...)` from patterns like `Filter(CrossJoin)`. -/// It intentionally does not register any physical plan rewrite rules. +/// This inserts rules at specific positions relative to DataFusion's built-in `PushDownFilter` +/// rule to ensure correct semantics for KNN joins: +/// +/// - `MergeSpatialProjectionIntoJoin` and `KnnJoinEarlyRewrite` are inserted *before* +/// `PushDownFilter` so that KNN joins are converted to `SpatialJoinPlanNode` extension nodes +/// before filter pushdown runs. Extension nodes naturally block filter pushdown via +/// `prevent_predicate_push_down_columns()`, preventing incorrect pushdown to the build side +/// of KNN joins. +/// +/// - `SpatialJoinLogicalRewrite` is appended at the end so that non-KNN spatial joins still +/// benefit from filter pushdown before being converted to extension nodes. pub(crate) fn register_spatial_join_logical_optimizer( - session_state_builder: SessionStateBuilder, + mut session_state_builder: SessionStateBuilder, ) -> SessionStateBuilder { + let optimizer = session_state_builder + .optimizer() + .get_or_insert_with(Optimizer::new); + + // Find PushDownFilter position by name + let push_down_pos = optimizer + .rules + .iter() + .position(|r| r.name() == "push_down_filter") + .expect("PushDownFilter rule not found in default optimizer rules"); + + // Insert KNN-specific rules BEFORE PushDownFilter. + // MergeSpatialProjectionIntoJoin must come first because it creates the Join(filter=...) + // nodes that KnnJoinEarlyRewrite then converts to SpatialJoinPlanNode. + optimizer + .rules + .insert(push_down_pos, Arc::new(KnnJoinEarlyRewrite)); + optimizer + .rules + .insert(push_down_pos, Arc::new(MergeSpatialProjectionIntoJoin)); + + // Append SpatialJoinLogicalRewrite at the end so non-KNN joins benefit from filter pushdown. + optimizer.rules.push(Arc::new(SpatialJoinLogicalRewrite)); + session_state_builder - .with_optimizer_rule(Arc::new(MergeSpatialProjectionIntoJoin)) - .with_optimizer_rule(Arc::new(SpatialJoinLogicalRewrite)) } -/// Logical optimizer rule that enables spatial join planning. +/// Early optimizer rule that converts KNN joins to `SpatialJoinPlanNode` extension nodes +/// *before* DataFusion's `PushDownFilter` runs. +/// +/// This prevents `PushDownFilter` from pushing filters to the build (object) side of KNN joins, +/// which would change which objects are the K nearest neighbors and produce incorrect results. /// -/// This rule turns eligible `Join(filter=...)` nodes into a `SpatialJoinPlanNode` extension. +/// Extension nodes naturally block filter pushdown because their default +/// `prevent_predicate_push_down_columns()` returns all columns. +/// +/// Handles two patterns that DataFusion's SQL planner creates: +/// +/// 1. `Join(filter=ST_KNN(...))` — when the ON clause has only the spatial predicate +/// 2. `Filter(ST_KNN(...), Join(on=[...]))` — when the ON clause also has equi-join conditions, +/// the SQL planner separates equi-conditions into `on` and the spatial predicate into a Filter +#[derive(Default, Debug)] +struct KnnJoinEarlyRewrite; + +impl OptimizerRule for KnnJoinEarlyRewrite { + fn name(&self) -> &str { + "knn_join_early_rewrite" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::BottomUp) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result> { + let options = config.options(); + let Some(ext) = options.extensions.get::() else { + return Ok(Transformed::no(plan)); + }; + if !ext.spatial_join.enable { + return Ok(Transformed::no(plan)); + } + + // Case 1: Join(filter=ST_KNN(...)) + if let LogicalPlan::Join(join) = &plan { + if let Some(filter) = join.filter.as_ref() { + let names = collect_spatial_predicate_names(filter); + if names.contains("st_knn") { + return rewrite_join_to_spatial_join_plan_node(join); + } + } + } + + // Case 2: Filter(ST_KNN(...), Join(on=[...])) + // When the ON clause has equi-conditions AND ST_KNN, DataFusion's SQL planner puts + // the equi-conditions in Join.on and the spatial predicate in a Filter above. + if let LogicalPlan::Filter(filter_node) = &plan { + let names = collect_spatial_predicate_names(&filter_node.predicate); + if names.contains("st_knn") { + if let LogicalPlan::Join(join) = filter_node.input.as_ref() { + return rewrite_filter_join_to_spatial_join_plan_node( + &filter_node.predicate, + join, + ); + } + } + } + + Ok(Transformed::no(plan)) + } +} + +/// Logical optimizer rule that converts non-KNN spatial joins to `SpatialJoinPlanNode`. +/// +/// This rule runs *after* `PushDownFilter` so that non-KNN spatial joins benefit from +/// filter pushdown before being converted to extension nodes. +/// +/// KNN joins are skipped here because they are already handled by `KnnJoinEarlyRewrite`. #[derive(Default, Debug)] struct SpatialJoinLogicalRewrite; @@ -86,46 +192,110 @@ impl OptimizerRule for SpatialJoinLogicalRewrite { return Ok(Transformed::no(plan)); } - // Join with with equi-join condition and spatial join condition. Only handle it - // when the join condition contains ST_KNN. KNN join is not a regular join and - // ST_KNN is also not a regular predicate. It must be handled by our spatial join exec. - if !join.on.is_empty() && !spatial_predicate_names.contains("st_knn") { + // KNN joins are already handled by KnnJoinEarlyRewrite — skip them here. + if spatial_predicate_names.contains("st_knn") { return Ok(Transformed::no(plan)); } - // Build new filter expression including equi-join conditions - let filter = filter.clone(); - let eq_op = if join.null_equality == NullEquality::NullEqualsNothing { - Operator::Eq - } else { - Operator::IsNotDistinctFrom - }; - let filter = join.on.iter().fold(filter, |acc, (l, r)| { - let eq_expr = Expr::BinaryExpr(BinaryExpr::new( - Box::new(l.clone()), - eq_op, - Box::new(r.clone()), - )); - Expr::and(acc, eq_expr) - }); - - let schema = Arc::clone(&join.schema); - let node = SpatialJoinPlanNode { - left: join.left.as_ref().clone(), - right: join.right.as_ref().clone(), - join_type: join.join_type, - filter, - schema, - join_constraint: join.join_constraint, - null_equality: join.null_equality, - }; + // Non-KNN spatial joins with equi-join conditions are not supported. + if !join.on.is_empty() { + return Ok(Transformed::no(plan)); + } - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))) + rewrite_join_to_spatial_join_plan_node(join) } } +/// Shared helper: convert a `Join` node (with spatial predicate in `filter`) to a +/// `SpatialJoinPlanNode`, folding any equi-join `on` conditions into the filter expression. +fn rewrite_join_to_spatial_join_plan_node(join: &Join) -> Result> { + let filter = join + .filter + .as_ref() + .expect("join filter must be present") + .clone(); + + let eq_op = if join.null_equality == NullEquality::NullEqualsNothing { + Operator::Eq + } else { + Operator::IsNotDistinctFrom + }; + let filter = join.on.iter().fold(filter, |acc, (l, r)| { + let eq_expr = Expr::BinaryExpr(BinaryExpr::new( + Box::new(l.clone()), + eq_op, + Box::new(r.clone()), + )); + Expr::and(acc, eq_expr) + }); + + let schema = Arc::clone(&join.schema); + let node = SpatialJoinPlanNode { + left: join.left.as_ref().clone(), + right: join.right.as_ref().clone(), + join_type: join.join_type, + filter, + schema, + join_constraint: join.join_constraint, + null_equality: join.null_equality, + }; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(node), + }))) +} + +/// Convert a `Filter(predicate, Join(on=[...]))` pattern to a `SpatialJoinPlanNode`. +/// +/// This handles the case where DataFusion's SQL planner separates the ON clause into: +/// - equi-join conditions in `Join.on` +/// - the spatial predicate (+ any remaining conjuncts) in a `Filter` above the `Join` +/// +/// All predicates (filter predicate + join's existing filter + equi-join on conditions) are +/// combined into a single filter expression in the `SpatialJoinPlanNode`. +fn rewrite_filter_join_to_spatial_join_plan_node( + filter_predicate: &Expr, + join: &Join, +) -> Result> { + // Start with the filter predicate from the Filter node above + let mut combined = filter_predicate.clone(); + + // Add any existing join filter + if let Some(join_filter) = &join.filter { + combined = Expr::and(combined, join_filter.clone()); + } + + // Fold equi-join on conditions + let eq_op = if join.null_equality == NullEquality::NullEqualsNothing { + Operator::Eq + } else { + Operator::IsNotDistinctFrom + }; + combined = join.on.iter().fold(combined, |acc, (l, r)| { + let eq_expr = Expr::BinaryExpr(BinaryExpr::new( + Box::new(l.clone()), + eq_op, + Box::new(r.clone()), + )); + Expr::and(acc, eq_expr) + }); + + let schema = Arc::clone(&join.schema); + let node = SpatialJoinPlanNode { + left: join.left.as_ref().clone(), + right: join.right.as_ref().clone(), + join_type: join.join_type, + filter: combined, + schema, + join_constraint: join.join_constraint, + null_equality: join.null_equality, + }; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(node), + }))) +} + /// Logical optimizer rule that enables spatial join planning. /// /// This rule turns eligible `Filter(Join(filter=...))` nodes into a `Join(filter=...)` node, diff --git a/rust/sedona-spatial-join/tests/spatial_join_integration.rs b/rust/sedona-spatial-join/tests/spatial_join_integration.rs index 671ce2c804..e888bc54d1 100644 --- a/rust/sedona-spatial-join/tests/spatial_join_integration.rs +++ b/rust/sedona-spatial-join/tests/spatial_join_integration.rs @@ -27,6 +27,7 @@ use datafusion::{ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::Result; use datafusion_expr::{ColumnarValue, JoinType}; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::NestedLoopJoinExec; use datafusion_physical_plan::ExecutionPlan; use geo::{Distance, Euclidean}; @@ -1368,3 +1369,107 @@ async fn test_knn_join_include_tie_breakers( Ok(()) } + +/// Recursively check whether any node in the physical plan tree is a `FilterExec`. +fn subtree_contains_filter_exec(plan: &Arc) -> bool { + let mut found = false; + plan.apply(|node| { + if node.as_any().downcast_ref::().is_some() { + found = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("failed to walk plan"); + found +} + +/// Create a session context with two small tables for filter-pushdown tests. +/// +/// L(id INT, x DOUBLE) and R(id INT, x DOUBLE) each with 10 rows. +/// Geometry is constructed in SQL via ST_Point so no geometry column exists on the table itself. +async fn plan_for_filter_pushdown_test(sql: &str) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("x", DataType::Float64, false), + ])); + + let ids: Vec = (0..10).collect(); + let xs: Vec = (0..10).map(|i| i as f64).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow_array::Int32Array::from(ids)), + Arc::new(arrow_array::Float64Array::from(xs)), + ], + )?; + + let options = SpatialJoinOptions::default(); + let ctx = setup_context(Some(options), 100)?; + let mem_l: Arc = Arc::new(MemTable::try_new( + schema.clone(), + vec![vec![batch.clone()]], + )?); + let mem_r: Arc = + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch]])?); + ctx.register_table("L", mem_l)?; + ctx.register_table("R", mem_r)?; + + let df = ctx.sql(sql).await?; + df.create_physical_plan().await +} + +/// Verify that a filter on the *object* (build / right) side of a KNN join is NOT pushed down +/// into the build side subtree. +/// +/// If `PushDownFilter` incorrectly pushes `R.id > 5` below the spatial join, the set of objects +/// considered for KNN changes, yielding wrong nearest-neighbor results. +#[tokio::test] +async fn test_knn_join_object_side_filter_not_pushed_down() -> Result<()> { + let sql = "SELECT L.id, R.id \ + FROM L JOIN R ON ST_KNN(ST_Point(L.x, 0), ST_Point(R.x, 1), 3, false) \ + WHERE R.id > 5"; + let plan = plan_for_filter_pushdown_test(sql).await?; + + let spatial_joins = collect_spatial_join_exec(&plan)?; + assert_eq!( + spatial_joins.len(), + 1, + "expected exactly one SpatialJoinExec" + ); + let sj = spatial_joins[0]; + + // The build (right / object) side must NOT have a FilterExec pushed into it. + assert!( + !subtree_contains_filter_exec(&sj.right), + "FilterExec should NOT be pushed into the object (right/build) side of a KNN join" + ); + + Ok(()) +} + +/// Verify that for a *non-KNN* spatial join, a filter on the build side IS pushed down +/// (the normal, desirable behaviour). +#[tokio::test] +async fn test_non_knn_join_object_side_filter_is_pushed_down() -> Result<()> { + let sql = "SELECT L.id, R.id \ + FROM L JOIN R ON ST_Intersects(ST_Buffer(ST_Point(L.x, 0), 1.5), ST_Point(R.x, 1)) \ + WHERE R.id > 5"; + let plan = plan_for_filter_pushdown_test(sql).await?; + + let spatial_joins = collect_spatial_join_exec(&plan)?; + assert_eq!( + spatial_joins.len(), + 1, + "expected exactly one SpatialJoinExec" + ); + let sj = spatial_joins[0]; + + // For non-KNN joins, the filter SHOULD be pushed down to the build side. + assert!( + subtree_contains_filter_exec(&sj.right), + "FilterExec should be pushed into the object (right/build) side of a non-KNN spatial join" + ); + + Ok(()) +} From d6bc0d777db222bc68121d231aaf21bb81dc6307 Mon Sep 17 00:00:00 2001 From: Kristin Cowalcijk Date: Fri, 13 Feb 2026 21:47:18 +0800 Subject: [PATCH 2/5] Manually tuned the implementation --- .../src/planner/optimizer.rs | 106 ++------- .../src/planner/spatial_expr_utils.rs | 40 ++-- .../tests/spatial_join_integration.rs | 223 +++++++++--------- 3 files changed, 156 insertions(+), 213 deletions(-) diff --git a/rust/sedona-spatial-join/src/planner/optimizer.rs b/rust/sedona-spatial-join/src/planner/optimizer.rs index cce68f5724..b4721728c6 100644 --- a/rust/sedona-spatial-join/src/planner/optimizer.rs +++ b/rust/sedona-spatial-join/src/planner/optimizer.rs @@ -18,7 +18,6 @@ use std::sync::Arc; use crate::planner::logical_plan_node::SpatialJoinPlanNode; use crate::planner::spatial_expr_utils::collect_spatial_predicate_names; -use crate::planner::spatial_expr_utils::is_spatial_predicate; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::optimizer::{ApplyOrder, Optimizer, OptimizerConfig, OptimizerRule}; use datafusion_common::tree_node::Transformed; @@ -34,7 +33,7 @@ use sedona_common::option::SedonaOptions; /// This inserts rules at specific positions relative to DataFusion's built-in `PushDownFilter` /// rule to ensure correct semantics for KNN joins: /// -/// - `MergeSpatialProjectionIntoJoin` and `KnnJoinEarlyRewrite` are inserted *before* +/// - `MergeSpatialFilterIntoJoin` and `KnnJoinEarlyRewrite` are inserted *before* /// `PushDownFilter` so that KNN joins are converted to `SpatialJoinPlanNode` extension nodes /// before filter pushdown runs. Extension nodes naturally block filter pushdown via /// `prevent_predicate_push_down_columns()`, preventing incorrect pushdown to the build side @@ -57,20 +56,21 @@ pub(crate) fn register_spatial_join_logical_optimizer( .expect("PushDownFilter rule not found in default optimizer rules"); // Insert KNN-specific rules BEFORE PushDownFilter. - // MergeSpatialProjectionIntoJoin must come first because it creates the Join(filter=...) + // MergeSpatialFilterIntoJoin must come first because it creates the Join(filter=...) // nodes that KnnJoinEarlyRewrite then converts to SpatialJoinPlanNode. optimizer .rules .insert(push_down_pos, Arc::new(KnnJoinEarlyRewrite)); optimizer .rules - .insert(push_down_pos, Arc::new(MergeSpatialProjectionIntoJoin)); + .insert(push_down_pos, Arc::new(MergeSpatialFilterIntoJoin)); // Append SpatialJoinLogicalRewrite at the end so non-KNN joins benefit from filter pushdown. optimizer.rules.push(Arc::new(SpatialJoinLogicalRewrite)); session_state_builder } + /// Early optimizer rule that converts KNN joins to `SpatialJoinPlanNode` extension nodes /// *before* DataFusion's `PushDownFilter` runs. /// @@ -114,7 +114,7 @@ impl OptimizerRule for KnnJoinEarlyRewrite { return Ok(Transformed::no(plan)); } - // Case 1: Join(filter=ST_KNN(...)) + // Join(filter=ST_KNN(...)) if let LogicalPlan::Join(join) = &plan { if let Some(filter) = join.filter.as_ref() { let names = collect_spatial_predicate_names(filter); @@ -124,21 +124,6 @@ impl OptimizerRule for KnnJoinEarlyRewrite { } } - // Case 2: Filter(ST_KNN(...), Join(on=[...])) - // When the ON clause has equi-conditions AND ST_KNN, DataFusion's SQL planner puts - // the equi-conditions in Join.on and the spatial predicate in a Filter above. - if let LogicalPlan::Filter(filter_node) = &plan { - let names = collect_spatial_predicate_names(&filter_node.predicate); - if names.contains("st_knn") { - if let LogicalPlan::Join(join) = filter_node.input.as_ref() { - return rewrite_filter_join_to_spatial_join_plan_node( - &filter_node.predicate, - join, - ); - } - } - } - Ok(Transformed::no(plan)) } } @@ -192,13 +177,10 @@ impl OptimizerRule for SpatialJoinLogicalRewrite { return Ok(Transformed::no(plan)); } - // KNN joins are already handled by KnnJoinEarlyRewrite — skip them here. - if spatial_predicate_names.contains("st_knn") { - return Ok(Transformed::no(plan)); - } - - // Non-KNN spatial joins with equi-join conditions are not supported. - if !join.on.is_empty() { + // Join with with equi-join condition and spatial join condition. Only handle it + // when the join condition contains ST_KNN. KNN join is not a regular join and + // ST_KNN is also not a regular predicate. It must be handled by our spatial join exec. + if !join.on.is_empty() && !spatial_predicate_names.contains("st_knn") { return Ok(Transformed::no(plan)); } @@ -245,65 +227,14 @@ fn rewrite_join_to_spatial_join_plan_node(join: &Join) -> Result Result> { - // Start with the filter predicate from the Filter node above - let mut combined = filter_predicate.clone(); - - // Add any existing join filter - if let Some(join_filter) = &join.filter { - combined = Expr::and(combined, join_filter.clone()); - } - - // Fold equi-join on conditions - let eq_op = if join.null_equality == NullEquality::NullEqualsNothing { - Operator::Eq - } else { - Operator::IsNotDistinctFrom - }; - combined = join.on.iter().fold(combined, |acc, (l, r)| { - let eq_expr = Expr::BinaryExpr(BinaryExpr::new( - Box::new(l.clone()), - eq_op, - Box::new(r.clone()), - )); - Expr::and(acc, eq_expr) - }); - - let schema = Arc::clone(&join.schema); - let node = SpatialJoinPlanNode { - left: join.left.as_ref().clone(), - right: join.right.as_ref().clone(), - join_type: join.join_type, - filter: combined, - schema, - join_constraint: join.join_constraint, - null_equality: join.null_equality, - }; - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))) -} - /// Logical optimizer rule that enables spatial join planning. /// /// This rule turns eligible `Filter(Join(filter=...))` nodes into a `Join(filter=...)` node, /// so that the spatial join can be rewritten later by [SpatialJoinLogicalRewrite]. #[derive(Debug, Default)] -struct MergeSpatialProjectionIntoJoin; +struct MergeSpatialFilterIntoJoin; -impl OptimizerRule for MergeSpatialProjectionIntoJoin { +impl OptimizerRule for MergeSpatialFilterIntoJoin { fn name(&self) -> &str { "merge_spatial_filter_into_join" } @@ -358,7 +289,9 @@ impl OptimizerRule for MergeSpatialProjectionIntoJoin { else { return Ok(Transformed::no(plan)); }; - if !is_spatial_predicate(predicate) { + + let spatial_predicates = collect_spatial_predicate_names(predicate); + if spatial_predicates.is_empty() { return Ok(Transformed::no(plan)); } @@ -377,20 +310,25 @@ impl OptimizerRule for MergeSpatialProjectionIntoJoin { }; // Check if this is a suitable join for rewriting + let is_equi_join = !on.is_empty() && !spatial_predicates.contains("st_knn"); if !matches!( join_type, JoinType::Inner | JoinType::Left | JoinType::Right - ) || !on.is_empty() - || filter.is_some() + ) || is_equi_join { return Ok(Transformed::no(plan)); } + let new_filter = match filter { + Some(existing_filter) => Expr::and(predicate.clone(), existing_filter.clone()), + None => predicate.clone(), + }; + let rewritten_plan = Join::try_new( Arc::clone(left), Arc::clone(right), on.clone(), - Some(predicate.clone()), + Some(new_filter), JoinType::Inner, *join_constraint, *null_equality, diff --git a/rust/sedona-spatial-join/src/planner/spatial_expr_utils.rs b/rust/sedona-spatial-join/src/planner/spatial_expr_utils.rs index 5fd85a50e3..1d858d45ec 100644 --- a/rust/sedona-spatial-join/src/planner/spatial_expr_utils.rs +++ b/rust/sedona-spatial-join/src/planner/spatial_expr_utils.rs @@ -98,13 +98,6 @@ pub(crate) fn collect_spatial_predicate_names(expr: &Expr) -> HashSet { acc } -/// Check if a given logical expression contains a spatial predicate component or not. We assume that the given -/// `expr` evaluates to a boolean value and originates from a filter logical node. -pub(crate) fn is_spatial_predicate(expr: &Expr) -> bool { - let pred_names = collect_spatial_predicate_names(expr); - !pred_names.is_empty() -} - /// Transform the join filter to a spatial predicate and a remainder. /// /// * The spatial predicate is a spatial predicate that is extracted from the join filter. @@ -2244,16 +2237,17 @@ mod tests { } #[test] - fn test_is_spatial_predicate() { - // Test 1: ST_ functions should return true + fn test_collect_spatial_predicate_names() { + // ST_Intersects should be collected let st_intersects_udf = create_dummy_st_intersects_udf(); let st_intersects_expr = Expr::ScalarFunction(datafusion_expr::expr::ScalarFunction { func: st_intersects_udf, args: vec![col("geom1"), col("geom2")], }); - assert!(is_spatial_predicate(&st_intersects_expr)); + let names = collect_spatial_predicate_names(&st_intersects_expr); + assert_eq!(names, HashSet::from(["st_intersects".to_string()])); - // ST_Distance(geom1, geom2) < 100 should return true + // ST_Distance(geom1, geom2) < 100 should be collected as st_dwithin let st_distance_udf = create_dummy_st_distance_udf(); let st_distance_expr = Expr::ScalarFunction(datafusion_expr::expr::ScalarFunction { func: st_distance_udf, @@ -2264,29 +2258,33 @@ mod tests { op: Operator::Lt, right: Box::new(lit(100.0)), }); - assert!(is_spatial_predicate(&distance_lt_expr)); + let names = collect_spatial_predicate_names(&distance_lt_expr); + assert_eq!(names, HashSet::from(["st_dwithin".to_string()])); - // ST_Distance(geom1, geom2) > 100 should return false + // ST_Distance(geom1, geom2) > 100 should not be collected (wrong comparison direction) let distance_gt_expr = Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr { left: Box::new(st_distance_expr.clone()), op: Operator::Gt, right: Box::new(lit(100.0)), }); - assert!(!is_spatial_predicate(&distance_gt_expr)); + let names = collect_spatial_predicate_names(&distance_gt_expr); + assert!(names.is_empty()); - // AND expressions with spatial predicates should return true + // AND expression: spatial predicate should be collected through conjunction let and_expr = Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr { left: Box::new(st_intersects_expr.clone()), op: Operator::And, right: Box::new(col("id").eq(lit(1))), }); - assert!(is_spatial_predicate(&and_expr)); + let names = collect_spatial_predicate_names(&and_expr); + assert_eq!(names, HashSet::from(["st_intersects".to_string()])); - // Non-spatial expressions should return false + // Non-spatial expressions should return empty set // Simple column comparison let non_spatial_expr = col("id").eq(lit(1)); - assert!(!is_spatial_predicate(&non_spatial_expr)); + let names = collect_spatial_predicate_names(&non_spatial_expr); + assert!(names.is_empty()); // Not a spatial relationship function let non_st_func = Expr::ScalarFunction(datafusion_expr::expr::ScalarFunction { @@ -2299,7 +2297,8 @@ mod tests { ))), args: vec![col("id")], }); - assert!(!is_spatial_predicate(&non_st_func)); + let names = collect_spatial_predicate_names(&non_st_func); + assert!(names.is_empty()); // AND expression with no spatial predicates let non_spatial_and = Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr { @@ -2307,6 +2306,7 @@ mod tests { op: Operator::And, right: Box::new(col("name").eq(lit("test"))), }); - assert!(!is_spatial_predicate(&non_spatial_and)); + let names = collect_spatial_predicate_names(&non_spatial_and); + assert!(names.is_empty()); } } diff --git a/rust/sedona-spatial-join/tests/spatial_join_integration.rs b/rust/sedona-spatial-join/tests/spatial_join_integration.rs index e888bc54d1..785c968255 100644 --- a/rust/sedona-spatial-join/tests/spatial_join_integration.rs +++ b/rust/sedona-spatial-join/tests/spatial_join_integration.rs @@ -21,6 +21,7 @@ use arrow_array::{Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::{ catalog::{MemTable, TableProvider}, + datasource::empty::EmptyTable, execution::SessionStateBuilder, prelude::{SessionConfig, SessionContext}, }; @@ -1089,72 +1090,90 @@ async fn test_knn_join_with_filter_correctness( }; let k = 3; - let sql = format!( - "SELECT L.id AS l_id, R.id AS r_id FROM L JOIN R ON ST_KNN(L.geometry, R.geometry, {}, false) AND (L.id % 7) = (R.id % 7)", - k - ); + let sqls = [ + format!( + "SELECT L.id AS l_id, R.id AS r_id FROM L JOIN R ON ST_KNN(L.geometry, R.geometry, {}, false) AND (L.id % 7) = (R.id % 7)", + k + ), + format!( + "SELECT L.id AS l_id, R.id AS r_id FROM L JOIN R ON ST_KNN(L.geometry, R.geometry, {}, false) AND L.id % 7 = 0", + k + ), + format!( + "SELECT L.id AS l_id, R.id AS r_id FROM L JOIN R ON ST_KNN(L.geometry, R.geometry, {}, false) AND R.id % 7 = 0", + k + ), + ]; - let batches = run_spatial_join_query( - &left_schema, - &right_schema, - left_partitions.clone(), - right_partitions.clone(), - Some(options), - max_batch_size, - &sql, - ) - .await?; + for (idx, sql) in sqls.iter().enumerate() { + let batches = run_spatial_join_query( + &left_schema, + &right_schema, + left_partitions.clone(), + right_partitions.clone(), + Some(options.clone()), + max_batch_size, + sql, + ) + .await?; - let mut actual_results = Vec::new(); - let combined_batch = arrow::compute::concat_batches(&batches.schema(), &[batches])?; - let l_ids = combined_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let r_ids = combined_batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); + let mut actual_results = Vec::new(); + let combined_batch = arrow::compute::concat_batches(&batches.schema(), &[batches])?; + let l_ids = combined_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let r_ids = combined_batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); - for i in 0..combined_batch.num_rows() { - actual_results.push((l_ids.value(i), r_ids.value(i))); - } - actual_results.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1))); + for i in 0..combined_batch.num_rows() { + actual_results.push((l_ids.value(i), r_ids.value(i))); + } + actual_results.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1))); - // Prove the test actually exercises the "< K rows after filtering" case. - // Build a list of all probe-side IDs and count how many results each has. - let all_left_ids: Vec = extract_geoms_and_ids(&left_partitions) - .into_iter() - .map(|(id, _)| id) - .collect(); - let mut per_left_counts: std::collections::HashMap = - std::collections::HashMap::new(); - for (l_id, _) in &actual_results { - *per_left_counts.entry(*l_id).or_default() += 1; - } - let min_count = all_left_ids - .iter() - .map(|l_id| *per_left_counts.get(l_id).unwrap_or(&0)) - .min() - .unwrap_or(0); - assert!( - min_count < k, - "expected at least one probe row to produce < K rows after filtering; min_count={min_count}, k={k}" - ); + // Prove the test actually exercises the "< K rows after filtering" case. + // Build a list of all probe-side IDs and count how many results each has. + let all_left_ids: Vec = extract_geoms_and_ids(&left_partitions) + .into_iter() + .map(|(id, _)| id) + .collect(); + let mut per_left_counts: std::collections::HashMap = + std::collections::HashMap::new(); + for (l_id, _) in &actual_results { + *per_left_counts.entry(*l_id).or_default() += 1; + } + let min_count = all_left_ids + .iter() + .map(|l_id| *per_left_counts.get(l_id).unwrap_or(&0)) + .min() + .unwrap_or(0); + assert!( + min_count < k, + "expected at least one probe row to produce < K rows after filtering; min_count={min_count}, k={k}" + ); - let expected_results = compute_knn_ground_truth_with_pair_filter( - &left_partitions, - &right_partitions, - k, - |l_id, r_id| (l_id.rem_euclid(7)) == (r_id.rem_euclid(7)), - ) - .into_iter() - .map(|(l, r, _)| (l, r)) - .collect::>(); + let filter_closure = match idx { + 0 => |l_id: i32, r_id: i32| (l_id.rem_euclid(7)) == (r_id.rem_euclid(7)), + 1 => |l_id: i32, _r_id: i32| l_id.rem_euclid(7) == 0, + 2 => |_l_id: i32, r_id: i32| r_id.rem_euclid(7) == 0, + _ => unreachable!(), + }; + let expected_results = compute_knn_ground_truth_with_pair_filter( + &left_partitions, + &right_partitions, + k, + filter_closure, + ) + .into_iter() + .map(|(l, r, _)| (l, r)) + .collect::>(); - assert_eq!(actual_results, expected_results); + assert_eq!(actual_results, expected_results); + } Ok(()) } @@ -1370,55 +1389,6 @@ async fn test_knn_join_include_tie_breakers( Ok(()) } -/// Recursively check whether any node in the physical plan tree is a `FilterExec`. -fn subtree_contains_filter_exec(plan: &Arc) -> bool { - let mut found = false; - plan.apply(|node| { - if node.as_any().downcast_ref::().is_some() { - found = true; - return Ok(TreeNodeRecursion::Stop); - } - Ok(TreeNodeRecursion::Continue) - }) - .expect("failed to walk plan"); - found -} - -/// Create a session context with two small tables for filter-pushdown tests. -/// -/// L(id INT, x DOUBLE) and R(id INT, x DOUBLE) each with 10 rows. -/// Geometry is constructed in SQL via ST_Point so no geometry column exists on the table itself. -async fn plan_for_filter_pushdown_test(sql: &str) -> Result> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("x", DataType::Float64, false), - ])); - - let ids: Vec = (0..10).collect(); - let xs: Vec = (0..10).map(|i| i as f64).collect(); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int32Array::from(ids)), - Arc::new(arrow_array::Float64Array::from(xs)), - ], - )?; - - let options = SpatialJoinOptions::default(); - let ctx = setup_context(Some(options), 100)?; - let mem_l: Arc = Arc::new(MemTable::try_new( - schema.clone(), - vec![vec![batch.clone()]], - )?); - let mem_r: Arc = - Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch]])?); - ctx.register_table("L", mem_l)?; - ctx.register_table("R", mem_r)?; - - let df = ctx.sql(sql).await?; - df.create_physical_plan().await -} - /// Verify that a filter on the *object* (build / right) side of a KNN join is NOT pushed down /// into the build side subtree. /// @@ -1473,3 +1443,38 @@ async fn test_non_knn_join_object_side_filter_is_pushed_down() -> Result<()> { Ok(()) } + +/// Recursively check whether any node in the physical plan tree is a `FilterExec`. +fn subtree_contains_filter_exec(plan: &Arc) -> bool { + let mut found = false; + plan.apply(|node| { + if node.as_any().downcast_ref::().is_some() { + found = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("failed to walk plan"); + found +} + +/// Create a session context with two small tables for filter-pushdown tests. +/// +/// L(id INT, x DOUBLE) and R(id INT, x DOUBLE) each with 10 rows. +/// Geometry is constructed in SQL via ST_Point so no geometry column exists on the table itself. +async fn plan_for_filter_pushdown_test(sql: &str) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("x", DataType::Float64, false), + ])); + + let options = SpatialJoinOptions::default(); + let ctx = setup_context(Some(options), 100)?; + let empty_l: Arc = Arc::new(EmptyTable::new(schema.clone())); + let empty_r: Arc = Arc::new(EmptyTable::new(schema.clone())); + ctx.register_table("L", empty_l)?; + ctx.register_table("R", empty_r)?; + + let df = ctx.sql(sql).await?; + df.create_physical_plan().await +} From ce913b7154f1a8000e158433ca0cb69f8972f53b Mon Sep 17 00:00:00 2001 From: Kristin Cowalcijk Date: Sat, 14 Feb 2026 01:50:39 +0800 Subject: [PATCH 3/5] Implement necessary_children_exprs to make optimize_projections work properly with extension node --- .../sedona-spatial-join/src/planner/logical_plan_node.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/rust/sedona-spatial-join/src/planner/logical_plan_node.rs b/rust/sedona-spatial-join/src/planner/logical_plan_node.rs index e93e228dd6..3247695587 100644 --- a/rust/sedona-spatial-join/src/planner/logical_plan_node.rs +++ b/rust/sedona-spatial-join/src/planner/logical_plan_node.rs @@ -106,6 +106,15 @@ impl UserDefinedLogicalNodeCore for SpatialJoinPlanNode { ) } + fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { + // Request all columns from both children. This ensures the optimizer + // recurses into children while preserving all columns needed by the + // join filter and output schema. + let left_indices: Vec = (0..self.left.schema().fields().len()).collect(); + let right_indices: Vec = (0..self.right.schema().fields().len()).collect(); + Some(vec![left_indices, right_indices]) + } + fn with_exprs_and_inputs( &self, mut exprs: Vec, From c9ee76c86f3a597b69217179122ec2598987bc9d Mon Sep 17 00:00:00 2001 From: Kristin Cowalcijk Date: Tue, 17 Feb 2026 14:31:45 +0800 Subject: [PATCH 4/5] Update the comment to include the link to the DataFusion PR --- rust/sedona-spatial-join/src/planner/logical_plan_node.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/rust/sedona-spatial-join/src/planner/logical_plan_node.rs b/rust/sedona-spatial-join/src/planner/logical_plan_node.rs index 3247695587..23f35d09b3 100644 --- a/rust/sedona-spatial-join/src/planner/logical_plan_node.rs +++ b/rust/sedona-spatial-join/src/planner/logical_plan_node.rs @@ -107,9 +107,10 @@ impl UserDefinedLogicalNodeCore for SpatialJoinPlanNode { } fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { - // Request all columns from both children. This ensures the optimizer - // recurses into children while preserving all columns needed by the - // join filter and output schema. + // Request all columns from both children. The default implementation returns None, which + // should also be fine, but we need to return the columns indices explicitly to workaround + // a bug in DataFusion's handling of None projection indices in FFI table provider. + // See https://github.com/apache/datafusion/pull/20393 let left_indices: Vec = (0..self.left.schema().fields().len()).collect(); let right_indices: Vec = (0..self.right.schema().fields().len()).collect(); Some(vec![left_indices, right_indices]) From b4381c327489401f0061876903838917a4ca5d83 Mon Sep 17 00:00:00 2001 From: Kristin Cowalcijk Date: Tue, 17 Feb 2026 15:27:05 +0800 Subject: [PATCH 5/5] Address review comments --- rust/sedona-spatial-join/src/planner.rs | 9 +++-- .../src/planner/optimizer.rs | 33 ++++++++++++++----- .../tests/spatial_join_integration.rs | 7 ++-- rust/sedona/src/context.rs | 2 +- 4 files changed, 35 insertions(+), 16 deletions(-) diff --git a/rust/sedona-spatial-join/src/planner.rs b/rust/sedona-spatial-join/src/planner.rs index 11312aeceb..e40303bfe1 100644 --- a/rust/sedona-spatial-join/src/planner.rs +++ b/rust/sedona-spatial-join/src/planner.rs @@ -21,6 +21,7 @@ //! can produce `SpatialJoinExec`. use datafusion::execution::SessionStateBuilder; +use datafusion_common::Result; mod logical_plan_node; mod optimizer; @@ -34,10 +35,12 @@ mod spatial_expr_utils; /// implementation provided by this crate and ensures joins created by SQL or using /// a DataFrame API that meet certain conditions (e.g. contain a spatial predicate as /// a join condition) are executed using the `SpatialJoinExec`. -pub fn register_planner(state_builder: SessionStateBuilder) -> SessionStateBuilder { +pub fn register_planner(state_builder: SessionStateBuilder) -> Result { // Enable the logical rewrite that turns Filter(CrossJoin) into Join(filter=...) - let state_builder = optimizer::register_spatial_join_logical_optimizer(state_builder); + let state_builder = optimizer::register_spatial_join_logical_optimizer(state_builder)?; // Enable planning SpatialJoinExec via an extension node during logical->physical planning. - physical_planner::register_spatial_join_planner(state_builder) + Ok(physical_planner::register_spatial_join_planner( + state_builder, + )) } diff --git a/rust/sedona-spatial-join/src/planner/optimizer.rs b/rust/sedona-spatial-join/src/planner/optimizer.rs index b4721728c6..811c866b2d 100644 --- a/rust/sedona-spatial-join/src/planner/optimizer.rs +++ b/rust/sedona-spatial-join/src/planner/optimizer.rs @@ -27,6 +27,7 @@ use datafusion_expr::logical_plan::Extension; use datafusion_expr::{BinaryExpr, Expr, Operator}; use datafusion_expr::{Filter, Join, JoinType, LogicalPlan}; use sedona_common::option::SedonaOptions; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err}; /// Register the logical spatial join optimizer rules. /// @@ -43,7 +44,7 @@ use sedona_common::option::SedonaOptions; /// benefit from filter pushdown before being converted to extension nodes. pub(crate) fn register_spatial_join_logical_optimizer( mut session_state_builder: SessionStateBuilder, -) -> SessionStateBuilder { +) -> Result { let optimizer = session_state_builder .optimizer() .get_or_insert_with(Optimizer::new); @@ -53,7 +54,11 @@ pub(crate) fn register_spatial_join_logical_optimizer( .rules .iter() .position(|r| r.name() == "push_down_filter") - .expect("PushDownFilter rule not found in default optimizer rules"); + .ok_or_else(|| { + sedona_internal_datafusion_err!( + "PushDownFilter rule not found in default optimizer rules" + ) + })?; // Insert KNN-specific rules BEFORE PushDownFilter. // MergeSpatialFilterIntoJoin must come first because it creates the Join(filter=...) @@ -68,7 +73,7 @@ pub(crate) fn register_spatial_join_logical_optimizer( // Append SpatialJoinLogicalRewrite at the end so non-KNN joins benefit from filter pushdown. optimizer.rules.push(Arc::new(SpatialJoinLogicalRewrite)); - session_state_builder + Ok(session_state_builder) } /// Early optimizer rule that converts KNN joins to `SpatialJoinPlanNode` extension nodes @@ -133,7 +138,7 @@ impl OptimizerRule for KnnJoinEarlyRewrite { /// This rule runs *after* `PushDownFilter` so that non-KNN spatial joins benefit from /// filter pushdown before being converted to extension nodes. /// -/// KNN joins are skipped here because they are already handled by `KnnJoinEarlyRewrite`. +/// KNN joins are skipped here because they are already handled by [[KnnJoinEarlyRewrite]]. #[derive(Default, Debug)] struct SpatialJoinLogicalRewrite; @@ -177,10 +182,16 @@ impl OptimizerRule for SpatialJoinLogicalRewrite { return Ok(Transformed::no(plan)); } - // Join with with equi-join condition and spatial join condition. Only handle it - // when the join condition contains ST_KNN. KNN join is not a regular join and - // ST_KNN is also not a regular predicate. It must be handled by our spatial join exec. - if !join.on.is_empty() && !spatial_predicate_names.contains("st_knn") { + if spatial_predicate_names.contains("st_knn") { + // KNN joins should have already been rewritten by KnnJoinEarlyRewrite, so we shouldn't + // see them here. + return sedona_internal_err!( + "Found KNN predicate in SpatialJoinLogicalRewrite, which should have been handled by KnnJoinEarlyRewrite"); + } + + // Join with with equi-join condition should be planned as a regular HashJoin + // or SortMergeJoin. + if !join.on.is_empty() { return Ok(Transformed::no(plan)); } @@ -194,7 +205,11 @@ fn rewrite_join_to_spatial_join_plan_node(join: &Join) -> Result, batch_size: usize) -> Resu session_config = add_sedona_option_extension(session_config); let mut state_builder = SessionStateBuilder::new(); if let Some(options) = options { - state_builder = register_planner(state_builder); + state_builder = register_planner(state_builder)?; let opts = session_config .options_mut() .extensions @@ -1393,7 +1393,7 @@ async fn test_knn_join_include_tie_breakers( /// into the build side subtree. /// /// If `PushDownFilter` incorrectly pushes `R.id > 5` below the spatial join, the set of objects -/// considered for KNN changes, yielding wrong nearest-neighbor results. +/// considered for the KNN search changes, yielding wrong nearest-neighbor results. #[tokio::test] async fn test_knn_join_object_side_filter_not_pushed_down() -> Result<()> { let sql = "SELECT L.id, R.id \ @@ -1460,7 +1460,8 @@ fn subtree_contains_filter_exec(plan: &Arc) -> bool { /// Create a session context with two small tables for filter-pushdown tests. /// -/// L(id INT, x DOUBLE) and R(id INT, x DOUBLE) each with 10 rows. +/// L(id INT, x DOUBLE) and R(id INT, x DOUBLE) are all empty, this is just for exercising the +/// plan optimizer and physical planner. /// Geometry is constructed in SQL via ST_Point so no geometry column exists on the table itself. async fn plan_for_filter_pushdown_test(sql: &str) -> Result> { let schema = Arc::new(Schema::new(vec![ diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs index ee2c963559..dfd77913e4 100644 --- a/rust/sedona/src/context.rs +++ b/rust/sedona/src/context.rs @@ -117,7 +117,7 @@ impl SedonaContext { // Register the spatial join planner extension #[cfg(feature = "spatial-join")] { - state_builder = sedona_spatial_join::register_planner(state_builder); + state_builder = sedona_spatial_join::register_planner(state_builder)?; } let mut state = state_builder.build();