Move children isolator union insertion before distribution (#565)

## What changed

- Add a pre-pass that replaces `UnionExec` with
`ChildrenIsolatorUnionExec` before annotation and network-boundary
distribution.
- Keep the pre-inserted placeholder single-node-correct by mapping all
children into one default task.
- Make annotation react to pre-inserted `ChildrenIsolatorUnionExec` for
child task-count propagation without constructing the node there.
- Finalize the per-task child map during `distribute_plan`, where final
task counts are available.

## Why

Fixes #539. This keeps union-isolator insertion out of the
network-boundary distribution recursion while preserving the existing
per-task assignment behavior.

## Validation

- `cargo fmt`
- `cargo test --lib`
- `cargo test --lib insert_children_isolator_union`
- `cargo test --features integration

Author: alexanderbianchi

src/distributed_planner/distributed_query_planner.rs

@@ -4,6 +4,7 @@ use crate::distributed_planner::inject_network_boundaries::{
     CardinalityBasedNetworkBoundaryBuilder, inject_network_boundaries,
 };
 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::partial_reduce_below_network_shuffles::partial_reduce_below_network_shuffles;
 use crate::distributed_planner::prepare_network_boundaries::prepare_network_boundaries;
 use crate::distributed_planner::push_fetch_into_network_coalesce::push_fetch_into_network_coalesce;
@@ -28,6 +29,9 @@ use std::sync::Arc;
 ///    partition-collecting parent and injects a `NetworkCoalesceExec` above its child). Then
 ///    [insert_broadcast_execs] adds `BroadcastExec` nodes on the build side of `CollectLeft`
 ///    hash joins so those build sides can later be wrapped in `NetworkBroadcastExec`.
+///    [insert_children_isolator_unions] replaces `UnionExec` nodes with single-node-correct
+///    `ChildrenIsolatorUnionExec` placeholders whose task maps are finalized during boundary
+///    injection.
 ///
 /// 2. **Boundary injection.** [inject_network_boundaries] walks the plan, computes a task count
 ///    for each node, and inserts `NetworkShuffleExec` / `NetworkBroadcastExec` /
@@ -109,6 +113,7 @@ impl QueryPlanner for DistributedQueryPlanner {
         let cfg = session_state.config_options();
 
         plan = insert_broadcast_execs(plan, cfg)?;
+        plan = insert_children_isolator_unions(plan, cfg)?;
 
         if d_cfg.dynamic_task_count {
             // The task count will be decided dynamically at execution time.

src/distributed_planner/inject_network_boundaries.rs

@@ -17,7 +17,6 @@ use datafusion::physical_plan::execution_plan::CardinalityEffect;
 use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
 use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use datafusion::physical_plan::union::UnionExec;
 use datafusion::physical_plan::{ExecutionPlan, PlanProperties};
 use datafusion::prelude::SessionConfig;
 use std::any::TypeId;
@@ -259,9 +258,9 @@ async fn _inject_network_boundaries(
     let mut task_count = estimator
         .task_estimation(&plan, nb_ctx.cfg)
         .map_or(Desired(1), |v| v.task_count);
-    if nb_ctx.d_cfg.children_isolator_unions && plan.is::<UnionExec>() {
-        // Unions have the chance to decide how many tasks they should run on. If there's a union
-        // with a bunch of children, the user might want to increase parallelism and increase the
+    if plan.is::<ChildrenIsolatorUnionExec>() {
+        // Isolating unions have the chance to decide how many tasks they should run on. If there
+        // is a union with a bunch of children, the user might want to increase parallelism and the
         // task count for the stage running that.
         let mut count = 0;
         for processed_child in processed_children.iter() {
@@ -404,10 +403,8 @@ async fn _inject_network_boundaries(
 ///   stage). Instead, rescale the boundary's input via [network_boundary_scale_input] using the
 ///   *consumer* partition and task counts of this side of the boundary, and stitch the rescaled
 ///   stage back in via [NetworkBoundary::with_input_stage].
-/// - **Eligible `UnionExec`s** (when `children_isolator_unions` is on): rewrite to
-///   [ChildrenIsolatorUnionExec] and recurse into each child with the per-child task count
-///   chosen by [ChildrenIsolatorUnionExec::from_children_and_task_counts] — each child runs
-///   isolated in its own subset of tasks.
+/// - **`ChildrenIsolatorUnionExec`s**: finalize the placeholder's task map and recurse into each
+///   child with the allocated task count. Each child runs isolated in its own subset of tasks.
 /// - **Everything else**: recurse into children with the same `task_count`, then rebuild the
 ///   node with the rebuilt children.
 impl InjectNetworkBoundaryContext<'_> {
@@ -440,8 +437,8 @@ impl InjectNetworkBoundaryContext<'_> {
             // Just annotate the network boundary and stop recursion here.
             Ok(self.plan_with_task_count(Arc::clone(plan), task_count))
 
-        // Handle ChildrenIsolatorUnionExec.
-        } else if self.d_cfg.children_isolator_unions && plan.is::<UnionExec>() {
+        // Finalize a pre-inserted ChildrenIsolatorUnionExec.
+        } else if plan.is::<ChildrenIsolatorUnionExec>() {
             // Propagating through ChildrenIsolatorUnionExec is not that easy, each child will
             // be executed in its own task, and therefore, they will act as if they were in executing
             // in a non-distributed context. The ChildrenIsolatorUnionExec itself will make sure to
@@ -629,6 +626,7 @@ impl NetworkBoundaryBuilder for CardinalityBasedNetworkBoundaryBuilder {
 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::test_utils::plans::{BuildSideOneTaskEstimator, TestPlanBuilder};
     use crate::{TaskEstimation, TaskEstimator, assert_snapshot};
     use datafusion::config::ConfigOptions;
@@ -1264,8 +1262,10 @@ mod tests {
         let plan = test_plan.physical_plan(query).await;
         let session_config = test_plan.get_ctx().copied_config();
 
-        let plan_w_broadcast = insert_broadcast_execs(plan, session_config.options())
+        let plan = insert_broadcast_execs(plan, session_config.options())
             .expect("failed to insert broadcasts");
+        let plan = insert_children_isolator_unions(plan, session_config.options())
+            .expect("failed to insert children isolator unions");
         let network_boundaries_ctx = InjectNetworkBoundaryContext {
             cfg: session_config.options(),
             d_cfg: DistributedConfig::from_config_options(session_config.options()).unwrap(),
@@ -1277,7 +1277,7 @@ mod tests {
             nb_builder: &CardinalityBasedNetworkBoundaryBuilder,
         };
 
-        let annotated = _inject_network_boundaries(plan_w_broadcast, None, &network_boundaries_ctx)
+        let annotated = _inject_network_boundaries(plan, None, &network_boundaries_ctx)
             .await
             .expect("failed to annotate plan");
         debug_annotated(&annotated, 0, &network_boundaries_ctx)

src/distributed_planner/insert_children_isolator_union.rs

@@ -0,0 +1,129 @@
+use std::sync::Arc;
+
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::error::DataFusionError;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::union::UnionExec;
+
+use crate::execution_plans::ChildrenIsolatorUnionExec;
+
+use super::DistributedConfig;
+
+/// Replaces every [`UnionExec`] with a single-node-correct [`ChildrenIsolatorUnionExec`].
+///
+/// The placeholder maps every child to one default task. Network-boundary injection later
+/// replaces that mapping with the final per-child task allocation.
+pub(super) fn insert_children_isolator_unions(
+    plan: Arc<dyn ExecutionPlan>,
+    cfg: &ConfigOptions,
+) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+    let d_cfg = DistributedConfig::from_config_options(cfg)?;
+    if !d_cfg.children_isolator_unions {
+        return Ok(plan);
+    }
+
+    plan.transform_down(|node| {
+        if !node.is::<UnionExec>() {
+            return Ok(Transformed::no(node));
+        }
+        Ok(Transformed::yes(Arc::new(
+            ChildrenIsolatorUnionExec::new_single_task(node.children().into_iter().cloned())?,
+        )))
+    })
+    .map(|transformed| transformed.data)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::distributed_ext::DistributedExt;
+    use datafusion::arrow::array::{ArrayRef, Int32Array};
+    use datafusion::arrow::record_batch::RecordBatch;
+    use datafusion::datasource::memory::MemorySourceConfig;
+    use datafusion::physical_plan::ExecutionPlanProperties;
+    use datafusion::physical_plan::common::collect;
+    use datafusion::prelude::{SessionConfig, SessionContext};
+
+    #[tokio::test]
+    async fn inserted_union_is_single_node_correct() {
+        let original = union_plan();
+        let expected_schema = original.schema();
+        let expected_partition_count = original.output_partitioning().partition_count();
+        let expected_rows = collect_rows(Arc::clone(&original)).await;
+
+        let plan = insert_children_isolator_unions(original, config(true).options())
+            .expect("insert children isolator union");
+
+        assert!(plan.is::<ChildrenIsolatorUnionExec>());
+        assert_eq!(plan.schema(), expected_schema);
+        assert_eq!(
+            plan.output_partitioning().partition_count(),
+            expected_partition_count
+        );
+        assert_eq!(collect_rows(plan).await, expected_rows);
+    }
+
+    #[test]
+    fn replaces_nested_unions() {
+        let inner = union_plan();
+        let plan =
+            UnionExec::try_new(vec![inner, memory_plan(&[5, 6])]).expect("create outer union");
+
+        let plan = insert_children_isolator_unions(plan, config(true).options())
+            .expect("insert children isolator unions");
+
+        assert!(plan.is::<ChildrenIsolatorUnionExec>());
+        assert!(plan.children()[0].is::<ChildrenIsolatorUnionExec>());
+    }
+
+    #[test]
+    fn leaves_union_when_disabled() {
+        let plan = insert_children_isolator_unions(union_plan(), config(false).options())
+            .expect("skip children isolator union");
+
+        assert!(plan.is::<UnionExec>());
+    }
+
+    fn union_plan() -> Arc<dyn ExecutionPlan> {
+        UnionExec::try_new(vec![memory_plan(&[1, 2]), memory_plan(&[3, 4])]).expect("create union")
+    }
+
+    fn memory_plan(values: &[i32]) -> Arc<dyn ExecutionPlan> {
+        let values: ArrayRef = Arc::new(Int32Array::from(values.to_vec()));
+        let batch = RecordBatch::try_from_iter([("value", values)]).expect("create record batch");
+        MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None)
+            .expect("create memory plan")
+    }
+
+    async fn collect_rows(plan: Arc<dyn ExecutionPlan>) -> Vec<i32> {
+        let context = SessionContext::new().task_ctx();
+        let mut rows = Vec::new();
+        for partition in 0..plan.output_partitioning().partition_count() {
+            let stream = plan
+                .execute(partition, Arc::clone(&context))
+                .expect("execute partition");
+            let batches = collect(stream).await.expect("collect partition");
+            for batch in batches {
+                let values = batch
+                    .column(0)
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .expect("int32 column");
+                for row in 0..values.len() {
+                    rows.push(values.value(row));
+                }
+            }
+        }
+        rows
+    }
+
+    fn config(children_isolator_unions: bool) -> SessionConfig {
+        let mut config = SessionConfig::new();
+        config.set_distributed_option_extension(DistributedConfig {
+            children_isolator_unions,
+            ..Default::default()
+        });
+        config
+    }
+}

src/distributed_planner/mod.rs

@@ -2,6 +2,7 @@ mod distributed_config;
 mod distributed_query_planner;
 mod inject_network_boundaries;
 mod insert_broadcast;
+mod insert_children_isolator_union;
 mod network_boundary;
 mod partial_reduce_below_network_shuffles;
 mod prepare_network_boundaries;

src/execution_plans/children_isolator_union.rs

@@ -129,6 +129,15 @@ impl ChildWeight {
 }
 
 impl ChildrenIsolatorUnionExec {
+    /// Creates a single-node union placeholder with every child assigned to the default task.
+    pub(crate) fn new_single_task(
+        children: impl IntoIterator<Item = Arc<dyn ExecutionPlan>>,
+    ) -> Result<Self, DataFusionError> {
+        let children = children.into_iter().collect_vec();
+        let child_count = children.len();
+        Self::from_children_and_weights(children, vec![ChildWeight::desired(1.0); child_count], 1)
+    }
+
     pub(crate) fn from_children_and_weights(
         children: impl IntoIterator<Item = Arc<dyn ExecutionPlan>>,
         children_weights: impl IntoIterator<Item = ChildWeight>,