From 2d4327313ed8a3e202f080fc751bfd6dc2da7153 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 16 Jul 2026 15:23:35 -0600 Subject: [PATCH] fix: map a task's partition through UnionExec when restricting file scans (#2070) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A task executes one partition of its stage, so each file scan beneath it is pinned to the single file group that partition reads — otherwise DataFusion 54's shared work-queue lets a lone `execute(p)` drain the queue and scan the whole table. That pinning applied the stage's partition number directly to every scan, which is only right when the partition passes straight through. A `UnionExec` concatenates its children's partitions: with N-partition children, stage partition `N + k` is the right child's local partition `k`. Applying the stage number there addressed the wrong group, and for partitions past a scan's group count it addressed none at all, leaving that scan unrestricted and free to read its entire table. The result was a silently inflated answer — no error, just several times too much data. Walk the plan from the stage root instead, mirroring `UnionExec::execute`'s own partition arithmetic, so each scan is pinned to the group its local partition reads. Children a partition never reaches are left untouched; this task does not execute them. Minimal case, SF1, executor `-c 8`: select round(sum(p),2) from (select ws_ext_sales_price p from web_sales union all select cs_ext_sales_price p from catalog_sales); -- expected 5492796521.50 -- was 49389026433.70 at --partitions 16 (~9x, varying run to run) Correct now at every partition count tried (4, 5, 6, 8, 16, 20, 32); previously wrong whenever the union stage exceeded one wave of executor task slots. TPC-DS q2 and q5 now match single-process DataFusion and join the correctness gate. Both were skipped as "non-deterministic (ORDER BY ties)"; that diagnosis was wrong — with a total order imposed and LIMIT removed they still mismatched, because they were hitting this bug. q31 and q71 are genuinely ordering-only and stay skipped. --- ballista/executor/src/execution_engine.rs | 141 +++++++++++++++++++++- 1 file changed, 137 insertions(+), 4 deletions(-) diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index e579621c9e..ea61bd5361 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -36,6 +36,7 @@ use datafusion::error::{DataFusionError, Result}; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::union::UnionExec; use datafusion::prelude::SessionConfig; use log::warn; use std::any::Any; @@ -165,6 +166,65 @@ fn restrict_scan_to_partition( Some(DataSourceExec::from_data_source(config)) } +/// Restrict every file scan in `plan` to the file group the task running +/// `partition` will actually poll. +/// +/// A task executes exactly one partition of its stage, so each scan beneath it +/// must be pinned to the one file group that partition reads (see +/// [`restrict_scan_to_partition`]). Which group that is depends on the path from +/// the stage root: most operators pass a partition straight through to their +/// children, but a `UnionExec` concatenates its children's partitions, so its +/// child `i` sees `partition` minus the partition counts of the children before +/// it. Applying the stage's partition number directly to a scan under a union +/// therefore addresses the wrong group, and for partitions past the scan's group +/// count it addresses none at all — leaving that scan unrestricted and free to +/// read the whole table. +/// +/// Children a partition never reaches are left untouched: this task will not +/// execute them. +fn restrict_scans_for_partition( + plan: &Arc, + partition: usize, +) -> Result> { + if plan.downcast_ref::().is_some() { + // `UnionExec::execute(p)` walks its children subtracting each one's + // partition count until `p` lands inside a child, then executes that + // child alone. Mirror that walk so the scan under the serving child is + // pinned to the group that child's local partition reads. + let mut local = Some(partition); + let mut children = Vec::with_capacity(plan.children().len()); + for child in plan.children() { + let count = child.properties().output_partitioning().partition_count(); + match local { + Some(p) if p < count => { + children.push(restrict_scans_for_partition(child, p)?); + local = None; + } + Some(p) => { + local = Some(p - count); + children.push(Arc::clone(child)); + } + None => children.push(Arc::clone(child)), + } + } + return Arc::clone(plan).with_new_children(children); + } + + if let Some(rewritten) = restrict_scan_to_partition(plan, partition) { + return Ok(rewritten); + } + + let children = plan + .children() + .into_iter() + .map(|child| restrict_scans_for_partition(child, partition)) + .collect::>>()?; + if children.is_empty() { + return Ok(Arc::clone(plan)); + } + Arc::clone(plan).with_new_children(children) +} + impl ExecutionEngine for DefaultExecutionEngine { fn create_query_stage_exec( &self, @@ -188,15 +248,12 @@ impl ExecutionEngine for DefaultExecutionEngine { reader.with_work_dir(work_dir.to_string()), ))), } - } else if let Some(rewritten) = - restrict_scan_to_partition(&p, partition_id) - { - Ok(Transformed::yes(rewritten)) } else { Ok(Transformed::no(p)) } })? .data; + let plan = restrict_scans_for_partition(&plan, partition_id)?; // the query plan created by the scheduler always starts with a shuffle writer // (either ShuffleWriterExec or SortShuffleWriterExec) @@ -383,4 +440,80 @@ mod tests { let plan: Arc = Arc::new(EmptyExec::new(schema)); assert!(restrict_scan_to_partition(&plan, 0).is_none()); } + + /// Find the file-group layout of each scan under a union, left child first. + fn union_group_counts(plan: &Arc) -> Vec> { + plan.children().into_iter().map(group_file_counts).collect() + } + + /// A `UnionExec` concatenates its children's partitions, so stage partition + /// `left_count + k` is the right child's local partition `k`. The scan under + /// the serving child must be pinned to that local group. + #[test] + fn union_partition_maps_to_the_right_childs_local_group() { + let union: Arc = + UnionExec::try_new(vec![scan_with_file_groups(4), scan_with_file_groups(4)]) + .unwrap(); + assert_eq!( + union.properties().output_partitioning().partition_count(), + 8 + ); + + // Partition 1 is served by the left child's local partition 1. + let restricted = restrict_scans_for_partition(&union, 1).unwrap(); + assert_eq!( + union_group_counts(&restricted), + vec![vec![0, 1, 0, 0], vec![1, 1, 1, 1]], + "left scan pinned to group 1; right scan is never polled for this partition" + ); + + // Partition 6 is served by the right child's local partition 2 (6 - 4). + // Before the mapping existed this addressed no group at all and left the + // scan free to read the whole table. + let restricted = restrict_scans_for_partition(&union, 6).unwrap(); + assert_eq!( + union_group_counts(&restricted), + vec![vec![1, 1, 1, 1], vec![0, 0, 1, 0]], + "right scan pinned to local group 2" + ); + } + + /// Every partition of a union must pin exactly one group in exactly one + /// child, so across the whole stage each file group is read exactly once. + #[test] + fn every_union_partition_pins_exactly_one_group() { + let union: Arc = + UnionExec::try_new(vec![scan_with_file_groups(3), scan_with_file_groups(3)]) + .unwrap(); + for partition in 0..6 { + let restricted = restrict_scans_for_partition(&union, partition).unwrap(); + let counts = union_group_counts(&restricted); + let (serving, idle) = if partition < 3 { (0, 1) } else { (1, 0) }; + let expected_local = partition % 3; + assert_eq!( + counts[serving].iter().sum::(), + 1, + "partition {partition}: serving child must keep exactly one group" + ); + assert_eq!( + counts[serving][expected_local], 1, + "partition {partition}: expected local group {expected_local}, got {:?}", + counts[serving] + ); + assert_eq!( + counts[idle], + vec![1, 1, 1], + "partition {partition}: the child it never polls is left untouched" + ); + } + } + + /// Without a union a partition passes straight through, so the behaviour is + /// unchanged from restricting the scan directly. + #[test] + fn scan_without_union_is_pinned_to_its_own_partition() { + let plan = scan_with_file_groups(4); + let restricted = restrict_scans_for_partition(&plan, 2).unwrap(); + assert_eq!(group_file_counts(&restricted), vec![0, 0, 1, 0]); + } }