diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 57d7dba29526f..d003692267463 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -66,6 +66,9 @@ use datafusion_physical_plan::{ aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}, coalesce_partitions::CoalescePartitionsExec, collect, + execution_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, plan_contains_expression_id, + }, filter::{FilterExec, FilterExecBuilder}, joins::{HashJoinExec, PartitionMode}, projection::ProjectionExec, @@ -2964,29 +2967,12 @@ async fn test_hashjoin_hash_table_pushdown_collect_left() { ); } -// Not portable to sqllogictest: verifies whether the optimized probe-side plan -// retains the HashJoinExec's dynamic filter expression. The with_support(false) -// branch has no SQL analog because parquet supports filter pushdown. +// Not portable to sqllogictest: verifies that the HashJoinExec only keeps its +// dynamic filter when the optimized probe-side plan retains the expression, i.e. +// when there is something to consume it. The with_support(false) branch has no +// SQL analog because parquet supports filter pushdown. #[test] -fn test_hashjoin_dynamic_filter_pushdown_is_used() { - fn contains_expression_id(plan: &Arc, expression_id: u64) -> bool { - let mut found = false; - plan.apply(|node| { - node.apply_expressions(&mut |root| { - root.apply(|expr| { - if expr.expression_id() == Some(expression_id) { - found = true; - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - }) - }) - }) - .unwrap(); - found - } - +fn test_hashjoin_dynamic_filter_requires_probe_consumer() { for (probe_supports_pushdown, expected_consumer) in [(false, false), (true, true)] { let build_side_schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Utf8, false), @@ -3050,20 +3036,199 @@ fn test_hashjoin_dynamic_filter_pushdown_is_used() { .downcast_ref::() .expect("Plan should be HashJoinExec"); let dynamic_filters = hash_join.dynamic_expressions_produced(); - let expression_id = dynamic_filters - .first() - .expect("Dynamic filter should be created") - .expression_id() - .expect("Dynamic filters always have an expression ID"); + // The join keeps a dynamic filter only if the pushdown left a consumer for it + // in the probe subtree. Otherwise it is dropped at planning time so that + // `execute` skips build side bounds accumulation entirely. assert_eq!( - contains_expression_id(hash_join.right(), expression_id), - expected_consumer, - "probe consumer should be {expected_consumer} when pushdown support is {probe_supports_pushdown}" + dynamic_filters.len(), + usize::from(expected_consumer), + "dynamic filter should {}be produced when pushdown support is {probe_supports_pushdown}", + if expected_consumer { "" } else { "not " } ); + + if let Some(dynamic_filter) = dynamic_filters.first() { + let expression_id = dynamic_filter + .expression_id() + .expect("Dynamic filters always have an expression ID"); + assert!( + plan_contains_expression_id(hash_join.right(), expression_id).unwrap(), + "probe subtree should contain the dynamic filter it accepted" + ); + } } } +// Not portable to sqllogictest: stands in for a distributed planner that splits +// an already-optimized plan into stages, leaving the HashJoinExec and the scan +// consuming its dynamic filter on different workers. Whether to produce the +// filter is decided during pushdown, so it has to survive the probe subtree +// being swapped out afterwards. +#[tokio::test] +async fn test_hashjoin_dynamic_filter_survives_probe_subtree_replacement() { + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); + + let on = vec![ + ( + col("a", &build_side_schema).unwrap(), + col("a", &probe_side_schema).unwrap(), + ), + ( + col("b", &build_side_schema).unwrap(), + col("b", &probe_side_schema).unwrap(), + ), + ]; + let plan = Arc::new( + HashJoinExec::try_new( + Arc::clone(&build_scan), + probe_scan, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_dynamic_filter_pushdown = true; + let plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + + // The probe scan accepted the filter, so the join kept it. + assert_eq!(plan.dynamic_expressions_produced().len(), 1); + + // Swap the probe subtree for one that does not hold the filter: in a real + // deployment the consumer ends up in another stage, on another worker, and is + // not reachable from this node at all. + let detached_probe = TestScanBuilder::new(Arc::clone(&probe_side_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!( + ("a", Utf8, ["aa", "ab", "ac", "ad"]), + ("b", Utf8, ["ba", "bb", "bc", "bd"]), + ("e", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]) + .build(); + let plan = plan + .replace_children( + vec![build_scan, detached_probe], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + assert_eq!(plan.dynamic_expressions_produced().len(), 1); + + let session_ctx = + SessionContext::new_with_config(SessionConfig::from(config).with_batch_size(10)); + session_ctx.register_object_store( + ObjectStoreUrl::parse("test://").unwrap().as_ref(), + Arc::new(InMemory::new()), + ); + collect(Arc::clone(&plan), session_ctx.state().task_ctx()) + .await + .unwrap(); + + // The build side bounds were still computed and published, so whatever holds + // the other end of this filter sees them. + let dynamic_filter = plan + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("dynamic filter should be retained"); + insta::assert_snapshot!( + format!("{dynamic_filter}"), + @"DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ]", + ); +} + +// Not portable to sqllogictest: custom `PhysicalOptimizerRule`s are appended +// after the built in ones, so a user rule runs after the Post phase +// `FilterPushdown` that decides whether to keep the join's dynamic filter. Such a +// rule can bring a consumer back by re-running the pushdown; there is nothing to +// undo first, because a join with no consumer holds no dynamic filter. +// +// Re-running the Post phase is an already supported mode rather than something +// this test invents: see `post_phase_is_idempotent_on_hash_join` below, added by +// apache/datafusion#22523 because AQE (datafusion-ballista#1359) re-runs the +// optimizer chain after every completed stage. +#[test] +fn test_hashjoin_dynamic_filter_recreated_when_pushdown_reruns() { + let (build_side_schema, build_scan, probe_side_schema, _) = hashjoin_pushdown_scans(); + + let probe_batches = vec![ + record_batch!( + ("a", Utf8, ["aa", "ab", "ac", "ad"]), + ("b", Utf8, ["ba", "bb", "bc", "bd"]), + ("e", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]; + let unsupported_probe = TestScanBuilder::new(Arc::clone(&probe_side_schema)) + .with_support(false) + .with_batches(probe_batches.clone()) + .build(); + + let on = vec![ + ( + col("a", &build_side_schema).unwrap(), + col("a", &probe_side_schema).unwrap(), + ), + ( + col("b", &build_side_schema).unwrap(), + col("b", &probe_side_schema).unwrap(), + ), + ]; + let plan = Arc::new( + HashJoinExec::try_new( + Arc::clone(&build_scan), + unsupported_probe, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_dynamic_filter_pushdown = true; + + // Nothing in the probe side accepts the filter, so the join drops it. + let plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + assert_eq!(plan.dynamic_expressions_produced().len(), 0); + + // A later rule swaps in a probe side that does accept filters and re-runs the + // pushdown. The join creates a fresh filter and finds its consumer. + let supported_probe = TestScanBuilder::new(Arc::clone(&probe_side_schema)) + .with_support(true) + .with_batches(probe_batches) + .build(); + let plan = plan + .replace_children( + vec![build_scan, supported_probe], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + let plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + assert_eq!(plan.dynamic_expressions_produced().len(), 1); +} + /// Regression test for https://github.com/apache/datafusion/issues/20109. /// /// Not portable to sqllogictest: the regression specifically targets the diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index a4d081b3d9e75..c0fa4c6bede41 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1122,7 +1122,12 @@ where /// /// This traverses both the execution plan and the children of each expression root /// reported by [`ExecutionPlan::apply_expressions`]. -pub(crate) fn plan_contains_expression_id( +/// +/// Producers of dynamic filters use this to find out whether anything downstream +/// holds the filter they pushed, since a node that replies +/// [`PushedDown::No`](crate::filter_pushdown::PushedDown::No) may still retain it +/// for statistics pruning. +pub fn plan_contains_expression_id( plan: &Arc, expression_id: u64, ) -> Result { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 977a36578da79..420dff4419922 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -999,6 +999,11 @@ impl HashJoinExec { /// Set the dynamic filter on this hash join. /// + /// Setting a dynamic filter is what makes the join compute and publish + /// build-side bounds during execution. [`Self::handle_child_pushdown_result`] + /// sets one after finding a consumer for it in the probe side. Callers wiring a + /// filter up by hand take on that check themselves. + /// /// Resets any internal state that depends on any existing dynamic filter. /// /// Validates that the filter's children reference valid columns in @@ -1449,20 +1454,9 @@ impl ExecutionPlan for HashJoinExec { consider using CoalescePartitionsExec or the EnforceDistribution rule" ); - // Only compute a dynamic filter when the probe subtree contains a consumer. - // Searching from `self` would always find the producer expression owned by this join. - let enable_dynamic_filter_pushdown = if self + let enable_dynamic_filter_pushdown = self .allow_join_dynamic_filter_pushdown(context.session_config().options()) - { - self.dynamic_filter - .as_ref() - .and_then(|df| df.filter.expression_id()) - .map(|id| plan_contains_expression_id(&self.right, id)) - .transpose()? - .unwrap_or(false) - } else { - false - }; + && self.dynamic_filter.is_some(); let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); @@ -1814,21 +1808,40 @@ impl ExecutionPlan for HashJoinExec { let right_child_self_filters = &child_pushdown_result.self_filters[1]; // We only push down filters to the right child // We expect 0 or 1 self filters if let Some(filter) = right_child_self_filters.first() { - // Note that we don't check PushdDownPredicate::discrimnant because even if nothing said - // "yes, I can fully evaluate this filter" things might still use it for statistics -> it's worth updating let predicate = Arc::clone(&filter.predicate); if let Ok(dynamic_filter) = Arc::downcast::(predicate) { - // We successfully pushed down our self filter - we need to make a new node with the dynamic filter - let new_node = self - .builder() - .with_dynamic_filter(Some(HashJoinExecDynamicFilter { - filter: dynamic_filter, - build_accumulator: OnceLock::new(), - })) - .build_exec()?; - result = result.with_updated_node(new_node); + // Note that we don't check `PushedDownPredicate::discriminant`: a node + // that replies `PushedDown::No` may still retain the filter for + // statistics pruning, so the reply does not tell us whether anyone will + // actually read the filter. Instead we look for a consumer holding the + // expression in the probe subtree. + // + // `self` here is the join with its post-pushdown children, so anything + // that accepted the filter is already wired into `self.right`. Searching + // from `self` would always find the producer expression pushed by this + // join. This is the last chance to make the decision: the Post-phase + // `FilterPushdown` rule is the final rule that mutates the plan. + let has_consumer = dynamic_filter + .expression_id() + .map(|id| plan_contains_expression_id(&self.right, id)) + .transpose()? + .unwrap_or(false); + if has_consumer { + // Our self filter reached a consumer: rebuild the node holding onto + // the dynamic filter so that `execute` populates it from the build + // side. If it did not, we leave `dynamic_filter` as `None` and skip + // the (not cheap) bounds accumulation entirely. + let new_node = self + .builder() + .with_dynamic_filter(Some(HashJoinExecDynamicFilter { + filter: dynamic_filter, + build_accumulator: OnceLock::new(), + })) + .build_exec()?; + result = result.with_updated_node(new_node); + } } } Ok(result) diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 9e50a93b2163f..5b0cb720ac6ac 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -49,8 +49,8 @@ pub use crate::execution_plan::{ AsPhysicalExprRef, ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, apply_expression_roots, collect, collect_partitioned, displayable, execute_input_stream, execute_stream, - execute_stream_partitioned, get_plan_string, replace_children_if_necessary, - with_new_children_if_necessary, + execute_stream_partitioned, get_plan_string, plan_contains_expression_id, + replace_children_if_necessary, with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode;