fix union recomputing child properties (datafusion-contrib#423)

The `ChildrenIsolatorUnionExec` when replacing children would keep the
old cached properties (things like `output_partitioning()`) rather than
recomputing them via the replaced children. This would cause stale
properties to remain.

This is evident in a scenario like this:

```text
ChildrenIsolatorUnionExec -> advertises output partitions = 2
    child 0: 2 partitions
    child 1: 1 partition
```

This is ok because:
```
children say: max child/task output = 2
ChildrenIsolatorUnionExec properties: output_partitioning = 2
```

But say an optimizer rule replaces child 0 with a new plan that has more
partitions:

```text
ChildrenIsolatorUnionExec -> advertises output partitions = 2  (now stale)
    child 0: 8 partitions (updated)
    child 1: 1 partition
```

Now `ChildrenIsolatorUntionExec is incorrect about its output.

**If `ChildrenIsolatorUntionExec` advertises too few partitions:**

```text
actual child partitions: 0 1 2 3
advertised ChildrenIsolatorUntionExec partitions:  0 1

scheduled downstream work:  0 1
missing work: 2 3
```

This can cause correctness issues if those extra partitions might never
be executed.

**If `ChildrenIsolatorUntionExec` advertises too many partitions:**

```text
actual child partitions: 0 1
advertised CIU partitions:  0 1 2 3
```

Then downstream operators might schedule extra empty work.

Author: gene-bordegaray

src/execution_plans/children_isolator_union.rs

@@ -148,6 +148,18 @@ impl ChildrenIsolatorUnionExec {
             task_idx_map,
         })
     }
+
+    fn child_task_counts(&self) -> Vec<usize> {
+        // Preserve the task assignment in task_idx_map and allow child plans to be
+        // replaced and properties to be recomputed from these new children.
+        let mut counts = vec![0; self.children.len()];
+        for children_in_task in &self.task_idx_map {
+            for (child_idx, child_task_ctx) in children_in_task {
+                counts[*child_idx] = counts[*child_idx].max(child_task_ctx.task_count);
+            }
+        }
+        counts
+    }
 }
 
 impl DisplayAs for ChildrenIsolatorUnionExec {
@@ -205,9 +217,11 @@ impl ExecutionPlan for ChildrenIsolatorUnionExec {
                 self.children.len()
             );
         }
-        let mut clone = self.as_ref().clone();
-        clone.children = children;
-        Ok(Arc::new(clone))
+        Ok(Arc::new(Self::from_children_and_task_counts(
+            children,
+            self.child_task_counts(),
+            self.task_idx_map.len(),
+        )?))
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
@@ -452,6 +466,11 @@ fn split_children(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::physical_plan::ExecutionPlan;
+    use datafusion::physical_plan::empty::EmptyExec;
+    use datafusion::physical_plan::repartition::RepartitionExec;
+    use std::sync::Arc;
 
     #[test]
     fn children_split_all_1_task() -> Result<(), Box<dyn std::error::Error>> {
@@ -531,4 +550,40 @@ mod tests {
             task_count,
         }
     }
+
+    fn empty_partitions(partitions: usize) -> Arc<dyn ExecutionPlan> {
+        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)]));
+        Arc::new(EmptyExec::new(schema).with_partitions(partitions))
+    }
+
+    #[test]
+    fn with_new_children_recomputes_partitioning() -> Result<(), Box<dyn std::error::Error>> {
+        let child0 = empty_partitions(2);
+        let child1 = empty_partitions(1);
+
+        let union = Arc::new(ChildrenIsolatorUnionExec::from_children_and_task_counts(
+            vec![Arc::clone(&child0), Arc::clone(&child1)],
+            vec![2, 1],
+            3,
+        )?);
+        assert_eq!(
+            union.properties().output_partitioning().partition_count(),
+            2
+        );
+
+        // Rewrites can wrap a child in a partition-changing operator. The cached union
+        // properties must be part of the replacement. If not, the downstream planning can skip
+        // partitions that the new child produces.
+        let repartitioned_child0: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new(
+            child0,
+            Partitioning::RoundRobinBatch(8),
+        )?);
+        let rebuilt = union.with_new_children(vec![repartitioned_child0, child1])?;
+        assert_eq!(
+            rebuilt.properties().output_partitioning().partition_count(),
+            8
+        );
+
+        Ok(())
+    }
 }

tests/work_unit_feed.rs

@@ -376,9 +376,9 @@ mod tests {
         assert_snapshot!(plan + &results, @r"
         ┌───── DistributedExec ── Tasks: t0:[p0]
         │ SortPreservingMergeExec: [tag@0 ASC NULLS LAST, task@1 ASC NULLS LAST, partition@2 ASC NULLS LAST, letter@3 ASC NULLS LAST]
-        │   [Stage 1] => NetworkCoalesceExec: output_partitions=12, input_tasks=3
+        │   [Stage 1] => NetworkCoalesceExec: output_partitions=6, input_tasks=3
         └──────────────────────────────────────────────────
-          ┌───── Stage 1 ── Tasks: t0:[p0..p3] t1:[p4..p7] t2:[p8..p11]
+          ┌───── Stage 1 ── Tasks: t0:[p0..p1] t1:[p2..p3] t2:[p4..p5]
           │ DistributedUnionExec: t0:[c0(0/2)] t1:[c0(1/2)] t2:[c1]
           │   SortExec: expr=[tag@0 ASC NULLS LAST, task@1 ASC NULLS LAST, partition@2 ASC NULLS LAST, letter@3 ASC NULLS LAST], preserve_partitioning=[true]
           │     RowGeneratorExec: tag=feed, tasks=2, rows_per_partition=[[2], [1], [1], [2]]