diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 7bd4191e5a..c114a72e9a 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 c3f841adf6..26efc0491f 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> = 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 5dd4b15095..9d40958f96 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 eb9b43896d..cb4689c68d 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 9f9efa8f1d..31708a701e 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, } #[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 045b303f01..213a600a78 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 = 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 = 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::() + .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 b9e44f6dc1..83f0351db5 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 c2daa312b2..b9f3d5d7f1 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)