feat: Improve TPC-H performance with AQE on

ballista/core/src/config.rs

@@ -111,6 +111,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.
@@ -264,6 +271,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,
@@ -615,6 +628,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)
@@ -884,4 +902,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/writer.rs

@@ -62,7 +62,7 @@ use datafusion::physical_plan::{
     SendableRecordBatchStream, Statistics, displayable,
 };
 use futures::{StreamExt, TryStreamExt};
-use log::{debug, warn};
+use log::{debug, info, warn};
 
 /// Result of finalizing shuffle output: (data_path, index_path, partition_write_stats)
 /// where partition_write_stats is (partition_id, num_batches, num_rows, num_bytes)
@@ -439,8 +439,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;
             // A limit of 0 disables the per-task budget, leaving the runtime
             // `MemoryPool` as the sole spill trigger.
@@ -567,7 +568,7 @@ impl SortShuffleWriterExec {
                     write_time,
                 );
             } else {
-                debug!(
+                info!(
                     "Sort shuffle write for partition {} completed. \
                      Output: {:?}, Index: {:?}, Rows: {}, \
                      repart_time={:?} spill_time={:?} write_time={:?}, \

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/mod.rs

@@ -1063,6 +1063,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();