fix: Enforce data distribution requirements and fix plans that break them

src/coordinator/prepare_dynamic_plan.rs

@@ -3,7 +3,7 @@ use crate::coordinator::distributed::PreparedPlan;
 use crate::coordinator::query_coordinator::QueryCoordinator;
 use crate::distributed_planner::{
     InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost,
-    inject_network_boundaries,
+    inject_network_boundaries, validate_stage_plan,
 };
 use crate::events::TaskCountAnnotation::{Desired, Maximum};
 use crate::execution_plans::SamplerExec;
@@ -75,6 +75,9 @@ pub(super) async fn prepare_dynamic_plan(
             input_stage.plan = nb_ctx
                 .propagate_task_count_until_network_boundaries(&input_stage.plan, task_count)?;
             input_stage.tasks = task_count.as_usize();
+            // The task count is now final; refuse to dispatch a stage whose plan cannot run
+            // correctly on that many tasks.
+            validate_stage_plan(&input_stage.plan, input_stage.tasks, nb_ctx.cfg)?;
             // In order to infer the compute the cost of the stage above this one, here a sampler
             // is injected to gather runtime statistics.
             input_stage.plan = ProducerHead::insert_sampler(input_stage.plan)?;

src/distributed_planner/inject_network_boundaries.rs

@@ -12,11 +12,13 @@ use crate::{
 };
 use async_trait::async_trait;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
-use datafusion::common::{HashMap, Result, plan_err};
+use datafusion::common::{HashMap, JoinType, Result, plan_err};
 use datafusion::physical_expr::Partitioning;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::execution_plan::CardinalityEffect;
-use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
+use datafusion::physical_plan::joins::{
+    CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
+};
 use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::{ExecutionPlan, PlanProperties};
@@ -159,7 +161,7 @@ pub(crate) async fn inject_network_boundaries(
 pub(crate) struct InjectNetworkBoundaryContext<'a> {
     pub(crate) d_cfg: &'a DistributedConfig,
 
-    cfg: &'a SessionConfig,
+    pub(crate) cfg: &'a SessionConfig,
     worker_resolver: Arc<WorkerResolverExtension>,
     nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync),
     task_counts: &'a Mutex<HashMap<usize, TaskCountAnnotation>>,
@@ -273,10 +275,28 @@ async fn _inject_network_boundaries(
         task_count = Desired(count);
     } else if let Some(node) = plan.downcast_ref::<HashJoinExec>()
         && node.mode == PartitionMode::CollectLeft
-        && !broadcast_joins_enabled
+        && (!broadcast_joins_enabled || node.null_aware)
+    {
+        // A CollectLeft join collects its entire build side in every task, so it can only run in
+        // a multi-task stage when [insert_broadcast_execs] broadcast the build side. That is not
+        // possible when broadcast joins are disabled, nor for null-aware anti joins, whose
+        // NULL-existence checks live in process-local shared state that cannot span tasks in any
+        // orientation. Other build-side-emitting join types were already rewritten to Partitioned
+        // by [normalize_collect_joins], so they never reach this arm.
+        task_count = Maximum(1);
+    } else if let Some(node) = plan.downcast_ref::<NestedLoopJoinExec>()
+        && (!broadcast_joins_enabled || node.join_type() == &JoinType::Full)
     {
-        // Only distribute CollectLeft HashJoins after we broadcast more intelligently or when it
-        // is explicitly enabled.
+        // A NestedLoopJoin always collects its entire left side in every task, so it also needs
+        // its build side broadcast to run in a multi-task stage. Full joins emit unmatched rows
+        // from both sides, which needs global match knowledge that cannot span tasks, so they
+        // always run in a single task. Other build-side-emitting join types were already swapped
+        // to probe-side-emitting ones by [normalize_collect_joins].
+        task_count = Maximum(1);
+    } else if plan.is::<CrossJoinExec>() && !broadcast_joins_enabled {
+        // A CrossJoin also collects its entire left side in every task. It is always safe to
+        // broadcast (it emits only pair rows), so it is only restricted to a single task when
+        // broadcasts are unavailable.
         task_count = Maximum(1);
     } else {
         // The task count for this plan is decided by the biggest task count from the children; unless
@@ -632,6 +652,7 @@ mod tests {
     use super::*;
     use crate::distributed_planner::insert_broadcast::insert_broadcast_execs;
     use crate::distributed_planner::insert_children_isolator_union::insert_children_isolator_unions;
+    use crate::distributed_planner::normalize_collect_joins::normalize_collect_joins;
     use crate::test_utils::plans::{TestPlanBuilder, build_side_one_desired_task_count_handler};
     use crate::{DesiredTaskCountEvent, DesiredTaskCountEventResponse, assert_snapshot};
     use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
@@ -713,11 +734,14 @@ mod tests {
             .distributed_planner(false)
             .broadcast_joins(false);
         let annotated = annotate_test_plan(test_plan_builder, query).await;
-        assert_snapshot!(annotated, @r"
-        HashJoinExec: task_count=Maximum(1)
-          CoalescePartitionsExec: task_count=Maximum(1)
-            DistributedLeafExec: task_count=Maximum(1)
-          DistributedLeafExec: task_count=Maximum(1)
+        assert_snapshot!(annotated, @"
+        HashJoinExec: task_count=Desired(4)
+          NetworkShuffleExec: task_count=Desired(4)
+            RepartitionExec: task_count=Desired(4)
+              DistributedLeafExec: task_count=Desired(4)
+          NetworkShuffleExec: task_count=Desired(4)
+            RepartitionExec: task_count=Desired(4)
+              DistributedLeafExec: task_count=Desired(4)
         ")
     }
 
@@ -753,10 +777,20 @@ mod tests {
             .distributed_planner(false)
             .broadcast_joins(false);
         let annotated = annotate_test_plan(test_plan_builder, query).await;
-        assert_snapshot!(annotated, @r"
-        HashJoinExec: task_count=Maximum(1)
-          CoalescePartitionsExec: task_count=Maximum(1)
-            NetworkCoalesceExec: task_count=Maximum(1)
+        assert_snapshot!(annotated, @"
+        HashJoinExec: task_count=Desired(2)
+          NetworkShuffleExec: task_count=Desired(2)
+            RepartitionExec: task_count=Desired(2)
+              ProjectionExec: task_count=Desired(2)
+                AggregateExec: task_count=Desired(2)
+                  NetworkShuffleExec: task_count=Desired(2)
+                    RepartitionExec: task_count=Desired(4)
+                      AggregateExec: task_count=Desired(4)
+                        FilterExec: task_count=Desired(4)
+                          RepartitionExec: task_count=Desired(4)
+                            DistributedLeafExec: task_count=Desired(4)
+          NetworkShuffleExec: task_count=Desired(2)
+            RepartitionExec: task_count=Desired(2)
               ProjectionExec: task_count=Desired(2)
                 AggregateExec: task_count=Desired(2)
                   NetworkShuffleExec: task_count=Desired(2)
@@ -765,14 +799,6 @@ mod tests {
                         FilterExec: task_count=Desired(4)
                           RepartitionExec: task_count=Desired(4)
                             DistributedLeafExec: task_count=Desired(4)
-          ProjectionExec: task_count=Maximum(1)
-            AggregateExec: task_count=Maximum(1)
-              NetworkShuffleExec: task_count=Maximum(1)
-                RepartitionExec: task_count=Desired(4)
-                  AggregateExec: task_count=Desired(4)
-                    FilterExec: task_count=Desired(4)
-                      RepartitionExec: task_count=Desired(4)
-                        DistributedLeafExec: task_count=Desired(4)
         ")
     }
 
@@ -1085,6 +1111,32 @@ mod tests {
         ");
     }
 
+    #[tokio::test]
+    async fn test_nested_loop_inner_join_broadcast() {
+        // Inner joins emit only probe-driven rows, so the left side can be broadcast and
+        // the join can run in a multi-task stage.
+        let query = r#"
+        SELECT a."MinTemp", b."MaxTemp"
+        FROM weather a INNER JOIN weather b
+        ON a."MinTemp" < b."MaxTemp"
+        "#;
+        let test_plan_builder = TestPlanBuilder::new()
+            .target_partitions(4)
+            .num_workers(4)
+            // annotate_test_plan wants this as false so its s a single node plan
+            .distributed_planner(false)
+            .broadcast_joins(true);
+        let annotated = annotate_test_plan(test_plan_builder, query).await;
+        assert_snapshot!(annotated, @"
+        NestedLoopJoinExec: task_count=Desired(4)
+          CoalescePartitionsExec: task_count=Desired(4)
+            NetworkBroadcastExec: task_count=Desired(4)
+              BroadcastExec: task_count=Desired(4)
+                DistributedLeafExec: task_count=Desired(4)
+          DistributedLeafExec: task_count=Desired(4)
+        ");
+    }
+
     #[tokio::test]
     async fn test_broadcast_disabled_default() {
         let query = r#"
@@ -1210,6 +1262,8 @@ mod tests {
         let plan = test_plan.physical_plan(query).await;
         let session_config = test_plan.get_ctx().copied_config();
 
+        let plan = normalize_collect_joins(plan, session_config.options())
+            .expect("failed to normalize collect joins");
         let plan = insert_broadcast_execs(plan, session_config.options())
             .expect("failed to insert broadcasts");
         let plan = insert_children_isolator_unions(plan, session_config.options())

src/distributed_planner/insert_broadcast.rs

@@ -165,7 +165,7 @@ fn can_broadcast_left_input(plan: &dyn ExecutionPlan) -> bool {
     plan.downcast_ref::<CrossJoinExec>().is_some()
 }
 
-fn is_left_broadcast_safe(join_type: &JoinType) -> bool {
+pub(super) fn is_left_broadcast_safe(join_type: &JoinType) -> bool {
     matches!(
         join_type,
         JoinType::Inner

src/distributed_planner/mod.rs

@@ -4,11 +4,13 @@ mod inject_network_boundaries;
 mod insert_broadcast;
 mod insert_children_isolator_union;
 mod network_boundary;
+mod normalize_collect_joins;
 mod partial_reduce_below_network_shuffles;
 mod prepare_network_boundaries;
 mod push_fetch_into_network_coalesce;
 mod session_state_builder_ext;
 mod statistics;
+mod validate_stages;
 
 pub use distributed_config::DistributedConfig;
 pub(crate) use inject_network_boundaries::{
@@ -18,3 +20,4 @@ pub(crate) use network_boundary::ProducerHead;
 pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt};
 pub use session_state_builder_ext::SessionStateBuilderExt;
 pub(crate) use statistics::calculate_cost;
+pub(crate) use validate_stages::validate_stage_plan;