Kontinuation commented on code in PR #611:
URL: https://github.com/apache/sedona-db/pull/611#discussion_r2815388592
##########
rust/sedona-spatial-join/src/planner/optimizer.rs:
##########
@@ -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));
Review Comment:
Changed the plan rule registration methods to return Result.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
