fix: propagate sort shuffle memory_limit_per_task_bytes to executors

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 {

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(),

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,
 }
 

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

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 {

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();

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