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/logical_plan_node.rs b/rust/sedona-spatial-join/src/planner/logical_plan_node.rs index e93e228dd6..23f35d09b3 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,16 @@ impl UserDefinedLogicalNodeCore for SpatialJoinPlanNode { ) } + fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { + // 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]) + } + fn with_exprs_and_inputs( &self, mut exprs: Vec, diff --git a/rust/sedona-spatial-join/src/planner/optimizer.rs b/rust/sedona-spatial-join/src/planner/optimizer.rs index 1edbc2d3da..811c866b2d 100644 --- a/rust/sedona-spatial-join/src/planner/optimizer.rs +++ b/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; @@ -28,21 +27,118 @@ 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 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, -) -> SessionStateBuilder { - session_state_builder - .with_optimizer_rule(Arc::new(MergeSpatialProjectionIntoJoin)) - .with_optimizer_rule(Arc::new(SpatialJoinLogicalRewrite)) + mut session_state_builder: SessionStateBuilder, +) -> Result { + 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") + .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=...) + // 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)); + + Ok(session_state_builder) } -/// 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 rule turns eligible `Join(filter=...)` nodes into a `SpatialJoinPlanNode` extension. +/// 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: +/// +/// 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)); + } + + // 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; @@ -86,54 +182,74 @@ 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") { - return Ok(Transformed::no(plan)); + 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"); } - // 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, - }; + // 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)); + } - 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() + .ok_or_else(|| { + datafusion_common::DataFusionError::Internal( + "join filter must be present for spatial join rewrite".to_string(), + ) + })? + .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 +304,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 +325,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 671ce2c804..2add149d38 100644 --- a/rust/sedona-spatial-join/tests/spatial_join_integration.rs +++ b/rust/sedona-spatial-join/tests/spatial_join_integration.rs @@ -21,12 +21,14 @@ 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}, }; 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}; @@ -142,7 +144,7 @@ fn setup_context(options: Option, 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 @@ -1088,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(()) } @@ -1368,3 +1388,94 @@ async fn test_knn_join_include_tie_breakers( Ok(()) } + +/// 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 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 \ + 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(()) +} + +/// 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) 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![ + 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 +} 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();