This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch branch-54
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/branch-54 by this push:
new aa0b62622 fix: map a task's partition through UnionExec when
restricting file scans (#2070) (#2131)
aa0b62622 is described below
commit aa0b62622e94baf829c6d6686656091d7bac8d73
Author: Andy Grove <[email protected]>
AuthorDate: Wed Jul 22 06:40:15 2026 -0600
fix: map a task's partition through UnionExec when restricting file scans
(#2070) (#2131)
---
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 e579621c9..ea61bd536 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<dyn ExecutionPlan>,
+ partition: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ if plan.downcast_ref::<UnionExec>().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::<Result<Vec<_>>>()?;
+ 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan>) -> Vec<Vec<usize>> {
+ 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<dyn ExecutionPlan> =
+ 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<dyn ExecutionPlan> =
+ 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::<usize>(),
+ 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]);
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
