fix: Various fixes for AQE join selection

ballista/core/src/config.rs

@@ -110,8 +110,10 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str =
 
 /// 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).
+/// exceeds this, the join falls back to SortMergeJoin (spillable). Defaults to
+/// 64 MiB. `0` disables the check (hash join is used regardless of build size);
+/// because AQE does not consult `datafusion.optimizer.prefer_hash_join`, that
+/// also makes SortMergeJoin unreachable for a Partitioned join.
 pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str =
     "ballista.optimizer.hash_join_max_build_partition_bytes";
 
@@ -267,9 +269,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
         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(),
+                         SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, \
+                         which makes AQE use a hash join regardless of build size.".to_string(),
                          DataType::UInt64,
-                         Some("0".to_string())),
+                         Some((64 * 1024 * 1024).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,
@@ -888,11 +891,15 @@ mod tests {
         Ok(())
     }
 
+    // The default must stay non-zero: `0` disables the fit check, and since AQE
+    // no longer consults `prefer_hash_join`, a disabled check would leave the
+    // Partitioned arm unconditionally on hash join and make SortMergeJoin — the
+    // spillable strategy — unreachable.
     #[test]
-    fn hash_join_max_build_partition_bytes_defaults_to_zero() {
+    fn hash_join_max_build_partition_bytes_defaults_to_64_mib() {
         assert_eq!(
             BallistaConfig::default().hash_join_max_build_partition_bytes(),
-            0
+            64 * 1024 * 1024
         );
     }
 }

ballista/core/src/extension.rs

@@ -20,9 +20,10 @@ use crate::config::{
     BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE, BALLISTA_CLIENT_USE_TLS,
     BALLISTA_COALESCE_ENABLED, BALLISTA_COALESCE_MERGED_PARTITION_FACTOR,
     BALLISTA_COALESCE_SMALL_PARTITION_FACTOR, BALLISTA_COALESCE_TARGET_PARTITION_BYTES,
-    BALLISTA_JOB_NAME, BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
-    BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
-    BALLISTA_STANDALONE_PARALLELISM, BallistaConfig,
+    BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, BALLISTA_JOB_NAME,
+    BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
+    BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, BALLISTA_STANDALONE_PARALLELISM,
+    BallistaConfig,
 };
 use crate::planner::BallistaQueryPlanner;
 use crate::serde::protobuf::KeyValuePair;
@@ -208,6 +209,11 @@ pub trait SessionConfigExt {
     /// falling back to SortMergeJoin under AQE. `0` disables the check.
     fn ballista_hash_join_max_build_partition_bytes(&self) -> usize;
 
+    /// Sets the maximum per-partition hash-join build-side bytes before falling
+    /// back to SortMergeJoin under AQE. Setting `0` disables the check, which
+    /// leaves AQE on a hash join whatever the build size.
+    fn with_ballista_hash_join_max_build_partition_bytes(self, max_bytes: usize) -> Self;
+
     /// retrieves grpc client max message size
     fn ballista_grpc_client_max_message_size(&self) -> usize;
 
@@ -524,6 +530,15 @@ impl SessionConfigExt for SessionConfig {
             })
     }
 
+    fn with_ballista_hash_join_max_build_partition_bytes(self, max_bytes: usize) -> Self {
+        if self.options().extensions.get::<BallistaConfig>().is_some() {
+            self.set_usize(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, max_bytes)
+        } else {
+            self.with_option_extension(BallistaConfig::default())
+                .set_usize(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES, max_bytes)
+        }
+    }
+
     fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize {
         self.options()
             .extensions