This is an automated email from the ASF dual-hosted git repository.
andygrove 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 752d4de13 fix: propagate sort shuffle memory_limit_per_task_bytes to
executors (#2090)
752d4de13 is described below
commit 752d4de136e81418f217a133beef51119bf5592e
Author: Andy Grove <[email protected]>
AuthorDate: Sun Jul 19 20:53:55 2026 -0600
fix: propagate sort shuffle memory_limit_per_task_bytes to executors (#2090)
---
ballista/core/proto/ballista.proto | 5 ++
ballista/core/src/config.rs | 8 +++-
.../src/execution_plans/sort_shuffle/config.rs | 4 +-
.../src/execution_plans/sort_shuffle/writer.rs | 40 +++++++++++++++-
ballista/core/src/serde/generated/ballista.rs | 6 +++
ballista/core/src/serde/mod.rs | 56 +++++++++++++++++++++-
docs/source/user-guide/configs.md | 18 ++++---
docs/source/user-guide/tuning-guide.md | 49 ++++++++++++++-----
8 files changed, 158 insertions(+), 28 deletions(-)
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index 7bd4191e5..c114a72e9 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -77,6 +77,11 @@ message SortShuffleWriterExecNode {
datafusion.PhysicalHashRepartition output_partitioning = 4;
// Target batch size in rows when materializing buffered shuffle data.
uint64 batch_size = 8;
+ // Per-task buffered-bytes budget at which the writer spills its in-memory
+ // batches to disk. A value of 0 disables the per-task budget so spilling is
+ // driven solely by runtime memory-pool pressure. Optional so a plan produced
+ // before this field existed decodes to the built-in default rather than 0.
+ optional uint64 memory_limit_per_task_bytes = 9;
}
message UnresolvedShuffleExecNode {
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index bef5556fa..716597ff6 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -94,7 +94,8 @@ pub const BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE: &str =
pub const BALLISTA_SHUFFLE_WRITER_CHANNEL_CAPACITY: &str =
"ballista.shuffle.writer_channel_capacity";
/// Configuration key for the per-task buffered-bytes budget at which the
-/// sort shuffle writer spills its in-memory batches to disk.
+/// sort shuffle writer spills its in-memory batches to disk. Set to 0 to
+/// disable the per-task budget and spill only under memory-pool pressure.
pub const BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str =
"ballista.shuffle.sort_based.memory_limit_per_task_bytes";
/// Configuration key for the byte-size threshold below which a hash join's
@@ -241,7 +242,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
"Per-task buffered-bytes budget at which the sort
shuffle writer spills its \
in-memory batches to disk. Counted independently of
the runtime memory pool, so \
spilling kicks in even when the pool is unbounded.
Total worst-case sort shuffle \
- memory per executor is approximately vcores * this
value.".to_string(),
+ memory per executor is approximately vcores * this
value. Set to 0 to disable the \
+ per-task budget and rely solely on runtime
memory-pool pressure to trigger spilling; \
+ this is safe only with a bounded memory pool,
otherwise the writer never spills and \
+ may run out of memory.".to_string(),
DataType::UInt64,
Some((256 * 1024 * 1024).to_string())),
ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES.to_string(),
diff --git a/ballista/core/src/execution_plans/sort_shuffle/config.rs
b/ballista/core/src/execution_plans/sort_shuffle/config.rs
index 5dd4b1509..9d40958f9 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/config.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/config.rs
@@ -27,7 +27,9 @@ pub struct SortShuffleConfig {
pub batch_size: usize,
/// Per-task buffered-bytes budget at which the writer spills its in-memory
/// batches to disk. Counted independently of the runtime `MemoryPool`, so
- /// spilling kicks in even when the pool is unbounded.
+ /// spilling kicks in even when the pool is unbounded. A value of 0
disables
+ /// the per-task budget, leaving the runtime `MemoryPool` as the sole spill
+ /// trigger.
pub memory_limit_per_task_bytes: usize,
}
diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
index eb9b43896..cb4689c68 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
@@ -443,7 +443,10 @@ impl SortShuffleWriterExec {
// `MemoryPool`. Drives spill decisions so the writer bounds its
// RSS even when the pool is unbounded.
let mut buffered_bytes: usize = 0;
+ // A limit of 0 disables the per-task budget, leaving the runtime
+ // `MemoryPool` as the sole spill trigger.
let memory_limit = config.memory_limit_per_task_bytes;
+ let per_task_budget_enabled = memory_limit > 0;
while let Some(result) = stream.next().await {
let input_batch = result?;
@@ -476,7 +479,9 @@ impl SortShuffleWriterExec {
let pool_rejected_growth =
reservation.try_grow(growth).is_err();
buffered_bytes = buffered_bytes.saturating_add(growth);
- if pool_rejected_growth || buffered_bytes >= memory_limit {
+ if pool_rejected_growth
+ || (per_task_budget_enabled && buffered_bytes >=
memory_limit)
+ {
let spill_timer = metrics.spill_time.timer();
let (event_batches, event_bytes) = spill_all_partitions(
&mut buffered,
@@ -1350,6 +1355,39 @@ mod tests {
.await
}
+ #[tokio::test]
+ async fn zero_budget_disables_per_task_spill_trigger() -> Result<()> {
+ // A per-task budget of 0 disables trigger 2 entirely. With a pool that
+ // never rejects, neither trigger can fire, so the whole payload stays
in
+ // memory and round-trips without spilling — even though the same
payload
+ // would spill under any small non-zero budget.
+ run_round_trip(
+ /* num_batches */ 10,
+ /* rows_per_batch */ 8192,
+ /* num_partitions */ 4,
+ /* sort_shuffle_memory_limit_bytes */ 0,
+ /* pool_bytes */ POOL_NEVER_REJECTS,
+ /* expect_spills */ false,
+ )
+ .await
+ }
+
+ #[tokio::test]
+ async fn zero_budget_still_spills_under_memory_pool_pressure() ->
Result<()> {
+ // With the per-task budget disabled (0), the runtime memory pool is
the
+ // sole spill trigger. A 64 KiB pool cannot admit a ~128 KiB batch, so
+ // pool rejection still forces spilling.
+ run_round_trip(
+ /* num_batches */ 8,
+ /* rows_per_batch */ 8192,
+ /* num_partitions */ 2,
+ /* sort_shuffle_memory_limit_bytes */ 0,
+ /* pool_bytes */ 64 * 1024,
+ /* expect_spills */ true,
+ )
+ .await
+ }
+
#[tokio::test]
async fn spills_when_memory_pool_rejects_growth() -> Result<()> {
// Trigger 1 only: a 64 KiB pool cannot admit even one ~128 KiB batch,
so
diff --git a/ballista/core/src/serde/generated/ballista.rs
b/ballista/core/src/serde/generated/ballista.rs
index 9f9efa8f1..31708a701 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -93,6 +93,12 @@ pub struct SortShuffleWriterExecNode {
/// Target batch size in rows when materializing buffered shuffle data.
#[prost(uint64, tag = "8")]
pub batch_size: u64,
+ /// Per-task buffered-bytes budget at which the writer spills its in-memory
+ /// batches to disk. A value of 0 disables the per-task budget so spilling
is
+ /// driven solely by runtime memory-pool pressure. Optional so a plan
produced
+ /// before this field existed decodes to the built-in default rather than
0.
+ #[prost(uint64, optional, tag = "9")]
+ pub memory_limit_per_task_bytes: ::core::option::Option<u64>,
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct UnresolvedShuffleExecNode {
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index 045b303f0..213a600a7 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -423,7 +423,14 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
} else {
8192
};
- let config = SortShuffleConfig::new(true, batch_size);
+ let mut config = SortShuffleConfig::new(true, batch_size);
+ // Absent (legacy plan) keeps the built-in default; a present
+ // value — including 0, which disables the per-task budget —
+ // is applied verbatim so the session override reaches the
+ // executor where the writer actually runs.
+ if let Some(bytes) =
sort_shuffle_writer.memory_limit_per_task_bytes {
+ config = config.with_memory_limit_per_task_bytes(bytes as
usize);
+ }
Ok(Arc::new(SortShuffleWriterExec::try_new(
sort_shuffle_writer.job_id.clone().into(),
@@ -622,6 +629,9 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
input: None,
output_partitioning,
batch_size: config.batch_size as u64,
+ memory_limit_per_task_bytes: Some(
+ config.memory_limit_per_task_bytes as u64,
+ ),
},
)),
};
@@ -964,6 +974,50 @@ mod test {
);
}
+ /// The sort shuffle writer's per-task memory budget must survive the
+ /// protobuf round trip so a session override reaches the executor where
the
+ /// writer actually runs (issue #2089). A non-default value and the special
+ /// `0` (per-task budget disabled) are both checked.
+ #[tokio::test]
+ async fn test_sort_shuffle_writer_exec_roundtrip_preserves_memory_limit() {
+ use datafusion::physical_plan::empty::EmptyExec;
+
+ for memory_limit in [7 * 1024 * 1024_usize, 0] {
+ let schema = create_test_schema();
+ let input: Arc<dyn ExecutionPlan> =
Arc::new(EmptyExec::new(schema.clone()));
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()],
4);
+
+ let original_exec = SortShuffleWriterExec::try_new(
+ "job1".into(),
+ 1,
+ input.clone(),
+ String::new(),
+ partitioning.clone(),
+ SortShuffleConfig::new(true, 4096)
+ .with_memory_limit_per_task_bytes(memory_limit),
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec.try_encode(Arc::new(original_exec), &mut buf).unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap();
+ let decoded_exec = decoded_plan
+ .downcast_ref::<SortShuffleWriterExec>()
+ .expect("Expected SortShuffleWriterExec");
+
+ assert_eq!(
+ decoded_exec.config().memory_limit_per_task_bytes,
+ memory_limit,
+ "memory_limit_per_task_bytes must survive the wire, including
0"
+ );
+ assert_eq!(decoded_exec.config().batch_size, 4096);
+ }
+ }
+
#[tokio::test]
async fn test_shuffle_reader_exec_coalesced_roundtrip_single_group() {
let schema = create_test_schema();
diff --git a/docs/source/user-guide/configs.md
b/docs/source/user-guide/configs.md
index b9e44f6dc..83f0351db 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -94,16 +94,14 @@ The following session-level keys control Ballista's shuffle
behavior. See
the [tuning guide](tuning-guide.md#shuffle-implementation) for an
explanation of the sort-based (default) and hash-based shuffle writers.
-| key | type | default |
description
|
-| --------------------------------------------- | ------- | --------- |
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
-| ballista.shuffle.max_concurrent_read_requests | UInt64 | 64 |
Maximum number of concurrent fetch requests the shuffle reader will issue.
|
-| ballista.shuffle.force_remote_read | Boolean | false | Forces
the shuffle reader to fetch every partition through Arrow Flight, even when the
data is local. Intended for testing. |
-| ballista.shuffle.remote_read_prefer_flight | Boolean | false | For
remote reads, prefer the Arrow Flight reader over the block reader. The block
reader is generally faster.
|
-| ballista.shuffle.sort_based.enabled | Boolean | true |
Enables the sort-based shuffle writer (consolidated data file per input
partition with an index, instead of one file per (input partition, output
partition) pair). |
-| ballista.shuffle.sort_based.buffer_size | UInt64 | 1048576 |
Per-partition buffer size in bytes for the sort-based writer (1 MiB default).
|
-| ballista.shuffle.sort_based.memory_limit | UInt64 | 268435456 | Total
in-memory budget across all output-partition buffers for the sort-based writer
(256 MiB default). |
-| ballista.shuffle.sort_based.spill_threshold | Utf8 | "0.8" |
Fraction of `memory_limit` at which the largest buffers spill to disk. Must be
in the range 0–1.
|
-| ballista.shuffle.sort_based.batch_size | UInt64 | 8192 | Target
row count when coalescing buffered batches before they are written or spilled.
|
+| key | type | default
| description
|
+| ------------------------------------------------------- | ------- |
--------- |
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
+| ballista.shuffle.max_concurrent_read_requests | UInt64 | 64
| Maximum number of concurrent fetch requests the shuffle reader will issue.
|
+| ballista.shuffle.force_remote_read | Boolean | false
| Forces the shuffle reader to fetch every partition through Arrow Flight,
even when the data is local. Intended for testing.
|
+| ballista.shuffle.remote_read_prefer_flight | Boolean | false
| For remote reads, prefer the Arrow Flight reader over the block reader. The
block reader is generally faster.
|
+| ballista.shuffle.sort_based.enabled | Boolean | true
| Enables the sort-based shuffle writer (consolidated data file per input
partition with an index, instead of one file per (input partition, output
partition) pair).
|
+| ballista.shuffle.sort_based.batch_size | UInt64 | 8192
| Target row count when coalescing buffered batches before they are written or
spilled.
|
+| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 |
268435456 | Per-task buffered-bytes budget at which the sort-based writer
spills to disk (256 MiB default), counted independently of the runtime memory
pool. Set to `0` to disable the per-task budget and spill only under
memory-pool pressure — safe only with a bounded memory pool, otherwise the
writer never spills and may run out of memory. |
## Ballista Scheduler Configuration Settings
diff --git a/docs/source/user-guide/tuning-guide.md
b/docs/source/user-guide/tuning-guide.md
index c2daa312b..b9f3d5d7f 100644
--- a/docs/source/user-guide/tuning-guide.md
+++ b/docs/source/user-guide/tuning-guide.md
@@ -133,17 +133,42 @@ write latency.
### Sort-based shuffle (default)
-The sort-based writer accumulates batches in a per-output-partition
-in-memory buffer. When the total buffered size crosses a threshold, the
-largest buffers are spilled to disk. After the input stream finishes, the
-remaining in-memory data and any spilled batches are merged and written
-into a single consolidated Arrow IPC file per input partition, alongside
-an index file that lets readers seek directly to a given output partition.
+The sort-based writer accumulates incoming batches in memory, tracking each
+row's output partition. It spills the buffered batches to disk when either of
+two independent triggers fires:
+
+- the runtime memory pool rejects a growth request (the executor is under
+ memory pressure — see [Configuring Executor Memory
Pool](#configuring-executor-memory-pool)), or
+- the per-task buffered-bytes budget
+ (`ballista.shuffle.sort_based.memory_limit_per_task_bytes`) is reached. This
+ budget is counted independently of the memory pool, so it bounds the writer's
+ memory even when the pool is unbounded (the default). Setting it to `0`
+ disables this trigger, leaving memory-pool pressure as the sole spill signal.
+
+**Warning:** setting `memory_limit_per_task_bytes` to `0` while the executor
+uses the default unbounded memory pool disables both spill triggers. The writer
+then buffers the entire task's shuffle output in memory and never spills, which
+can exhaust the executor and cause an out-of-memory failure on large inputs.
+Only use `0` together with a bounded memory pool (see
+[Configuring Executor Memory Pool](#configuring-executor-memory-pool)), so
+pool pressure still forces spilling.
+
+After the input stream finishes, the remaining in-memory data and any spilled
+batches are written into a single consolidated Arrow IPC file per input
+partition, alongside an index file that lets readers seek directly to a given
+output partition.
This produces `2 × N` files instead of `N × M`, coalesces small batches
to a target size before writing, and bounds shuffle memory use via
spilling — at the cost of higher write latency than the hash writer.
+Worst-case sort-shuffle memory per executor is approximately
+`vcores × memory_limit_per_task_bytes`, since one writer task can run per
+core. Lower the per-task budget on memory-constrained executors to spill
+sooner, or raise it to keep more data in memory and reduce spill I/O. Setting
it
+to `0` removes the budget entirely and is safe only with a bounded memory pool
+(see the warning above).
+
### Hash-based shuffle (opt-in)
The hash-based writer hashes each incoming `RecordBatch` and immediately
@@ -164,13 +189,11 @@ let session_config = SessionConfig::new_with_ballista()
The following session-level keys tune its behavior:
-| key | type | default |
description
|
-| ------------------------------------------- | ------- | --------- |
---------------------------------------------------------------------------------------------------------
|
-| ballista.shuffle.sort_based.enabled | Boolean | true | Enables
the sort-based shuffle writer.
|
-| ballista.shuffle.sort_based.buffer_size | UInt64 | 1048576 |
Per-partition buffer size in bytes (1 MiB default).
|
-| ballista.shuffle.sort_based.memory_limit | UInt64 | 268435456 | Total
in-memory budget across all output-partition buffers (256 MiB default).
|
-| ballista.shuffle.sort_based.spill_threshold | Utf8 | "0.8" | Fraction
of `memory_limit` at which the largest buffers begin spilling to disk. Must be
in the range 0–1. |
-| ballista.shuffle.sort_based.batch_size | UInt64 | 8192 | Target
row count when coalescing buffered batches before they are written or spilled.
|
+| key | type | default
| description
|
+| ------------------------------------------------------- | ------- |
--------- |
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
+| ballista.shuffle.sort_based.enabled | Boolean | true
| Enables the sort-based shuffle writer.
|
+| ballista.shuffle.sort_based.batch_size | UInt64 | 8192
| Target row count when coalescing buffered batches before they are written or
spilled.
|
+| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 |
268435456 | Per-task buffered-bytes budget at which the writer spills to disk
(256 MiB default). Counted independently of the runtime memory pool. Set to `0`
to spill only under memory pressure — safe only with a bounded memory pool,
otherwise the writer never spills and may run out of memory. |
## Adaptive Query Execution (Experimental)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
