feat: spilling hash join operator (opt-in, inner joins)

ballista/core/Cargo.toml

@@ -64,6 +64,7 @@ prost = { workspace = true }
 prost-types = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+smallvec = "1.15"
 tokio = { workspace = true, features = ["rt-multi-thread"] }
 tokio-stream = { workspace = true, features = ["net"] }
 tonic = { workspace = true }

ballista/core/proto/ballista.proto

@@ -51,6 +51,7 @@ message BallistaPhysicalPlanNode {
     UnresolvedShuffleExecNode unresolved_shuffle = 3;
     SortShuffleWriterExecNode sort_shuffle_writer = 4;
     ChaosExecNode chaos_exec = 5;
+    SpillingHashJoinExecNode spilling_hash_join = 6;
   }
 }
 
@@ -60,6 +61,19 @@ message ChaosExecNode {
   uint64 seed = 3;
 }
 
+// Left/right children are not carried here: `datafusion-proto`'s
+// `PhysicalPlanNode` decodes them itself and passes the results to
+// `PhysicalExtensionCodec::try_decode` via its `inputs` parameter, matching
+// the ChaosExecNode convention. This message serializes only the join's own
+// state.
+message SpillingHashJoinExecNode {
+  repeated datafusion.PhysicalExprNode left_keys = 1;
+  repeated datafusion.PhysicalExprNode right_keys = 2;
+  // 0 = Partitioned, 1 = CollectLeft (v1 only emits 0)
+  uint32 partition_mode = 3;
+  uint64 num_sub_partitions = 4;
+}
+
 message ShuffleWriterExecNode {
   //TODO it seems redundant to provide job and stage id here since we also have them
   // in the TaskDefinition that wraps this plan

ballista/core/src/config.rs

@@ -154,6 +154,15 @@ pub const BALLISTA_SHUFFLE_COMPRESSION_CODEC: &str = "ballista.shuffle.compressi
 pub const BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK: &str =
     "ballista.scheduler.max_partitions_per_task";
 
+/// Enables substituting eligible Partitioned inner `HashJoinExec` nodes with the
+/// spilling hash join operator. Disabled by default.
+pub const BALLISTA_SPILLING_HASH_JOIN_ENABLED: &str =
+    "ballista.execution.spilling_hash_join.enabled";
+/// Number of in-memory sub-partitions the spilling hash join splits each build
+/// side into.
+pub const BALLISTA_SPILLING_HASH_JOIN_PARTITIONS: &str =
+    "ballista.execution.spilling_hash_join.partitions";
+
 /// Result type for configuration parsing operations.
 pub type ParseResult<T> = result::Result<T, String>;
 use std::sync::LazyLock;
@@ -373,6 +382,18 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
             DataType::UInt64,
             Some(1.to_string()),
         ),
+        ConfigEntry::new(
+            BALLISTA_SPILLING_HASH_JOIN_ENABLED.to_string(),
+            "Enable the spilling hash join operator".to_string(),
+            DataType::Boolean,
+            Some(false.to_string()),
+        ),
+        ConfigEntry::new(
+            BALLISTA_SPILLING_HASH_JOIN_PARTITIONS.to_string(),
+            "Sub-partitions per build side in the spilling hash join".to_string(),
+            DataType::UInt64,
+            Some(16.to_string()),
+        ),
     ];
     entries
         .into_iter()
@@ -615,6 +636,18 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS)
     }
 
+    /// Returns whether eligible Partitioned inner `HashJoinExec` nodes are
+    /// substituted with the spilling hash join operator.
+    pub fn spilling_hash_join_enabled(&self) -> bool {
+        self.get_bool_setting(BALLISTA_SPILLING_HASH_JOIN_ENABLED)
+    }
+
+    /// Returns the number of in-memory sub-partitions the spilling hash join
+    /// splits each build side into.
+    pub fn spilling_hash_join_partitions(&self) -> usize {
+        self.get_usize_setting(BALLISTA_SPILLING_HASH_JOIN_PARTITIONS)
+    }
+
     /// 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 +917,11 @@ mod tests {
         assert_eq!(16777216, config.grpc_client_max_message_size());
         Ok(())
     }
+
+    #[test]
+    fn spilling_hash_join_defaults() {
+        let config = BallistaConfig::default();
+        assert!(!config.spilling_hash_join_enabled());
+        assert_eq!(16, config.spilling_hash_join_partitions());
+    }
 }

ballista/core/src/execution_plans/mod.rs

@@ -25,6 +25,8 @@ mod shuffle_reader;
 mod shuffle_writer;
 mod shuffle_writer_trait;
 pub mod sort_shuffle;
+/// Hash join operator whose build side can spill sub-partitions to disk.
+pub mod spilling_hash_join;
 mod unresolved_shuffle;
 
 use std::path::{Path, PathBuf};
@@ -40,6 +42,7 @@ pub use shuffle_writer::ShuffleWriterExec;
 pub use shuffle_writer::compute_global_output_partition_ids;
 pub use shuffle_writer_trait::ShuffleWriter;
 pub use sort_shuffle::SortShuffleWriterExec;
+pub use spilling_hash_join::SpillingHashJoinExec;
 pub use unresolved_shuffle::UnresolvedShuffleExec;
 
 use crate::JobId;