feat: Improve TPC-H performance [WIP]

ballista/core/proto/ballista.proto

@@ -77,6 +77,9 @@ 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 to disk. 0 means
+  // no cap (spill is driven solely by memory-pool pressure).
+  uint64 memory_limit_per_task_bytes = 9;
 }
 
 message UnresolvedShuffleExecNode {

ballista/core/src/config.rs

@@ -110,6 +110,13 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
 pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str =
     "ballista.optimizer.broadcast_join_threshold_rows";
 
+/// Configuration key for the maximum per-partition hash-join build-side bytes
+/// permitted for a Partitioned hash join under AQE. When a build partition
+/// exceeds this, the join falls back to SortMergeJoin (spillable). `0` disables
+/// the check (hash join is used regardless of build size).
+pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str =
+    "ballista.optimizer.hash_join_max_build_partition_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.
@@ -235,11 +242,12 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          Some(DEFAULT_SHUFFLE_CHANNEL_CAPACITY.to_string())),
         ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES.to_string(),
                          "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(),
+                         in-memory batches to disk, counted independently of the runtime memory pool. A \
+                         value of 0 (the default) disables this cap so spilling is driven solely by \
+                         memory-pool pressure. When non-zero, total worst-case sort shuffle memory per \
+                         executor is approximately vcores * this value.".to_string(),
                          DataType::UInt64,
-                         Some((256 * 1024 * 1024).to_string())),
+                         Some("0".to_string())),
         ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES.to_string(),
                          "Byte-size threshold below which a hash join's smaller side is \
                           promoted to CollectLeft and lowered via the broadcast pattern. \
@@ -256,6 +264,12 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                           promotion via the row-count path.".to_string(),
                          DataType::UInt64,
                          Some((1_000_000).to_string())),
+        ConfigEntry::new(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES.to_string(),
+                         "Maximum per-partition hash-join build-side bytes for a Partitioned \
+                         hash join under AQE. A build partition larger than this falls back to \
+                         SortMergeJoin (spillable). 0 (the default) disables the check.".to_string(),
+                         DataType::UInt64,
+                         Some("0".to_string())),
         ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(),
                          "Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients.".to_string(),
                          DataType::Boolean,
@@ -594,6 +608,11 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS)
     }
 
+    /// Maximum per-partition hash-join build-side bytes before falling back to SMJ.
+    pub fn hash_join_max_build_partition_bytes(&self) -> usize {
+        self.get_usize_setting(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES)
+    }
+
     /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
     pub fn coalesce_enabled(&self) -> bool {
         self.get_bool_setting(BALLISTA_COALESCE_ENABLED)
@@ -858,4 +877,12 @@ mod tests {
         assert_eq!(16777216, config.grpc_client_max_message_size());
         Ok(())
     }
+
+    #[test]
+    fn hash_join_max_build_partition_bytes_defaults_to_zero() {
+        assert_eq!(
+            BallistaConfig::default().hash_join_max_build_partition_bytes(),
+            0
+        );
+    }
 }

ballista/core/src/execution_plans/sort_shuffle/config.rs

@@ -26,8 +26,10 @@ pub struct SortShuffleConfig {
     /// `interleave_record_batch` (default: 8192).
     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.
+    /// batches to disk, counted independently of the runtime `MemoryPool`. A
+    /// value of 0 (the default) disables this cap so spilling is driven solely
+    /// by memory-pool pressure; a non-zero value adds a second spill trigger
+    /// that bounds buffered memory even when the pool is unbounded.
     pub memory_limit_per_task_bytes: usize,
 }
 
@@ -36,13 +38,14 @@ impl Default for SortShuffleConfig {
         Self {
             enabled: false,
             batch_size: 8192,
-            memory_limit_per_task_bytes: 256 * 1024 * 1024,
+            memory_limit_per_task_bytes: 0,
         }
     }
 }
 
 impl SortShuffleConfig {
-    /// Creates a new configuration with the default per-task memory limit.
+    /// Creates a new configuration with the default (uncapped) per-task memory
+    /// limit.
     pub fn new(enabled: bool, batch_size: usize) -> Self {
         Self {
             enabled,
@@ -67,15 +70,15 @@ mod tests {
         let config = SortShuffleConfig::default();
         assert!(!config.enabled);
         assert_eq!(config.batch_size, 8192);
-        assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024);
+        assert_eq!(config.memory_limit_per_task_bytes, 0);
     }
 
     #[test]
     fn test_new() {
         let config = SortShuffleConfig::new(true, 4096);
         assert!(config.enabled);
         assert_eq!(config.batch_size, 4096);
-        assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024);
+        assert_eq!(config.memory_limit_per_task_bytes, 0);
     }
 
     #[test]

ballista/core/src/execution_plans/sort_shuffle/writer.rs

@@ -65,7 +65,7 @@ use datafusion::physical_plan::{
     SendableRecordBatchStream, Statistics, displayable,
 };
 use futures::{StreamExt, TryFutureExt, TryStreamExt};
-use log::debug;
+use log::{debug, info, warn};
 
 use crate::serde::scheduler::PartitionStats;
 
@@ -246,8 +246,9 @@ impl SortShuffleWriterExec {
             let mut hash_buffer: Vec<u64> = Vec::new();
             let mut spill_events: u64 = 0;
             // Absolute buffered-bytes counter, independent of the runtime
-            // `MemoryPool`. Drives spill decisions so the writer bounds its
-            // RSS even when the pool is unbounded.
+            // `MemoryPool`. When `memory_limit` is non-zero it caps this counter
+            // as a second spill trigger; a `memory_limit` of 0 disables the cap
+            // so spilling is driven solely by memory-pool pressure.
             let mut buffered_bytes: usize = 0;
             let memory_limit = config.memory_limit_per_task_bytes;
 
@@ -282,7 +283,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
+                    || (memory_limit > 0 && buffered_bytes >= memory_limit)
+                {
                     let spill_timer = metrics.spill_time.timer();
                     let (event_batches, event_bytes) = spill_all_partitions(
                         &mut buffered,
@@ -350,18 +353,35 @@ impl SortShuffleWriterExec {
             // Reservation drops naturally; nothing left to free.
             drop(reservation);
 
-            debug!(
-                "Sort shuffle write for partition {} completed in {} seconds. \
-                 Output: {:?}, Index: {:?}, Spill events: {}, Spill batches: {}, \
-                 Spill bytes: {}",
-                input_partition,
-                now.elapsed().as_secs(),
-                data_path,
-                index_path,
-                spill_events,
-                total_spilled_batches,
-                total_bytes_spilled
-            );
+            let elapsed_secs = now.elapsed().as_secs();
+            if total_bytes_spilled > 0 {
+                warn!(
+                    "Sort shuffle spill: job={} stage={} input_partition={} \
+                     spilled {} bytes in {} batches ({} events) under memory \
+                     pressure; write completed in {} seconds",
+                    job_id,
+                    stage_id,
+                    input_partition,
+                    total_bytes_spilled,
+                    total_spilled_batches,
+                    spill_events,
+                    elapsed_secs,
+                );
+            } else {
+                info!(
+                    "Sort shuffle write for partition {} completed in {} seconds. \
+                     Output: {:?}, Index: {:?}, Rows: {}, Spill events: {}, \
+                     Spill batches: {}, Spill bytes: {}",
+                    input_partition,
+                    elapsed_secs,
+                    data_path,
+                    index_path,
+                    total_rows,
+                    spill_events,
+                    total_spilled_batches,
+                    total_bytes_spilled
+                );
+            }
 
             let mut results = Vec::new();
             for (part_id, num_batches, num_rows, num_bytes) in partition_stats {

ballista/core/src/extension.rs

@@ -204,6 +204,10 @@ pub trait SessionConfigExt {
     /// disables promotion via the row-count path.
     fn with_ballista_broadcast_join_threshold_rows(self, threshold_rows: usize) -> Self;
 
+    /// Returns the maximum per-partition hash-join build-side bytes before
+    /// falling back to SortMergeJoin under AQE. `0` disables the check.
+    fn ballista_hash_join_max_build_partition_bytes(&self) -> usize;
+
     /// retrieves grpc client max message size
     fn ballista_grpc_client_max_message_size(&self) -> usize;
 
@@ -513,6 +517,16 @@ impl SessionConfigExt for SessionConfig {
         }
     }
 
+    fn ballista_hash_join_max_build_partition_bytes(&self) -> usize {
+        self.options()
+            .extensions
+            .get::<BallistaConfig>()
+            .map(|c| c.hash_join_max_build_partition_bytes())
+            .unwrap_or_else(|| {
+                BallistaConfig::default().hash_join_max_build_partition_bytes()
+            })
+    }
+
     fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize {
         self.options()
             .extensions

ballista/core/src/serde/generated/ballista.rs

@@ -93,6 +93,10 @@ 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 to disk. 0 means
+    /// no cap (spill is driven solely by memory-pool pressure).
+    #[prost(uint64, tag = "9")]
+    pub memory_limit_per_task_bytes: u64,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct UnresolvedShuffleExecNode {

ballista/core/src/serde/mod.rs

@@ -423,7 +423,10 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                 } else {
                     8192
                 };
-                let config = SortShuffleConfig::new(true, batch_size);
+                let config = SortShuffleConfig::new(true, batch_size)
+                    .with_memory_limit_per_task_bytes(
+                        sort_shuffle_writer.memory_limit_per_task_bytes as usize,
+                    );
 
                 Ok(Arc::new(SortShuffleWriterExec::try_new(
                     sort_shuffle_writer.job_id.clone().into(),
@@ -622,6 +625,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                         input: None,
                         output_partitioning,
                         batch_size: config.batch_size as u64,
+                        memory_limit_per_task_bytes: config.memory_limit_per_task_bytes
+                            as u64,
                     },
                 )),
             };
@@ -1009,6 +1014,45 @@ mod test {
         assert_eq!(stored.groups[0].upstream_indices, vec![0, 1, 2, 3]);
     }
 
+    #[tokio::test]
+    async fn sort_shuffle_writer_memory_limit_survives_roundtrip() {
+        use datafusion::physical_plan::empty::EmptyExec;
+
+        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 config = SortShuffleConfig::new(true, 4096)
+            .with_memory_limit_per_task_bytes(1024 * 1024 * 1024);
+        let original = SortShuffleWriterExec::try_new(
+            "job-1".to_string().into(),
+            3,
+            input.clone(),
+            String::new(),
+            partitioning,
+            config,
+        )
+        .unwrap();
+
+        let codec = BallistaPhysicalExtensionCodec::default();
+        let mut buf: Vec<u8> = vec![];
+        codec.try_encode(Arc::new(original), &mut buf).unwrap();
+
+        let ctx = SessionContext::new().task_ctx();
+        let decoded = codec.try_decode(&buf, &[input], &ctx).unwrap();
+        let decoded = decoded
+            .downcast_ref::<SortShuffleWriterExec>()
+            .expect("Expected SortShuffleWriterExec");
+
+        assert_eq!(
+            decoded.config().memory_limit_per_task_bytes,
+            1024 * 1024 * 1024,
+            "memory limit override must survive serialization to the executor"
+        );
+        assert_eq!(decoded.config().batch_size, 4096);
+    }
+
     #[tokio::test]
     async fn test_shuffle_reader_exec_coalesced_roundtrip_multi_group_mixed_sizes() {
         let schema = create_test_schema();