This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 083f7a2cb feat(aqe): CoalescePartitionsRule — shuffle-partition
coalescing on resolved stats (#1684)
083f7a2cb is described below
commit 083f7a2cb528805524ee97cc9d8f2059390fba71
Author: mete <[email protected]>
AuthorDate: Mon May 18 14:18:46 2026 +0100
feat(aqe): CoalescePartitionsRule — shuffle-partition coalescing on
resolved stats (#1684)
* feat: AQE rule to coalesce shuffle partitions (Spark-port)
After an upstream shuffle stage finalizes with per-partition byte stats,
this rule rewrites the downstream stage to read K < M logical partitions
when several upstream partitions are near-empty. Mirrors Spark's
CoalesceShufflePartitions for distributed query execution.
The rule is opt-in (`ballista.coalesce.enabled` defaults to false). The
only tunable is `ballista.coalesce.target_partition_bytes` (advisory,
64 MB default) — Spark's `advisoryPartitionSizeInBytes`.
## How it plugs into the AQE planner
The rule is invoked **per stage** in `actionable_stages()`, right
before each newly-runnable stage is handed to `BallistaAdapter`. It is
deliberately NOT registered in `default_optimizers()` because that
chain runs on the entire residual plan tree on every `replan_stages()`
call, which causes two problems:
1. Cross-stage alignment groups. The walker collects "topmost
ExchangeExecs", which spans multiple future stages' inputs into
one alignment group. A K decision driven by stage N's byte
distribution gets stamped onto stage N+1's input exchanges too.
2. Stale state across fires. `set_coalesce` uses interior mutability
on a shared `Arc<ExchangeExec>`. A fire that bails on K=1
degenerate doesn't invalidate the K=5 a previous fire wrote on
the same exchange — the next stage then sees mismatched leg Ks.
Per-stage invocation fixes both: each call receives one stage's plan
as its root, the walker descends to that stage's input exchanges only,
and the K decision is local to that stage. The "unresolved leaf" bail
path becomes unreachable by construction (a stage becomes runnable
precisely when all its inputs are resolved).
## Example: TPC-H Q8 stage 11, both legs aligned
After stage 10 finalizes (writes `Hash(o_custkey, 16)`), the rule
fires on stage 11's plan:
AdaptiveDatafusionExec
ProjectionExec
SortMergeJoinExec on (o_custkey, c_custkey)
SortExec
ExchangeExec Hash([o_custkey], 16) plan_id=11 <- leaf #1,
resolved
...stage 10 writer output...
SortExec
ExchangeExec Hash([c_custkey], 16) plan_id=4 <- leaf #2,
resolved
...customer table scan...
Walker collects exactly 2 leaves: {plan_id=11, plan_id=4}. Per-leaf
bytes are ~19 MB each across 16 partitions, summed [38M×16]. Bin-pack
at 64 MB target → K=5. Both leaves get the SAME `CoalescePlan` (K=5,
M=16, identical group mapping) → the SMJ runs with 5 partitions on
each side, hash buckets stay aligned.
## Example: Q8 stage 12, three resolved siblings, K=1 degenerate
When stage 11 finalizes and stage 12 is surfaced, the walker sees:
leaves = [plan_id=12 (c_nationkey from stage 11),
plan_id=5 (n_nationkey from stage 5),
plan_id=6 (n_nationkey from stage 6),
plan_id=7 (r_regionkey from stage 7)]
All four resolved, but summed bytes total only ~76 MB (the multi-join
filtered hard). Bin-pack returns K=1 → degenerate → the rule no-ops
on this stage. The SMJ runs at native M=16 across all four legs.
No partition-count mismatch.
## Components
- Proto: `CoalescePlan` / `PartitionGroup` on `ShuffleReaderExecNode`
and `UnresolvedShuffleExecNode` (field numbers 7/8 — non-clashing
with PR #1647's `broadcast` / `upstream_partition_count` fields 5/6).
- Reader: `ShuffleReaderExec::try_new_coalesced` builds the K-shape
pre-concatenated reader; legacy `try_new` path byte-for-byte
unchanged when no `CoalescePlan` is attached. Coexists with
`try_new_broadcast` from PR #1647.
- Algorithm: `split_size_list_by_target_size` is a verbatim Rust port
of Spark's legacy `splitSizeListByTargetSize` — merged-factor early
flush and small-tail post-loop merge preserved.
- Rule: `CoalescePartitionsRule` (unit struct), invoked per-stage in
`actionable_stages()`. ExchangeExec's display conditionally appends
`coalesce=K of M` only when attached, so existing AQE snapshots that
don't involve coalesce are unchanged.
## Behavior preservation
When `coalesce.enabled=false`, the rule short-circuits as the first
statement of `optimize()` and returns the input Arc verbatim. The
reader path with no `CoalescePlan` attached is byte-for-byte identical
to today's reader. Orthogonal to PR #1647 (broadcast hash join,
merged) and PR #1649 (lazy AQE planner, merged). All three use
DataFusion's standard `PhysicalOptimizerRule` trait.
## Validation
- `cargo test --workspace`: ~600 tests pass, 0 failures.
- `cargo clippy --workspace --all-targets --tests`: 0 warnings.
- TPC-H SF100 sanity, 22 queries × 2 join variants (hash-pref and
sort-pref), `coalesce.enabled=true`: 44/44 rc=0, row counts
identical across variants, 9 queries get meaningful K reductions
(K=2..K=13 across nation/region/supplier/customer joins).
* refactor(coalesce): address PR review — namespace keys, Float64 storage
- Rename ballista.coalesce.* → ballista.planner.coalesce.* to match the
existing ballista.planner.* convention used for adaptive/broadcast knobs
- Add DataType::Float64 to BallistaConfig::parse_value
- Switch small/merged partition factors from Utf8 to Float64 storage
- Extract get_float_setting helper mirroring get_usize_setting; drop the
manual f64::from_str + unwrap_or fallback at the call sites
- Document the neighboring-partitions-only grouping rule in the module doc
* fix(coalesce): bail on heterogeneous M to avoid Q22 panic
Guard against the byte-sum loop indexing past `summed.len()` when an
alignment group contains leaf Exchanges with different partition_counts
(e.g. Q22's scalar avg subquery at M=14 alongside hash-join legs at M=48).
* style: cargo fmt the Q22 guard
* refactor(coalesce): address PR review nits
- Remove dead ShuffleReaderExec branch in BallistaAdapter — the coalesce
rule never produces one (it only calls set_coalesce on the Exchange),
so transform_children only ever sees ExchangeExec.
- Restore PartitionStats.num_bytes to pub(crate); rule reads via the
existing num_bytes() getter instead of direct field access.
---
.gitignore | 2 +-
ballista/core/proto/ballista.proto | 19 +
ballista/core/src/config.rs | 82 +++-
ballista/core/src/execution_plans/mod.rs | 2 +-
.../core/src/execution_plans/shuffle_reader.rs | 273 ++++++++++++-
.../core/src/execution_plans/unresolved_shuffle.rs | 102 ++++-
ballista/core/src/extension.rs | 103 ++++-
ballista/core/src/serde/generated/ballista.rs | 24 ++
ballista/core/src/serde/mod.rs | 338 +++++++++++++++-
ballista/core/src/serde/scheduler/mod.rs | 8 +
ballista/scheduler/src/state/aqe/adapter.rs | 43 +-
.../scheduler/src/state/aqe/coalesce/algorithm.rs | 241 ++++++++++++
.../state/aqe/{optimizer_rule => coalesce}/mod.rs | 17 +-
ballista/scheduler/src/state/aqe/execution_plan.rs | 58 ++-
ballista/scheduler/src/state/aqe/mod.rs | 3 +-
.../aqe/optimizer_rule/coalesce_partitions.rs | 311 +++++++++++++++
.../scheduler/src/state/aqe/optimizer_rule/mod.rs | 2 +
ballista/scheduler/src/state/aqe/planner.rs | 17 +-
.../scheduler/src/state/aqe/test/coalesce_rule.rs | 434 +++++++++++++++++++++
ballista/scheduler/src/state/aqe/test/mod.rs | 2 +
20 files changed, 2049 insertions(+), 32 deletions(-)
diff --git a/.gitignore b/.gitignore
index 20c28224d..42f5e3ebe 100644
--- a/.gitignore
+++ b/.gitignore
@@ -115,4 +115,4 @@ CLAUDE.md
# git worktrees (local only)
.worktrees/
# ignore insta captures
-*-snap
\ No newline at end of file
+*-snap
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index bdd68ba88..cb209b55e 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -78,6 +78,8 @@ message UnresolvedShuffleExecNode {
datafusion.Partitioning partitioning = 5;
bool broadcast = 6;
uint32 upstream_partition_count = 7;
+ // Optional coalesce metadata. Absent means "no coalesce" (legacy one-to-one
read behavior).
+ optional CoalescePlan coalesce = 8;
}
message ShuffleReaderExecNode {
@@ -88,6 +90,8 @@ message ShuffleReaderExecNode {
datafusion.Partitioning partitioning = 4;
bool broadcast = 5;
uint32 upstream_partition_count = 6;
+ // Optional coalesce metadata. Absent means "no coalesce" (legacy one-to-one
read behavior).
+ optional CoalescePlan coalesce = 7;
}
message ShuffleReaderPartition {
@@ -95,6 +99,21 @@ message ShuffleReaderPartition {
repeated PartitionLocation location = 1;
}
+// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
+// Empty when no coalesce is applied (the optional field on the parent message
is absent).
+message CoalescePlan {
+ // Original number of upstream partitions (M) before coalescing. Required
for EXPLAIN's "K of M" rendering.
+ uint32 upstream_partition_count = 1;
+ // Coalesced output groups. Length is K (the post-coalesce partition count).
+ repeated PartitionGroup groups = 2;
+}
+
+// One coalesced output partition's source list: a set of upstream partition
indices in [0, upstream_partition_count).
+// Default algorithm produces only contiguous indices, but proto allows
arbitrary index sets for future strategies.
+message PartitionGroup {
+ repeated uint32 upstream_indices = 1;
+}
+
///////////////////////////////////////////////////////////////////////////////////////////////////
// Ballista Scheduling
///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 334d5b991..fbd0de708 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -86,6 +86,21 @@ pub const
BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str =
pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
"ballista.optimizer.broadcast_join_threshold_bytes";
+/// Configuration key to enable AQE coalesce-shuffle-partitions rule.
+/// Disabled by default — opt in when the workload benefits from larger
+/// downstream tasks more than from preserved parallelism.
+pub const BALLISTA_COALESCE_ENABLED: &str =
"ballista.planner.coalesce.enabled";
+/// Configuration key for the target post-coalesce partition byte size (bytes).
+/// Mirrors Spark's `spark.sql.adaptive.advisoryPartitionSizeInBytes`.
+pub const BALLISTA_COALESCE_TARGET_PARTITION_BYTES: &str =
+ "ballista.planner.coalesce.target_partition_bytes";
+/// Configuration key for the small-partition merge factor (Spark legacy
semantics).
+pub const BALLISTA_COALESCE_SMALL_PARTITION_FACTOR: &str =
+ "ballista.planner.coalesce.small_partition_factor";
+/// Configuration key for the merged-partition early-flush factor (Spark
legacy semantics).
+pub const BALLISTA_COALESCE_MERGED_PARTITION_FACTOR: &str =
+ "ballista.planner.coalesce.merged_partition_factor";
+
/// Result type for configuration parsing operations.
pub type ParseResult<T> = result::Result<T, String>;
use std::sync::LazyLock;
@@ -178,7 +193,33 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
ConfigEntry::new(BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS.to_string(),
"Wait time in milliseconds between IO retries in the
Ballista client.".to_string(),
DataType::UInt64,
- Some(3000.to_string()))
+ Some(3000.to_string())),
+ ConfigEntry::new(BALLISTA_COALESCE_ENABLED.to_string(),
+ "Enables the AQE coalesce-shuffle-partitions rule. \
+ Disabled by default — opt in when fewer/larger \
+ downstream tasks matter more than
parallelism.".to_string(),
+ DataType::Boolean,
+ Some(false.to_string())),
+ ConfigEntry::new(
+ BALLISTA_COALESCE_TARGET_PARTITION_BYTES.to_string(),
+ "Target post-coalesce partition byte size in bytes. Mirrors
Spark's \
+ advisoryPartitionSizeInBytes."
+ .to_string(),
+ DataType::UInt64,
+ Some((64 * 1024 * 1024_usize).to_string()),
+ ),
+ ConfigEntry::new(
+ BALLISTA_COALESCE_SMALL_PARTITION_FACTOR.to_string(),
+ "Small-partition merge factor (Spark legacy).".to_string(),
+ DataType::Float64,
+ Some("0.2".to_string()),
+ ),
+ ConfigEntry::new(
+ BALLISTA_COALESCE_MERGED_PARTITION_FACTOR.to_string(),
+ "Merged-partition early-flush factor (Spark legacy).".to_string(),
+ DataType::Float64,
+ Some("1.2".to_string()),
+ ),
];
entries
.into_iter()
@@ -272,6 +313,11 @@ impl BallistaConfig {
DataType::Utf8 => {
val.to_string();
}
+ DataType::Float64 => {
+ val.to_string()
+ .parse::<f64>()
+ .map_err(|e| format!("{e:?}"))?;
+ }
_ => {
return Err(format!("not support data type: {data_type}"));
}
@@ -383,6 +429,27 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES)
}
+ /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
+ pub fn coalesce_enabled(&self) -> bool {
+ self.get_bool_setting(BALLISTA_COALESCE_ENABLED)
+ }
+
+ /// Returns the target post-coalesce partition byte size in bytes
+ /// (Spark's `advisoryPartitionSizeInBytes`).
+ pub fn coalesce_target_partition_bytes(&self) -> u64 {
+ self.get_usize_setting(BALLISTA_COALESCE_TARGET_PARTITION_BYTES) as u64
+ }
+
+ /// Returns the small-partition merge factor (Spark legacy).
+ pub fn coalesce_small_partition_factor(&self) -> f64 {
+ self.get_float_setting(BALLISTA_COALESCE_SMALL_PARTITION_FACTOR)
+ }
+
+ /// Returns the merged-partition early-flush factor (Spark legacy).
+ pub fn coalesce_merged_partition_factor(&self) -> f64 {
+ self.get_float_setting(BALLISTA_COALESCE_MERGED_PARTITION_FACTOR)
+ }
+
/// Should client employ pull or push job tracking strategy
pub fn client_pull(&self) -> bool {
self.get_bool_setting(BALLISTA_CLIENT_PULL)
@@ -439,6 +506,19 @@ impl BallistaConfig {
v.to_string()
}
}
+
+ #[allow(dead_code)]
+ fn get_float_setting(&self, key: &str) -> f64 {
+ if let Some(v) = self.settings.get(key) {
+ // infallible because we validate all configs in the constructor
+ v.parse::<f64>().unwrap()
+ } else {
+ let entries = Self::valid_entries();
+ // infallible because we validate all configs in the constructor
+ let v = entries.get(key).unwrap().default_value.as_ref().unwrap();
+ v.parse::<f64>().unwrap()
+ }
+ }
}
impl datafusion::config::ExtensionOptions for BallistaConfig {
diff --git a/ballista/core/src/execution_plans/mod.rs
b/ballista/core/src/execution_plans/mod.rs
index d94fa4476..ae46fad68 100644
--- a/ballista/core/src/execution_plans/mod.rs
+++ b/ballista/core/src/execution_plans/mod.rs
@@ -31,7 +31,7 @@ use std::path::{Path, PathBuf};
use datafusion::common::exec_err;
pub use distributed_explain_analyze::DistributedExplainAnalyzeExec;
pub use distributed_query::DistributedQueryExec;
-pub use shuffle_reader::ShuffleReaderExec;
+pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec};
pub use shuffle_reader::{stats_for_partition, stats_for_partitions};
pub use shuffle_writer::DEFAULT_SHUFFLE_CHANNEL_CAPACITY;
pub use shuffle_writer::ShuffleWriterExec;
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs
b/ballista/core/src/execution_plans/shuffle_reader.rs
index ee75917c5..fff578350 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -60,6 +60,39 @@ use std::task::{Context, Poll};
use tokio::sync::{Semaphore, mpsc};
use tokio_stream::wrappers::ReceiverStream;
+/// Coalesce plan attached to a `ShuffleReaderExec` or `UnresolvedShuffleExec`.
+///
+/// Produced by the AQE `CoalescePartitionsRule` and round-tripped through
+/// proto so it survives stage retries. Absent (`None` on the parent operator)
+/// means "no coalesce" — the existing one-to-one read behavior.
+///
+/// `K = self.groups.len()` is the post-coalesce partition count.
+/// `M = self.upstream_partition_count` is the original upstream partition
count.
+/// EXPLAIN renders this as `coalesce: K of M` (see `DisplayAs::fmt_as`).
+///
+/// Note: `Default` is intentionally NOT derived. Callers must construct
explicitly
+/// to keep "absent coalesce" (`Option::None`) semantically distinct from
"empty plan".
+#[derive(Debug, Clone, PartialEq)]
+pub struct CoalescePlan {
+ /// Original upstream partition count (M) before coalescing.
+ pub upstream_partition_count: u32,
+ /// Output partition groups. Length is K (the post-coalesce partition
count).
+ pub groups: Vec<PartitionGroup>,
+}
+
+/// One output partition's upstream-index list.
+///
+/// Each value is an index into the M-shape `Vec<Vec<PartitionLocation>>`
produced by
+/// the upstream `ShuffleWriterExec` (or `SortShuffleWriterExec`). The default
+/// `split_size_list_by_target_size` algorithm produces only contiguous ranges,
+/// but proto permits arbitrary index sets for future strategies.
+#[derive(Debug, Clone, PartialEq)]
+pub struct PartitionGroup {
+ /// Indices into the upstream `Vec<Vec<PartitionLocation>>` that this
output
+ /// partition concatenates.
+ pub upstream_indices: Vec<u32>,
+}
+
/// ShuffleReaderExec reads partitions that have already been materialized by
a ShuffleWriterExec
/// being executed by an executor
#[derive(Debug, Clone)]
@@ -76,8 +109,14 @@ pub struct ShuffleReaderExec {
pub broadcast: bool,
/// Number of shuffle output partitions on the upstream stage. Useful for
/// metrics and EXPLAIN output. For non-broadcast readers this equals
- /// `partition.len()`.
+ /// `partition.len()` (or, when coalesced,
`coalesce.upstream_partition_count`).
pub upstream_partition_count: usize,
+ /// Optional coalesce metadata. `None` means the reader behaves
identically to
+ /// the legacy one-to-one read (no coalescing). When `Some`,
`partition.len()` equals
+ /// `coalesce.groups.len()` (= K, the post-coalesce partition count); the
rule
+ /// is responsible for pre-concatenating the M-shape upstream
+ /// `Vec<Vec<PartitionLocation>>` into K-shape before invoking
`try_new_coalesced`.
+ pub coalesce: Option<CoalescePlan>,
/// Execution metrics
metrics: ExecutionPlanMetricsSet,
properties: Arc<PlanProperties>,
@@ -107,6 +146,7 @@ impl ShuffleReaderExec {
partition,
broadcast: false,
upstream_partition_count,
+ coalesce: None,
metrics: ExecutionPlanMetricsSet::new(),
properties,
work_dir: None, // to be updated at the executor side
@@ -136,6 +176,58 @@ impl ShuffleReaderExec {
partition: vec![all_locations],
broadcast: true,
upstream_partition_count,
+ coalesce: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ properties,
+ work_dir: None, // to be updated at the executor side
+ client_pool: None, // to be updated at the executor side
+ })
+ }
+
+ /// Create a new coalesced ShuffleReaderExec.
+ ///
+ /// `partition` MUST be the K-shape, pre-concatenated
`Vec<Vec<PartitionLocation>>`
+ /// produced by the AQE rule: each output index `idx` in `0..K` holds the
+ /// concatenation of the upstream `Vec<PartitionLocation>`s named by
+ /// `coalesce.groups[idx].upstream_indices`. `partitioning` MUST
+ /// be `Partitioning::Hash(keys, K)` (or another `Partitioning` of width
K) so
+ /// `partition_count() == K` and `Partitioning::Hash` co-partitioning is
preserved
+ /// across joins.
+ ///
+ /// In debug builds this constructor asserts `partition.len() ==
coalesce.groups.len()`
+ /// and `partitioning.partition_count() == coalesce.groups.len()` to catch
+ /// rule-side mistakes early; release builds skip the check.
+ pub fn try_new_coalesced(
+ stage_id: usize,
+ partition: Vec<Vec<PartitionLocation>>,
+ coalesce: CoalescePlan,
+ schema: SchemaRef,
+ partitioning: Partitioning,
+ ) -> Result<Self> {
+ debug_assert_eq!(
+ partition.len(),
+ coalesce.groups.len(),
+ "K-shape partition vector length must equal coalesce.groups.len()",
+ );
+ debug_assert_eq!(
+ partitioning.partition_count(),
+ coalesce.groups.len(),
+ "partitioning.partition_count() must equal coalesce.groups.len()
(= K)",
+ );
+ let upstream_partition_count = coalesce.upstream_partition_count as
usize;
+ let properties = Arc::new(PlanProperties::new(
+
datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
+ partitioning,
+
datafusion::physical_plan::execution_plan::EmissionType::Incremental,
+ datafusion::physical_plan::execution_plan::Boundedness::Bounded,
+ ));
+ Ok(Self {
+ stage_id,
+ schema,
+ partition,
+ broadcast: false,
+ upstream_partition_count,
+ coalesce: Some(coalesce),
metrics: ExecutionPlanMetricsSet::new(),
properties,
work_dir: None, // to be updated at the executor side
@@ -151,6 +243,7 @@ impl ShuffleReaderExec {
partition: self.partition.clone(),
broadcast: self.broadcast,
upstream_partition_count: self.upstream_partition_count,
+ coalesce: self.coalesce.clone(),
metrics: self.metrics.clone(),
properties: self.properties.clone(),
work_dir: Some(work_dir),
@@ -165,6 +258,7 @@ impl ShuffleReaderExec {
partition: self.partition.clone(),
broadcast: self.broadcast,
upstream_partition_count: self.upstream_partition_count,
+ coalesce: self.coalesce.clone(),
metrics: self.metrics.clone(),
properties: self.properties.clone(),
work_dir: self.work_dir.clone(),
@@ -192,7 +286,16 @@ impl DisplayAs for ShuffleReaderExec {
f,
"ShuffleReaderExec: partitioning: {}",
self.properties.partitioning,
- )
+ )?;
+ if let Some(c) = &self.coalesce {
+ write!(
+ f,
+ ", coalesce: {} of {}",
+ c.groups.len(),
+ c.upstream_partition_count,
+ )?;
+ }
+ Ok(())
}
}
DisplayFormatType::TreeRender => {
@@ -233,6 +336,7 @@ impl ExecutionPlan for ShuffleReaderExec {
partition: self.partition.clone(),
broadcast: self.broadcast,
upstream_partition_count: self.upstream_partition_count,
+ coalesce: self.coalesce.clone(),
metrics: ExecutionPlanMetricsSet::new(),
properties: self.properties.clone(),
work_dir: self.work_dir.clone(),
@@ -351,8 +455,19 @@ impl ExecutionPlan for ShuffleReaderExec {
partition_count
);
}
- let stat_for_partition =
- stats_for_partition(idx, self.schema.fields().len(),
&self.partition);
+ // K-shape (coalesced): self.partition[idx] is the inner Vec
holding
+ // the concatenated upstream PartitionLocations for output
partition
+ // `idx`. Sum across that inner Vec.
+ // M-shape (legacy): outer = replicas, inner = partition index.
+ // Use the existing axis-flipped helper.
+ let stat_for_partition = if self.coalesce.is_some() {
+ Ok(stats_for_partitions(
+ self.schema.fields().len(),
+ self.partition[idx].iter().map(|loc| loc.partition_stats),
+ ))
+ } else {
+ stats_for_partition(idx, self.schema.fields().len(),
&self.partition)
+ };
trace!(
"shuffle reader at stage: {} and partition {} returned
statistics: {:?}",
@@ -881,6 +996,67 @@ mod tests {
use datafusion::prelude::SessionContext;
use tempfile::{TempDir, tempdir};
+ /// Build an M-shape upstream `Vec<Vec<PartitionLocation>>` with
per-partition
+ /// `num_bytes` and `num_rows` taken from parallel slices.
+ ///
+ /// `bytes_per_partition.len()` and `rows_per_partition.len()` define M
and must
+ /// be equal. Used by tests that require distinct per-partition stats so
the
+ /// test cannot accidentally pass under a wrong-axis aggregation.
+ fn make_upstream_partitions_nonuniform(
+ stage_id: usize,
+ bytes_per_partition: &[u64],
+ rows_per_partition: &[u64],
+ ) -> Vec<Vec<PartitionLocation>> {
+ assert_eq!(bytes_per_partition.len(), rows_per_partition.len());
+ let job_id = "test_job_coalesce_nonuniform";
+ bytes_per_partition
+ .iter()
+ .zip(rows_per_partition.iter())
+ .enumerate()
+ .map(|(i, (&bytes, &rows))| {
+ vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: PartitionId {
+ job_id: job_id.to_string(),
+ stage_id,
+ partition_id: i,
+ },
+ executor_meta: ExecutorMetadata {
+ id: "executor_1".to_string(),
+ host: "executor_1".to_string(),
+ port: 7070,
+ grpc_port: 8080,
+ specification: ExecutorSpecification::default()
+ .with_task_slots(1),
+ os_info:
ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats {
+ num_rows: Some(rows),
+ num_batches: None,
+ num_bytes: Some(bytes),
+ },
+ file_id: None,
+ is_sort_shuffle: false,
+ }]
+ })
+ .collect()
+ }
+
+ /// Concatenate selected upstream M-shape inner-Vecs into a K-shape
inner-Vec.
+ ///
+ /// Used by the coalesce tests to mirror what the rule does at
+ /// construction time.
+ fn coalesce_upstream(
+ upstream: &[Vec<PartitionLocation>],
+ indices: &[u32],
+ ) -> Vec<PartitionLocation> {
+ let mut out = Vec::new();
+ for &i in indices {
+ out.extend(upstream[i as usize].iter().cloned());
+ }
+ out
+ }
+
#[tokio::test]
async fn test_stats_for_partitions_empty() {
let result = stats_for_partitions(0, std::iter::empty());
@@ -1582,4 +1758,93 @@ mod tests {
"unexpected error message: {msg}"
);
}
+
+ #[tokio::test]
+ async fn test_shuffle_reader_exec_display_with_coalesce_renders_k_of_m()
-> Result<()>
+ {
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, false)]));
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 8,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2],
+ },
+ PartitionGroup {
+ upstream_indices: vec![3, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![5, 6, 7],
+ },
+ ],
+ };
+ let exec = ShuffleReaderExec::try_new_coalesced(
+ 1,
+ vec![vec![], vec![], vec![]],
+ coalesce,
+ schema,
+ Partitioning::UnknownPartitioning(3),
+ )?;
+ // Exercise propagation through with_work_dir to verify the field
survives a
+ // builder chain (Self-literal propagation pitfall).
+ let exec = exec.with_work_dir("/tmp".to_string());
+ let s = format!(
+ "{}",
+ datafusion::physical_plan::displayable(&exec).indent(false)
+ );
+ assert!(
+ s.contains(", coalesce: 3 of 8"),
+ "expected ', coalesce: 3 of 8' annotation; got: {s}"
+ );
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn
test_coalesced_reader_partition_statistics_sums_concatenated_bytes()
+ -> Result<()> {
+ // Non-uniform group sizes [3,2] over M=5, with non-uniform
per-partition
+ // byte counts [10,20,30,40,50]. Distinct expected totals (60 vs 90)
ensure
+ // the wrong axis cannot accidentally produce the right result.
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, false)]));
+ let stage_id = 16;
+ let bytes = [10u64, 20, 30, 40, 50];
+ let rows_per_partition = [1u64, 2, 3, 4, 5];
+ let m = bytes.len();
+ let upstream =
+ make_upstream_partitions_nonuniform(stage_id, &bytes,
&rows_per_partition);
+ let groups = vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2],
+ },
+ PartitionGroup {
+ upstream_indices: vec![3, 4],
+ },
+ ];
+ let k = groups.len();
+ let coalesce = CoalescePlan {
+ upstream_partition_count: m as u32,
+ groups: groups.clone(),
+ };
+ let k_shape: Vec<Vec<PartitionLocation>> = groups
+ .iter()
+ .map(|g| coalesce_upstream(&upstream, &g.upstream_indices))
+ .collect();
+
+ let exec = ShuffleReaderExec::try_new_coalesced(
+ stage_id,
+ k_shape,
+ coalesce,
+ schema,
+ Partitioning::UnknownPartitioning(k),
+ )?;
+
+ // partition[0] = upstream [0,1,2] -> 10+20+30 = 60 bytes, 1+2+3 = 6
rows
+ let stats0 = exec.partition_statistics(Some(0))?;
+ assert_eq!(60, *stats0.total_byte_size.get_value().unwrap());
+ assert_eq!(6, *stats0.num_rows.get_value().unwrap());
+ // partition[1] = upstream [3,4] -> 40+50 = 90 bytes, 4+5 = 9 rows
+ let stats1 = exec.partition_statistics(Some(1))?;
+ assert_eq!(90, *stats1.total_byte_size.get_value().unwrap());
+ assert_eq!(9, *stats1.num_rows.get_value().unwrap());
+ Ok(())
+ }
}
diff --git a/ballista/core/src/execution_plans/unresolved_shuffle.rs
b/ballista/core/src/execution_plans/unresolved_shuffle.rs
index 06157b75a..3f3567b6a 100644
--- a/ballista/core/src/execution_plans/unresolved_shuffle.rs
+++ b/ballista/core/src/execution_plans/unresolved_shuffle.rs
@@ -26,6 +26,8 @@ use datafusion::physical_plan::{
SendableRecordBatchStream,
};
+use crate::execution_plans::CoalescePlan;
+
/// UnresolvedShuffleExec represents a dependency on the results of a
ShuffleWriterExec node which hasn't computed yet.
///
/// An ExecutionPlan that contains an UnresolvedShuffleExec isn't ready for
execution. The presence of this ExecutionPlan
@@ -42,15 +44,22 @@ pub struct UnresolvedShuffleExec {
pub output_partition_count: usize,
/// The number of shuffle output partitions on the upstream stage. For
- /// non-broadcast readers this equals `output_partition_count`. For
- /// broadcast readers this is M (one logical output partition fans in
- /// all M upstream partition files).
+ /// non-broadcast, non-coalesced readers this equals
`output_partition_count`.
+ /// For broadcast readers this is M (one logical output partition fans in
+ /// all M upstream partition files). For coalesced readers this is M
+ /// (= `coalesce.upstream_partition_count`).
pub upstream_partition_count: usize,
/// When true, the resolved `ShuffleReaderExec` reads *all* upstream
/// partition files into its single output partition (broadcast pattern).
pub broadcast: bool,
+ /// Optional coalesce metadata. `None` means the unresolved placeholder
will
+ /// resolve to a non-coalesced ShuffleReaderExec (legacy behavior). When
`Some`,
+ /// the value is forwarded at resolution time so the resulting reader is
+ /// constructed via `ShuffleReaderExec::try_new_coalesced`.
+ pub coalesce: Option<CoalescePlan>,
+
properties: Arc<PlanProperties>,
}
@@ -71,6 +80,7 @@ impl UnresolvedShuffleExec {
output_partition_count: partition_count,
upstream_partition_count: partition_count,
broadcast: false,
+ coalesce: None,
properties,
}
}
@@ -95,6 +105,41 @@ impl UnresolvedShuffleExec {
output_partition_count: 1,
upstream_partition_count,
broadcast: true,
+ coalesce: None,
+ properties,
+ }
+ }
+
+ /// Create a new coalesce-aware UnresolvedShuffleExec.
+ ///
+ /// `partitioning` MUST already reflect the post-coalesce K (e.g.
`Partitioning::Hash(keys, K)`)
+ /// so `output_partition_count == K` matches `coalesce.groups.len()`. The
forwarding
+ /// to `ShuffleReaderExec::try_new_coalesced` happens at stage resolution
time.
+ pub fn new_coalesced(
+ stage_id: usize,
+ schema: SchemaRef,
+ partitioning: Partitioning,
+ coalesce: CoalescePlan,
+ ) -> Self {
+ debug_assert_eq!(
+ partitioning.partition_count(),
+ coalesce.groups.len(),
+ "partitioning.partition_count() must equal coalesce.groups.len()
(= K)",
+ );
+ let upstream_partition_count = coalesce.upstream_partition_count as
usize;
+ let properties = Arc::new(PlanProperties::new(
+
datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
+ partitioning,
+
datafusion::physical_plan::execution_plan::EmissionType::Incremental,
+ datafusion::physical_plan::execution_plan::Boundedness::Bounded,
+ ));
+ Self {
+ stage_id,
+ schema,
+ output_partition_count: properties.partitioning.partition_count(),
+ upstream_partition_count,
+ broadcast: false,
+ coalesce: Some(coalesce),
properties,
}
}
@@ -119,7 +164,16 @@ impl DisplayAs for UnresolvedShuffleExec {
f,
"UnresolvedShuffleExec: partitioning: {}",
self.properties().output_partitioning()
- )
+ )?;
+ if let Some(c) = &self.coalesce {
+ write!(
+ f,
+ ", coalesce: {} of {}",
+ c.groups.len(),
+ c.upstream_partition_count,
+ )?;
+ }
+ Ok(())
}
}
DisplayFormatType::TreeRender => {
@@ -178,3 +232,43 @@ impl ExecutionPlan for UnresolvedShuffleExec {
))
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::execution_plans::PartitionGroup;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+
+ #[tokio::test]
+ async fn
test_unresolved_shuffle_exec_display_with_coalesce_renders_k_of_m() {
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, false)]));
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 8,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2],
+ },
+ PartitionGroup {
+ upstream_indices: vec![3, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![5, 6, 7],
+ },
+ ],
+ };
+ let exec = UnresolvedShuffleExec::new_coalesced(
+ 1,
+ schema,
+ Partitioning::UnknownPartitioning(3),
+ coalesce,
+ );
+ let s = format!(
+ "{}",
+ datafusion::physical_plan::displayable(&exec).indent(false)
+ );
+ assert!(
+ s.contains(", coalesce: 3 of 8"),
+ "expected ', coalesce: 3 of 8' annotation; got: {s}"
+ );
+ }
+}
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index 08040b4e7..a9da57c78 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -17,7 +17,9 @@
use crate::config::{
BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES,
BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE,
- BALLISTA_CLIENT_USE_TLS, BALLISTA_JOB_NAME,
+ BALLISTA_CLIENT_USE_TLS, BALLISTA_COALESCE_ENABLED,
+ BALLISTA_COALESCE_MERGED_PARTITION_FACTOR,
BALLISTA_COALESCE_SMALL_PARTITION_FACTOR,
+ BALLISTA_COALESCE_TARGET_PARTITION_BYTES, BALLISTA_JOB_NAME,
BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
BALLISTA_STANDALONE_PARALLELISM,
BallistaConfig,
@@ -245,6 +247,26 @@ pub trait SessionConfigExt {
/// Is short shuffle used
fn ballista_sort_shuffle_enabled(&self) -> bool;
+
+ /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
+ fn ballista_coalesce_enabled(&self) -> bool;
+ /// Sets whether the AQE coalesce-shuffle-partitions rule is enabled.
+ fn with_ballista_coalesce_enabled(self, enabled: bool) -> Self;
+
+ /// Returns the target post-coalesce partition byte size in bytes.
+ fn ballista_coalesce_target_partition_bytes(&self) -> u64;
+ /// Sets the target post-coalesce partition byte size in bytes.
+ fn with_ballista_coalesce_target_partition_bytes(self, bytes: u64) -> Self;
+
+ /// Returns the small-partition merge factor (Spark legacy).
+ fn ballista_coalesce_small_partition_factor(&self) -> f64;
+ /// Sets the small-partition merge factor (Spark legacy).
+ fn with_ballista_coalesce_small_partition_factor(self, factor: f64) ->
Self;
+
+ /// Returns the merged-partition early-flush factor (Spark legacy).
+ fn ballista_coalesce_merged_partition_factor(&self) -> f64;
+ /// Sets the merged-partition early-flush factor (Spark legacy).
+ fn with_ballista_coalesce_merged_partition_factor(self, factor: f64) ->
Self;
}
/// [SessionConfigHelperExt] is set of [SessionConfig] extension methods
@@ -581,6 +603,85 @@ impl SessionConfigExt for SessionConfig {
.map(|c| c.client_use_tls())
.unwrap_or_else(|| BallistaConfig::default().client_use_tls())
}
+
+ fn ballista_coalesce_enabled(&self) -> bool {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.coalesce_enabled())
+ .unwrap_or_else(|| BallistaConfig::default().coalesce_enabled())
+ }
+
+ fn with_ballista_coalesce_enabled(self, enabled: bool) -> Self {
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_bool(BALLISTA_COALESCE_ENABLED, enabled)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_bool(BALLISTA_COALESCE_ENABLED, enabled)
+ }
+ }
+
+ fn ballista_coalesce_target_partition_bytes(&self) -> u64 {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.coalesce_target_partition_bytes())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().coalesce_target_partition_bytes()
+ })
+ }
+
+ fn with_ballista_coalesce_target_partition_bytes(self, bytes: u64) -> Self
{
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_usize(BALLISTA_COALESCE_TARGET_PARTITION_BYTES, bytes as
usize)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_usize(BALLISTA_COALESCE_TARGET_PARTITION_BYTES, bytes as
usize)
+ }
+ }
+
+ fn ballista_coalesce_small_partition_factor(&self) -> f64 {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.coalesce_small_partition_factor())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().coalesce_small_partition_factor()
+ })
+ }
+
+ // f64 setter — uses set_str because SessionConfig has no set_f64 in this
+ // workspace; the stored string is round-tripped via f64::to_string() /
+ // f64::from_str(), mirroring the `with_ballista_job_name` set_str pattern.
+ fn with_ballista_coalesce_small_partition_factor(self, factor: f64) ->
Self {
+ let s = factor.to_string();
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_str(BALLISTA_COALESCE_SMALL_PARTITION_FACTOR, &s)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_str(BALLISTA_COALESCE_SMALL_PARTITION_FACTOR, &s)
+ }
+ }
+
+ fn ballista_coalesce_merged_partition_factor(&self) -> f64 {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.coalesce_merged_partition_factor())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().coalesce_merged_partition_factor()
+ })
+ }
+
+ fn with_ballista_coalesce_merged_partition_factor(self, factor: f64) ->
Self {
+ let s = factor.to_string();
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_str(BALLISTA_COALESCE_MERGED_PARTITION_FACTOR, &s)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_str(BALLISTA_COALESCE_MERGED_PARTITION_FACTOR, &s)
+ }
+ }
}
impl SessionConfigHelperExt for SessionConfig {
diff --git a/ballista/core/src/serde/generated/ballista.rs
b/ballista/core/src/serde/generated/ballista.rs
index 897326fdd..774c53ad3 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -95,6 +95,9 @@ pub struct UnresolvedShuffleExecNode {
pub broadcast: bool,
#[prost(uint32, tag = "7")]
pub upstream_partition_count: u32,
+ /// Optional coalesce metadata. Absent means "no coalesce" (legacy
one-to-one read behavior).
+ #[prost(message, optional, tag = "8")]
+ pub coalesce: ::core::option::Option<CoalescePlan>,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ShuffleReaderExecNode {
@@ -111,6 +114,9 @@ pub struct ShuffleReaderExecNode {
pub broadcast: bool,
#[prost(uint32, tag = "6")]
pub upstream_partition_count: u32,
+ /// Optional coalesce metadata. Absent means "no coalesce" (legacy
one-to-one read behavior).
+ #[prost(message, optional, tag = "7")]
+ pub coalesce: ::core::option::Option<CoalescePlan>,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ShuffleReaderPartition {
@@ -118,6 +124,24 @@ pub struct ShuffleReaderPartition {
#[prost(message, repeated, tag = "1")]
pub location: ::prost::alloc::vec::Vec<PartitionLocation>,
}
+/// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
+/// Empty when no coalesce is applied (the optional field on the parent
message is absent).
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct CoalescePlan {
+ /// Original number of upstream partitions (M) before coalescing. Required
for EXPLAIN's "K of M" rendering.
+ #[prost(uint32, tag = "1")]
+ pub upstream_partition_count: u32,
+ /// Coalesced output groups. Length is K (the post-coalesce partition
count).
+ #[prost(message, repeated, tag = "2")]
+ pub groups: ::prost::alloc::vec::Vec<PartitionGroup>,
+}
+/// One coalesced output partition's source list: a set of upstream partition
indices in \[0, upstream_partition_count).
+/// Default algorithm produces only contiguous indices, but proto allows
arbitrary index sets for future strategies.
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct PartitionGroup {
+ #[prost(uint32, repeated, tag = "1")]
+ pub upstream_indices: ::prost::alloc::vec::Vec<u32>,
+}
///
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Ballista Scheduling
///
/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index d24398201..b90867433 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -53,7 +53,8 @@ use std::{convert::TryInto, io::Cursor};
use crate::execution_plans::sort_shuffle::SortShuffleConfig;
use crate::execution_plans::{
- ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec,
UnresolvedShuffleExec,
+ CoalescePlan, PartitionGroup, ShuffleReaderExec, ShuffleWriterExec,
+ SortShuffleWriterExec, UnresolvedShuffleExec,
};
use crate::serde::protobuf::{
ballista_logical_plan_node::LogicalPlanType,
@@ -67,6 +68,47 @@ pub mod generated;
/// Scheduler-specific serialization types and conversions.
pub mod scheduler;
+// ============================ CoalescePlan codec ============================
+//
+// Native ↔ proto conversions for `CoalescePlan` and `PartitionGroup`. Borrow-
+// based on the encode side because the call site only has a borrow
+// (`exec.coalesce.as_ref()`); the `Vec<u32>` clone is intentional and cheap
+// for typical K (small post-coalesce partition counts).
+
+impl From<&protobuf::PartitionGroup> for PartitionGroup {
+ fn from(p: &protobuf::PartitionGroup) -> Self {
+ Self {
+ upstream_indices: p.upstream_indices.clone(),
+ }
+ }
+}
+
+impl From<&PartitionGroup> for protobuf::PartitionGroup {
+ fn from(p: &PartitionGroup) -> Self {
+ Self {
+ upstream_indices: p.upstream_indices.clone(),
+ }
+ }
+}
+
+impl From<&protobuf::CoalescePlan> for CoalescePlan {
+ fn from(p: &protobuf::CoalescePlan) -> Self {
+ Self {
+ upstream_partition_count: p.upstream_partition_count,
+ groups: p.groups.iter().map(PartitionGroup::from).collect(),
+ }
+ }
+}
+
+impl From<&CoalescePlan> for protobuf::CoalescePlan {
+ fn from(p: &CoalescePlan) -> Self {
+ Self {
+ upstream_partition_count: p.upstream_partition_count,
+ groups: p.groups.iter().map(Into::into).collect(),
+ }
+ }
+}
+
impl ProstMessageExt for protobuf::Action {
fn type_url() -> &'static str {
"type.googleapis.com/arrow.flight.protocol.sql.Action"
@@ -425,7 +467,15 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
)?;
let partitioning = partitioning
.ok_or_else(|| proto_error("missing required partitioning
field"))?;
- let exec = if shuffle_reader.broadcast {
+ let exec = if let Some(c) = shuffle_reader.coalesce.as_ref() {
+ ShuffleReaderExec::try_new_coalesced(
+ stage_id,
+ partition_location,
+ CoalescePlan::from(c),
+ schema,
+ partitioning,
+ )?
+ } else if shuffle_reader.broadcast {
let all_locations = partition_location
.into_iter()
.next()
@@ -460,7 +510,14 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
)?;
let partitioning = partitioning
.ok_or_else(|| proto_error("missing required partitioning
field"))?;
- let exec = if unresolved_shuffle.broadcast {
+ let exec = if let Some(c) =
unresolved_shuffle.coalesce.as_ref() {
+ UnresolvedShuffleExec::new_coalesced(
+ unresolved_shuffle.stage_id as usize,
+ schema,
+ partitioning,
+ CoalescePlan::from(c),
+ )
+ } else if unresolved_shuffle.broadcast {
UnresolvedShuffleExec::new_broadcast(
unresolved_shuffle.stage_id as usize,
schema,
@@ -597,6 +654,7 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
partitioning: Some(partitioning),
broadcast: exec.broadcast,
upstream_partition_count:
exec.upstream_partition_count as u32,
+ coalesce: exec.coalesce.as_ref().map(|c| c.into()),
},
)),
};
@@ -622,6 +680,7 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
partitioning: Some(partitioning),
broadcast: exec.broadcast,
upstream_partition_count:
exec.upstream_partition_count as u32,
+ coalesce: exec.coalesce.as_ref().map(|c| c.into()),
},
)),
};
@@ -662,6 +721,7 @@ struct FileFormatProto {
#[cfg(test)]
mod test {
use super::*;
+ use crate::execution_plans::PartitionGroup;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::physical_plan::Partitioning;
use datafusion::physical_plan::expressions::col;
@@ -747,6 +807,10 @@ mod test {
assert_eq!(decoded_exec.stage_id, 1);
assert_eq!(decoded_exec.schema().as_ref(), schema.as_ref());
assert_eq!(&decoded_exec.properties().partitioning, &partitioning);
+ assert!(
+ decoded_exec.coalesce.is_none(),
+ "absent coalesce field must decode to None (codec inertness)"
+ );
}
#[tokio::test]
@@ -780,6 +844,274 @@ mod test {
assert_eq!(decoded_exec.stage_id, 1);
assert_eq!(decoded_exec.schema().as_ref(), schema.as_ref());
assert_eq!(&decoded_exec.properties().partitioning, &partitioning);
+ assert!(
+ decoded_exec.coalesce.is_none(),
+ "absent coalesce field must decode to None (codec inertness)"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_shuffle_reader_exec_coalesced_roundtrip_single_group() {
+ let schema = create_test_schema();
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()], 1);
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 4,
+ groups: vec![PartitionGroup {
+ upstream_indices: vec![0, 1, 2, 3],
+ }],
+ };
+
+ let original_exec = ShuffleReaderExec::try_new_coalesced(
+ 7,
+ vec![vec![]; 1], // K-shape: 1 output partition
+ coalesce.clone(),
+ schema.clone(),
+ partitioning.clone(),
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<ShuffleReaderExec>()
+ .expect("Expected ShuffleReaderExec");
+
+ assert_eq!(decoded_exec.stage_id, 7);
+ assert_eq!(&decoded_exec.properties().partitioning, &partitioning);
+ let stored = decoded_exec
+ .coalesce
+ .as_ref()
+ .expect("coalesce must round-trip");
+ assert_eq!(stored, &coalesce);
+ assert_eq!(stored.upstream_partition_count, 4);
+ assert_eq!(stored.groups.len(), 1);
+ assert_eq!(stored.groups[0].upstream_indices, vec![0, 1, 2, 3]);
+ }
+
+ #[tokio::test]
+ async fn
test_shuffle_reader_exec_coalesced_roundtrip_multi_group_mixed_sizes() {
+ let schema = create_test_schema();
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()], 3);
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 8,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2],
+ },
+ PartitionGroup {
+ upstream_indices: vec![3, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![5, 6, 7],
+ },
+ ],
+ };
+
+ let original_exec = ShuffleReaderExec::try_new_coalesced(
+ 1,
+ vec![vec![]; 3], // K-shape: 3 output partitions
+ coalesce.clone(),
+ schema.clone(),
+ partitioning.clone(),
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<ShuffleReaderExec>()
+ .expect("Expected ShuffleReaderExec");
+
+ let stored = decoded_exec
+ .coalesce
+ .as_ref()
+ .expect("coalesce must round-trip");
+ assert_eq!(stored, &coalesce);
+ assert_eq!(stored.upstream_partition_count, 8);
+ assert_eq!(stored.groups.len(), 3);
+ }
+
+ #[tokio::test]
+ async fn test_unresolved_shuffle_exec_coalesced_roundtrip_multi_index() {
+ let schema = create_test_schema();
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()], 2);
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 5,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2, 3],
+ },
+ PartitionGroup {
+ upstream_indices: vec![4],
+ },
+ ],
+ };
+
+ let original_exec = UnresolvedShuffleExec::new_coalesced(
+ 9,
+ schema.clone(),
+ partitioning.clone(),
+ coalesce.clone(),
+ );
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<UnresolvedShuffleExec>()
+ .expect("Expected UnresolvedShuffleExec");
+
+ assert_eq!(decoded_exec.stage_id, 9);
+ let stored = decoded_exec
+ .coalesce
+ .as_ref()
+ .expect("coalesce must round-trip");
+ assert_eq!(stored, &coalesce);
+ assert_eq!(stored.upstream_partition_count, 5);
+ assert_eq!(stored.groups[0].upstream_indices.len(), 4);
+ }
+
+ #[tokio::test]
+ async fn
test_shuffle_reader_exec_coalesced_roundtrip_non_contiguous_indices() {
+ // Proto allows arbitrary upstream_indices sets even though the default
+ // algorithm only emits contiguous ranges.
+ let schema = create_test_schema();
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()], 2);
+ let coalesce = CoalescePlan {
+ upstream_partition_count: 6,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 2, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![1, 3, 5],
+ },
+ ],
+ };
+
+ let original_exec = ShuffleReaderExec::try_new_coalesced(
+ 3,
+ vec![vec![]; 2],
+ coalesce.clone(),
+ schema.clone(),
+ partitioning.clone(),
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec
+ .try_encode(Arc::new(original_exec.clone()), &mut buf)
+ .unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+ let decoded_exec = decoded_plan
+ .as_any()
+ .downcast_ref::<ShuffleReaderExec>()
+ .expect("Expected ShuffleReaderExec");
+
+ let stored = decoded_exec
+ .coalesce
+ .as_ref()
+ .expect("coalesce must round-trip");
+ assert_eq!(stored, &coalesce);
+ // Non-contiguous indices preserved bit-for-bit:
+ assert_eq!(stored.groups[0].upstream_indices, vec![0, 2, 4]);
+ assert_eq!(stored.groups[1].upstream_indices, vec![1, 3, 5]);
+ }
+
+ // ---- CoalescePlan native ↔ proto direct conversion round-trips ----
+
+ #[test]
+ fn coalesce_plan_native_to_proto_roundtrip_empty() {
+ let native = CoalescePlan {
+ upstream_partition_count: 0,
+ groups: vec![],
+ };
+ let proto: protobuf::CoalescePlan = (&native).into();
+ let back: CoalescePlan = (&proto).into();
+ assert_eq!(native, back);
+ }
+
+ #[test]
+ fn coalesce_plan_native_to_proto_roundtrip_single_group() {
+ let native = CoalescePlan {
+ upstream_partition_count: 4,
+ groups: vec![PartitionGroup {
+ upstream_indices: vec![0, 1, 2, 3],
+ }],
+ };
+ let proto: protobuf::CoalescePlan = (&native).into();
+ let back: CoalescePlan = (&proto).into();
+ assert_eq!(native, back);
+ }
+
+ #[test]
+ fn coalesce_plan_native_to_proto_roundtrip_multi_group_mixed_sizes() {
+ let native = CoalescePlan {
+ upstream_partition_count: 8,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 1, 2],
+ },
+ PartitionGroup {
+ upstream_indices: vec![3, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![5, 6, 7],
+ },
+ ],
+ };
+ let proto: protobuf::CoalescePlan = (&native).into();
+ let back: CoalescePlan = (&proto).into();
+ assert_eq!(native, back);
+ assert_eq!(back.groups.len(), 3);
+ assert_eq!(back.upstream_partition_count, 8);
+ }
+
+ #[test]
+ fn coalesce_plan_native_to_proto_roundtrip_non_contiguous_indices() {
+ // Proto allows arbitrary index sets even though the default algorithm
+ // only produces contiguous ranges.
+ let native = CoalescePlan {
+ upstream_partition_count: 6,
+ groups: vec![
+ PartitionGroup {
+ upstream_indices: vec![0, 2, 4],
+ },
+ PartitionGroup {
+ upstream_indices: vec![1, 3, 5],
+ },
+ ],
+ };
+ let proto: protobuf::CoalescePlan = (&native).into();
+ let back: CoalescePlan = (&proto).into();
+ assert_eq!(native, back);
}
#[tokio::test]
diff --git a/ballista/core/src/serde/scheduler/mod.rs
b/ballista/core/src/serde/scheduler/mod.rs
index 9e80d91b6..eca6700d9 100644
--- a/ballista/core/src/serde/scheduler/mod.rs
+++ b/ballista/core/src/serde/scheduler/mod.rs
@@ -270,6 +270,9 @@ pub struct ExecutorDataChange {
pub struct PartitionStats {
pub(crate) num_rows: Option<u64>,
pub(crate) num_batches: Option<u64>,
+ /// Per-partition byte size reported by the shuffle writer. Read by the
+ /// AQE coalesce rule (in `ballista-scheduler`) to bin-pack alignment
+ /// groups, so this field is `pub` rather than `pub(crate)`.
pub(crate) num_bytes: Option<u64>,
}
@@ -297,6 +300,11 @@ impl PartitionStats {
}
}
+ /// Returns the per-partition byte size, if populated by the writer.
+ pub fn num_bytes(&self) -> Option<u64> {
+ self.num_bytes
+ }
+
/// Returns the Arrow struct field representation of these statistics.
pub fn arrow_struct_repr(self) -> Field {
Field::new(
diff --git a/ballista/scheduler/src/state/aqe/adapter.rs
b/ballista/scheduler/src/state/aqe/adapter.rs
index 111f036b6..dae2a2944 100644
--- a/ballista/scheduler/src/state/aqe/adapter.rs
+++ b/ballista/scheduler/src/state/aqe/adapter.rs
@@ -22,6 +22,7 @@ use ballista_core::execution_plans::ShuffleReaderExec;
use datafusion::common::exec_err;
use datafusion::config::ConfigOptions;
use datafusion::error::DataFusionError;
+use datafusion::physical_plan::Partitioning;
use datafusion::{
common::tree_node::{Transformed, TreeNode},
physical_plan::ExecutionPlan,
@@ -57,10 +58,46 @@ impl BallistaAdapter {
})?;
self.inputs.push(stage_id);
let partitioning = exchange.properties().partitioning.clone();
- let shuffle_read =
- ShuffleReaderExec::try_new(stage_id, partitions, schema,
partitioning)?;
- Ok(Transformed::yes(Arc::new(shuffle_read)))
+ let reader = match exchange.coalesce() {
+ Some(cp) => {
+ // Concatenate M-shape locations into K-shape per
CoalescePlan.groups.
+ let k_shape: Vec<Vec<_>> = cp
+ .groups
+ .iter()
+ .map(|pg| {
+ let mut concat = Vec::new();
+ for &idx in &pg.upstream_indices {
+ if let Some(inner) = partitions.get(idx as
usize) {
+ concat.extend_from_slice(inner);
+ }
+ }
+ concat
+ })
+ .collect();
+ let new_partitioning = match &partitioning {
+ Partitioning::Hash(keys, _m) => {
+ Partitioning::Hash(keys.clone(), cp.groups.len())
+ }
+ _ =>
Partitioning::UnknownPartitioning(cp.groups.len()),
+ };
+ ShuffleReaderExec::try_new_coalesced(
+ stage_id,
+ k_shape,
+ (*cp).clone(),
+ schema,
+ new_partitioning,
+ )?
+ }
+ None => ShuffleReaderExec::try_new(
+ stage_id,
+ partitions,
+ schema,
+ partitioning,
+ )?,
+ };
+
+ Ok(Transformed::yes(Arc::new(reader)))
} else {
Ok(Transformed::no(plan))
}
diff --git a/ballista/scheduler/src/state/aqe/coalesce/algorithm.rs
b/ballista/scheduler/src/state/aqe/coalesce/algorithm.rs
new file mode 100644
index 000000000..1e7a1c739
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/coalesce/algorithm.rs
@@ -0,0 +1,241 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Bin-packing helpers that turn per-partition byte sizes into coalesce
+//! decisions.
+//!
+//! `split_size_list_by_target_size` walks the size list left-to-right,
+//! accumulating into a bucket and flushing when adding the next size would
+//! overshoot `target`. Two refinements smooth out pathological shapes that
+//! pure overshoot-flushing produces:
+//!
+//! - **merged-factor early flush** — when a small flushed bucket sits next
+//! to another small one, fold them together rather than leaving two tiny
+//! downstream tasks.
+//! - **small-tail folding** — the post-loop pass folds a small final bucket
+//! into its predecessor.
+//!
+//! The implementation derives from Spark's `ShufflePartitionsUtil`; both
+//! refinements are load-bearing on bursty workloads and aren't optional
+//! polish.
+//!
+//! Float arithmetic is intentional: every comparison casts `u64 → f64` so
+//! `target * factor` matches the original semantics exactly. Don't
+//! pre-compute integer thresholds.
+
+use ballista_core::execution_plans::PartitionGroup;
+
+/// Pack `sizes` into bins whose total bytes approach `target` and return
+/// each bin's start index.
+///
+/// Output `starts` defines bins as: bin `k` covers
+/// `sizes[starts[k]..starts.get(k+1).unwrap_or(&sizes.len())]`. Always
+/// returns at least `vec![0]`, including for an empty input.
+///
+/// `small_factor` (default 0.2) and `merged_factor` (default 1.2) tune the
+/// merge-on-flush refinement: a bucket is merged back into its predecessor
+/// when their combined size is below `target * merged_factor`, or when
+/// either is below `target * small_factor`. The rule wires these from
+/// `ConfigOptions`; tests pass them inline.
+pub fn split_size_list_by_target_size(
+ sizes: &[u64],
+ target: u64,
+ small_factor: f64,
+ merged_factor: f64,
+) -> Vec<usize> {
+ let mut starts: Vec<usize> = vec![0];
+ let mut current: u64 = 0;
+ // Last flushed bucket's size, or None before the first flush.
+ let mut last: Option<u64> = None;
+
+ for (i, &size) in sizes.iter().enumerate() {
+ // Strict `>`: if the next size would push current PAST target, flush.
+ if i > 0 && current + size > target {
+ try_merge_partitions(
+ &mut starts,
+ current,
+ &mut last,
+ target,
+ small_factor,
+ merged_factor,
+ );
+ starts.push(i);
+ current = size;
+ } else {
+ current += size;
+ }
+ }
+ // Unconditional post-loop merge so the small tail bucket has a chance to
+ // fold back into its predecessor.
+ try_merge_partitions(
+ &mut starts,
+ current,
+ &mut last,
+ target,
+ small_factor,
+ merged_factor,
+ );
+ starts
+}
+
+// Decide whether to fold `current` back into the previous bucket. Only
+// callable from `split_size_list_by_target_size`. Float casts are
+// intentional — see module docs.
+fn try_merge_partitions(
+ starts: &mut Vec<usize>,
+ current: u64,
+ last: &mut Option<u64>,
+ target: u64,
+ small_factor: f64,
+ merged_factor: f64,
+) {
+ // Skipped on the first flush, when there's nothing to merge into.
+ let should_merge = match *last {
+ None => false,
+ Some(l) => {
+ let combined = (current + l) as f64;
+ let curr_f = current as f64;
+ let last_f = l as f64;
+ let target_f = target as f64;
+ combined < target_f * merged_factor
+ || curr_f < target_f * small_factor
+ || last_f < target_f * small_factor
+ }
+ };
+ if should_merge {
+ // Pop the last start — merging the current bucket back into the
previous.
+ starts.pop();
+ // Safe because should_merge implies last.is_some().
+ *last = Some(last.expect("should_merge guarantees Some") + current);
+ } else {
+ *last = Some(current);
+ }
+}
+
+/// Expand the start-index output of [`split_size_list_by_target_size`]
+/// into `PartitionGroup`s.
+///
+/// Group `k` covers `[starts[k], starts.get(k+1).copied().unwrap_or(n))`.
+/// Produces only contiguous ranges; non-contiguous groups stay representable
+/// in proto for future strategies but this algorithm never emits them.
+///
+/// `n` is the total number of upstream partitions (M) — the same value the
+/// `CoalescePlan::upstream_partition_count` field will carry.
+pub fn start_indices_to_partition_groups(
+ starts: &[usize],
+ n: usize,
+) -> Vec<PartitionGroup> {
+ starts
+ .iter()
+ .enumerate()
+ .map(|(k, &start)| {
+ let end = starts.get(k + 1).copied().unwrap_or(n);
+ let upstream_indices: Vec<u32> = (start..end).map(|i| i as
u32).collect();
+ PartitionGroup { upstream_indices }
+ })
+ .collect()
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ // Defaults the rule passes in real life. Pulled out so each test below
+ // reads as "this input → this start-indices vector at the production
+ // configuration", and so a reader doesn't have to wonder whether a
+ // tweaked factor is what's driving the assertion.
+ const TARGET: u64 = 1024;
+ const SMALL: f64 = 0.2;
+ const MERGED: f64 = 1.2;
+
+ fn pack(sizes: &[u64]) -> Vec<usize> {
+ split_size_list_by_target_size(sizes, TARGET, SMALL, MERGED)
+ }
+
+ #[test]
+ fn empty_input_returns_single_bucket() {
+ // No partitions → one (empty) bucket. The `[0]` seed survives the
+ // post-loop merge because `last=None` skips it.
+ assert_eq!(pack(&[]), vec![0]);
+ }
+
+ #[test]
+ fn single_partition_returns_single_bucket() {
+ // Single 2048-byte partition (≥ target). The `i > 0` guard never
+ // fires so no flush happens mid-loop; one bucket.
+ assert_eq!(pack(&[2048]), vec![0]);
+ }
+
+ #[test]
+ fn all_zero_sizes_collapse_to_one_bucket() {
+ // 0 + 0 + 0 + 0 never exceeds target; never flushes.
+ assert_eq!(pack(&[0, 0, 0, 0]), vec![0]);
+ }
+
+ #[test]
+ fn sum_exactly_target_does_not_split() {
+ // 512 + 512 = 1024. The comparison is strict `>`, so 1024 > 1024
+ // is false; the bucket fills to target without flushing.
+ assert_eq!(pack(&[512, 512]), vec![0]);
+ }
+
+ #[test]
+ fn one_byte_over_target_flushes_then_post_merges() {
+ // 512 + 513 = 1025 > 1024 → mid-loop flush. The post-loop merge
+ // then folds the second bucket back in (combined 1025 < 1024*1.2),
+ // so the final result is one bucket.
+ assert_eq!(pack(&[512, 513]), vec![0]);
+ }
+
+ #[test]
+ fn small_tail_folds_into_predecessor() {
+ // 1000 fills bucket 0, 50 starts bucket 1, post-loop sees a
+ // 50-byte tail next to a 1000-byte previous. Combined 1050 <
+ // 1024 * 1.2 = 1228.8 → merge. One bucket.
+ assert_eq!(pack(&[1000, 50]), vec![0]);
+ }
+
+ #[test]
+ fn alternating_pattern_keeps_buckets_separate() {
+ // 800 + 100 = 900 in bucket 0. Adding the third 800 would push to
+ // 1700 > 1024, so bucket 0 closes and bucket 1 starts with 800;
+ // the trailing 100 fills bucket 1 to 900. Post-loop merge is
+ // rejected (1800 ≥ 1228.8; both buckets ≥ small threshold of 205).
+ assert_eq!(pack(&[800, 100, 800, 100]), vec![0, 2]);
+ }
+
+ #[test]
+ fn factor_zero_disables_merging() {
+ // With small=0 and merged=1, only the strict `>` overshoot drives
+ // boundaries; the merging refinement is fully suppressed. Each
+ // 600-byte partition gets its own bucket.
+ let starts = split_size_list_by_target_size(&[600, 600, 600], 1024,
0.0, 1.0);
+ assert_eq!(starts, vec![0, 1, 2]);
+ }
+
+ #[test]
+ fn start_indices_expand_into_contiguous_partition_groups() {
+ // [0, 3, 5] over n=8 means: group 0 covers indices [0,1,2],
+ // group 1 covers [3,4], group 2 covers [5,6,7]. Pure unpacking,
+ // no algorithm logic.
+ let groups = start_indices_to_partition_groups(&[0, 3, 5], 8);
+ assert_eq!(groups.len(), 3);
+ assert_eq!(groups[0].upstream_indices, vec![0, 1, 2]);
+ assert_eq!(groups[1].upstream_indices, vec![3, 4]);
+ assert_eq!(groups[2].upstream_indices, vec![5, 6, 7]);
+ }
+}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
b/ballista/scheduler/src/state/aqe/coalesce/mod.rs
similarity index 63%
copy from ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
copy to ballista/scheduler/src/state/aqe/coalesce/mod.rs
index af23d2a7b..50e9aa592 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
+++ b/ballista/scheduler/src/state/aqe/coalesce/mod.rs
@@ -15,10 +15,15 @@
// specific language governing permissions and limitations
// under the License.
-pub mod datafusion_patch;
-pub mod distributed_exchange;
-pub mod propagate_empty;
+//! AQE coalesce-shuffle-partitions helpers.
+//!
+//! This submodule packages the pure-CPU helpers `CoalescePartitionsRule`
+//! consumes:
+//!
+//! - [`split_size_list_by_target_size`] — bin-packs per-partition byte sizes
+//! into bucket boundaries near a target size.
+//! - [`start_indices_to_partition_groups`] — expands those boundaries into
+//! `PartitionGroup`s attached to a `CoalescePlan`.
-pub use datafusion_patch::*;
-pub use distributed_exchange::*;
-pub use propagate_empty::*;
+pub(crate) mod algorithm;
+pub(crate) use algorithm::*;
diff --git a/ballista/scheduler/src/state/aqe/execution_plan.rs
b/ballista/scheduler/src/state/aqe/execution_plan.rs
index 3fc23de49..ac37328f6 100644
--- a/ballista/scheduler/src/state/aqe/execution_plan.rs
+++ b/ballista/scheduler/src/state/aqe/execution_plan.rs
@@ -33,7 +33,9 @@
//! adaptive and to carry mutable state such as `is_final` and resolved
//! shuffle metadata.
-use ballista_core::execution_plans::{stats_for_partition,
stats_for_partitions};
+use ballista_core::execution_plans::{
+ CoalescePlan, stats_for_partition, stats_for_partitions,
+};
use ballista_core::serde::scheduler::PartitionLocation;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_plan::Statistics;
@@ -64,7 +66,7 @@ use std::sync::{Arc, atomic::AtomicI64};
/// Note: this type implements DataFusion's `ExecutionPlan` trait but returns
/// an error from `execute` because it is not directly runnable.
#[derive(Debug)]
-pub(crate) struct ExchangeExec {
+pub struct ExchangeExec {
input: Arc<dyn ExecutionPlan>,
properties: Arc<PlanProperties>,
pub(crate) partitioning: Option<Partitioning>,
@@ -81,6 +83,19 @@ pub(crate) struct ExchangeExec {
/// can not be assumed.
shuffle_partitions: Arc<Mutex<Option<Vec<Vec<PartitionLocation>>>>>,
+ /// Per-stage coalesce decision attached to this Exchange by
+ /// `CoalescePartitionsRule` before adapter conversion.
+ ///
+ /// `None` means: build the SR with `try_new` (M-partition, no coalesce).
+ /// `Some(cp)` means: build the SR with `try_new_coalesced(cp)` so the
+ /// reader exposes K = `cp.groups.len()` partitions, each backed by the
+ /// upstream-index range described by the corresponding `PartitionGroup`.
+ ///
+ /// Wrapped in `Arc<Mutex<…>>` so `with_new_children` can clone the slot
+ /// alongside the Exchange, keeping rule decisions in sync across
+ /// transform-rebuilt parent chains. Same pattern as `shuffle_partitions`.
+ coalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>,
+
/// this disables stage from running even it would be suitable to run.
///
/// the main reason for this property this is to allow rules to override
@@ -90,6 +105,9 @@ pub(crate) struct ExchangeExec {
}
impl ExchangeExec {
+ /// Creates a new `ExchangeExec` with default stage ID (-1) and empty
+ /// partition set. The stage ID and partitions should be resolved
+ /// before the exchange participates in AQE rules.
pub fn new(
input: Arc<dyn ExecutionPlan>,
partitioning: Option<Partitioning>,
@@ -104,6 +122,9 @@ impl ExchangeExec {
)
}
+ /// Creates a new `ExchangeExec` with explicitly-provided stage ID and
+ /// partition storage. Used by the AQE rule infrastructure to construct
+ /// exchanges that share atomic state with the enclosing `AdaptivePlanner`.
pub fn new_with_details(
input: Arc<dyn ExecutionPlan>,
partitioning: Option<Partitioning>,
@@ -130,6 +151,7 @@ impl ExchangeExec {
stage_id,
shuffle_partitions: stage_partitions,
partitioning,
+ coalesce: Arc::new(Mutex::new(None)),
inactive_stage: false,
}
}
@@ -176,6 +198,8 @@ impl ExchangeExec {
.store(id as i64, std::sync::atomic::Ordering::Relaxed);
}
+ /// Returns the stage ID assigned to this exchange, or `None` if the
+ /// stage has not yet been resolved (initial value -1).
pub fn stage_id(&self) -> Option<usize> {
let stage_id =
self.stage_id.load(std::sync::atomic::Ordering::Relaxed);
@@ -186,9 +210,23 @@ impl ExchangeExec {
}
}
+ /// Returns a reference to the input (child) execution plan.
pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
&self.input
}
+
+ /// Attaches a `CoalescePlan` to this Exchange. The adapter consumes the
+ /// plan when converting Exchange → ShuffleReader: a Some value triggers
+ /// `try_new_coalesced` (K-partition reader); None uses `try_new`
+ /// (M-partition reader). Idempotent overwrite.
+ pub fn set_coalesce(&self, cp: Arc<CoalescePlan>) {
+ self.coalesce.lock().replace(cp);
+ }
+
+ /// Returns the attached `CoalescePlan`, if `set_coalesce` was called.
+ pub fn coalesce(&self) -> Option<Arc<CoalescePlan>> {
+ self.coalesce.lock().clone()
+ }
}
impl DisplayAs for ExchangeExec {
@@ -210,8 +248,17 @@ impl DisplayAs for ExchangeExec {
self.stage_id()
.map(|stage_id| format!("{}", stage_id))
.unwrap_or_else(|| "pending".to_string()),
- self.shuffle_created()
- )
+ self.shuffle_created(),
+ )?;
+ if let Some(cp) = self.coalesce.lock().as_ref() {
+ write!(
+ f,
+ ", coalesce={} of {}",
+ cp.groups.len(),
+ cp.upstream_partition_count,
+ )?;
+ }
+ Ok(())
}
DisplayFormatType::TreeRender => {
writeln!(
@@ -273,6 +320,9 @@ impl ExecutionPlan for ExchangeExec {
self.shuffle_partitions.clone(),
);
new_exec.inactive_stage = self.inactive_stage;
+ // Carry the coalesce slot so a transform-rebuilt parent chain
+ // doesn't lose the rule's decision.
+ new_exec.coalesce = self.coalesce.clone();
Ok(Arc::new(new_exec))
} else {
diff --git a/ballista/scheduler/src/state/aqe/mod.rs
b/ballista/scheduler/src/state/aqe/mod.rs
index 6b1903dc9..79e6d6595 100644
--- a/ballista/scheduler/src/state/aqe/mod.rs
+++ b/ballista/scheduler/src/state/aqe/mod.rs
@@ -50,7 +50,8 @@ use std::vec;
// an AQE optimizer rule in a follow-up PR.
mod adapter;
-mod execution_plan;
+pub(crate) mod coalesce;
+pub(crate) mod execution_plan;
pub mod optimizer_rule;
pub mod planner;
#[cfg(test)]
diff --git
a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs
new file mode 100644
index 000000000..6efd6b332
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs
@@ -0,0 +1,311 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! AQE rule that coalesces shuffle partitions after upstream stages finalize.
+//!
+//! [`CoalescePartitionsRule`] runs once per `replan_stages()` pass on a
+//! stage subtree whose root is either an [`ExchangeExec`] (intermediate
+//! stage) or an [`AdaptiveDatafusionExec`] (final stage). The rule walks the
+//! subtree, collects every leaf [`ExchangeExec`] — the resolved upstream
+//! shuffles feeding this stage — and decides whether to coalesce.
+//!
+//! # The alignment group
+//!
+//! Every leaf `ExchangeExec` in a single stage subtree forms one **alignment
+//! group**. Why one group, not one decision per leaf?
+//!
+//! - Hash-partitioned joins (`HashJoinExec(Partitioned)`, `SortMergeJoinExec`)
+//! require their two inputs to have the *same partition count* and to be
+//! hash-partitioned on the join key. If we coalesced left to `K=4` and
+//! right to `K=2`, DataFusion's `EnforceDistribution` would either reject
+//! the plan or insert remediation repartitions that undo the optimization.
+//! - Both join legs read shuffle output from upstream stages that wrote
+//! `M` partitions using the *same* hash function on the *same* key
+//! (that's what made them joinable in the first place). So upstream
+//! partition `i` of the left and upstream partition `i` of the right
+//! hold rows that must meet at downstream partition `f(i)`. Coalescing
+//! them with the *same* mapping `i → group(i)` keeps that meeting point
+//! consistent; coalescing them with different mappings scatters it.
+//!
+//! Practically: we treat all leaf Exchanges as a single workload, sum their
+//! per-partition byte counts element-wise, bin-pack the summed sizes once,
+//! and attach the *same* `CoalescePlan` to every leaf. Joins with two leaves
+//! and chains of joins with three or more leaves all go through the same
+//! code path — there is no per-leaf decision.
+//!
+//! Concretely for `[25; 8]` bytes per partition on both sides of a join:
+//! summed `[50; 8]`, bin-pack at target `200` produces `K=2` (4 upstream
+//! partitions per group), both leaves get `coalesce=2 of 8`, the downstream
+//! join runs with 2 partitions on each side, hash buckets stay aligned.
+//!
+//! # Default off
+//!
+//! `ballista.planner.coalesce.enabled` is `false` by default. The rule is an
opt-in
+//! trade — coalescing reduces task overhead and IPC cost, but at the price
+//! of less downstream parallelism. Users who want the trade explicitly turn
+//! the rule on. When off, the rule short-circuits at the first statement
+//! of `optimize()` and the plan flows through untouched.
+//!
+//! Conceptually:
+//! - `coalesce.enabled=false` (default) ≈ Spark's `parallelismFirst=true`
+//! outcome — partitions preserved, no packing.
+//! - `coalesce.enabled=true` (opt-in) ≈ Spark's `parallelismFirst=false`
+//! outcome — pack toward the advisory target, accept fewer/larger tasks.
+//!
+//! # Algorithm
+//!
+//! 1. Find leaf `ExchangeExec`s — the alignment group. If empty, this
+//! stage reads from scans and has nothing to coalesce.
+//! 2. All leaves share the upstream partition count `M` (the writer side).
+//! 3. Sum per-partition byte sizes element-wise across the group to get
+//! combined work per upstream index.
+//! 4. Bin-pack the summed sizes into `K` buckets near
+//! `target_partition_bytes` (Spark's `advisoryPartitionSizeInBytes`,
+//! 64 MB by default) using `split_size_list_by_target_size`.
+//! 5. If `K >= M` or `K <= 1`, the rewrite is degenerate and is skipped.
+//! 6. Otherwise, attach a shared [`CoalescePlan`] (with `K` partition
+//! groups) to every leaf `ExchangeExec` via `set_coalesce(..)`. The
+//! adapter consumes that decision when it builds the downstream
+//! `ShuffleReaderExec`s.
+//!
+//! # Carrier semantics
+//!
+//! The `CoalescePlan` lives on the upstream `ExchangeExec`; the rule does
+//! not rewrite the plan tree. Idempotency is structural — `set_coalesce`
+//! overwrites the slot with an equivalent plan on re-entry, and the
+//! bin-pack is a pure function of the resolved byte sizes, so the second
+//! pass produces the same decision.
+//!
+//! # Grouping discipline
+//!
+//! The bin-pack groups **neighboring** upstream partitions only — each
+//! output partition `k` covers a contiguous index range `[start_k,
+//! start_k+1)` over the input. Non-adjacent partitions are never folded
+//! together, even when that would yield a tighter byte fit. This matches
+//! Spark's `CoalesceShufflePartitions` and is what keeps hash
+//! co-partitioning intact across the rewrite: a hash bucket that used to
+//! live at index `i` still lives in the single output group that covers
+//! `i`, on every leaf of the alignment group.
+//!
+//! # Behavior preservation
+//!
+//! When `ballista.planner.coalesce.enabled=false`, when the subtree has no
+//! leaf Exchanges, or when the bin-pack returns a degenerate `K`, the rule
+//! is a no-op and returns the input `Arc` verbatim (preserving
+//! `Arc::ptr_eq`).
+
+use std::sync::Arc;
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::CoalescePlan;
+use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use log::debug;
+
+use crate::state::aqe::coalesce::{
+ split_size_list_by_target_size, start_indices_to_partition_groups,
+};
+use crate::state::aqe::execution_plan::AdaptiveDatafusionExec;
+use crate::state::aqe::execution_plan::ExchangeExec;
+
+/// AQE rule that attaches a coalesce decision to every leaf `ExchangeExec`
+/// feeding the current stage, so the downstream reader exposes `K < M`
+/// partitions.
+///
+/// See module docs for design intent.
+#[derive(Debug, Default)]
+pub struct CoalescePartitionsRule;
+
+impl PhysicalOptimizerRule for CoalescePartitionsRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ let bc = config
+ .extensions
+ .get::<BallistaConfig>()
+ .cloned()
+ .unwrap_or_default();
+ if !bc.coalesce_enabled() {
+ return Ok(plan);
+ }
+ let target = bc.coalesce_target_partition_bytes();
+ let small = bc.coalesce_small_partition_factor();
+ let merged = bc.coalesce_merged_partition_factor();
+
+ debug!(
+ "[coalesce-rule] fire: target_partition_bytes={target}
small_factor={small} merged_factor={merged}",
+ );
+
+ // Get the subtree below the root. Two root kinds, same outcome.
+ let input = if let Some(ex) =
plan.as_any().downcast_ref::<ExchangeExec>() {
+ debug!(
+ "[coalesce-rule] root=ExchangeExec plan_id={} stage_id={:?}
stage_resolved={}",
+ ex.plan_id,
+ ex.stage_id(),
+ ex.shuffle_partitions().is_some(),
+ );
+ ex.input().clone()
+ } else if let Some(adp) =
plan.as_any().downcast_ref::<AdaptiveDatafusionExec>() {
+ debug!(
+ "[coalesce-rule] root=AdaptiveDatafusionExec stage_id={:?}",
+ adp.stage_id(),
+ );
+ adp.input().clone()
+ } else {
+ debug!(
+ "[coalesce-rule] root is neither ExchangeExec nor
AdaptiveDatafusionExec; bail"
+ );
+ return Ok(plan); // unexpected root — adapter will fail anyway,
just bail
+ };
+
+ // Collect the alignment group: every leaf `ExchangeExec` feeding
+ // this stage. `Jump` after each hit stops the walk from descending
+ // into the upstream stage's compute — those nodes aren't part of
+ // *this* stage's group, they belong to whatever stage wrote them.
+ let mut leaves: Vec<Arc<dyn ExecutionPlan>> = Vec::new();
+ input.apply(|node| {
+ if node.as_any().is::<ExchangeExec>() {
+ leaves.push(node.clone());
+ Ok(TreeNodeRecursion::Jump)
+ } else {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ })?;
+
+ // Helper: downcast each Arc back to &ExchangeExec.
+ fn as_exchange(arc: &Arc<dyn ExecutionPlan>) -> &ExchangeExec {
+ arc.as_any()
+ .downcast_ref::<ExchangeExec>()
+ .expect("filtered to ExchangeExec above")
+ }
+
+ debug!(
+ "[coalesce-rule] collected {} leaf ExchangeExec(s)",
+ leaves.len()
+ );
+ for arc in &leaves {
+ let ex = as_exchange(arc);
+ debug!(
+ "[coalesce-rule] leaf: plan_id={} stage_id={:?}
partitioning={} M={} resolved={} existing_coalesce={:?}",
+ ex.plan_id,
+ ex.stage_id(),
+ ex.properties().partitioning,
+ ex.properties().partitioning.partition_count(),
+ ex.shuffle_partitions().is_some(),
+ ex.coalesce()
+ .as_ref()
+ .map(|cp| (cp.groups.len(), cp.upstream_partition_count)),
+ );
+ }
+
+ // Leaf-scan stage with no upstream Exchanges → nothing to coalesce.
+ if leaves.is_empty() {
+ debug!("[coalesce-rule] no leaves; bail");
+ return Ok(plan);
+ }
+
+ // The alignment-group invariant assumes a shared `M`. In every plan
+ // shape we currently produce, all leaves of one stage subtree are
+ // hash-partitioned by the same target_partitions setting upstream,
+ // so reading `M` from leaf 0 is sufficient.
+ let m = as_exchange(&leaves[0])
+ .properties()
+ .partitioning
+ .partition_count();
+
+ // TODO: per-M subgrouping; for now bail on heterogeneous M (Q22 panic
guard).
+ if leaves
+ .iter()
+ .any(|arc|
as_exchange(arc).properties().partitioning.partition_count() != m)
+ {
+ return Ok(plan);
+ }
+
+ // Sum byte sizes element-wise across the alignment group. Upstream
+ // partition `i` is the same logical hash bucket on every leaf, so
+ // `summed[i]` is the total downstream work for that bucket. If any
+ // leaf is still unresolved we bail — early `replan_stages()` passes
+ // run before all upstream stages finalize, and the rule reruns on
+ // every later pass anyway, so the no-op is free.
+ let mut summed = vec![0u64; m];
+ for arc in &leaves {
+ let ex = as_exchange(arc);
+ let Some(parts) = ex.shuffle_partitions() else {
+ debug!(
+ "[coalesce-rule] leaf plan_id={} unresolved; bail entire
group",
+ ex.plan_id
+ );
+ return Ok(plan);
+ };
+ for (i, locs) in parts.iter().enumerate() {
+ summed[i] += locs
+ .iter()
+ .filter_map(|l| l.partition_stats.num_bytes())
+ .sum::<u64>();
+ }
+ }
+ debug!("[coalesce-rule] summed bytes per upstream partition:
{summed:?}");
+
+ // One bin-pack decision for the whole alignment group, packing toward
+ // `target_partition_bytes` (Spark's `advisoryPartitionSizeInBytes`).
+ // The rule is opt-in (`coalesce.enabled=false` by default), so users
+ // get parallelism preservation unless they explicitly trade it for
+ // larger tasks. This corresponds to Spark's
+ // `parallelismFirst=false` mode — direct advisory-driven packing.
+ let starts = split_size_list_by_target_size(&summed, target, small,
merged);
+ let k = starts.len();
+ debug!("[coalesce-rule] bin-pack result: K={k} M={m}");
+ if k >= m || k <= 1 {
+ debug!(
+ "[coalesce-rule] K degenerate (K>=M or K<=1); bail without
setting coalesce"
+ );
+ return Ok(plan);
+ }
+
+ // Attach the same `CoalescePlan` to every member of the alignment
+ // group. Sharing the plan (not just the K value) keeps the upstream
+ // index → group mapping identical across leaves — hash buckets that
+ // were aligned at M stay aligned at K, and the join's
+ // partition-count requirement still holds after the rewrite.
+ let cp = Arc::new(CoalescePlan {
+ upstream_partition_count: m as u32,
+ groups: start_indices_to_partition_groups(&starts, m),
+ });
+ for arc in &leaves {
+ let ex = as_exchange(arc);
+ debug!(
+ "[coalesce-rule] set_coalesce(K={k}) on plan_id={} (was {:?})",
+ ex.plan_id,
+ ex.coalesce().as_ref().map(|cp| cp.groups.len()),
+ );
+ ex.set_coalesce(cp.clone());
+ }
+ Ok(plan)
+ }
+
+ fn name(&self) -> &str {
+ "CoalescePartitionsRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
index af23d2a7b..348158bb4 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
@@ -15,10 +15,12 @@
// specific language governing permissions and limitations
// under the License.
+pub mod coalesce_partitions;
pub mod datafusion_patch;
pub mod distributed_exchange;
pub mod propagate_empty;
+pub use coalesce_partitions::*;
pub use datafusion_patch::*;
pub use distributed_exchange::*;
pub use propagate_empty::*;
diff --git a/ballista/scheduler/src/state/aqe/planner.rs
b/ballista/scheduler/src/state/aqe/planner.rs
index 2346bb380..03361dc3c 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -17,7 +17,8 @@
use crate::state::aqe::adapter::BallistaAdapter;
use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec};
use crate::state::aqe::optimizer_rule::{
- DistributedExchangeRule, PropagateEmptyExecRule, WarnOnDuplicateExecRule,
+ CoalescePartitionsRule, DistributedExchangeRule, PropagateEmptyExecRule,
+ WarnOnDuplicateExecRule,
};
use crate::state::execution_stage::StageOutput;
@@ -26,6 +27,7 @@ use ballista_core::serde::scheduler::PartitionLocation;
use datafusion::common;
use datafusion::common::{HashMap, exec_err};
use datafusion::execution::{SessionState, SessionStateBuilder};
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::optimizer::PhysicalOptimizer;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties,
displayable};
use datafusion::physical_planner::DefaultPhysicalPlanner;
@@ -296,8 +298,13 @@ impl AdaptivePlanner {
let (stage_ids, shuffle_writers) = stages
.into_iter()
.map(|plan| {
- // TODO: we need to find input stages for given stage
- // thus result should change
+ // Run the coalesce rule per-stage: the root of `plan`
is
+ // the stage's wrapper exchange, so the rule's walker
sees
+ // only this stage's input exchanges as the alignment
+ // group. This avoids cross-stage gluing and stale
state
+ // that would arise if the rule walked the entire
residual
+ // plan in `default_optimizers()`.
+ let plan = CoalescePartitionsRule.optimize(plan,
config)?;
BallistaAdapter::adapt_to_ballista(
plan,
self.job_name.as_str(),
@@ -431,6 +438,10 @@ impl AdaptivePlanner {
// rule does not mutate plan hance it can go after
`DistributedExchangeRule`
physical_optimizers.push(Arc::new(WarnOnDuplicateExecRule::default()));
+ // `CoalescePartitionsRule` is invoked per-stage in
`actionable_stages()`
+ // rather than registered here, so each invocation sees only one
stage's
+ // plan and forms an alignment group scoped to that stage's inputs.
+
physical_optimizers
}
/// Creates a session state with the given configuration and optimizer
rules.
diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
new file mode 100644
index 000000000..d055f1ea9
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
@@ -0,0 +1,434 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Functional tests for [`CoalescePartitionsRule`]: drive a query through
+//! `AdaptivePlanner`, finalize the upstream stage with synthetic per-partition
+//! byte stats, and snapshot the displayed plan tree so the rule's effect on
+//! the leaf `ExchangeExec` is visible at the `coalesce=K of M` field.
+//!
+//! Each test uses small synthetic byte sizes paired with a small
+//! `coalesce_target_partition_bytes` so the bin-pack outcome is hand-traceable
+//! against `split_size_list_by_target_size`.
+
+use crate::assert_plan;
+use crate::state::aqe::planner::AdaptivePlanner;
+use crate::state::aqe::test::{mock_batch, mock_schema};
+use ballista_core::extension::SessionConfigExt;
+use ballista_core::serde::scheduler::{
+ ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
+ PartitionId, PartitionLocation, PartitionStats,
+};
+use datafusion::datasource::MemTable;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use std::sync::Arc;
+
+/// Build a session context with the Ballista config extension installed and
+/// the coalesce-relevant knobs forced to specific values. The rule packs
+/// directly toward `target_partition_bytes` (here 200 for test scale), so
+/// every scenario below traces from inputs alone.
+fn coalesce_context(target_partitions: usize, enabled: bool) -> SessionContext
{
+ let config = SessionConfig::new_with_ballista()
+ .with_target_partitions(target_partitions)
+ .with_round_robin_repartition(false)
+ .with_ballista_coalesce_enabled(enabled)
+ .with_ballista_coalesce_target_partition_bytes(200);
+
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build();
+
+ SessionContext::new_with_state(state)
+}
+
+/// Register a MemTable with `n_partitions` partitions, each holding one copy
+/// of `mock_batch`. Multi-partition sources force DataFusion's
+/// `EnforceDistribution` to insert a hash repartition before partitioned
+/// joins — without that, two 1-partition inputs satisfy `Partitioned` on a
+/// single partition and no `ExchangeExec` shows up in the final plan.
+fn register_partitioned_table(
+ ctx: &SessionContext,
+ name: &str,
+ n_partitions: usize,
+) -> datafusion::error::Result<()> {
+ let data = (0..n_partitions)
+ .map(|_| Ok(vec![mock_batch()?]))
+ .collect::<datafusion::error::Result<Vec<_>>>()?;
+ let table = MemTable::try_new(mock_schema(), data)?;
+ ctx.register_table(name, Arc::new(table))?;
+ Ok(())
+}
+
+/// Build a `Vec<Vec<PartitionLocation>>` of length `per_partition_bytes.len()`
+/// where each upstream partition reports the given byte size. The rule sums
+/// `partition_stats.num_bytes` across leaves before bin-packing — that's the
+/// only field these tests need to vary.
+fn partitions_with_byte_sizes(
+ per_partition_bytes: &[u64],
+) -> Vec<Vec<PartitionLocation>> {
+ per_partition_bytes
+ .iter()
+ .enumerate()
+ .map(|(idx, &bytes)| {
+ vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: PartitionId {
+ job_id: "".to_string(),
+ stage_id: 0,
+ partition_id: idx,
+ },
+ executor_meta: ExecutorMetadata {
+ id: "".to_string(),
+ host: "".to_string(),
+ port: 0,
+ grpc_port: 0,
+ specification:
ExecutorSpecification::default().with_task_slots(0),
+ os_info: ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats::new(Some(1), None,
Some(bytes)),
+ file_id: None,
+ is_sort_shuffle: false,
+ }]
+ })
+ .collect()
+}
+
+/// Happy path: M=8 upstream partitions @ 50 bytes each, target=200.
+/// Bin-pack trace (small_factor=0.2 → 40, merged_factor=1.2 → 240):
+/// i=0..3 accumulate into bucket=200; i=4 overshoots, flush, start new;
+/// i=5..7 accumulate into bucket=200; post-loop merge is rejected
+/// (200 + 200 = 400, not below 240). Result: K=2.
+/// Plan tree therefore shows `coalesce=2 of 8` on the leaf Exchange after
+/// stage 0 finalizes.
+#[tokio::test]
+async fn should_attach_coalesce_when_partitions_pack_below_m()
+-> datafusion::error::Result<()> {
+ let ctx = coalesce_context(8, true);
+ ctx.register_batch("t", mock_batch()?)?;
+
+ let plan = ctx
+ .sql("select min(a) as c0, c as c2 from t group by c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ // Before any stage finalizes the leaves are unresolved, so the rule
+ // no-ops: `coalesce=none`.
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending,
stage_resolved=false
+ ProjectionExec: expr=[min(t.a)@1 as c0, c@0 as c2]
+ AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)]
+ ExchangeExec: partitioning=Hash([c@0], 8), plan_id=0,
stage_id=pending, stage_resolved=false
+ AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
+ DataSourceExec: partitions=1, partition_sizes=[1]
+ ");
+
+ // Surface the runnable stage so its id is registered before we finalize.
+ let _ = planner.runnable_stages()?.unwrap();
+
+ // Finalize stage 0 with 8 partitions of 50 bytes each (total = 400,
target = 200).
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[50; 8]))?;
+
+ // Surface the next runnable stage — this is where `CoalescePartitionsRule`
+ // fires per-stage on the downstream consumer and attaches the
+ // `CoalescePlan` to plan_id=0.
+ let _ = planner.runnable_stages()?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1,
stage_resolved=false
+ ProjectionExec: expr=[min(t.a)@1 as c0, c@0 as c2]
+ AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)]
+ ExchangeExec: partitioning=Hash([c@0], 8), plan_id=0, stage_id=0,
stage_resolved=true, coalesce=2 of 8
+ AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
+ DataSourceExec: partitions=1, partition_sizes=[1]
+ ");
+
+ Ok(())
+}
+
+/// Disabled path: same inputs as above but
`ballista.planner.coalesce.enabled=false`.
+/// The rule short-circuits at the first statement of `optimize()` and returns
+/// the plan untouched, so the leaf Exchange's coalesce slot stays None.
+#[tokio::test]
+async fn should_skip_coalesce_when_rule_disabled() ->
datafusion::error::Result<()> {
+ let ctx = coalesce_context(8, false);
+ ctx.register_batch("t", mock_batch()?)?;
+
+ let plan = ctx
+ .sql("select min(a) as c0, c as c2 from t group by c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ let _ = planner.runnable_stages()?.unwrap();
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[50; 8]))?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending,
stage_resolved=false
+ ProjectionExec: expr=[min(t.a)@1 as c0, c@0 as c2]
+ AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)]
+ ExchangeExec: partitioning=Hash([c@0], 8), plan_id=0, stage_id=0,
stage_resolved=true
+ AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
+ DataSourceExec: partitions=1, partition_sizes=[1]
+ ");
+
+ Ok(())
+}
+
+/// Degenerate K=M path: every partition is already at target. Bin-pack
+/// flushes after each one (adding the next would exceed target=200), and
+/// the post-flush merge is rejected (each bucket = 300, neither small nor
+/// combinable). K = M = 8 → the rule treats it as no work and returns the
+/// plan as-is.
+#[tokio::test]
+async fn should_skip_coalesce_when_partitions_are_full() ->
datafusion::error::Result<()>
+{
+ let ctx = coalesce_context(8, true);
+ ctx.register_batch("t", mock_batch()?)?;
+
+ let plan = ctx
+ .sql("select min(a) as c0, c as c2 from t group by c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ let _ = planner.runnable_stages()?.unwrap();
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[300; 8]))?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending,
stage_resolved=false
+ ProjectionExec: expr=[min(t.a)@1 as c0, c@0 as c2]
+ AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)]
+ ExchangeExec: partitioning=Hash([c@0], 8), plan_id=0, stage_id=0,
stage_resolved=true
+ AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
+ DataSourceExec: partitions=1, partition_sizes=[1]
+ ");
+
+ Ok(())
+}
+
+/// Partitioned hash join: both sides shuffled by the join key, both leaf
+/// Exchanges live in the same final-stage subtree. Per-leaf bytes
+/// `[25; 8]` × 2 leaves → summed `[50; 8]` → bin-pack at target=200
+/// collapses to K=2 (same trace as the happy-path test). Both leaves get
+/// the SAME `CoalescePlan` so the join's partition-count requirement holds
+/// across the rewrite.
+#[tokio::test]
+async fn should_attach_coalesce_to_both_sides_of_hash_join()
+-> datafusion::error::Result<()> {
+ let ctx = coalesce_context(8, true);
+ register_partitioned_table(&ctx, "t1", 8)?;
+ register_partitioned_table(&ctx, "t2", 8)?;
+
+ let plan = ctx
+ .sql("select t1.a, t2.b from t1 join t2 on t1.c = t2.c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ let stages = planner.runnable_stages()?.unwrap();
+ assert_eq!(2, stages.len());
+
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[25; 8]))?;
+ planner.finalise_stage_internal(1, partitions_with_byte_sizes(&[25; 8]))?;
+
+ // Surface the join stage so `CoalescePartitionsRule` fires per-stage and
+ // attaches the shared `CoalescePlan` to both leaf Exchanges.
+ let _ = planner.runnable_stages()?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2,
stage_resolved=false
+ ProjectionExec: expr=[a@0 as a, b@2 as b]
+ SortMergeJoinExec: join_type=Inner, on=[(c@1, c@1)]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=0, stage_id=0,
stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1,
1, 1]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=1, stage_id=1,
stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1,
1, 1]
+ ");
+
+ Ok(())
+}
+
+/// Two hash joins in one final stage: 3 upstream Exchanges feed the join
+/// chain (t1 ⋈ t2 ⋈ t3 on a shared key). Per-leaf bytes `[16; 8]` × 3
+/// leaves → summed `[48; 8]`. Bin-pack at target=200: 4 partitions fill
+/// bucket to 192, 5th overshoots (240 > 200) and flushes; next bucket fills
+/// remaining 4 to 192; post-loop merge rejected. K=2.
+/// All three leaves get the same `CoalescePlan`.
+#[tokio::test]
+async fn should_attach_coalesce_to_all_three_legs_of_two_hash_joins()
+-> datafusion::error::Result<()> {
+ let ctx = coalesce_context(8, true);
+ register_partitioned_table(&ctx, "t1", 8)?;
+ register_partitioned_table(&ctx, "t2", 8)?;
+ register_partitioned_table(&ctx, "t3", 8)?;
+
+ let plan = ctx
+ .sql(
+ "select t1.a, t2.b, t3.c from t1 join t2 on t1.c = t2.c \
+ join t3 on t1.c = t3.c",
+ )
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ let stages = planner.runnable_stages()?.unwrap();
+ assert_eq!(3, stages.len());
+
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[16; 8]))?;
+ planner.finalise_stage_internal(1, partitions_with_byte_sizes(&[16; 8]))?;
+ planner.finalise_stage_internal(2, partitions_with_byte_sizes(&[16; 8]))?;
+
+ // Surface the join stage so `CoalescePartitionsRule` fires per-stage and
+ // attaches the same `CoalescePlan` to all three leaf Exchanges.
+ let _ = planner.runnable_stages()?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=true, plan_id=3, stage_id=3,
stage_resolved=false
+ ProjectionExec: expr=[a@0 as a, b@2 as b, c@3 as c]
+ SortMergeJoinExec: join_type=Inner, on=[(c@1, c@0)]
+ ProjectionExec: expr=[a@0 as a, c@1 as c, b@2 as b]
+ SortMergeJoinExec: join_type=Inner, on=[(c@1, c@1)]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=0,
stage_id=0, stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1,
1, 1, 1, 1]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=1,
stage_id=1, stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1,
1, 1, 1, 1]
+ SortExec: expr=[c@0 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@0], 8), plan_id=2, stage_id=2,
stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1,
1, 1]
+ ");
+
+ Ok(())
+}
+
+/// Sort-merge join: same shuffle structure as the hash-join case, but
+/// DataFusion picks `SortMergeJoinExec` when `prefer_hash_join=false`. The
+/// rule is structural — it walks down to the leaf Exchanges and attaches
+/// `CoalescePlan` regardless of the parent join kind. Per-leaf bytes
+/// `[25; 8]` × 2 leaves trace identically to the hash-join case → K=2.
+#[tokio::test]
+async fn should_attach_coalesce_to_both_sides_of_sort_merge_join()
+-> datafusion::error::Result<()> {
+ let config = SessionConfig::new_with_ballista()
+ .with_target_partitions(8)
+ .with_round_robin_repartition(false)
+ .with_ballista_coalesce_enabled(true)
+ .with_ballista_coalesce_target_partition_bytes(200)
+ .set_bool("datafusion.optimizer.prefer_hash_join", false);
+
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build();
+ let ctx = SessionContext::new_with_state(state);
+ register_partitioned_table(&ctx, "t1", 8)?;
+ register_partitioned_table(&ctx, "t2", 8)?;
+
+ let plan = ctx
+ .sql("select t1.a, t2.b from t1 join t2 on t1.c = t2.c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ let stages = planner.runnable_stages()?.unwrap();
+ assert_eq!(2, stages.len());
+
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[25; 8]))?;
+ planner.finalise_stage_internal(1, partitions_with_byte_sizes(&[25; 8]))?;
+
+ // Surface the join stage so `CoalescePartitionsRule` fires per-stage.
+ let _ = planner.runnable_stages()?;
+
+ assert_plan!(planner.current_plan(), @ "
+ AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2,
stage_resolved=false
+ ProjectionExec: expr=[a@0 as a, b@2 as b]
+ SortMergeJoinExec: join_type=Inner, on=[(c@1, c@1)]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=0, stage_id=0,
stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1,
1, 1]
+ SortExec: expr=[c@1 ASC], preserve_partitioning=[true]
+ ExchangeExec: partitioning=Hash([c@1], 8), plan_id=1, stage_id=1,
stage_resolved=true, coalesce=2 of 8
+ DataSourceExec: partitions=8, partition_sizes=[1, 1, 1, 1, 1, 1,
1, 1]
+ ");
+
+ Ok(())
+}
+
+/// End-to-end: after the rule attaches `coalesce=K of M` to a leaf Exchange,
+/// the adapter must build the downstream `ShuffleReaderExec` with `K`
+/// partitions instead of `M`. The next runnable stage's plan tree is the
+/// proof — its `ShuffleReaderExec: partitioning: Hash([c@0], 2)` shows the
+/// rule's decision has flowed through the adapter into the runnable plan,
+/// not just sitting on the Exchange as metadata.
+#[tokio::test]
+async fn shuffle_reader_uses_coalesced_k_when_rule_fires() ->
datafusion::error::Result<()>
+{
+ let ctx = coalesce_context(8, true);
+ ctx.register_batch("t", mock_batch()?)?;
+
+ let plan = ctx
+ .sql("select min(a) as c0, max(b) as c1, c as c2 from t group by c")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let mut planner =
+ AdaptivePlanner::try_new(ctx.state().config(), plan,
"test_job".to_string())?;
+
+ // Stage 0 is the upstream shuffle writer, partitioning by `c` into M=8.
+ let stages = planner.runnable_stages()?.unwrap();
+ assert_eq!(1, stages.len());
+ assert_plan!(stages[0].plan.as_ref(), @ r"
+ SortShuffleWriterExec: partitioning=Hash([c@0], 8)
+ AggregateExec: mode=Partial, gby=[c@2 as c], aggr=[min(t.a), max(t.b)]
+ DataSourceExec: partitions=1, partition_sizes=[1]
+ ");
+
+ // Finalize stage 0 with 8 partitions × 50 bytes. Bin-pack at target=200
+ // yields K=2 (same trace as the happy-path test).
+ planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[50; 8]))?;
+
+ // Stage 1 is the final stage. Its `ShuffleReaderExec` exposes K=2
+ // partitions — the rule's coalesce decision is now baked into the
+ // adapter's reader-construction path.
+ let stages = planner.runnable_stages()?.unwrap();
+ assert_eq!(1, stages.len());
+ assert_plan!(stages[0].plan.as_ref(), @ r"
+ ShuffleWriterExec: partitioning: None
+ ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
+ AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a),
max(t.b)]
+ ShuffleReaderExec: partitioning: Hash([c@0], 2), coalesce: 2 of 8
+ ");
+
+ Ok(())
+}
diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs
b/ballista/scheduler/src/state/aqe/test/mod.rs
index a60cd038a..831f25c1c 100644
--- a/ballista/scheduler/src/state/aqe/test/mod.rs
+++ b/ballista/scheduler/src/state/aqe/test/mod.rs
@@ -17,6 +17,8 @@
/// Test if stages can be added or removed
mod alter_stages;
+/// Functional tests for the CoalescePartitionsRule end-to-end through the
planner
+mod coalesce_rule;
/// Tests if plan is going to be split to stages correctly
mod plan_to_stages;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
