fix(rust/sedona-spatial-join): prevent filter pushdown past KNN joins

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<Vec<Vec<usize>>> {
+        // 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<usize> = (0..self.left.schema().fields().len()).collect();
+        let right_indices: Vec<usize> = (0..self.right.schema().fields().len()).collect();
+        Some(vec![left_indices, right_indices])
+    }
+
     fn with_exprs_and_inputs(
         &self,
         mut exprs: Vec<Expr>,

rust/sedona-spatial-join/src/planner/optimizer.rs

@@ -18,9 +18,8 @@ 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, 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 +28,112 @@ 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:
+///
+/// - `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
+///   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.
+    // 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(MergeSpatialFilterIntoJoin));
+
+    // 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.
+///
+/// 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:
 ///
-/// This rule turns eligible `Join(filter=...)` nodes into a `SpatialJoinPlanNode` extension.
+/// 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<ApplyOrder> {
+        Some(ApplyOrder::BottomUp)
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
+        let options = config.options();
+        let Some(ext) = options.extensions.get::<SedonaOptions>() else {
+            return Ok(Transformed::no(plan));
+        };
+        if !ext.spatial_join.enable {
+            return Ok(Transformed::no(plan));
+        }
+
+        // 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);
+                }
+            }
+        }
+
+        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;
 
@@ -93,47 +184,57 @@ impl OptimizerRule for SpatialJoinLogicalRewrite {
             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,
-        };
-
-        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<Transformed<LogicalPlan>> {
+    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),
+    })))
+}
+
 /// 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"
     }
@@ -188,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));
         }
 
@@ -207,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,

rust/sedona-spatial-join/src/planner/spatial_expr_utils.rs

@@ -98,13 +98,6 @@ pub(crate) fn collect_spatial_predicate_names(expr: &Expr) -> HashSet<String> {
     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,14 +2297,16 @@ 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 {
             left: Box::new(col("id").eq(lit(1))),
             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());
     }
 }