diff --git a/src/distributed_planner/insert_broadcast.rs b/src/distributed_planner/insert_broadcast.rs index 6986223b..852a940f 100644 --- a/src/distributed_planner/insert_broadcast.rs +++ b/src/distributed_planner/insert_broadcast.rs @@ -21,8 +21,9 @@ use super::DistributedConfig; /// The pass searches for joins whose left input can be broadcast without duplicating output rows: /// CollectLeft [HashJoinExec]s, [NestedLoopJoinExec]s, and [CrossJoinExec]s. Then it does one of /// two things: -/// 1. If the build child is a [CoalescePartitionsExec] -> Insert a [BroadcastExec] directly -/// below it. +/// 1. If the build child is a fetch-less [CoalescePartitionsExec] -> Insert a +/// [BroadcastExec] directly below it. A fetch-bearing coalesce stays below the broadcast +/// so its global limit is applied before the rows are replicated to consumers. /// 2. Otherwise (means it is already single partitioned going into the join) -> Insert a /// [BroadcastExec] -> [CoalescePartitionsExec] below the join but above its /// original build child. @@ -133,17 +134,18 @@ pub(super) fn insert_broadcast_execs( return Ok(Transformed::no(node)); }; - let (broadcast_input, coalesce_fetch) = build_child + let broadcast_input = build_child .downcast_ref::() + .filter(|coalesce| coalesce.fetch().is_none()) .map_or_else( - || (Arc::clone(build_child), None), - |coalesce| (Arc::clone(coalesce.input()), coalesce.fetch()), + || Arc::clone(build_child), + |coalesce| Arc::clone(coalesce.input()), ); // consumer_task_count=1 is a placeholder and will be corrected during optimizer rule. let broadcast: Arc = Arc::new(BroadcastExec::new(broadcast_input, 1)); let new_build_child: Arc = - Arc::new(CoalescePartitionsExec::new(broadcast).with_fetch(coalesce_fetch)); + Arc::new(CoalescePartitionsExec::new(broadcast)); let mut new_children: Vec> = children.into_iter().cloned().collect(); new_children[0] = new_build_child; diff --git a/tests/multi_task_collect_join_repros.rs b/tests/multi_task_collect_join_repros.rs index 91e0d5dc..c0c8ebf3 100644 --- a/tests/multi_task_collect_join_repros.rs +++ b/tests/multi_task_collect_join_repros.rs @@ -344,8 +344,10 @@ mod tests { } /// A build-side `LIMIT` is carried by the `CoalescePartitionsExec` that a - /// broadcast rewrite replaces. The replacement must preserve that fetch or - /// the join observes every build row instead of the requested 50. + /// broadcast rewrite replaces. It must run before broadcasting so the same + /// 50 rows reach every consumer, rather than being applied independently after fan-out. + /// Every build id has 50 probe matches, so the count is 2500 regardless of which + /// unordered 50 ids survive. #[tokio::test] async fn build_side_fetch_is_preserved_by_broadcast() { assert_distributed_matches_single_node(