From f9d253491daff802fd88ea82003276da5c6fb6b2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 10:18:24 -0600 Subject: [PATCH 01/15] feat(core): add spilling hash join config keys --- ballista/core/src/config.rs | 40 +++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 716597ff6b..8677c2b51b 100644 --- a/ballista/core/src/config.rs +++ b/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 = result::Result; use std::sync::LazyLock; @@ -373,6 +382,18 @@ static CONFIG_ENTRIES: LazyLock> = 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()); + } } From ffe43222c92f06e20e427a8f5b21fb6532d828e2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 10:28:23 -0600 Subject: [PATCH 02/15] feat(core): scaffold SpillingHashJoinExec operator --- ballista/core/src/execution_plans/mod.rs | 3 + .../spilling_hash_join/exec.rs | 268 ++++++++++++++++++ .../execution_plans/spilling_hash_join/mod.rs | 20 ++ 3 files changed, 291 insertions(+) create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/exec.rs create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/mod.rs diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index f863fd4776..3efe51f200 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/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; diff --git a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs new file mode 100644 index 0000000000..1853f82f46 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -0,0 +1,268 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// SpillingHashJoinExec is a physical execution plan node for a hash join +// whose build side spills to disk instead of requiring the entire hash +// table to fit in memory. This is currently a scaffold only: the struct, +// its constructor, and its accessors exist so later work can wire up +// planning/serde/substitution, but `execute` is not yet implemented. +// +// v1 invariants (enforced by the constructor, which always builds an +// inner join with no residual filter and no output projection): +// - join_type is always JoinType::Inner +// - filter is always None +// - projection is always None + +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{JoinType, NullEquality, Result, internal_err, not_impl_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; + +/// Physical execution plan node for a hash join whose build side can spill +/// sub-partitions to disk rather than requiring the whole build side to fit +/// in memory at once. +/// +/// This is currently a scaffold: `execute` is not yet implemented (see +/// `SpillingHashJoinExec::execute` below). +#[derive(Debug)] +pub struct SpillingHashJoinExec { + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + partition_mode: PartitionMode, + num_sub_partitions: usize, + schema: SchemaRef, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl SpillingHashJoinExec { + /// Creates a new `SpillingHashJoinExec` joining `left` and `right` on the + /// equijoin pairs in `on`. + /// + /// This is a v1 scaffold: the join is always `JoinType::Inner`, with no + /// residual filter and no output projection. A throwaway `HashJoinExec` + /// is constructed purely to borrow its output `schema()` and + /// `properties()` (partitioning/equivalence info) — it is never + /// executed. + pub fn try_new( + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + partition_mode: PartitionMode, + num_sub_partitions: usize, + ) -> Result { + let hj = HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + None, + &JoinType::Inner, + None, + partition_mode, + NullEquality::NullEqualsNothing, + false, + )?; + let schema = hj.schema(); + let properties = hj.properties().clone(); + + Ok(Self { + left, + right, + on, + partition_mode, + num_sub_partitions, + schema, + properties, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + /// Returns the equijoin column pairs `(left_expr, right_expr)`. + pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] { + &self.on + } + + /// Returns the configured partition mode. + pub fn partition_mode(&self) -> PartitionMode { + self.partition_mode + } + + /// Returns the number of sub-partitions the build side is split into + /// for spilling. + pub fn num_sub_partitions(&self) -> usize { + self.num_sub_partitions + } + + /// Returns the left (build) child. + pub fn left(&self) -> &Arc { + &self.left + } + + /// Returns the right (probe) child. + pub fn right(&self) -> &Arc { + &self.right + } +} + +impl ExecutionPlan for SpillingHashJoinExec { + fn name(&self) -> &str { + "SpillingHashJoinExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 2 { + return internal_err!( + "SpillingHashJoinExec expected two children, got {}", + children.len() + ); + } + let right = children.pop().unwrap(); + let left = children.pop().unwrap(); + Ok(Arc::new(Self::try_new( + left, + right, + self.on.clone(), + self.partition_mode, + self.num_sub_partitions, + )?)) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + not_impl_err!("SpillingHashJoinExec::execute") + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } +} + +impl DisplayAs for SpillingHashJoinExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let on = self + .on + .iter() + .map(|(l, r)| format!("({l}, {r})")) + .collect::>() + .join(", "); + write!( + f, + "SpillingHashJoinExec: on=[{on}], mode={:?}", + self.partition_mode + ) + } + DisplayFormatType::TreeRender => { + writeln!(f, "mode={:?}", self.partition_mode) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::expressions::Column; + + /// Returns `(left, right, on)`: two single-partition-schema `DataSourceExec` + /// plans with schema `(k Int32, v Int32)`, each split into 4 partitions, + /// joined on `on = [(k_left, k_right)]`. + #[allow(clippy::type_complexity)] + fn two_col_inputs() -> ( + Arc, + Arc, + Vec<(PhysicalExprRef, PhysicalExprRef)>, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let partitions: Vec> = (0..4).map(|_| vec![]).collect(); + + let left_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("left MemorySourceConfig"); + let right_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("right MemorySourceConfig"); + + let left: Arc = + Arc::new(DataSourceExec::new(Arc::new(left_source))); + let right: Arc = + Arc::new(DataSourceExec::new(Arc::new(right_source))); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = + vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k", 0)))]; + + (left, right, on) + } + + #[test] + fn scaffold_schema_and_name() { + let (left, right, on) = two_col_inputs(); + let exec = SpillingHashJoinExec::try_new( + left, + right, + on, + PartitionMode::Partitioned, + 16, + ) + .unwrap(); + assert_eq!(exec.name(), "SpillingHashJoinExec"); + // Output schema is left ++ right columns. + assert_eq!(exec.schema().fields().len(), 4); + let rendered = + format!("{}", displayable(&exec as &dyn ExecutionPlan).indent(false)); + assert!(rendered.contains("SpillingHashJoinExec"), "{rendered}"); + } +} diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs new file mode 100644 index 0000000000..03575faa3f --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod exec; + +pub use exec::SpillingHashJoinExec; From 1a1d13e809c23ac1f05ce5c885d287a212bb4075 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 10:34:18 -0600 Subject: [PATCH 03/15] feat(core): add hash row partitioner for spilling hash join --- .../execution_plans/spilling_hash_join/mod.rs | 2 + .../spilling_hash_join/partitioner.rs | 205 ++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs index 03575faa3f..09eb3434bc 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -16,5 +16,7 @@ // under the License. mod exec; +mod partitioner; pub use exec::SpillingHashJoinExec; +pub use partitioner::{PartitionedBatch, RowPartitioner}; diff --git a/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs b/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs new file mode 100644 index 0000000000..1cd3e37aa0 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs @@ -0,0 +1,205 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// RowPartitioner hashes the join key columns of a `RecordBatch` and splits +// its rows into a fixed number of sub-partitions ("buckets"). It is used for +// both the build side and the probe side of `SpillingHashJoinExec` so that +// rows with equal join keys always land in the same sub-partition, whether +// they are processed in memory or after being spilled and re-read. +// +// The random state used for hashing is fixed and distinct from DataFusion's +// default so that bucket assignment here is independent of any hash used +// upstream (e.g. shuffle partitioning). + +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, RecordBatch, UInt32Array}; +use datafusion::arrow::compute::take; +use datafusion::common::Result; +use datafusion::common::hash_utils::{RandomState, create_hashes}; +use datafusion::physical_expr::PhysicalExprRef; + +/// A fixed, non-default seed for [`RandomState`] so that `RowPartitioner` +/// bucketing is independent of the hash used for shuffle partitioning +/// upstream. `RandomState::with_seed(0)` is defined to match +/// `RandomState::default()`, so any non-zero constant here is sufficient to +/// diverge from it; the exact value carries no other meaning. +const ROW_PARTITIONER_SEED: u64 = 0x5350_4a5f_484a_3121; + +/// Splits the rows of a `RecordBatch` into `num_sub` sub-partitions by +/// hashing a set of join key expressions, so that rows with equal keys +/// always land in the same sub-partition. +pub struct RowPartitioner { + keys: Vec, + num_sub: usize, + random_state: RandomState, +} + +/// One non-empty bucket produced by [`RowPartitioner::partition`]: the +/// bucket index, the sub-batch of rows routed to it, and the per-row hash +/// values for those rows (in the same order as the sub-batch's rows) so +/// callers can reuse them without re-hashing. +pub struct PartitionedBatch { + /// Index of the sub-partition this batch was routed to, in + /// `0..num_sub`. + pub bucket: usize, + /// The subset of rows from the input batch routed to `bucket`. + pub batch: RecordBatch, + /// The hash value of each row in `batch`, in the same row order. + pub hashes: Vec, +} + +impl RowPartitioner { + /// Creates a new `RowPartitioner` that hashes `keys` and routes rows + /// into `num_sub` sub-partitions using a fixed, non-default random + /// state. + pub fn new(keys: Vec, num_sub: usize) -> Self { + Self { + keys, + num_sub, + random_state: RandomState::with_seed(ROW_PARTITIONER_SEED), + } + } + + /// Hashes the key columns of `batch` and splits its rows into + /// sub-partitions. Bucket assignment is `(hash >> 32) % num_sub`. Only + /// non-empty buckets are returned, each carrying the hash values for the + /// rows it contains. + pub fn partition(&self, batch: &RecordBatch) -> Result> { + let num_rows = batch.num_rows(); + + let key_arrays: Vec = self + .keys + .iter() + .map(|expr| expr.evaluate(batch)?.into_array(num_rows)) + .collect::>()?; + + let mut hashes = vec![0u64; num_rows]; + create_hashes(&key_arrays, &self.random_state, &mut hashes)?; + + let mut indices: Vec> = vec![Vec::new(); self.num_sub]; + let mut bucket_hashes: Vec> = vec![Vec::new(); self.num_sub]; + for (row, &hash) in hashes.iter().enumerate() { + let bucket = ((hash >> 32) as usize) % self.num_sub; + indices[bucket].push(row as u32); + bucket_hashes[bucket].push(hash); + } + + let schema = batch.schema(); + let mut out = Vec::new(); + for (bucket, row_indices) in indices.into_iter().enumerate() { + if row_indices.is_empty() { + continue; + } + let idx = UInt32Array::from(row_indices); + let taken_cols = batch + .columns() + .iter() + .map(|col| take(col, &idx, None)) + .collect::, _>>()?; + let sub_batch = RecordBatch::try_new(Arc::clone(&schema), taken_cols)?; + out.push(PartitionedBatch { + bucket, + batch: sub_batch, + hashes: std::mem::take(&mut bucket_hashes[bucket]), + }); + } + + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + + fn batch_with_keys(keys: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(keys))]).unwrap() + } + + fn partitioner(num_sub: usize) -> RowPartitioner { + let keys: Vec = vec![Arc::new(Column::new("k", 0))]; + RowPartitioner::new(keys, num_sub) + } + + #[test] + fn partitions_cover_all_rows_and_are_stable() { + let num_sub = 4; + let p = partitioner(num_sub); + + let batch_a = batch_with_keys(vec![1, 2, 3, 1, 2]); + let batch_b = batch_with_keys(vec![2, 1]); + + let out_a = p.partition(&batch_a).unwrap(); + let out_b = p.partition(&batch_b).unwrap(); + + // All rows accounted for, per batch. + let rows_a: usize = out_a.iter().map(|pb| pb.batch.num_rows()).sum(); + let rows_b: usize = out_b.iter().map(|pb| pb.batch.num_rows()).sum(); + assert_eq!(rows_a, batch_a.num_rows()); + assert_eq!(rows_b, batch_b.num_rows()); + + // No empty buckets returned, and every bucket index is in range. + for pb in out_a.iter().chain(out_b.iter()) { + assert!(pb.bucket < num_sub); + assert!(pb.batch.num_rows() > 0); + assert_eq!(pb.batch.num_rows(), pb.hashes.len()); + } + + // Stability: key value 1 must land in the same bucket in both batches. + let bucket_for_key = |out: &[PartitionedBatch], key: i32| -> usize { + for pb in out { + let col = pb + .batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + if col.iter().any(|v| v == Some(key)) { + return pb.bucket; + } + } + panic!("key {key} not found in any bucket"); + }; + + assert_eq!(bucket_for_key(&out_a, 1), bucket_for_key(&out_b, 1)); + assert_eq!(bucket_for_key(&out_a, 2), bucket_for_key(&out_b, 2)); + } + + #[test] + fn stable_across_recreated_partitioner() { + // Later tasks recompute hashes on spilled data with a brand-new + // `RowPartitioner`, so the fixed seed must make bucket assignment + // (and the hash values themselves) reproducible across instances, + // not just across calls on the same instance. + let num_sub = 4; + let batch = batch_with_keys(vec![1, 2, 3, 1, 2]); + + let out_1 = partitioner(num_sub).partition(&batch).unwrap(); + let out_2 = partitioner(num_sub).partition(&batch).unwrap(); + + assert_eq!(out_1.len(), out_2.len()); + for (a, b) in out_1.iter().zip(out_2.iter()) { + assert_eq!(a.bucket, b.bucket); + assert_eq!(a.hashes, b.hashes); + } + } +} From 4f6742bb68e59c8d0c7110f3be13560eec1a857b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 10:44:24 -0600 Subject: [PATCH 04/15] feat(core): add in-memory probe table for spilling hash join Add ProbeTable, the correctness core of the spilling hash join: builds an inner-join hash table over one resident build-side bucket and probes it against a probe-side batch, returning matched (build_row, probe_row) index pairs. Also adds assemble_output to materialize the joined RecordBatch from those index pairs. Key equality is resolved via Arrow's row encoding (RowConverter/Rows) rather than a per-DataType match, so Int32/Int64/Decimal128/Date32/Utf8 (and other row-encodable types) all work without dedicated dispatch code. A shared per-row null-key mask enforces NullEquality::NullEqualsNothing semantics, since Arrow's row encoding otherwise treats two NULLs as equal. --- Cargo.lock | 1 + ballista/core/Cargo.toml | 1 + .../spilling_hash_join/hash_table.rs | 400 ++++++++++++++++++ .../execution_plans/spilling_hash_join/mod.rs | 2 + 4 files changed, 404 insertions(+) create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/hash_table.rs diff --git a/Cargo.lock b/Cargo.lock index d9b89761b0..6ae374bf75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1105,6 +1105,7 @@ dependencies = [ "rand 0.10.2", "rustc_version", "serde", + "smallvec", "tempfile", "tokio", "tokio-stream", diff --git a/ballista/core/Cargo.toml b/ballista/core/Cargo.toml index 7359183906..bbf1ba279d 100644 --- a/ballista/core/Cargo.toml +++ b/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 } diff --git a/ballista/core/src/execution_plans/spilling_hash_join/hash_table.rs b/ballista/core/src/execution_plans/spilling_hash_join/hash_table.rs new file mode 100644 index 0000000000..6c2ba5837e --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/hash_table.rs @@ -0,0 +1,400 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// ProbeTable is the in-memory correctness core of the spilling hash join: an +// inner-join hash table built over one resident build-side bucket, probed +// with a probe-side batch. It is deliberately agnostic to spilling — the +// stream (a later task) is responsible for choosing which build bucket is +// resident and feeding it, along with matching probe batches, to this table. +// +// Equality between build and probe keys is resolved via Arrow's row format +// (`arrow::row::RowConverter`) rather than a per-`DataType` match, so any key +// type combination Arrow can encode into row format works without dedicated +// dispatch code here. The one correctness wrinkle this creates: Arrow's row +// encoding treats two nulls as equal, which is wrong for +// `NullEquality::NullEqualsNothing` (the only null semantics this join +// supports; see `SpillingHashJoinExec`). This table therefore tracks, per +// side, which rows have a null in any key column, and excludes those rows +// from both hash-table insertion (build side) and matching (probe side). + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, ArrayRef, RecordBatch, UInt32Array}; +use datafusion::arrow::compute::take; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::row::{RowConverter, Rows, SortField}; +use datafusion::common::Result; +use datafusion::physical_expr::PhysicalExprRef; +use smallvec::SmallVec; + +/// An in-memory hash table over one resident build-side bucket, used to +/// probe matching rows from a probe-side batch for an inner join with +/// `NullEquality::NullEqualsNothing` semantics. +pub struct ProbeTable { + /// The build-side bucket this table was built from, kept so + /// `assemble_output` can later `take` its columns by matched row index. + build: RecordBatch, + /// Row-encoded build key columns, indexed the same as `build`. Shares a + /// `RowConverter` with the probe side (see `probe`) so encodings are + /// directly comparable. + build_rows: Rows, + /// The row converter used to encode `build_rows`, reused to encode + /// probe keys so both sides use identical row-format parameters. + converter: RowConverter, + /// `hash -> build row indices`, populated only for build rows with no + /// null in any key column. + map: HashMap>, +} + +impl ProbeTable { + /// Builds a table from one resident build bucket (already concatenated) + /// and the per-row hash values computed for `keys` on that batch (e.g. + /// by `RowPartitioner`). + pub fn build( + build: RecordBatch, + build_hashes: &[u64], + keys: &[PhysicalExprRef], + ) -> Result { + let key_arrays = evaluate_keys(keys, &build)?; + let has_null_key = has_null_key_mask(&key_arrays); + + let converter = RowConverter::new(sort_fields(&key_arrays))?; + let build_rows = converter.convert_columns(&key_arrays)?; + + let mut map: HashMap> = HashMap::new(); + for (i, &is_null_key) in has_null_key.iter().enumerate() { + // A null key never matches anything under NullEquality::NullEqualsNothing, + // so it must never become a matchable build-side entry. + if is_null_key { + continue; + } + map.entry(build_hashes[i]).or_default().push(i as u32); + } + + Ok(Self { + build, + build_rows, + converter, + map, + }) + } + + /// Probes one probe-side batch against this table, returning the + /// matched `(build_row, probe_row)` index pairs as two parallel + /// `UInt32Array`s. Order within a matched key group is unspecified. + pub fn probe( + &self, + probe: &RecordBatch, + probe_hashes: &[u64], + keys: &[PhysicalExprRef], + ) -> Result<(UInt32Array, UInt32Array)> { + let probe_key_arrays = evaluate_keys(keys, probe)?; + let has_null_key = has_null_key_mask(&probe_key_arrays); + + // Reuse the build side's converter so both sides use identical row + // encoding parameters and their `Rows` are directly comparable. + let probe_rows = self.converter.convert_columns(&probe_key_arrays)?; + + let mut build_out: Vec = Vec::new(); + let mut probe_out: Vec = Vec::new(); + + for (j, &is_null_key) in has_null_key.iter().enumerate() { + // A null key never matches anything under NullEquality::NullEqualsNothing. + if is_null_key { + continue; + } + let Some(candidates) = self.map.get(&probe_hashes[j]) else { + continue; + }; + let probe_row = probe_rows.row(j); + for &i in candidates { + // The hash lookup above is only a pre-filter: distinct keys can + // share a hash value, so equality must still be resolved on the + // actual encoded row values before counting a match. + if self.build_rows.row(i as usize) == probe_row { + build_out.push(i); + probe_out.push(j as u32); + } + } + } + + Ok((UInt32Array::from(build_out), UInt32Array::from(probe_out))) + } + + /// Returns the resident build-side batch this table was built from, so + /// callers (e.g. the join stream) can pass it to `assemble_output` + /// without keeping a separate copy. + pub fn build_batch(&self) -> &RecordBatch { + &self.build + } +} + +/// Evaluates `keys` against `batch`, returning one array per key expression. +fn evaluate_keys(keys: &[PhysicalExprRef], batch: &RecordBatch) -> Result> { + let num_rows = batch.num_rows(); + keys.iter() + .map(|expr| expr.evaluate(batch)?.into_array(num_rows)) + .collect() +} + +/// Returns one `SortField` per key array, describing its data type to +/// `RowConverter`. Sort options are irrelevant here: rows are only ever +/// compared for equality, never ordered. +fn sort_fields(key_arrays: &[ArrayRef]) -> Vec { + key_arrays + .iter() + .map(|a| SortField::new(a.data_type().clone())) + .collect() +} + +/// Returns a per-row mask that is `true` iff the row is null in at least one +/// of `key_arrays`. Used to exclude such rows from matching under +/// `NullEquality::NullEqualsNothing`, since Arrow's row encoding otherwise +/// treats two nulls as equal. +fn has_null_key_mask(key_arrays: &[ArrayRef]) -> Vec { + let num_rows = key_arrays.first().map(|a| a.len()).unwrap_or(0); + let mut mask = vec![false; num_rows]; + for arr in key_arrays { + if arr.null_count() == 0 { + continue; + } + for (i, is_null_key) in mask.iter_mut().enumerate() { + *is_null_key = *is_null_key || arr.is_null(i); + } + } + mask +} + +/// Assembles a joined output batch from matched `(build_row, probe_row)` +/// index pairs: `take`s `build`'s columns by `build_rows`, `take`s `probe`'s +/// columns by `probe_rows`, and concatenates them left-then-right into +/// `schema`. +pub fn assemble_output( + schema: &SchemaRef, + build: &RecordBatch, + probe: &RecordBatch, + build_rows: &UInt32Array, + probe_rows: &UInt32Array, +) -> Result { + let bcols = build + .columns() + .iter() + .map(|c| take(c, build_rows, None)) + .collect::, _>>()?; + let pcols = probe + .columns() + .iter() + .map(|c| take(c, probe_rows, None)) + .collect::, _>>()?; + + Ok(RecordBatch::try_new( + Arc::clone(schema), + [bcols, pcols].concat(), + )?) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::hash_utils::{RandomState, create_hashes}; + use datafusion::physical_expr::expressions::Column; + + /// Same fixed seed `RowPartitioner` uses, so tests hash build/probe + /// batches the way the real pipeline would. The exact value doesn't + /// matter for these tests (only that build/probe use the same one) but + /// reusing it documents the intended caller contract. + const SEED: u64 = 0x5350_4a5f_484a_3121; + + fn batch_i32(name: &str, values: Vec>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, true)])); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + fn batch_two_i32( + names: (&str, &str), + a: Vec>, + b: Vec>, + ) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new(names.0, DataType::Int32, true), + Field::new(names.1, DataType::Int32, true), + ])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(a)), Arc::new(Int32Array::from(b))], + ) + .unwrap() + } + + fn batch_utf8(name: &str, values: Vec>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Utf8, true)])); + RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(values))]).unwrap() + } + + fn hashes(batch: &RecordBatch, keys: &[PhysicalExprRef]) -> Vec { + let arrays = evaluate_keys(keys, batch).unwrap(); + let mut out = vec![0u64; batch.num_rows()]; + create_hashes(&arrays, &RandomState::with_seed(SEED), &mut out).unwrap(); + out + } + + fn sorted_pairs( + build_rows: &UInt32Array, + probe_rows: &UInt32Array, + ) -> Vec<(u32, u32)> { + let mut pairs: Vec<(u32, u32)> = build_rows + .values() + .iter() + .copied() + .zip(probe_rows.values().iter().copied()) + .collect(); + pairs.sort_unstable(); + pairs + } + + fn single_key(name: &str) -> Vec { + vec![Arc::new(Column::new(name, 0))] + } + + #[test] + fn inner_join_single_key_matches() { + let keys = single_key("k"); + // build k=[1,2,2,3], probe k=[2,3,4] + let build = batch_i32("k", vec![Some(1), Some(2), Some(2), Some(3)]); + let probe = batch_i32("k", vec![Some(2), Some(3), Some(4)]); + + let build_hashes = hashes(&build, &keys); + let probe_hashes = hashes(&probe, &keys); + + let table = ProbeTable::build(build, &build_hashes, &keys).unwrap(); + let (build_rows, probe_rows) = table.probe(&probe, &probe_hashes, &keys).unwrap(); + + // expect pairs: probe 2 -> build rows {1,2}; probe 3 -> build row {3}; probe 4 -> none + assert_eq!( + sorted_pairs(&build_rows, &probe_rows), + vec![(1, 0), (2, 0), (3, 1)] + ); + } + + #[test] + fn multi_key_and_nulls_never_match() { + let keys = vec![ + Arc::new(Column::new("a", 0)) as PhysicalExprRef, + Arc::new(Column::new("b", 1)) as PhysicalExprRef, + ]; + + // Build row 0: (1,10) - a valid, matchable row. + // Build row 1: (NULL,20) - null in `a`. + // Build row 2: (1,NULL) - null in `b`. + let build = batch_two_i32( + ("a", "b"), + vec![Some(1), None, Some(1)], + vec![Some(10), Some(20), None], + ); + // Probe row 0: (1,10) - matches build row 0. + // Probe row 1: (NULL,20) - null in `a`; must NOT match build row 1 + // even though both are (NULL,20). + // Probe row 2: (1,NULL) - null in `b`; must NOT match build row 2 + // even though both are (1,NULL). + let probe = batch_two_i32( + ("a", "b"), + vec![Some(1), None, Some(1)], + vec![Some(10), Some(20), None], + ); + + let build_hashes = hashes(&build, &keys); + let probe_hashes = hashes(&probe, &keys); + + let table = ProbeTable::build(build, &build_hashes, &keys).unwrap(); + let (build_rows, probe_rows) = table.probe(&probe, &probe_hashes, &keys).unwrap(); + + assert_eq!(sorted_pairs(&build_rows, &probe_rows), vec![(0, 0)]); + } + + #[test] + fn hash_collision_resolved_by_key_equality() { + // Two build rows with distinct keys but forced to share a hash + // value; the probe key equal to only one of them must match just + // that row, proving equality is resolved on actual key values, not + // just the (colliding) hash. + let keys = single_key("k"); + let build = batch_i32("k", vec![Some(1), Some(2)]); + let probe = batch_i32("k", vec![Some(2)]); + + let build_hashes = vec![42u64, 42u64]; + let probe_hashes = vec![42u64]; + + let table = ProbeTable::build(build, &build_hashes, &keys).unwrap(); + let (build_rows, probe_rows) = table.probe(&probe, &probe_hashes, &keys).unwrap(); + + assert_eq!(sorted_pairs(&build_rows, &probe_rows), vec![(1, 0)]); + } + + #[test] + fn string_key_matches() { + let keys = single_key("k"); + let build = batch_utf8("k", vec![Some("a"), Some("b"), Some("b")]); + let probe = batch_utf8("k", vec![Some("b"), Some("c")]); + + let build_hashes = hashes(&build, &keys); + let probe_hashes = hashes(&probe, &keys); + + let table = ProbeTable::build(build, &build_hashes, &keys).unwrap(); + let (build_rows, probe_rows) = table.probe(&probe, &probe_hashes, &keys).unwrap(); + + assert_eq!(sorted_pairs(&build_rows, &probe_rows), vec![(1, 0), (2, 0)]); + } + + #[test] + fn assemble_output_concatenates_left_then_right() { + let build_schema = + Arc::new(Schema::new(vec![Field::new("bk", DataType::Int32, false)])); + let build = RecordBatch::try_new( + Arc::clone(&build_schema), + vec![Arc::new(Int32Array::from(vec![10, 20, 30]))], + ) + .unwrap(); + + let probe_schema = + Arc::new(Schema::new(vec![Field::new("pk", DataType::Int32, false)])); + let probe = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + let out_schema = Arc::new(Schema::new(vec![ + Field::new("bk", DataType::Int32, false), + Field::new("pk", DataType::Int32, false), + ])); + + let build_rows = UInt32Array::from(vec![2, 0]); + let probe_rows = UInt32Array::from(vec![1, 2]); + + let out = assemble_output(&out_schema, &build, &probe, &build_rows, &probe_rows) + .unwrap(); + + assert_eq!(out.num_rows(), 2); + let bk = out.column(0).as_any().downcast_ref::().unwrap(); + let pk = out.column(1).as_any().downcast_ref::().unwrap(); + assert_eq!(bk.values(), &[30, 10]); + assert_eq!(pk.values(), &[2, 3]); + } +} diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs index 09eb3434bc..c7a170dd88 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -16,7 +16,9 @@ // under the License. mod exec; +mod hash_table; mod partitioner; pub use exec::SpillingHashJoinExec; +pub use hash_table::{ProbeTable, assemble_output}; pub use partitioner::{PartitionedBatch, RowPartitioner}; From 19e833481f2a74ac113553ebec2499efc9bc8b11 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 10:56:13 -0600 Subject: [PATCH 05/15] feat(core): execute SpillingHashJoinExec in memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire SpillingHashJoinExec::execute() to a real join stream: drain the build (left) side fully into one resident ProbeTable per sub-partition bucket, then probe the right side against those tables as it streams in. All build buckets stay resident for now; spilling any of them to disk under memory pressure is left to a later change. The stream is built as a futures::stream::once future (drain + build) that resolves to the probe stream and is flattened with try_flatten, the same idiom SortShuffleWriterExec::execute already uses in this crate for "do async setup, then hand back a stream" — this keeps the build phase as plain async/await instead of a hand-rolled poll_next state machine. Adds an oracle test comparing SpillingHashJoinExec against DataFusion's own HashJoinExec (Inner, Partitioned) over the same hash-repartitioned inputs, so both joins see keys co-located the same way and the comparison actually exercises the Partitioned-mode contract instead of passing vacuously. --- .../spilling_hash_join/exec.rs | 229 +++++++++++++++++- .../execution_plans/spilling_hash_join/mod.rs | 1 + .../spilling_hash_join/stream.rs | 222 +++++++++++++++++ 3 files changed, 439 insertions(+), 13 deletions(-) create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/stream.rs diff --git a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs index 1853f82f46..37476bc559 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -16,10 +16,10 @@ // under the License. // SpillingHashJoinExec is a physical execution plan node for a hash join -// whose build side spills to disk instead of requiring the entire hash -// table to fit in memory. This is currently a scaffold only: the struct, -// its constructor, and its accessors exist so later work can wire up -// planning/serde/substitution, but `execute` is not yet implemented. +// whose build side is designed to spill to disk instead of requiring the +// entire hash table to fit in memory. `execute` currently keeps all build +// sub-partitions ("buckets") resident in memory — spilling any of them to +// disk under memory pressure is not yet implemented (see `stream.rs`). // // v1 invariants (enforced by the constructor, which always builds an // inner join with no residual filter and no output projection): @@ -30,7 +30,7 @@ use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{JoinType, NullEquality, Result, internal_err, not_impl_err}; +use datafusion::common::{JoinType, NullEquality, Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::PhysicalExprRef; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; @@ -40,12 +40,14 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; -/// Physical execution plan node for a hash join whose build side can spill -/// sub-partitions to disk rather than requiring the whole build side to fit -/// in memory at once. +use super::stream::execute_join; + +/// Physical execution plan node for a hash join whose build side is designed +/// to spill sub-partitions to disk rather than requiring the whole build +/// side to fit in memory at once. /// -/// This is currently a scaffold: `execute` is not yet implemented (see -/// `SpillingHashJoinExec::execute` below). +/// `execute` currently keeps every build-side sub-partition resident (no +/// spilling yet — see `stream.rs`). #[derive(Debug)] pub struct SpillingHashJoinExec { left: Arc, @@ -165,12 +167,32 @@ impl ExecutionPlan for SpillingHashJoinExec { )?)) } + /// Executes partition `partition` of both children (the Ballista shuffle + /// contract guarantees rows with equal join keys are co-located in the + /// same partition number on both sides — see `PartitionMode::Partitioned` + /// semantics) and returns a stream that drains the left (build) side + /// fully into resident per-sub-partition hash tables before probing the + /// right (probe) side against them. No spilling yet: all build buckets + /// stay in memory for the lifetime of the stream. fn execute( &self, - _partition: usize, - _context: Arc, + partition: usize, + context: Arc, ) -> Result { - not_impl_err!("SpillingHashJoinExec::execute") + let left = self.left.execute(partition, Arc::clone(&context))?; + let right = self.right.execute(partition, context)?; + + let (left_keys, right_keys): (Vec, Vec) = + self.on.iter().cloned().unzip(); + + Ok(execute_join( + self.schema(), + left, + right, + left_keys, + right_keys, + self.num_sub_partitions, + )) } fn metrics(&self) -> Option { @@ -208,6 +230,7 @@ impl DisplayAs for SpillingHashJoinExec { #[cfg(test)] mod tests { use super::*; + use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -265,4 +288,184 @@ mod tests { format!("{}", displayable(&exec as &dyn ExecutionPlan).indent(false)); assert!(rendered.contains("SpillingHashJoinExec"), "{rendered}"); } + + // --- oracle test: matches DataFusion's own Partitioned HashJoinExec --- + // + // `PartitionMode::Partitioned` carries a contract that is easy to break + // in a test without noticing: it assumes rows with equal join keys are + // already co-located in the same partition number on both sides (the + // Ballista shuffle contract). `HashJoinExec` (DataFusion's own + // Partitioned join) makes the exact same assumption — feed either join + // un-hash-partitioned input and both produce silently wrong results in + // the same way, so a naive comparison could pass for the wrong reason. + // + // To keep the comparison meaningful, both sides below are driven from + // data that has actually been hash-partitioned on the join key via + // `RepartitionExec`, so equal keys really are co-located per partition + // for both plans under test. + use datafusion::arrow::array::Int32Array; + use datafusion::physical_plan::Partitioning; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::repartition::RepartitionExec; + + /// Builds a single-input-partition `DataSourceExec` with schema + /// `(key_name Int32, val_name Int32)`, split into several batches of + /// `batch_size` rows (to exercise multi-batch accumulation on both the + /// build and probe sides), `num_rows` rows total. Row `i` gets key + /// `i % key_modulus` (so keys repeat, producing multi-row matches) and + /// value `i as i32 + val_offset`. + fn make_source( + key_name: &str, + val_name: &str, + num_rows: usize, + key_modulus: i32, + val_offset: i32, + batch_size: usize, + ) -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new(key_name, DataType::Int32, false), + Field::new(val_name, DataType::Int32, false), + ])); + + let mut batches = Vec::new(); + let mut start = 0usize; + while start < num_rows { + let end = (start + batch_size).min(num_rows); + let keys: Vec = (start..end).map(|r| (r as i32) % key_modulus).collect(); + let vals: Vec = (start..end).map(|r| r as i32 + val_offset).collect(); + batches.push( + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(vals)), + ], + ) + .unwrap(), + ); + start = end; + } + + let partitions = vec![batches]; + let source = MemorySourceConfig::try_new(&partitions, schema, None) + .expect("MemorySourceConfig"); + Arc::new(DataSourceExec::new(Arc::new(source))) + } + + /// Hash-repartitions `base` (a single-partition source) into `n` + /// partitions on the column at `key_index`, using a fresh + /// `RepartitionExec`. + /// + /// Deliberately builds a *new* `RepartitionExec` per call rather than + /// sharing one between the oracle and the exec under test: each output + /// partition of a `RepartitionExec` is single-consumer (its receiver is + /// taken on first `execute`), so one shared instance could not feed both + /// joins' full partition sweep. Hash bucketing is a pure, deterministic + /// function of the input batches, expressions, and partition count, so + /// two independent `RepartitionExec`s built over the same `base` data + /// reproduce byte-identical per-partition bucketing — which is all the + /// "same repartitioned input" fairness requirement actually needs. + fn hash_repartition( + base: Arc, + key_index: usize, + n: usize, + ) -> Arc { + let field = base.schema().field(key_index).clone(); + let expr: PhysicalExprRef = Arc::new(Column::new(field.name(), key_index)); + Arc::new( + RepartitionExec::try_new(base, Partitioning::Hash(vec![expr], n)) + .expect("RepartitionExec::try_new"), + ) + } + + /// Runs DataFusion's own `HashJoinExec` (Inner, `PartitionMode::Partitioned`) + /// over `left`/`right` and collects every output row across all output + /// partitions — the oracle `SpillingHashJoinExec` must match. + async fn oracle_join( + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + ctx: Arc, + ) -> Vec { + let hj = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + collect(Arc::new(hj), ctx).await.unwrap() + } + + /// Flattens output batches with schema `(lk, lv, rk, rv)` into a sorted + /// vector of row tuples, so two result sets can be compared independent + /// of batch boundaries or row order (both are unspecified across + /// partitions and within matched key groups). + fn sorted_rows(batches: &[RecordBatch]) -> Vec<(i32, i32, i32, i32)> { + let mut out = Vec::new(); + for batch in batches { + let col = |i: usize| { + batch + .column(i) + .as_any() + .downcast_ref::() + .unwrap() + }; + let (lk, lv, rk, rv) = (col(0), col(1), col(2), col(3)); + for i in 0..batch.num_rows() { + out.push((lk.value(i), lv.value(i), rk.value(i), rv.value(i))); + } + } + out.sort_unstable(); + out + } + + #[tokio::test] + async fn matches_datafusion_inner_partitioned_in_memory() { + let ctx = Arc::new(TaskContext::default()); // default = generous memory, nothing spills + let num_partitions = 4; + + // A few thousand rows per side; overlapping, repeating integer keys + // so both single- and multi-row matches occur on both sides. + let base_left = make_source("lk", "lv", 3_000, 500, 0, 777); + let base_right = make_source("rk", "rv", 3_500, 450, 1_000_000, 900); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![( + Arc::new(Column::new("lk", 0)), + Arc::new(Column::new("rk", 0)), + )]; + + // Independent (but content-identical, see `hash_repartition`) + // hash-partitioned inputs for the oracle vs. the exec under test. + let oracle_left = hash_repartition(Arc::clone(&base_left), 0, num_partitions); + let oracle_right = hash_repartition(Arc::clone(&base_right), 0, num_partitions); + let ours_left = hash_repartition(base_left, 0, num_partitions); + let ours_right = hash_repartition(base_right, 0, num_partitions); + + let expected = + oracle_join(oracle_left, oracle_right, on.clone(), Arc::clone(&ctx)).await; + + let exec = SpillingHashJoinExec::try_new( + ours_left, + ours_right, + on, + PartitionMode::Partitioned, + 16, + ) + .unwrap(); + let actual = collect(Arc::new(exec), ctx).await.unwrap(); + + let expected_rows = sorted_rows(&expected); + let actual_rows = sorted_rows(&actual); + assert!( + !expected_rows.is_empty(), + "test data should produce matching rows" + ); + assert_eq!(actual_rows, expected_rows); + } } diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs index c7a170dd88..759117b9ca 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -18,6 +18,7 @@ mod exec; mod hash_table; mod partitioner; +mod stream; pub use exec::SpillingHashJoinExec; pub use hash_table::{ProbeTable, assemble_output}; diff --git a/ballista/core/src/execution_plans/spilling_hash_join/stream.rs b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs new file mode 100644 index 0000000000..51fe34439c --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs @@ -0,0 +1,222 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// The in-memory (no-spill) execution stream for `SpillingHashJoinExec`. All +// build-side sub-partitions ("buckets") are kept resident for the lifetime of +// the stream — a later task adds spilling any of them to disk under memory +// pressure. Until then this is a plain, if internally sub-partitioned, hash +// join: build fully, then probe. +// +// The stream is built as a `futures::stream::once` future (drain the build +// side and construct one `ProbeTable` per bucket) that resolves to the probe +// stream, flattened via `TryStreamExt::try_flatten`. This mirrors the idiom +// `SortShuffleWriterExec::execute` already uses in this crate +// (`sort_shuffle/writer.rs`) for "do async setup, then hand back a stream" — +// it keeps the build phase as ordinary `async`/`await` code instead of a +// hand-rolled `Stream::poll_next` state machine, while still producing a +// single `SendableRecordBatchStream` that does no work until polled and +// never emits probe output before the build side is fully drained. + +use std::collections::VecDeque; +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::Result; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{Stream, StreamExt, TryStreamExt, stream}; + +use super::hash_table::{ProbeTable, assemble_output}; +use super::partitioner::RowPartitioner; + +/// Builds the `SendableRecordBatchStream` for `SpillingHashJoinExec::execute`: +/// drains `left` (the build side) fully into one resident `ProbeTable` per +/// sub-partition bucket, then probes `right` (the probe side) against those +/// tables, bucket-matched, as it arrives. +/// +/// `output_schema` is the join's output schema (left columns ++ right +/// columns), independent of `left`/`right`'s own schemas — it is what the +/// returned stream's `schema()` reports. +pub fn execute_join( + output_schema: SchemaRef, + left: SendableRecordBatchStream, + right: SendableRecordBatchStream, + left_keys: Vec, + right_keys: Vec, + num_sub_partitions: usize, +) -> SendableRecordBatchStream { + let adapter_schema = Arc::clone(&output_schema); + + // A single future that drains `left`, builds the resident hash tables, + // then resolves to the probe stream. `try_flatten` turns this + // `Stream>` of one element into the actual + // `Stream>` the adapter needs — nothing here + // polls `right` until the future above has completed, which is exactly + // the "fully build before any probe output" ordering the join requires. + let joined = stream::once(build_then_probe( + left, + right, + left_keys, + right_keys, + num_sub_partitions, + output_schema, + )) + .try_flatten(); + + Box::pin(RecordBatchStreamAdapter::new(adapter_schema, joined)) +} + +/// Drains `left` into resident hash tables, then returns a stream that +/// probes `right` against them. Split out of `execute_join` (rather than an +/// inline `async move` block) purely so its `Result>` return +/// type is written down once, instead of needing an explicit turbofish at +/// every call site to disambiguate the error type. +async fn build_then_probe( + left: SendableRecordBatchStream, + right: SendableRecordBatchStream, + left_keys: Vec, + right_keys: Vec, + num_sub_partitions: usize, + output_schema: SchemaRef, +) -> Result> + Send> { + let tables = build_tables(left, &left_keys, num_sub_partitions).await?; + Ok(probe_stream(right, right_keys, tables, output_schema)) +} + +/// Drains `left` fully, routing every row into one of `num_sub_partitions` +/// buckets by `left_keys` (via `RowPartitioner`), then builds one +/// `ProbeTable` per non-empty bucket. Buckets that never received a row are +/// `None`, so `probe_stream` can skip them (and later, spilling) without +/// treating "no rows" as an error. +async fn build_tables( + mut left: SendableRecordBatchStream, + left_keys: &[PhysicalExprRef], + num_sub_partitions: usize, +) -> Result>> { + let left_schema = left.schema(); + let partitioner = RowPartitioner::new(left_keys.to_vec(), num_sub_partitions); + + let mut bucket_batches: Vec> = vec![Vec::new(); num_sub_partitions]; + let mut bucket_hashes: Vec> = vec![Vec::new(); num_sub_partitions]; + + while let Some(batch) = left.next().await { + let batch = batch?; + for pb in partitioner.partition(&batch)? { + bucket_batches[pb.bucket].push(pb.batch); + bucket_hashes[pb.bucket].extend(pb.hashes); + } + } + + bucket_batches + .into_iter() + .zip(bucket_hashes) + .map(|(batches, hashes)| { + if batches.is_empty() { + return Ok(None); + } + let concatenated = concat_batches(&left_schema, batches.iter())?; + Ok(Some(ProbeTable::build(concatenated, &hashes, left_keys)?)) + }) + .collect() +} + +/// Mutable state threaded through `probe_stream`'s `futures::stream::unfold`. +struct ProbeState { + right: SendableRecordBatchStream, + partitioner: RowPartitioner, + right_keys: Vec, + tables: Vec>, + output_schema: SchemaRef, + /// Output batches assembled from the most recent `right` batch but not + /// yet yielded. A single `right` batch can fan out into multiple + /// sub-partition buckets, each producing its own output batch, so these + /// are queued and drained one at a time before pulling `right` again. + pending: VecDeque, +} + +/// Probes `right` against `tables` (one resident `ProbeTable` per bucket, +/// built from the fully-drained build side) as `right` batches arrive, +/// yielding one output batch per non-empty matched sub-partition. `right` +/// rows that hash to a bucket with no build-side table (`None`) have no +/// possible match and are dropped, matching inner-join semantics. +fn probe_stream( + right: SendableRecordBatchStream, + right_keys: Vec, + tables: Vec>, + output_schema: SchemaRef, +) -> impl Stream> + Send { + let num_sub_partitions = tables.len(); + let state = ProbeState { + right, + partitioner: RowPartitioner::new(right_keys.clone(), num_sub_partitions), + right_keys, + tables, + output_schema, + pending: VecDeque::new(), + }; + + stream::unfold(state, |mut state| async move { + loop { + if let Some(batch) = state.pending.pop_front() { + return Some((Ok(batch), state)); + } + + let next = state.right.next().await?; + let batch = match next { + Ok(batch) => batch, + Err(e) => return Some((Err(e), state)), + }; + + if let Err(e) = probe_batch(&batch, &mut state) { + return Some((Err(e), state)); + } + // Loop back around: this `right` batch may have produced no + // output (e.g. every bucket it touched was empty), in which case + // `pending` is still empty and we must pull the next `right` + // batch rather than returning `None` (which would end the + // stream early). + } + }) +} + +/// Partitions one `right` batch by `state.right_keys` and probes each +/// resulting sub-batch against its bucket's `ProbeTable` (if any), pushing +/// every non-empty assembled output batch onto `state.pending`. +fn probe_batch(batch: &RecordBatch, state: &mut ProbeState) -> Result<()> { + for pb in state.partitioner.partition(batch)? { + let Some(table) = &state.tables[pb.bucket] else { + continue; + }; + let (build_rows, probe_rows) = + table.probe(&pb.batch, &pb.hashes, &state.right_keys)?; + if build_rows.is_empty() { + continue; + } + let out = assemble_output( + &state.output_schema, + table.build_batch(), + &pb.batch, + &build_rows, + &probe_rows, + )?; + state.pending.push_back(out); + } + Ok(()) +} From d0c339fe43ef86030fd58a30e81e71c6dd85c331 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 11:12:32 -0600 Subject: [PATCH 06/15] feat(core): run spilling hash join in bounded memory Spill both sides of SpillingHashJoinExec to disk and drain spilled buckets one at a time so the join runs in memory bounded by the runtime MemoryPool. - Add per-bucket Arrow-IPC spill files (JoinSpillWriter/JoinSpillReader), with independent build-side and probe-side spill sets per task. - Build phase reserves resident bucket memory against the pool and, on a rejected try_grow, evicts the largest resident bucket to disk; once a bucket cannot fit it stays spilled. - Probe phase probes resident buckets and routes spilled buckets' probe rows to disk; the drain phase reads each spilled bucket back one at a time, recomputing hashes from the fixed partitioner seed, and drops each table before the next so peak drain memory is one bucket. - Publish spill_count/spilled_bytes metrics; gate with tiny-pool oracle tests asserting output identical to a generous-pool HashJoinExec. --- .../spilling_hash_join/exec.rs | 127 +++- .../execution_plans/spilling_hash_join/mod.rs | 1 + .../spilling_hash_join/spill.rs | 179 ++++++ .../spilling_hash_join/stream.rs | 567 ++++++++++++++---- 4 files changed, 733 insertions(+), 141 deletions(-) create mode 100644 ballista/core/src/execution_plans/spilling_hash_join/spill.rs diff --git a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs index 37476bc559..d11e747f81 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -16,10 +16,11 @@ // under the License. // SpillingHashJoinExec is a physical execution plan node for a hash join -// whose build side is designed to spill to disk instead of requiring the -// entire hash table to fit in memory. `execute` currently keeps all build -// sub-partitions ("buckets") resident in memory — spilling any of them to -// disk under memory pressure is not yet implemented (see `stream.rs`). +// whose build side spills to disk instead of requiring the entire hash table +// to fit in memory. `execute` runs in bounded memory: build sub-partitions +// ("buckets") are kept resident while the runtime `MemoryPool` grants space +// and spilled to disk under pressure, and spilled buckets are drained one at a +// time after the probe side ends (see `stream.rs`). // // v1 invariants (enforced by the constructor, which always builds an // inner join with no residual filter and no output projection): @@ -42,12 +43,12 @@ use datafusion::physical_plan::{ use super::stream::execute_join; -/// Physical execution plan node for a hash join whose build side is designed -/// to spill sub-partitions to disk rather than requiring the whole build -/// side to fit in memory at once. +/// Physical execution plan node for a hash join whose build side spills +/// sub-partitions to disk rather than requiring the whole build side to fit in +/// memory at once. /// -/// `execute` currently keeps every build-side sub-partition resident (no -/// spilling yet — see `stream.rs`). +/// `execute` runs in bounded memory, spilling buckets under `MemoryPool` +/// pressure and draining them one at a time (see `stream.rs`). #[derive(Debug)] pub struct SpillingHashJoinExec { left: Arc, @@ -170,17 +171,17 @@ impl ExecutionPlan for SpillingHashJoinExec { /// Executes partition `partition` of both children (the Ballista shuffle /// contract guarantees rows with equal join keys are co-located in the /// same partition number on both sides — see `PartitionMode::Partitioned` - /// semantics) and returns a stream that drains the left (build) side - /// fully into resident per-sub-partition hash tables before probing the - /// right (probe) side against them. No spilling yet: all build buckets - /// stay in memory for the lifetime of the stream. + /// semantics) and returns a stream that drains the left (build) side fully + /// into per-sub-partition hash tables before probing the right (probe) + /// side, spilling buckets to disk under memory pressure and draining any + /// spilled buckets one at a time after the probe side ends. fn execute( &self, partition: usize, context: Arc, ) -> Result { let left = self.left.execute(partition, Arc::clone(&context))?; - let right = self.right.execute(partition, context)?; + let right = self.right.execute(partition, Arc::clone(&context))?; let (left_keys, right_keys): (Vec, Vec) = self.on.iter().cloned().unzip(); @@ -192,6 +193,9 @@ impl ExecutionPlan for SpillingHashJoinExec { left_keys, right_keys, self.num_sub_partitions, + context, + &self.metrics, + partition, )) } @@ -468,4 +472,99 @@ mod tests { ); assert_eq!(actual_rows, expected_rows); } + + // --- spill oracle tests --- + // + // These reuse the same "match DataFusion's own Partitioned HashJoinExec" + // harness, but run `SpillingHashJoinExec` under a tiny `MemoryPool` so the + // build side cannot stay fully resident. Correctness under a tiny pool + // requires BOTH build-side spilling and probe-side routing plus a + // one-at-a-time drain of spilled buckets, so the two tests below gate the + // whole spilling path, not just a slice of it. + use datafusion::execution::TaskContext; + use datafusion::execution::memory_pool::FairSpillPool; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + + /// Builds a `TaskContext` whose `RuntimeEnv` has a `FairSpillPool` of + /// `pool_bytes` and a real (OS-temp) `DiskManager`, so `try_grow` fails + /// once the pool is exhausted and spilled batches have somewhere to go. + fn small_pool_ctx(pool_bytes: usize) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(pool_bytes))) + .build() + .expect("build RuntimeEnv"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) + } + + /// Sums the operator's `spill_count` metric across all partitions. + fn total_spill_count(plan: &Arc) -> usize { + plan.metrics() + .and_then(|m| m.sum_by_name("spill_count")) + .map(|v| v.as_usize()) + .unwrap_or(0) + } + + /// Shared body for the spill oracle tests: run the oracle under a generous + /// context and `SpillingHashJoinExec` under a `pool_bytes` pool, assert the + /// sorted outputs match and that at least `min_spills` spill events fired. + async fn assert_matches_under_pool(pool_bytes: usize, min_spills: usize) { + let generous = Arc::new(TaskContext::default()); + let small = small_pool_ctx(pool_bytes); + let num_partitions = 4; + + let base_left = make_source("lk", "lv", 3_000, 500, 0, 777); + let base_right = make_source("rk", "rv", 3_500, 450, 1_000_000, 900); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![( + Arc::new(Column::new("lk", 0)), + Arc::new(Column::new("rk", 0)), + )]; + + let oracle_left = hash_repartition(Arc::clone(&base_left), 0, num_partitions); + let oracle_right = hash_repartition(Arc::clone(&base_right), 0, num_partitions); + let ours_left = hash_repartition(base_left, 0, num_partitions); + let ours_right = hash_repartition(base_right, 0, num_partitions); + + let expected = oracle_join(oracle_left, oracle_right, on.clone(), generous).await; + + let exec: Arc = Arc::new( + SpillingHashJoinExec::try_new( + ours_left, + ours_right, + on, + PartitionMode::Partitioned, + 16, + ) + .unwrap(), + ); + let actual = collect(Arc::clone(&exec), small).await.unwrap(); + + let expected_rows = sorted_rows(&expected); + let actual_rows = sorted_rows(&actual); + assert!( + !expected_rows.is_empty(), + "test data should produce matching rows" + ); + assert_eq!(actual_rows, expected_rows); + + let spills = total_spill_count(&exec); + assert!( + spills >= min_spills, + "expected at least {min_spills} spill events, got {spills}" + ); + } + + #[tokio::test] + async fn matches_datafusion_with_forced_build_spill() { + // Pool small enough to force build-side spilling of the larger buckets + // while leaving room for some to stay resident. + assert_matches_under_pool(4 * 1024, 1).await; + } + + #[tokio::test] + async fn matches_datafusion_with_forced_two_sided_spill() { + // Pool tiny enough that many buckets spill, so their probe rows must be + // routed to disk and joined in the drain phase. + assert_matches_under_pool(512, 8).await; + } } diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs index 759117b9ca..8cd8534fc9 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -18,6 +18,7 @@ mod exec; mod hash_table; mod partitioner; +mod spill; mod stream; pub use exec::SpillingHashJoinExec; diff --git a/ballista/core/src/execution_plans/spilling_hash_join/spill.rs b/ballista/core/src/execution_plans/spilling_hash_join/spill.rs new file mode 100644 index 0000000000..d62271e218 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/spill.rs @@ -0,0 +1,179 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Per-bucket spill files for `SpillingHashJoinExec`. A `JoinSpillWriter` owns +// one Arrow-IPC stream file per sub-partition bucket, created lazily on the +// first `append` for that bucket and appended to on every subsequent one. The +// join uses two independent writers per task: one for build-side batches (build +// schema) and one for probe-side batches (probe schema). Only the batches are +// persisted — per-row hashes are recomputed on read-back from the same fixed +// seed used at partition time, so they never need to be stored. +// +// The IPC read/write pattern mirrors `sort_shuffle/spill.rs`. Files are backed +// by `DiskManager`-managed temp files (`RefCountedTempFile`), so they are +// cleaned up when the writer is dropped at the end of the task. + +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufReader, BufWriter}; +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::common::Result; +use datafusion::execution::disk_manager::{DiskManager, RefCountedTempFile}; + +/// Writes build- or probe-side batches to per-bucket Arrow-IPC spill files. +/// +/// One file is created per bucket on demand (the first `append` for that +/// bucket) and appended to thereafter. All batches written through a given +/// writer share `schema`. The temp files are owned by the writer via +/// `RefCountedTempFile`, so they outlive individual `File` handles and are +/// removed when the writer is dropped. +pub struct JoinSpillWriter { + schema: SchemaRef, + disk_manager: Arc, + files: HashMap>)>, +} + +impl JoinSpillWriter { + /// Creates a writer that will persist batches with `schema` to temp files + /// obtained from `disk_manager`. + pub fn new(schema: SchemaRef, disk_manager: Arc) -> Self { + Self { + schema, + disk_manager, + files: HashMap::new(), + } + } + + /// Appends `batch` to `bucket`'s spill file, creating the file on the first + /// call for that bucket. Empty batches are ignored so an all-empty bucket + /// never creates a file. + pub fn append(&mut self, bucket: usize, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + + if !self.files.contains_key(&bucket) { + let temp = self + .disk_manager + .create_tmp_file("SpillingHashJoin spill")?; + let file = File::create(temp.path())?; + let writer = StreamWriter::try_new(BufWriter::new(file), &self.schema)?; + self.files.insert(bucket, (temp, writer)); + } + + let (_, writer) = self.files.get_mut(&bucket).unwrap(); + writer.write(batch)?; + Ok(()) + } + + /// Flushes and finalizes every open writer so the spill files can be read. + /// Must be called before `reader`. + pub fn finish(&mut self) -> Result<()> { + for (_, writer) in self.files.values_mut() { + writer.finish()?; + } + Ok(()) + } + + /// Opens a streaming reader over `bucket`'s spill file, or `None` if + /// nothing was ever spilled for that bucket. `finish` must be called first. + pub fn reader(&self, bucket: usize) -> Result> { + match self.files.get(&bucket) { + Some((temp, _)) => { + let file = File::open(temp.path())?; + let inner = StreamReader::try_new(BufReader::new(file), None)?; + Ok(Some(JoinSpillReader { inner })) + } + None => Ok(None), + } + } +} + +/// A streaming reader over one bucket's spill file, yielding the batches back +/// in the order they were appended. +pub struct JoinSpillReader { + inner: StreamReader>, +} + +impl Iterator for JoinSpillReader { + type Item = Result; + + fn next(&mut self) -> Option { + self.inner.next().map(|r| r.map_err(Into::into)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::disk_manager::DiskManagerBuilder; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn batch(schema: &SchemaRef, values: Vec) -> RecordBatch { + RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(Int32Array::from(values))]) + .unwrap() + } + + #[test] + fn append_finish_read_roundtrip_per_bucket() { + let disk = Arc::new(DiskManagerBuilder::default().build().unwrap()); + let schema = schema(); + let mut writer = JoinSpillWriter::new(Arc::clone(&schema), disk); + + // Bucket 0 gets two batches; bucket 3 gets one; bucket 1 stays empty. + writer.append(0, &batch(&schema, vec![1, 2, 3])).unwrap(); + writer.append(0, &batch(&schema, vec![4, 5])).unwrap(); + writer.append(3, &batch(&schema, vec![6])).unwrap(); + // Empty batch must not create a file. + writer.append(2, &batch(&schema, vec![])).unwrap(); + + writer.finish().unwrap(); + + let b0: Vec = writer + .reader(0) + .unwrap() + .unwrap() + .collect::>() + .unwrap(); + assert_eq!(b0.len(), 2); + assert_eq!(b0[0].num_rows(), 3); + assert_eq!(b0[1].num_rows(), 2); + + let b3: Vec = writer + .reader(3) + .unwrap() + .unwrap() + .collect::>() + .unwrap(); + assert_eq!(b3.len(), 1); + assert_eq!(b3[0].num_rows(), 1); + + // Buckets that never received a non-empty batch have no reader. + assert!(writer.reader(1).unwrap().is_none()); + assert!(writer.reader(2).unwrap().is_none()); + } +} diff --git a/ballista/core/src/execution_plans/spilling_hash_join/stream.rs b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs index 51fe34439c..b36b582e38 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/stream.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs @@ -15,21 +15,40 @@ // specific language governing permissions and limitations // under the License. -// The in-memory (no-spill) execution stream for `SpillingHashJoinExec`. All -// build-side sub-partitions ("buckets") are kept resident for the lifetime of -// the stream — a later task adds spilling any of them to disk under memory -// pressure. Until then this is a plain, if internally sub-partitioned, hash -// join: build fully, then probe. +// The bounded-memory execution stream for `SpillingHashJoinExec`. Both sides +// are hash-partitioned into a fixed number of sub-partitions ("buckets"); a +// bucket is either kept resident or spilled to disk, so peak memory is bounded +// by the runtime `MemoryPool` during the build/probe phases and by a single +// bucket during the drain phase. // -// The stream is built as a `futures::stream::once` future (drain the build -// side and construct one `ProbeTable` per bucket) that resolves to the probe -// stream, flattened via `TryStreamExt::try_flatten`. This mirrors the idiom -// `SortShuffleWriterExec::execute` already uses in this crate -// (`sort_shuffle/writer.rs`) for "do async setup, then hand back a stream" — -// it keeps the build phase as ordinary `async`/`await` code instead of a -// hand-rolled `Stream::poll_next` state machine, while still producing a -// single `SendableRecordBatchStream` that does no work until polled and -// never emits probe output before the build side is fully drained. +// The stream runs in three phases: +// +// 1. Build. Drain the left (build) side. Each bucket's batches are buffered +// resident while the pool's `MemoryReservation` grants the space; when a +// `try_grow` is rejected, the largest resident not-yet-spilled bucket is +// evicted to a per-bucket build spill file (freeing its reservation), and +// the current batch is retried — it may itself become the victim and land +// on disk. Once the pool cannot hold a bucket, that whole bucket is +// spilled for the rest of the build. Resident buckets are finalized into +// `ProbeTable`s; spilled buckets stay on disk. +// +// 2. Probe. Stream the right (probe) side. A probe sub-batch whose bucket is +// resident is probed immediately and its output emitted; a probe sub-batch +// whose bucket is spilled is appended to a per-bucket probe spill file. +// +// 3. Drain. After the probe side ends, resident tables are dropped and the +// reservation freed, then each spilled bucket is processed one at a time: +// its build spill is read back and concatenated, its hashes recomputed by +// re-partitioning the concatenated build (the fixed seed guarantees all +// rows land in the one bucket, so the recomputed hashes match build time), +// a `ProbeTable` is built, and its probe spill is streamed through it. +// The table is dropped before the next bucket, so only one spilled +// bucket's build side is resident at a time. +// +// The stream is a single `futures::stream::once` future (do the async build +// phase) flattened via `try_flatten` into a `stream::unfold` that carries the +// probe and drain phases. Nothing polls the probe side until the build side is +// fully drained, which is exactly the ordering an inner hash join requires. use std::collections::VecDeque; use std::sync::Arc; @@ -38,22 +57,50 @@ use datafusion::arrow::array::RecordBatch; use datafusion::arrow::compute::concat_batches; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result; -use datafusion::execution::SendableRecordBatchStream; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::{Stream, StreamExt, TryStreamExt, stream}; use super::hash_table::{ProbeTable, assemble_output}; use super::partitioner::RowPartitioner; +use super::spill::{JoinSpillReader, JoinSpillWriter}; + +/// Spill counters for one `SpillingHashJoinExec` partition, published on the +/// operator's `ExecutionPlanMetricsSet`. +#[derive(Clone)] +struct SpillMetrics { + /// Number of batches appended to a build- or probe-side spill file. + spill_count: Count, + /// Sum of `RecordBatch::get_array_memory_size` over all spilled batches. + spilled_bytes: Count, +} + +impl SpillMetrics { + fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + spill_count: MetricBuilder::new(metrics).counter("spill_count", partition), + spilled_bytes: MetricBuilder::new(metrics) + .counter("spilled_bytes", partition), + } + } + + /// Records one spilled batch. + fn record(&self, batch: &RecordBatch) { + self.spill_count.add(1); + self.spilled_bytes.add(batch.get_array_memory_size()); + } +} -/// Builds the `SendableRecordBatchStream` for `SpillingHashJoinExec::execute`: -/// drains `left` (the build side) fully into one resident `ProbeTable` per -/// sub-partition bucket, then probes `right` (the probe side) against those -/// tables, bucket-matched, as it arrives. +/// Builds the `SendableRecordBatchStream` for `SpillingHashJoinExec::execute`. /// -/// `output_schema` is the join's output schema (left columns ++ right -/// columns), independent of `left`/`right`'s own schemas — it is what the -/// returned stream's `schema()` reports. +/// `output_schema` is the join's output schema (left columns ++ right columns), +/// independent of `left`/`right`'s own schemas — it is what the returned +/// stream's `schema()` reports. +#[allow(clippy::too_many_arguments)] pub fn execute_join( output_schema: SchemaRef, left: SendableRecordBatchStream, @@ -61,162 +108,428 @@ pub fn execute_join( left_keys: Vec, right_keys: Vec, num_sub_partitions: usize, + context: Arc, + metrics: &ExecutionPlanMetricsSet, + partition: usize, ) -> SendableRecordBatchStream { let adapter_schema = Arc::clone(&output_schema); + let spill_metrics = SpillMetrics::new(metrics, partition); - // A single future that drains `left`, builds the resident hash tables, - // then resolves to the probe stream. `try_flatten` turns this - // `Stream>` of one element into the actual - // `Stream>` the adapter needs — nothing here - // polls `right` until the future above has completed, which is exactly - // the "fully build before any probe output" ordering the join requires. - let joined = stream::once(build_then_probe( + let joined = stream::once(build_then_stream( left, right, left_keys, right_keys, num_sub_partitions, output_schema, + context, + spill_metrics, )) .try_flatten(); Box::pin(RecordBatchStreamAdapter::new(adapter_schema, joined)) } -/// Drains `left` into resident hash tables, then returns a stream that -/// probes `right` against them. Split out of `execute_join` (rather than an -/// inline `async move` block) purely so its `Result>` return -/// type is written down once, instead of needing an explicit turbofish at -/// every call site to disambiguate the error type. -async fn build_then_probe( +/// Runs the build phase (async, drains `left`), then hands back the stream +/// that carries the probe and drain phases. +#[allow(clippy::too_many_arguments)] +async fn build_then_stream( left: SendableRecordBatchStream, right: SendableRecordBatchStream, left_keys: Vec, right_keys: Vec, num_sub_partitions: usize, output_schema: SchemaRef, + context: Arc, + spill_metrics: SpillMetrics, ) -> Result> + Send> { - let tables = build_tables(left, &left_keys, num_sub_partitions).await?; - Ok(probe_stream(right, right_keys, tables, output_schema)) + let runtime = context.runtime_env(); + let mut reservation = MemoryConsumer::new(format!( + "SpillingHashJoin[{}]", + output_schema.fields().len() + )) + .with_can_spill(true) + .register(&runtime.memory_pool); + + let left_schema = left.schema(); + let right_schema = right.schema(); + + // Two independent spill sets: build-side batches (build schema) and + // probe-side batches (probe schema). + let mut build_spill = + JoinSpillWriter::new(Arc::clone(&left_schema), Arc::clone(&runtime.disk_manager)); + let probe_spill = JoinSpillWriter::new( + Arc::clone(&right_schema), + Arc::clone(&runtime.disk_manager), + ); + + let partitioner_left = RowPartitioner::new(left_keys.clone(), num_sub_partitions); + let partitioner_right = RowPartitioner::new(right_keys.clone(), num_sub_partitions); + + let (tables, spilled) = build_phase( + left, + &left_keys, + &left_schema, + num_sub_partitions, + &mut reservation, + &mut build_spill, + &spill_metrics, + &partitioner_left, + ) + .await?; + + let state = JoinState { + phase: Phase::Probe, + right, + left_keys, + right_keys, + left_schema, + output_schema, + tables, + spilled, + reservation, + build_spill, + probe_spill, + partitioner_left, + partitioner_right, + spill_metrics, + pending: VecDeque::new(), + drain_buckets: VecDeque::new(), + current: None, + }; + + Ok(join_stream(state)) } -/// Drains `left` fully, routing every row into one of `num_sub_partitions` -/// buckets by `left_keys` (via `RowPartitioner`), then builds one -/// `ProbeTable` per non-empty bucket. Buckets that never received a row are -/// `None`, so `probe_stream` can skip them (and later, spilling) without -/// treating "no rows" as an error. -async fn build_tables( +/// Drains `left` into resident `ProbeTable`s and a build spill, spilling the +/// largest resident bucket under memory pressure. +/// +/// Returns one entry per bucket: `Some(table)` for resident non-empty buckets, +/// `None` for empty or spilled buckets. `spilled[b]` is `true` iff bucket `b` +/// was spilled (its batches live in `build_spill`). +#[allow(clippy::too_many_arguments)] +async fn build_phase( mut left: SendableRecordBatchStream, left_keys: &[PhysicalExprRef], + left_schema: &SchemaRef, num_sub_partitions: usize, -) -> Result>> { - let left_schema = left.schema(); - let partitioner = RowPartitioner::new(left_keys.to_vec(), num_sub_partitions); - + reservation: &mut MemoryReservation, + build_spill: &mut JoinSpillWriter, + spill_metrics: &SpillMetrics, + partitioner_left: &RowPartitioner, +) -> Result<(Vec>, Vec)> { let mut bucket_batches: Vec> = vec![Vec::new(); num_sub_partitions]; let mut bucket_hashes: Vec> = vec![Vec::new(); num_sub_partitions]; + let mut resident_bytes: Vec = vec![0; num_sub_partitions]; + let mut spilled: Vec = vec![false; num_sub_partitions]; while let Some(batch) = left.next().await { let batch = batch?; - for pb in partitioner.partition(&batch)? { - bucket_batches[pb.bucket].push(pb.batch); - bucket_hashes[pb.bucket].extend(pb.hashes); + for pb in partitioner_left.partition(&batch)? { + let bucket = pb.bucket; + + // Already-spilled bucket: straight to disk, no reservation. + if spilled[bucket] { + build_spill.append(bucket, &pb.batch)?; + spill_metrics.record(&pb.batch); + continue; + } + + let size = pb.batch.get_array_memory_size(); + loop { + if reservation.try_grow(size).is_ok() { + bucket_batches[bucket].push(pb.batch); + bucket_hashes[bucket].extend(pb.hashes); + resident_bytes[bucket] += size; + break; + } + + // Pool rejected the grow: evict the largest resident + // not-yet-spilled bucket to free space, then retry. + match largest_resident(&resident_bytes, &spilled) { + Some(victim) => { + let freed = resident_bytes[victim]; + for b in bucket_batches[victim].drain(..) { + build_spill.append(victim, &b)?; + spill_metrics.record(&b); + } + bucket_hashes[victim].clear(); + reservation.shrink(freed); + resident_bytes[victim] = 0; + spilled[victim] = true; + + if victim == bucket { + // This batch's own bucket was the victim; the batch + // now belongs on disk with the rest of the bucket. + build_spill.append(bucket, &pb.batch)?; + spill_metrics.record(&pb.batch); + break; + } + // Otherwise loop and retry `try_grow` for this batch. + } + None => { + // Nothing resident to evict and the pool still won't + // grant the space: spill this batch directly and mark + // its bucket spilled for the remainder of the build. + build_spill.append(bucket, &pb.batch)?; + spill_metrics.record(&pb.batch); + spilled[bucket] = true; + break; + } + } + } } } - bucket_batches - .into_iter() - .zip(bucket_hashes) - .map(|(batches, hashes)| { - if batches.is_empty() { - return Ok(None); - } - let concatenated = concat_batches(&left_schema, batches.iter())?; - Ok(Some(ProbeTable::build(concatenated, &hashes, left_keys)?)) - }) - .collect() + // Finalize resident buckets into probe tables; spilled buckets stay on disk. + let mut tables: Vec> = Vec::with_capacity(num_sub_partitions); + for bucket in 0..num_sub_partitions { + if spilled[bucket] { + tables.push(None); + continue; + } + let batches = std::mem::take(&mut bucket_batches[bucket]); + if batches.is_empty() { + tables.push(None); + continue; + } + let hashes = std::mem::take(&mut bucket_hashes[bucket]); + let concatenated = concat_batches(left_schema, batches.iter())?; + tables.push(Some(ProbeTable::build(concatenated, &hashes, left_keys)?)); + } + + Ok((tables, spilled)) +} + +/// Returns the index of the largest resident (non-empty, not-yet-spilled) +/// bucket, or `None` if every bucket is either spilled or empty. +fn largest_resident(resident_bytes: &[usize], spilled: &[bool]) -> Option { + resident_bytes + .iter() + .enumerate() + .filter(|&(b, &bytes)| !spilled[b] && bytes > 0) + .max_by_key(|&(_, &bytes)| bytes) + .map(|(b, _)| b) +} + +/// Which phase `join_stream` is in. +enum Phase { + /// Streaming the probe side against resident tables, routing spilled + /// buckets' probe rows to disk. + Probe, + /// Draining spilled buckets one at a time. + Drain, +} + +/// The active spilled bucket being drained: its build-side `ProbeTable` and a +/// reader over its probe spill. +struct DrainBucket { + table: ProbeTable, + probe_reader: JoinSpillReader, } -/// Mutable state threaded through `probe_stream`'s `futures::stream::unfold`. -struct ProbeState { +/// All state threaded through `join_stream`'s `stream::unfold`. +struct JoinState { + phase: Phase, right: SendableRecordBatchStream, - partitioner: RowPartitioner, + left_keys: Vec, right_keys: Vec, - tables: Vec>, + left_schema: SchemaRef, output_schema: SchemaRef, - /// Output batches assembled from the most recent `right` batch but not - /// yet yielded. A single `right` batch can fan out into multiple - /// sub-partition buckets, each producing its own output batch, so these - /// are queued and drained one at a time before pulling `right` again. + /// Resident probe tables per bucket (`None` for empty or spilled buckets). + /// Dropped when the drain phase begins. + tables: Vec>, + spilled: Vec, + /// Kept alive so the pool sees this join's resident memory until drain. + reservation: MemoryReservation, + build_spill: JoinSpillWriter, + probe_spill: JoinSpillWriter, + partitioner_left: RowPartitioner, + partitioner_right: RowPartitioner, + spill_metrics: SpillMetrics, + /// Output batches assembled but not yet yielded. pending: VecDeque, + /// Spilled bucket indices still to be drained (populated at drain start). + drain_buckets: VecDeque, + /// The spilled bucket currently being drained, if any. + current: Option, } -/// Probes `right` against `tables` (one resident `ProbeTable` per bucket, -/// built from the fully-drained build side) as `right` batches arrive, -/// yielding one output batch per non-empty matched sub-partition. `right` -/// rows that hash to a bucket with no build-side table (`None`) have no -/// possible match and are dropped, matching inner-join semantics. -fn probe_stream( - right: SendableRecordBatchStream, - right_keys: Vec, - tables: Vec>, - output_schema: SchemaRef, -) -> impl Stream> + Send { - let num_sub_partitions = tables.len(); - let state = ProbeState { - right, - partitioner: RowPartitioner::new(right_keys.clone(), num_sub_partitions), - right_keys, - tables, - output_schema, - pending: VecDeque::new(), - }; +impl JoinState { + /// Partitions one probe batch and, per resulting bucket, either probes the + /// resident table (queuing output) or appends the sub-batch to the probe + /// spill for that bucket. + fn route_probe_batch(&mut self, batch: &RecordBatch) -> Result<()> { + for pb in self.partitioner_right.partition(batch)? { + if self.spilled[pb.bucket] { + self.probe_spill.append(pb.bucket, &pb.batch)?; + self.spill_metrics.record(&pb.batch); + continue; + } + let Some(table) = &self.tables[pb.bucket] else { + continue; + }; + let (build_rows, probe_rows) = + table.probe(&pb.batch, &pb.hashes, &self.right_keys)?; + if build_rows.is_empty() { + continue; + } + let out = assemble_output( + &self.output_schema, + table.build_batch(), + &pb.batch, + &build_rows, + &probe_rows, + )?; + self.pending.push_back(out); + } + Ok(()) + } - stream::unfold(state, |mut state| async move { + /// Ends the probe phase: finalizes both spill writers, releases the + /// resident tables and their reservation, and lists the spilled buckets to + /// drain. + fn begin_drain(&mut self) -> Result<()> { + self.build_spill.finish()?; + self.probe_spill.finish()?; + // Resident tables have been fully probed; drop them and free the pool + // so only one spilled bucket's build side is resident during drain. + self.tables = Vec::new(); + self.reservation.free(); + self.drain_buckets = (0..self.spilled.len()) + .filter(|&b| self.spilled[b]) + .collect(); + Ok(()) + } + + /// Advances the drain phase by one step. Queues output onto `pending` and + /// returns `Ok(true)` while there is more work, or `Ok(false)` once every + /// spilled bucket has been drained. + fn drain_step(&mut self) -> Result { loop { - if let Some(batch) = state.pending.pop_front() { - return Some((Ok(batch), state)); + if let Some(current) = &mut self.current { + // Pull the next probe batch for the active bucket. + match current.probe_reader.next() { + Some(Ok(probe_batch)) => { + // Recompute probe hashes from the same fixed seed so + // they match the build side's; all rows are in one + // bucket, so a single partitioned batch comes back. + for pb in self.partitioner_right.partition(&probe_batch)? { + let (build_rows, probe_rows) = current.table.probe( + &pb.batch, + &pb.hashes, + &self.right_keys, + )?; + if build_rows.is_empty() { + continue; + } + let out = assemble_output( + &self.output_schema, + current.table.build_batch(), + &pb.batch, + &build_rows, + &probe_rows, + )?; + self.pending.push_back(out); + } + if !self.pending.is_empty() { + return Ok(true); + } + // No output from this batch; keep pulling. + } + Some(Err(e)) => return Err(e), + None => { + // Bucket exhausted; drop its table before the next one. + self.current = None; + } + } + continue; } - let next = state.right.next().await?; - let batch = match next { - Ok(batch) => batch, - Err(e) => return Some((Err(e), state)), + // Load the next spilled bucket. + let Some(bucket) = self.drain_buckets.pop_front() else { + return Ok(false); }; - - if let Err(e) = probe_batch(&batch, &mut state) { - return Some((Err(e), state)); - } - // Loop back around: this `right` batch may have produced no - // output (e.g. every bucket it touched was empty), in which case - // `pending` is still empty and we must pull the next `right` - // batch rather than returning `None` (which would end the - // stream early). + let Some(table) = self.build_drain_table(bucket)? else { + // No build rows for this bucket: an inner join yields nothing, + // so skip it (and its probe spill, if any). + continue; + }; + let probe_reader = match self.probe_spill.reader(bucket)? { + Some(r) => r, + None => continue, // build rows but no probe rows -> no output + }; + self.current = Some(DrainBucket { + table, + probe_reader, + }); } - }) -} + } -/// Partitions one `right` batch by `state.right_keys` and probes each -/// resulting sub-batch against its bucket's `ProbeTable` (if any), pushing -/// every non-empty assembled output batch onto `state.pending`. -fn probe_batch(batch: &RecordBatch, state: &mut ProbeState) -> Result<()> { - for pb in state.partitioner.partition(batch)? { - let Some(table) = &state.tables[pb.bucket] else { - continue; + /// Reads a spilled bucket's build side back from disk, concatenates it, and + /// builds its `ProbeTable`, recomputing hashes by re-partitioning the + /// concatenated build. The fixed seed guarantees every row lands in the + /// single original bucket, so the recomputed hashes match build time. + /// Returns `None` if the bucket has no build rows. + fn build_drain_table(&self, bucket: usize) -> Result> { + let Some(reader) = self.build_spill.reader(bucket)? else { + return Ok(None); }; - let (build_rows, probe_rows) = - table.probe(&pb.batch, &pb.hashes, &state.right_keys)?; - if build_rows.is_empty() { - continue; + let batches: Vec = reader.collect::>()?; + if batches.is_empty() { + return Ok(None); } - let out = assemble_output( - &state.output_schema, - table.build_batch(), - &pb.batch, - &build_rows, - &probe_rows, - )?; - state.pending.push_back(out); + let concatenated = concat_batches(&self.left_schema, batches.iter())?; + + // Re-partition to recover per-row hashes consistent with build time. + // All rows share the one bucket, so exactly one partitioned batch is + // returned, carrying every row in the concatenated order. + let mut parts = self.partitioner_left.partition(&concatenated)?; + let Some(part) = parts.pop() else { + return Ok(None); + }; + Ok(Some(ProbeTable::build( + part.batch, + &part.hashes, + &self.left_keys, + )?)) } - Ok(()) +} + +/// Turns `JoinState` into the output stream: emit any queued output, then in +/// the probe phase pull the probe side (routing spilled buckets to disk), and +/// finally drain spilled buckets one at a time. +fn join_stream(state: JoinState) -> impl Stream> + Send { + stream::unfold(state, |mut state| async move { + loop { + if let Some(batch) = state.pending.pop_front() { + return Some((Ok(batch), state)); + } + + match state.phase { + Phase::Probe => match state.right.next().await { + Some(Ok(batch)) => { + if let Err(e) = state.route_probe_batch(&batch) { + return Some((Err(e), state)); + } + // Loop back: this batch may have produced no output. + } + Some(Err(e)) => return Some((Err(e), state)), + None => { + if let Err(e) = state.begin_drain() { + return Some((Err(e), state)); + } + state.phase = Phase::Drain; + } + }, + Phase::Drain => match state.drain_step() { + Ok(true) => { + // Loop back to drain queued output. + } + Ok(false) => return None, + Err(e) => return Some((Err(e), state)), + }, + } + } + }) } From 656958be7c13a6cf490c56d8039e208043e4976c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 11:44:38 -0600 Subject: [PATCH 07/15] feat(core): recursively repartition oversized spilled buckets Pool-track the drain-phase build side of SpillingHashJoinExec and, when a single spilled bucket's build side will not fit in memory, recursively re-partition both its build and probe spills into finer sub-buckets under a depth-varied hash seed (build and probe share the seed at each level so equal keys stay co-located). Only one (sub-)bucket's build table is resident at any depth. A build side for a single join key that exceeds the whole pool is reported as a clean error instead of an OOM; a best-effort infallible reservation covers residuals that fit the pool but lost a grow race with other resident buckets. A depth cap bounds the recursion. --- .../spilling_hash_join/exec.rs | 120 +++++- .../spilling_hash_join/partitioner.rs | 20 +- .../spilling_hash_join/stream.rs | 397 +++++++++++++++--- 3 files changed, 484 insertions(+), 53 deletions(-) diff --git a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs index d11e747f81..eda3670440 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -504,6 +504,15 @@ mod tests { .unwrap_or(0) } + /// Sums the operator's `repartition_count` metric (drain-phase recursive + /// re-partition events) across all partitions. + fn total_repartition_count(plan: &Arc) -> usize { + plan.metrics() + .and_then(|m| m.sum_by_name("repartition_count")) + .map(|v| v.as_usize()) + .unwrap_or(0) + } + /// Shared body for the spill oracle tests: run the oracle under a generous /// context and `SpillingHashJoinExec` under a `pool_bytes` pool, assert the /// sorted outputs match and that at least `min_spills` spill events fired. @@ -564,7 +573,114 @@ mod tests { #[tokio::test] async fn matches_datafusion_with_forced_two_sided_spill() { // Pool tiny enough that many buckets spill, so their probe rows must be - // routed to disk and joined in the drain phase. - assert_matches_under_pool(512, 8).await; + // routed to disk and joined in the drain phase. Kept at (just) one + // minimal Arrow batch so a spilled bucket's build side can still be held + // resident while it is drained. + assert_matches_under_pool(1024, 8).await; + } + + // --- recursive re-partition (skew fallback) tests --- + // + // When a single spilled bucket's build side still will not fit in memory + // during drain, the join re-partitions that bucket (build and probe) into + // finer sub-buckets with a depth-varied hash seed and recurses, keeping at + // most one sub-bucket's build table resident. If a bucket holds only one + // distinct join key it cannot be split, so the join returns a clean error + // rather than an OOM/panic/hang. + + #[tokio::test] + async fn matches_datafusion_with_forced_recursion() { + // Pool so small that even one spilled bucket's build side must be + // re-partitioned at least once during drain. Driven from a single + // partition so the tiny pool is exercised by one draining task, not + // shared across concurrent partitions. + let generous = Arc::new(TaskContext::default()); + let small = small_pool_ctx(2048); + let num_partitions = 1; + + // All-distinct keys, so any oversized bucket is guaranteed to hold more + // than one distinct join key and can therefore always be split. A small + // sub-partition count (2) packs many keys into each first-level bucket. + let base_left = make_source("lk", "lv", 600, 600, 0, 128); + let base_right = make_source("rk", "rv", 600, 600, 1_000_000, 128); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![( + Arc::new(Column::new("lk", 0)), + Arc::new(Column::new("rk", 0)), + )]; + + let oracle_left = hash_repartition(Arc::clone(&base_left), 0, num_partitions); + let oracle_right = hash_repartition(Arc::clone(&base_right), 0, num_partitions); + let ours_left = hash_repartition(base_left, 0, num_partitions); + let ours_right = hash_repartition(base_right, 0, num_partitions); + + let expected = oracle_join(oracle_left, oracle_right, on.clone(), generous).await; + + let exec: Arc = Arc::new( + SpillingHashJoinExec::try_new( + ours_left, + ours_right, + on, + PartitionMode::Partitioned, + 2, + ) + .unwrap(), + ); + let actual = collect(Arc::clone(&exec), small).await.unwrap(); + + let expected_rows = sorted_rows(&expected); + let actual_rows = sorted_rows(&actual); + assert!( + !expected_rows.is_empty(), + "test data should produce matching rows" + ); + assert_eq!(actual_rows, expected_rows); + assert!( + total_repartition_count(&exec) > 0, + "the recursive re-partition path must actually be taken" + ); + } + + #[tokio::test] + async fn single_key_too_large_errors_cleanly() { + // Every build row shares ONE join key, so the whole build side lands in + // a single bucket that cannot be split; under a tiny pool it cannot fit + // either, so draining it must return a clean, named error. + let ctx = small_pool_ctx(2048); + let num_partitions = 1; + + let base_left = make_source("lk", "lv", 2_000, 1, 0, 256); + let base_right = make_source("rk", "rv", 2_000, 1, 1_000_000, 256); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![( + Arc::new(Column::new("lk", 0)), + Arc::new(Column::new("rk", 0)), + )]; + + let ours_left = hash_repartition(base_left, 0, num_partitions); + let ours_right = hash_repartition(base_right, 0, num_partitions); + + let exec: Arc = Arc::new( + SpillingHashJoinExec::try_new( + ours_left, + ours_right, + on, + PartitionMode::Partitioned, + 2, + ) + .unwrap(), + ); + + let result = collect(exec, ctx).await; + let err = result.expect_err("single oversized join key must error, not hang/OOM"); + let msg = err.to_string(); + assert!( + msg.contains("SpillingHashJoinExec"), + "error should name the operator, got: {msg}" + ); + assert!( + msg.contains("single join key"), + "error should identify single-join-key skew, got: {msg}" + ); } } diff --git a/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs b/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs index 1cd3e37aa0..c7e825b498 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs @@ -40,6 +40,16 @@ use datafusion::physical_expr::PhysicalExprRef; /// diverge from it; the exact value carries no other meaning. const ROW_PARTITIONER_SEED: u64 = 0x5350_4a5f_484a_3121; +/// Derives the hashing seed for drain-phase recursion `depth`. +/// `seed_for_depth(0)` returns the base seed, so level-0 bucketing is +/// byte-identical to [`RowPartitioner::new`]; deeper levels perturb the seed so +/// re-partitioning a bucket that did not fit splits its keys differently than +/// the level that produced it. The multiplier is an arbitrary odd constant +/// (the golden-ratio mix) chosen only to spread successive depths apart. +pub fn seed_for_depth(depth: usize) -> u64 { + ROW_PARTITIONER_SEED ^ (depth as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) +} + /// Splits the rows of a `RecordBatch` into `num_sub` sub-partitions by /// hashing a set of join key expressions, so that rows with equal keys /// always land in the same sub-partition. @@ -68,10 +78,18 @@ impl RowPartitioner { /// into `num_sub` sub-partitions using a fixed, non-default random /// state. pub fn new(keys: Vec, num_sub: usize) -> Self { + Self::with_seed(keys, num_sub, ROW_PARTITIONER_SEED) + } + + /// Creates a `RowPartitioner` hashing with an explicit `seed`. Used by the + /// drain-phase recursion to split an oversized bucket differently than the + /// level that produced it (see [`seed_for_depth`]); `new` delegates here + /// with the fixed base seed. + pub fn with_seed(keys: Vec, num_sub: usize, seed: u64) -> Self { Self { keys, num_sub, - random_state: RandomState::with_seed(ROW_PARTITIONER_SEED), + random_state: RandomState::with_seed(seed), } } diff --git a/ballista/core/src/execution_plans/spilling_hash_join/stream.rs b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs index b36b582e38..e224e62776 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/stream.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs @@ -38,12 +38,18 @@ // // 3. Drain. After the probe side ends, resident tables are dropped and the // reservation freed, then each spilled bucket is processed one at a time: -// its build spill is read back and concatenated, its hashes recomputed by -// re-partitioning the concatenated build (the fixed seed guarantees all -// rows land in the one bucket, so the recomputed hashes match build time), -// a `ProbeTable` is built, and its probe spill is streamed through it. -// The table is dropped before the next bucket, so only one spilled -// bucket's build side is resident at a time. +// its build spill is read back and concatenated, and the reservation is +// grown for it. When it fits, its hashes are recomputed by re-partitioning +// the concatenated build (the fixed seed guarantees all rows land in the +// one bucket, so the recomputed hashes match build time), a `ProbeTable` +// is built, and its probe spill is streamed through it. When a bucket's +// build side will not fit, it is re-partitioned — build and probe together, +// under a depth-varied seed so equal keys still co-locate — into finer +// sub-buckets that are drained recursively, one resident table at a time. +// A build side for a single join key that exceeds the whole pool is +// irreducible skew and fails with a clean error rather than an OOM. The +// table is dropped before the next bucket, so only one (sub-)bucket's +// build side is resident at a time, at any recursion depth. // // The stream is a single `futures::stream::once` future (do the async build // phase) flattened via `try_flatten` into a `stream::unfold` that carries the @@ -53,12 +59,16 @@ use std::collections::VecDeque; use std::sync::Arc; -use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::{ArrayRef, RecordBatch}; use datafusion::arrow::compute::concat_batches; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{Result, exec_err}; use datafusion::execution::TaskContext; -use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::execution::disk_manager::DiskManager; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryReservation, +}; use datafusion::physical_expr::PhysicalExprRef; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}; @@ -66,9 +76,14 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::{Stream, StreamExt, TryStreamExt, stream}; use super::hash_table::{ProbeTable, assemble_output}; -use super::partitioner::RowPartitioner; +use super::partitioner::{RowPartitioner, seed_for_depth}; use super::spill::{JoinSpillReader, JoinSpillWriter}; +/// Maximum drain-phase re-partition depth. A bucket that still will not fit +/// after this many splits is treated as irreducible skew and reported as a +/// clean error rather than recursed on forever. +const MAX_DRAIN_DEPTH: usize = 8; + /// Spill counters for one `SpillingHashJoinExec` partition, published on the /// operator's `ExecutionPlanMetricsSet`. #[derive(Clone)] @@ -77,6 +92,9 @@ struct SpillMetrics { spill_count: Count, /// Sum of `RecordBatch::get_array_memory_size` over all spilled batches. spilled_bytes: Count, + /// Number of times a spilled bucket's build side did not fit in memory + /// during drain and was recursively re-partitioned into finer sub-buckets. + repartition_count: Count, } impl SpillMetrics { @@ -85,6 +103,8 @@ impl SpillMetrics { spill_count: MetricBuilder::new(metrics).counter("spill_count", partition), spilled_bytes: MetricBuilder::new(metrics) .counter("spilled_bytes", partition), + repartition_count: MetricBuilder::new(metrics) + .counter("repartition_count", partition), } } @@ -93,6 +113,11 @@ impl SpillMetrics { self.spill_count.add(1); self.spilled_bytes.add(batch.get_array_memory_size()); } + + /// Records one drain-phase recursive re-partition of an oversized bucket. + fn record_repartition(&self) { + self.repartition_count.add(1); + } } /// Builds the `SendableRecordBatchStream` for `SpillingHashJoinExec::execute`. @@ -166,6 +191,14 @@ async fn build_then_stream( let partitioner_left = RowPartitioner::new(left_keys.clone(), num_sub_partitions); let partitioner_right = RowPartitioner::new(right_keys.clone(), num_sub_partitions); + // The pool's total capacity, used during drain to tell irreducible + // single-key skew (a build side larger than the whole pool) apart from a + // grow that merely lost a race with other resident buckets. + let pool_limit = match runtime.memory_pool.memory_limit() { + MemoryLimit::Finite(bytes) => Some(bytes), + MemoryLimit::Infinite | MemoryLimit::Unknown => None, + }; + let (tables, spilled) = build_phase( left, &left_keys, @@ -184,7 +217,11 @@ async fn build_then_stream( left_keys, right_keys, left_schema, + right_schema, output_schema, + num_sub: num_sub_partitions, + pool_limit, + disk_manager: Arc::clone(&runtime.disk_manager), tables, spilled, reservation, @@ -321,11 +358,13 @@ enum Phase { Drain, } -/// The active spilled bucket being drained: its build-side `ProbeTable` and a -/// reader over its probe spill. +/// The active spilled bucket being drained: its build-side `ProbeTable`, a +/// reader over its probe spill, and the pool bytes reserved for the table +/// (released when the bucket is exhausted). struct DrainBucket { table: ProbeTable, probe_reader: JoinSpillReader, + reserved: usize, } /// All state threaded through `join_stream`'s `stream::unfold`. @@ -335,7 +374,15 @@ struct JoinState { left_keys: Vec, right_keys: Vec, left_schema: SchemaRef, + right_schema: SchemaRef, output_schema: SchemaRef, + /// Number of sub-partitions ("buckets"), reused as the fan-out for each + /// drain-phase re-partition level. + num_sub: usize, + /// The memory pool's total capacity, or `None` for an unbounded pool. + pool_limit: Option, + /// Source of temp files for spilling re-partitioned buckets during drain. + disk_manager: Arc, /// Resident probe tables per bucket (`None` for empty or spilled buckets). /// Dropped when the drain phase begins. tables: Vec>, @@ -439,7 +486,10 @@ impl JoinState { } Some(Err(e)) => return Err(e), None => { - // Bucket exhausted; drop its table before the next one. + // Bucket exhausted; free its build-table reservation and + // drop the table before loading the next one. + let reserved = current.reserved; + self.reservation.shrink(reserved); self.current = None; } } @@ -450,50 +500,297 @@ impl JoinState { let Some(bucket) = self.drain_buckets.pop_front() else { return Ok(false); }; - let Some(table) = self.build_drain_table(bucket)? else { - // No build rows for this bucket: an inner join yields nothing, - // so skip it (and its probe spill, if any). - continue; - }; - let probe_reader = match self.probe_spill.reader(bucket)? { - Some(r) => r, - None => continue, // build rows but no probe rows -> no output + + // Read the bucket's build side back from disk and concatenate it. + let build_batches: Vec = match self.build_spill.reader(bucket)? { + Some(r) => r.collect::>()?, + None => continue, // no build rows -> inner join yields nothing }; - self.current = Some(DrainBucket { - table, - probe_reader, - }); + if build_batches.is_empty() { + continue; + } + let concatenated = concat_batches(&self.left_schema, build_batches.iter())?; + let build_size = concatenated.get_array_memory_size(); + let probe_reader = self.probe_spill.reader(bucket)?; + + if self.reservation.try_grow(build_size).is_ok() { + // The build side fits: build its `ProbeTable` and stream the + // probe spill through it, one probe batch at a time. Re-partition + // to recover per-row hashes consistent with build time; all rows + // share the one bucket, so exactly one partitioned batch comes + // back, carrying every row in concatenated order. + let mut parts = self.partitioner_left.partition(&concatenated)?; + let Some(part) = parts.pop() else { + self.reservation.shrink(build_size); + continue; + }; + let table = ProbeTable::build(part.batch, &part.hashes, &self.left_keys)?; + match probe_reader { + Some(probe_reader) => { + self.current = Some(DrainBucket { + table, + probe_reader, + reserved: build_size, + }); + } + None => { + // Build rows but no probe rows -> no output. + self.reservation.shrink(build_size); + } + } + } else { + // The build side does not fit right now: re-partition this + // bucket into finer sub-buckets and drain them recursively (or + // fail cleanly on irreducible single-join-key skew). Output is + // materialized because this fallback is rare; boundedness is + // preserved by keeping only one sub-bucket's table resident. + let ctx = DrainCtx { + left_schema: &self.left_schema, + right_schema: &self.right_schema, + left_keys: &self.left_keys, + right_keys: &self.right_keys, + output_schema: &self.output_schema, + num_sub: self.num_sub, + pool_limit: self.pool_limit, + disk_manager: &self.disk_manager, + spill_metrics: &self.spill_metrics, + }; + let out = drain_overflowed_bucket( + &ctx, + &mut self.reservation, + concatenated, + probe_reader, + 0, + )?; + self.pending.extend(out); + if !self.pending.is_empty() { + // Surface this bucket's output before loading the next one, + // otherwise `join_stream` would end the drain on the next + // `Ok(false)` and discard the queued rows. + return Ok(true); + } + } } } +} - /// Reads a spilled bucket's build side back from disk, concatenates it, and - /// builds its `ProbeTable`, recomputing hashes by re-partitioning the - /// concatenated build. The fixed seed guarantees every row lands in the - /// single original bucket, so the recomputed hashes match build time. - /// Returns `None` if the bucket has no build rows. - fn build_drain_table(&self, bucket: usize) -> Result> { - let Some(reader) = self.build_spill.reader(bucket)? else { - return Ok(None); - }; - let batches: Vec = reader.collect::>()?; - if batches.is_empty() { - return Ok(None); +/// Shared, read-only context threaded through the drain-phase recursion. +struct DrainCtx<'a> { + left_schema: &'a SchemaRef, + right_schema: &'a SchemaRef, + left_keys: &'a [PhysicalExprRef], + right_keys: &'a [PhysicalExprRef], + output_schema: &'a SchemaRef, + num_sub: usize, + pool_limit: Option, + disk_manager: &'a Arc, + spill_metrics: &'a SpillMetrics, +} + +/// Drains one spilled (sub-)bucket whose rows were partitioned with +/// `seed_for_depth(depth)`. If the concatenated build side fits the pool it is +/// joined against its probe spill and the output returned; otherwise the bucket +/// is handled by [`drain_overflowed_bucket`] (re-partition, best-effort hold, +/// or clean error). Only one (sub-)bucket's build table is resident at any +/// instant, at any depth. +fn drain_bucket_recursive( + ctx: &DrainCtx, + reservation: &mut MemoryReservation, + build: RecordBatch, + probe_reader: Option, + depth: usize, +) -> Result> { + let build_size = build.get_array_memory_size(); + if reservation.try_grow(build_size).is_ok() { + let out = join_resident_bucket(ctx, &build, probe_reader, depth); + reservation.shrink(build_size); + return out; + } + drain_overflowed_bucket(ctx, reservation, build, probe_reader, depth) +} + +/// Handles a (sub-)bucket whose build side did not fit the pool via `try_grow`. +/// +/// - More than one distinct join key (and depth budget left) → re-partition +/// into finer sub-buckets and drain each recursively. +/// - A single distinct key (or the depth cap) whose build side is larger than +/// the whole pool → irreducible skew, returned as a clean error. +/// - Otherwise the build side fits the pool and the `try_grow` only lost a +/// race with other resident buckets → hold this one bucket anyway with an +/// infallible reservation (still at most one table resident). +fn drain_overflowed_bucket( + ctx: &DrainCtx, + reservation: &mut MemoryReservation, + build: RecordBatch, + probe_reader: Option, + depth: usize, +) -> Result> { + if depth < MAX_DRAIN_DEPTH && has_multiple_distinct_keys(&build, ctx.left_keys)? { + return repartition_and_recurse(ctx, reservation, build, probe_reader, depth); + } + + let build_size = build.get_array_memory_size(); + if depth >= MAX_DRAIN_DEPTH || exceeds_pool(build_size, ctx.pool_limit) { + return exec_err!( + "SpillingHashJoinExec: build side for a single join key exceeds \ + memory; enable SMJ fallback for this query" + ); + } + + reservation.grow(build_size); + let out = join_resident_bucket(ctx, &build, probe_reader, depth); + reservation.shrink(build_size); + out +} + +/// Returns `true` iff a bounded pool of total capacity `pool_limit` cannot hold +/// a build side of `size` even when otherwise empty (irreducible skew). An +/// unbounded pool (`None`) is never exceeded. +fn exceeds_pool(size: usize, pool_limit: Option) -> bool { + matches!(pool_limit, Some(limit) if size > limit) +} + +/// Joins one resident build bucket (already reserved by the caller) against its +/// probe spill, returning the assembled output batches. Hashes for both sides +/// are recomputed with `seed_for_depth(depth)` — the same seed the level that +/// produced this bucket used — so build and probe stay consistent. +fn join_resident_bucket( + ctx: &DrainCtx, + build: &RecordBatch, + probe_reader: Option, + depth: usize, +) -> Result> { + let seed = seed_for_depth(depth); + let partitioner_left = + RowPartitioner::with_seed(ctx.left_keys.to_vec(), ctx.num_sub, seed); + let partitioner_right = + RowPartitioner::with_seed(ctx.right_keys.to_vec(), ctx.num_sub, seed); + + // All build rows share the one bucket at this seed, so a single partitioned + // batch comes back carrying every row (with build-consistent hashes). + let mut parts = partitioner_left.partition(build)?; + let Some(part) = parts.pop() else { + return Ok(Vec::new()); + }; + let table = ProbeTable::build(part.batch, &part.hashes, ctx.left_keys)?; + + let mut out = Vec::new(); + let Some(reader) = probe_reader else { + return Ok(out); // build rows but no probe rows -> no output + }; + for probe_batch in reader { + let probe_batch = probe_batch?; + for pb in partitioner_right.partition(&probe_batch)? { + let (build_rows, probe_rows) = + table.probe(&pb.batch, &pb.hashes, ctx.right_keys)?; + if build_rows.is_empty() { + continue; + } + out.push(assemble_output( + ctx.output_schema, + table.build_batch(), + &pb.batch, + &build_rows, + &probe_rows, + )?); } - let concatenated = concat_batches(&self.left_schema, batches.iter())?; - - // Re-partition to recover per-row hashes consistent with build time. - // All rows share the one bucket, so exactly one partitioned batch is - // returned, carrying every row in the concatenated order. - let mut parts = self.partitioner_left.partition(&concatenated)?; - let Some(part) = parts.pop() else { - return Ok(None); + } + Ok(out) +} + +/// Re-partitions an oversized (sub-)bucket's build and probe spills into +/// `num_sub` finer sub-buckets using a deeper hash seed, then drains each +/// recursively. The caller (`drain_overflowed_bucket`) has already established +/// that the bucket has more than one distinct join key, so the split makes +/// progress. +fn repartition_and_recurse( + ctx: &DrainCtx, + reservation: &mut MemoryReservation, + build: RecordBatch, + probe_reader: Option, + depth: usize, +) -> Result> { + ctx.spill_metrics.record_repartition(); + + let child_depth = depth + 1; + let child_seed = seed_for_depth(child_depth); + let child_left = + RowPartitioner::with_seed(ctx.left_keys.to_vec(), ctx.num_sub, child_seed); + let child_right = + RowPartitioner::with_seed(ctx.right_keys.to_vec(), ctx.num_sub, child_seed); + + // Spill the build side into finer sub-buckets, then release it before + // recursing so only one sub-bucket's build side is ever resident. + let mut build_writer = + JoinSpillWriter::new(Arc::clone(ctx.left_schema), Arc::clone(ctx.disk_manager)); + for pb in child_left.partition(&build)? { + build_writer.append(pb.bucket, &pb.batch)?; + ctx.spill_metrics.record(&pb.batch); + } + build_writer.finish()?; + drop(build); + + // Spill the probe side into finer sub-buckets with the SAME seed, so rows + // with equal keys co-locate with their build rows. + let mut probe_writer = + JoinSpillWriter::new(Arc::clone(ctx.right_schema), Arc::clone(ctx.disk_manager)); + if let Some(reader) = probe_reader { + for probe_batch in reader { + let probe_batch = probe_batch?; + for pb in child_right.partition(&probe_batch)? { + probe_writer.append(pb.bucket, &pb.batch)?; + ctx.spill_metrics.record(&pb.batch); + } + } + } + probe_writer.finish()?; + + let mut out = Vec::new(); + for b in 0..ctx.num_sub { + let build_batches: Vec = match build_writer.reader(b)? { + Some(r) => r.collect::>()?, + None => continue, // no build rows -> inner join yields nothing }; - Ok(Some(ProbeTable::build( - part.batch, - &part.hashes, - &self.left_keys, - )?)) + if build_batches.is_empty() { + continue; + } + let sub_build = concat_batches(ctx.left_schema, build_batches.iter())?; + let sub_probe = probe_writer.reader(b)?; + out.extend(drain_bucket_recursive( + ctx, + reservation, + sub_build, + sub_probe, + child_depth, + )?); } + Ok(out) +} + +/// Returns `true` iff `build` contains more than one distinct join-key value. +/// Keys are compared via Arrow's row encoding (the same equality basis the +/// probe table uses), comparing every row to row 0; a bucket with a single +/// distinct key is irreducible and cannot be re-partitioned further. +fn has_multiple_distinct_keys( + build: &RecordBatch, + keys: &[PhysicalExprRef], +) -> Result { + let num_rows = build.num_rows(); + if num_rows <= 1 { + return Ok(false); + } + let key_arrays: Vec = keys + .iter() + .map(|expr| expr.evaluate(build)?.into_array(num_rows)) + .collect::>()?; + let fields: Vec = key_arrays + .iter() + .map(|a| SortField::new(a.data_type().clone())) + .collect(); + let converter = RowConverter::new(fields)?; + let rows = converter.convert_columns(&key_arrays)?; + let first = rows.row(0); + Ok((1..num_rows).any(|i| rows.row(i) != first)) } /// Turns `JoinState` into the output stream: emit any queued output, then in From 5dcab97f1201e41de2d5454ff8d2dae213b4ae44 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 11:59:26 -0600 Subject: [PATCH 08/15] feat(core): serialize SpillingHashJoinExec through the extension codec Add a SpillingHashJoinExecNode proto message and codec arms so SpillingHashJoinExec survives the scheduler->executor plan round trip. Following the ChaosExec convention, children are not embedded in the proto message; datafusion-proto decodes them and passes them through try_decode's inputs parameter. Join keys are serialized via the same physical-expr helpers already used for shuffle hash-repartition exprs. --- ballista/core/proto/ballista.proto | 14 ++ ballista/core/src/serde/generated/ballista.rs | 25 ++- ballista/core/src/serde/mod.rs | 172 +++++++++++++++++- 3 files changed, 209 insertions(+), 2 deletions(-) diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index c114a72e9a..6d178bc457 100644 --- a/ballista/core/proto/ballista.proto +++ b/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 diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 31708a701e..8a12faf053 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5" + tags = "1, 2, 3, 4, 5, 6" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -51,6 +51,8 @@ pub mod ballista_physical_plan_node { SortShuffleWriter(super::SortShuffleWriterExecNode), #[prost(message, tag = "5")] ChaosExec(super::ChaosExecNode), + #[prost(message, tag = "6")] + SpillingHashJoin(super::SpillingHashJoinExecNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -62,6 +64,27 @@ pub struct ChaosExecNode { #[prost(uint64, tag = "3")] pub seed: u64, } +/// 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. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SpillingHashJoinExecNode { + #[prost(message, repeated, tag = "1")] + pub left_keys: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalExprNode, + >, + #[prost(message, repeated, tag = "2")] + pub right_keys: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalExprNode, + >, + /// 0 = Partitioned, 1 = CollectLeft (v1 only emits 0) + #[prost(uint32, tag = "3")] + pub partition_mode: u32, + #[prost(uint64, tag = "4")] + pub num_sub_partitions: u64, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ShuffleWriterExecNode { /// TODO it seems redundant to provide job and stage id here since we also have them diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 213a600a78..3e11d1c9c5 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -26,14 +26,18 @@ use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; use datafusion::logical_expr::Extension; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_plan::joins::PartitionMode; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_proto::logical_plan::file_formats::{ ArrowLogicalExtensionCodec, AvroLogicalExtensionCodec, CsvLogicalExtensionCodec, JsonLogicalExtensionCodec, ParquetLogicalExtensionCodec, }; +use datafusion_proto::physical_plan::from_proto::parse_physical_exprs; use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning; use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; use datafusion_proto::physical_plan::to_proto::serialize_partitioning; +use datafusion_proto::physical_plan::to_proto::serialize_physical_exprs; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext, @@ -55,7 +59,7 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ ChaosExec, CoalescePlan, PartitionGroup, ShuffleReaderExec, ShuffleWriterExec, - SortShuffleWriterExec, UnresolvedShuffleExec, + SortShuffleWriterExec, SpillingHashJoinExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -550,6 +554,54 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { Some(chaos_exec.seed), )?)) } + PhysicalPlanType::SpillingHashJoin(spilling_hash_join) => { + let (left, right) = match inputs { + [left, right] => (left.clone(), right.clone()), + _ => { + return Err(DataFusionError::Internal(format!( + "SpillingHashJoinExec expects exactly 2 inputs, got {}", + inputs.len() + ))); + } + }; + let left_keys = parse_physical_exprs( + &spilling_hash_join.left_keys, + &decode_ctx, + left.schema().as_ref(), + &converter, + )?; + let right_keys = parse_physical_exprs( + &spilling_hash_join.right_keys, + &decode_ctx, + right.schema().as_ref(), + &converter, + )?; + if left_keys.len() != right_keys.len() { + return Err(DataFusionError::Internal(format!( + "SpillingHashJoinExec left_keys/right_keys length mismatch: {} vs {}", + left_keys.len(), + right_keys.len() + ))); + } + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = + left_keys.into_iter().zip(right_keys).collect(); + let partition_mode = match spilling_hash_join.partition_mode { + 0 => PartitionMode::Partitioned, + 1 => PartitionMode::CollectLeft, + other => { + return Err(DataFusionError::Internal(format!( + "SpillingHashJoinExec unknown partition_mode: {other}" + ))); + } + }; + Ok(Arc::new(SpillingHashJoinExec::try_new( + left, + right, + on, + partition_mode, + spilling_hash_join.num_sub_partitions as usize, + )?)) + } } } @@ -728,6 +780,45 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let converter = DefaultPhysicalProtoConverter {}; + let (left_exprs, right_exprs): (Vec, Vec) = + exec.on().iter().cloned().unzip(); + let left_keys = serialize_physical_exprs( + &left_exprs, + self.default_codec.as_ref(), + &converter, + )?; + let right_keys = serialize_physical_exprs( + &right_exprs, + self.default_codec.as_ref(), + &converter, + )?; + let partition_mode = match exec.partition_mode() { + PartitionMode::Partitioned => 0, + PartitionMode::CollectLeft => 1, + other => { + return Err(DataFusionError::Internal(format!( + "SpillingHashJoinExec unsupported partition_mode for serialization: {other:?}" + ))); + } + }; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::SpillingHashJoin( + protobuf::SpillingHashJoinExecNode { + left_keys, + right_keys, + partition_mode, + num_sub_partitions: exec.num_sub_partitions() as u64, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode spilling hash join execution plan: {e:?}" + )) + })?; + Ok(()) } else { Err(DataFusionError::Internal(format!( "Unsupported plan node, name: [{}] ", @@ -1327,4 +1418,83 @@ mod test { assert_eq!(decoded_exec.upstream_partition_count, 4); assert_eq!(decoded_exec.partition.len(), 1); } + + // Round-trips a `SpillingHashJoinExec` with a multi-key `on` (exercising + // key ordering/pairing) through `BallistaPhysicalExtensionCodec`. Children + // are not carried in the proto message (see `SpillingHashJoinExecNode`'s + // doc comment); they are supplied directly to `try_decode` via `inputs`, + // matching how `ChaosExec` and the other single/no-child nodes above are + // tested. The rendered plan string must be byte-identical before and + // after, which also proves `on`, `partition_mode`, and + // `num_sub_partitions` all survived the round trip. + #[tokio::test] + async fn spilling_hash_join_exec_roundtrip() { + use crate::execution_plans::SpillingHashJoinExec; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::PhysicalExprRef; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::expressions::Column; + use datafusion::physical_plan::joins::PartitionMode; + + let schema = Arc::new(Schema::new(vec![ + Field::new("k1", DataType::Int32, false), + Field::new("k2", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let partitions: Vec> = (0..4).map(|_| vec![]).collect(); + + let left_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("left MemorySourceConfig"); + let right_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("right MemorySourceConfig"); + + let left: Arc = + Arc::new(DataSourceExec::new(Arc::new(left_source))); + let right: Arc = + Arc::new(DataSourceExec::new(Arc::new(right_source))); + + // Multi-key `on`, deliberately not in column-index order, so a + // left/right key mismatch or a mis-zip would be caught. + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![ + ( + Arc::new(Column::new("k2", 1)), + Arc::new(Column::new("k2", 1)), + ), + ( + Arc::new(Column::new("k1", 0)), + Arc::new(Column::new("k1", 0)), + ), + ]; + + let exec = Arc::new( + SpillingHashJoinExec::try_new( + left.clone(), + right.clone(), + on, + PartitionMode::Partitioned, + 16, + ) + .unwrap(), + ); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec + .try_encode(exec.clone() as Arc, &mut buf) + .unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded = codec.try_decode(&buf, &[left, right], &ctx).unwrap(); + + let before = format!( + "{}", + displayable(exec.as_ref() as &dyn ExecutionPlan).indent(false) + ); + let after = format!("{}", displayable(decoded.as_ref()).indent(false)); + assert_eq!(before, after); + } } From f075a485111369d0c2f61bc9cee8d03341e3437d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 12:03:34 -0600 Subject: [PATCH 09/15] test: distinguish left/right schemas in spilling hash join codec roundtrip The roundtrip test previously built left and right inputs from the same schema with identical Column exprs on both sides of each on pair, so a bug swapping left_keys/right_keys through encode/decode would still produce byte-identical Display output. Give left and right distinct schemas (l_* vs r_* column names) so a side swap or mis-pairing now changes the rendered on=[...] string and fails the test. --- ballista/core/src/serde/mod.rs | 35 +++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 3e11d1c9c5..5424a4d568 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -1427,6 +1427,12 @@ mod test { // tested. The rendered plan string must be byte-identical before and // after, which also proves `on`, `partition_mode`, and // `num_sub_partitions` all survived the round trip. + // + // `left` and `right` deliberately use distinct schemas (`l_*` vs `r_*` + // column names) so that a left/right key-side swap through encode/decode + // (e.g. encoding right's exprs into `left_keys`, or decoding `right_keys` + // against `left.schema()`) changes the rendered `on=[...]` string instead + // of going unnoticed by symmetry. #[tokio::test] async fn spilling_hash_join_exec_roundtrip() { use crate::execution_plans::SpillingHashJoinExec; @@ -1438,18 +1444,23 @@ mod test { use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::joins::PartitionMode; - let schema = Arc::new(Schema::new(vec![ - Field::new("k1", DataType::Int32, false), - Field::new("k2", DataType::Int32, false), - Field::new("v", DataType::Int32, false), + let left_schema = Arc::new(Schema::new(vec![ + Field::new("l_k1", DataType::Int32, false), + Field::new("l_k2", DataType::Int32, false), + Field::new("l_v", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("r_k1", DataType::Int32, false), + Field::new("r_k2", DataType::Int32, false), + Field::new("r_v", DataType::Int32, false), ])); let partitions: Vec> = (0..4).map(|_| vec![]).collect(); let left_source = - MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + MemorySourceConfig::try_new(&partitions, Arc::clone(&left_schema), None) .expect("left MemorySourceConfig"); let right_source = - MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + MemorySourceConfig::try_new(&partitions, Arc::clone(&right_schema), None) .expect("right MemorySourceConfig"); let left: Arc = @@ -1458,15 +1469,17 @@ mod test { Arc::new(DataSourceExec::new(Arc::new(right_source))); // Multi-key `on`, deliberately not in column-index order, so a - // left/right key mismatch or a mis-zip would be caught. + // left/right key mismatch or a mis-zip would be caught. Left and + // right columns have distinct names, so a whole-side swap (not just + // intra-list reordering) also changes the `Display` output. let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = vec![ ( - Arc::new(Column::new("k2", 1)), - Arc::new(Column::new("k2", 1)), + Arc::new(Column::new("l_k2", 1)), + Arc::new(Column::new("r_k2", 1)), ), ( - Arc::new(Column::new("k1", 0)), - Arc::new(Column::new("k1", 0)), + Arc::new(Column::new("l_k1", 0)), + Arc::new(Column::new("r_k1", 0)), ), ]; From 0347aabebe6fa65132f3668d4ee569d9810d294e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 12:15:01 -0600 Subject: [PATCH 10/15] feat(scheduler): add rule substituting spilling hash join for eligible joins Add SpillingHashJoinRule, a PhysicalOptimizerRule that swaps eligible HashJoinExec nodes (Inner join, Partitioned mode, no projection, no residual filter) for SpillingHashJoinExec. The rule is gated by ballista.execution.spilling_hash_join.enabled read from the BallistaConfig extension on ConfigOptions, defaulting to disabled when the extension is absent, and uses ballista.execution.spilling_hash_join.partitions (default 16) for the number of spill sub-partitions. --- .../scheduler/src/physical_optimizer/mod.rs | 3 + .../physical_optimizer/spilling_hash_join.rs | 330 ++++++++++++++++++ 2 files changed, 333 insertions(+) create mode 100644 ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs diff --git a/ballista/scheduler/src/physical_optimizer/mod.rs b/ballista/scheduler/src/physical_optimizer/mod.rs index 4252d95492..a322950cd3 100644 --- a/ballista/scheduler/src/physical_optimizer/mod.rs +++ b/ballista/scheduler/src/physical_optimizer/mod.rs @@ -25,3 +25,6 @@ pub mod filter_pushdown; // specific cases. it has been used in static // execution graph only. pub mod join_selection; +// substitutes eligible `HashJoinExec` nodes with `SpillingHashJoinExec`, +// gated by `ballista.execution.spilling_hash_join.enabled`. +pub mod spilling_hash_join; diff --git a/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs new file mode 100644 index 0000000000..0701e69b60 --- /dev/null +++ b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs @@ -0,0 +1,330 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`SpillingHashJoinRule`] substitutes eligible `HashJoinExec` nodes with +//! [`SpillingHashJoinExec`], a hash join whose build side spills sub-partitions +//! to disk under memory pressure instead of requiring the whole build side to +//! fit in memory at once. +//! +//! # Eligibility (v1, strict) +//! +//! A `HashJoinExec` is substituted only when all of the following hold: +//! - `join_type() == JoinType::Inner` +//! - `*partition_mode() == PartitionMode::Partitioned` +//! - `projection` is `None` (no output projection folded into the join) +//! - `filter()` is `None` (no residual `JoinFilter`) +//! +//! `SpillingHashJoinExec` itself only supports this exact shape (see its +//! constructor), so the predicate here is not just a performance heuristic — +//! it is the full set of plans the replacement operator can represent. Any +//! join that does not match is left untouched. +//! +//! # Gating +//! +//! The rule is a no-op unless `ballista.execution.spilling_hash_join.enabled` +//! is set on the `BallistaConfig` extension carried by the `ConfigOptions` +//! passed to `optimize`. This mirrors how `CoalescePartitionsRule` (see +//! `crate::state::aqe::optimizer_rule::coalesce_partitions`) reads Ballista +//! options: `config.extensions.get::()`, defaulting to +//! `BallistaConfig::default()` (which has the flag off) when the extension is +//! absent. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use ballista_core::execution_plans::SpillingHashJoinExec; +use datafusion::common::JoinType; +use datafusion::common::Result; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::config::ConfigOptions; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; + +/// Physical optimizer rule that substitutes eligible `HashJoinExec` nodes with +/// `SpillingHashJoinExec`. +/// +/// See module docs for the eligibility predicate and the config-gating +/// approach. +#[derive(Debug, Default)] +pub struct SpillingHashJoinRule {} + +impl PhysicalOptimizerRule for SpillingHashJoinRule { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + let ballista_config = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + + if !ballista_config.spilling_hash_join_enabled() { + return Ok(plan); + } + + let num_sub_partitions = ballista_config.spilling_hash_join_partitions(); + + plan.transform_up(|node| { + let Some(hj) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + + if !is_eligible(hj) { + return Ok(Transformed::no(node)); + } + + let spilling = SpillingHashJoinExec::try_new( + Arc::clone(hj.left()), + Arc::clone(hj.right()), + hj.on().to_vec(), + PartitionMode::Partitioned, + num_sub_partitions, + )?; + + Ok(Transformed::yes( + Arc::new(spilling) as Arc + )) + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "SpillingHashJoinRule" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Whether `hj` matches the strict v1 shape `SpillingHashJoinExec` can +/// represent: inner, partitioned, no projection, no residual filter. +fn is_eligible(hj: &HashJoinExec) -> bool { + *hj.join_type() == JoinType::Inner + && *hj.partition_mode() == PartitionMode::Partitioned + && hj.projection.is_none() + && hj.filter().is_none() +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::NullEquality; + use datafusion::config::{ConfigOptions, ExtensionOptions}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::PhysicalExprRef; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::expressions::Column; + use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; + + /// Two 4-partition `DataSourceExec` inputs with schema `(k Int32, v + /// Int32)`, plus the equijoin pair `on = [(k_left, k_right)]`. + #[allow(clippy::type_complexity)] + fn two_inputs() -> ( + Arc, + Arc, + Vec<(PhysicalExprRef, PhysicalExprRef)>, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let partitions: Vec> = (0..4).map(|_| vec![]).collect(); + + let left_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("left MemorySourceConfig"); + let right_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("right MemorySourceConfig"); + + let left: Arc = + Arc::new(DataSourceExec::new(Arc::new(left_source))); + let right: Arc = + Arc::new(DataSourceExec::new(Arc::new(right_source))); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = + vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k", 0)))]; + + (left, right, on) + } + + fn hash_join( + join_type: JoinType, + partition_mode: PartitionMode, + filter: Option, + ) -> Arc { + let (left, right, on) = two_inputs(); + Arc::new( + HashJoinExec::try_new( + left, + right, + on, + filter, + &join_type, + None, + partition_mode, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) + } + + /// A trivial `JoinFilter` over `k_left > v_right`, just enough to make + /// `filter()` return `Some(..)`. + fn some_filter() -> JoinFilter { + let intermediate_schema = Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ]); + let expr = Arc::new(Column::new_with_schema("k", &intermediate_schema).unwrap()) + as PhysicalExprRef; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: datafusion::common::JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: datafusion::common::JoinSide::Right, + }, + ]; + JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)) + } + + fn config_with_spilling(enabled: bool) -> ConfigOptions { + let mut bc = BallistaConfig::default(); + bc.set( + "execution.spilling_hash_join.enabled", + if enabled { "true" } else { "false" }, + ) + .unwrap(); + let mut config = ConfigOptions::default(); + config.extensions.insert(bc); + config + } + + fn plan_string(plan: &Arc) -> String { + format!("{}", displayable(plan.as_ref()).indent(false)) + } + + /// Whether the rendered plan contains a bare `HashJoinExec` node (as + /// opposed to a `SpillingHashJoinExec` node, whose name has + /// `HashJoinExec` as a trailing substring and would otherwise produce a + /// false-positive `contains("HashJoinExec")` match). + fn has_bare_hash_join(s: &str) -> bool { + s.lines() + .any(|line| line.trim_start().starts_with("HashJoinExec")) + } + + #[test] + fn substitutes_eligible_inner_partitioned() { + let plan = hash_join(JoinType::Inner, PartitionMode::Partitioned, None); + let cfg = config_with_spilling(true); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(s.contains("SpillingHashJoinExec"), "{s}"); + assert!(!has_bare_hash_join(&s), "{s}"); + } + + #[test] + fn leaves_left_outer_untouched() { + let plan = hash_join(JoinType::Left, PartitionMode::Partitioned, None); + let cfg = config_with_spilling(true); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } + + #[test] + fn no_substitution_when_flag_off() { + let plan = hash_join(JoinType::Inner, PartitionMode::Partitioned, None); + let cfg = config_with_spilling(false); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } + + #[test] + fn leaves_filter_present_untouched() { + let plan = hash_join( + JoinType::Inner, + PartitionMode::Partitioned, + Some(some_filter()), + ); + let cfg = config_with_spilling(true); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } + + #[test] + fn leaves_collect_left_untouched() { + let plan = hash_join(JoinType::Inner, PartitionMode::CollectLeft, None); + let cfg = config_with_spilling(true); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } + + #[test] + fn config_extension_absent_defaults_to_off() { + // No `BallistaConfig` extension inserted at all: the rule must fall + // back to `BallistaConfig::default()`, whose flag is off, rather than + // panicking or substituting unconditionally. + let plan = hash_join(JoinType::Inner, PartitionMode::Partitioned, None); + let cfg = ConfigOptions::default(); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } +} From 72dce15359e245c38049beced96ddd0843d0a606 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 12:27:42 -0600 Subject: [PATCH 11/15] feat(scheduler): apply spilling hash join rule in default planner --- ballista/scheduler/src/planner.rs | 129 ++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 0128540c5b..e00b675d57 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -51,6 +51,7 @@ use log::debug; use crate::physical_optimizer::join_selection::{ collect_left_broadcast_safe, should_swap_join_order, }; +use crate::physical_optimizer::spilling_hash_join::SpillingHashJoinRule; type PartialQueryStageResult = (Arc, Vec>); @@ -112,6 +113,13 @@ impl DistributedPlanner for DefaultDistributedPlanner { config: &ConfigOptions, ) -> Result>> { debug!("Planning query stages for job: [{job_id}]"); + // Substitute eligible `HashJoinExec` nodes with `SpillingHashJoinExec` + // (no-op unless `ballista.execution.spilling_hash_join.enabled` is set + // on the `BallistaConfig` extension carried by `config`). Applied once, + // over the whole plan, before any stage splitting so the substituted + // node is what gets serialized to executors. + let execution_plan = + SpillingHashJoinRule::default().optimize(execution_plan, config)?; let (new_plan, mut stages) = self.plan_query_stages_internal(job_id, execution_plan, config)?; stages.push(create_shuffle_writer_with_config( @@ -744,10 +752,13 @@ mod test { use crate::assert_plan; use crate::planner::{DefaultDistributedPlanner, DistributedPlanner}; use crate::test_utils::datafusion_test_context; + use ballista_core::config::BallistaConfig; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{SortShuffleWriterExec, UnresolvedShuffleExec}; use ballista_core::serde::BallistaCodec; use datafusion::arrow::compute::SortOptions; + use datafusion::config::ConfigOptions; + use datafusion::physical_plan::joins::PartitionMode; use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::Column; @@ -1209,6 +1220,124 @@ order by Ok(()) } + /// An `Inner`/`Partitioned` `HashJoinExec` over two 2-partition inputs, + /// eligible for `SpillingHashJoinRule` substitution. Built directly + /// (bypassing SQL/`JoinSelection`) so the partition mode is deterministic. + fn inner_partitioned_hash_join_plan() -> Arc { + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinType, NullEquality}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::PhysicalExprRef; + use datafusion::physical_plan::expressions::Column; + + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let partitions: Vec> = (0..2).map(|_| vec![]).collect(); + + let left_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("left MemorySourceConfig"); + let right_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None) + .expect("right MemorySourceConfig"); + let left: Arc = + Arc::new(DataSourceExec::new(Arc::new(left_source))); + let right: Arc = + Arc::new(DataSourceExec::new(Arc::new(right_source))); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = + vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k", 0)))]; + + Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) + } + + /// A bare `ConfigOptions` carrying a `BallistaConfig` extension with + /// `ballista.execution.spilling_hash_join.enabled` set as requested, + /// mirroring how a real per-job `SessionConfig::options()` looks once + /// `BallistaConfig` has been upgraded onto it (see + /// `ballista_core::extension::SessionConfigExt`). + fn config_with_spilling_hash_join(enabled: bool) -> ConfigOptions { + use datafusion::config::ExtensionOptions; + + let mut bc = BallistaConfig::default(); + bc.set( + "execution.spilling_hash_join.enabled", + if enabled { "true" } else { "false" }, + ) + .unwrap(); + let mut config = ConfigOptions::default(); + config.extensions.insert(bc); + config + } + + /// Drives the actual `DefaultDistributedPlanner::plan_query_stages` entry + /// point (the real, production code path used for every job when AQE is + /// off) with an eligible join and the flag on: the resulting stage plan + /// must contain `SpillingHashJoinExec`. + #[tokio::test] + async fn planner_substitutes_spilling_hash_join_when_enabled() + -> Result<(), BallistaError> { + let plan = inner_partitioned_hash_join_plan(); + let config = config_with_spilling_hash_join(true); + + let mut planner = DefaultDistributedPlanner::new(); + let job_uuid = Uuid::new_v4(); + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &config)?; + + let rendered: String = stages + .iter() + .map(|s| format!("{}", displayable(s.as_ref()).indent(false))) + .collect(); + assert!(rendered.contains("SpillingHashJoinExec"), "{rendered}"); + + Ok(()) + } + + /// Same plan through the same entry point with the flag off: the plan + /// must be left as a plain `HashJoinExec`, with no substitution. + #[tokio::test] + async fn planner_leaves_hash_join_untouched_when_disabled() + -> Result<(), BallistaError> { + let plan = inner_partitioned_hash_join_plan(); + let config = config_with_spilling_hash_join(false); + + let mut planner = DefaultDistributedPlanner::new(); + let job_uuid = Uuid::new_v4(); + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &config)?; + + let rendered: String = stages + .iter() + .map(|s| format!("{}", displayable(s.as_ref()).indent(false))) + .collect(); + assert!(!rendered.contains("SpillingHashJoinExec"), "{rendered}"); + assert!( + rendered + .lines() + .any(|line| line.trim_start().starts_with("HashJoinExec")), + "{rendered}" + ); + + Ok(()) + } + // Across every expressible join type and both build-side orientations, any // CollectLeft join the static planner promotes must be broadcast-safe: a // CollectLeft join replicates the build (left) side to every probe task, so From 55894bccefb52309bdd3a40fa0da48a5e00c6285 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 12:34:45 -0600 Subject: [PATCH 12/15] refactor(scheduler): run spilling-hash-join substitution after broadcast promotion Broadcast promotion only matches HashJoinExec, so a join substituted to SpillingHashJoinExec up front could never be promoted to a shuffle-avoiding CollectLeft broadcast. Move the substitution to run per node immediately after broadcast promotion inside plan_query_stages_internal, so a small build side is offered a broadcast first and only joins left as Inner+Partitioned HashJoinExec are substituted. Factor the single-node eligibility and substitution into a shared maybe_substitute_spilling_hash_join function reused by both SpillingHashJoinRule and the distributed planner, keeping one definition of eligibility. --- .../physical_optimizer/spilling_hash_join.rs | 69 +++++++++---- ballista/scheduler/src/planner.rs | 96 +++++++++++++++++-- 2 files changed, 137 insertions(+), 28 deletions(-) diff --git a/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs index 0701e69b60..25b7d87da2 100644 --- a/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs +++ b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs @@ -79,28 +79,14 @@ impl PhysicalOptimizerRule for SpillingHashJoinRule { return Ok(plan); } - let num_sub_partitions = ballista_config.spilling_hash_join_partitions(); - plan.transform_up(|node| { - let Some(hj) = node.downcast_ref::() else { - return Ok(Transformed::no(node)); - }; - - if !is_eligible(hj) { - return Ok(Transformed::no(node)); + let substituted = + maybe_substitute_spilling_hash_join(Arc::clone(&node), config)?; + if Arc::ptr_eq(&substituted, &node) { + Ok(Transformed::no(node)) + } else { + Ok(Transformed::yes(substituted)) } - - let spilling = SpillingHashJoinExec::try_new( - Arc::clone(hj.left()), - Arc::clone(hj.right()), - hj.on().to_vec(), - PartitionMode::Partitioned, - num_sub_partitions, - )?; - - Ok(Transformed::yes( - Arc::new(spilling) as Arc - )) }) .map(|t| t.data) } @@ -114,6 +100,49 @@ impl PhysicalOptimizerRule for SpillingHashJoinRule { } } +/// Substitutes a single `HashJoinExec` node with a `SpillingHashJoinExec` when +/// it is eligible and the feature flag is on. Operates on the given node only +/// (non-recursive), returning `plan` unchanged when the flag is off, the node +/// is not a `HashJoinExec`, or the join does not match the eligibility +/// predicate. +/// +/// This is the single source of truth for spilling-join eligibility and +/// substitution. It is shared by [`SpillingHashJoinRule`] (which applies it +/// bottom-up over a whole tree) and the distributed planner (which applies it +/// per node, immediately after broadcast promotion). +pub(crate) fn maybe_substitute_spilling_hash_join( + plan: Arc, + config: &ConfigOptions, +) -> Result> { + let ballista_config = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + + if !ballista_config.spilling_hash_join_enabled() { + return Ok(plan); + } + + let Some(hj) = plan.downcast_ref::() else { + return Ok(plan); + }; + + if !is_eligible(hj) { + return Ok(plan); + } + + let spilling = SpillingHashJoinExec::try_new( + Arc::clone(hj.left()), + Arc::clone(hj.right()), + hj.on().to_vec(), + PartitionMode::Partitioned, + ballista_config.spilling_hash_join_partitions(), + )?; + + Ok(Arc::new(spilling) as Arc) +} + /// Whether `hj` matches the strict v1 shape `SpillingHashJoinExec` can /// represent: inner, partitioned, no projection, no residual filter. fn is_eligible(hj: &HashJoinExec) -> bool { diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index e00b675d57..4c37f745b5 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -51,7 +51,6 @@ use log::debug; use crate::physical_optimizer::join_selection::{ collect_left_broadcast_safe, should_swap_join_order, }; -use crate::physical_optimizer::spilling_hash_join::SpillingHashJoinRule; type PartialQueryStageResult = (Arc, Vec>); @@ -113,13 +112,6 @@ impl DistributedPlanner for DefaultDistributedPlanner { config: &ConfigOptions, ) -> Result>> { debug!("Planning query stages for job: [{job_id}]"); - // Substitute eligible `HashJoinExec` nodes with `SpillingHashJoinExec` - // (no-op unless `ballista.execution.spilling_hash_join.enabled` is set - // on the `BallistaConfig` extension carried by `config`). Applied once, - // over the whole plan, before any stage splitting so the substituted - // node is what gets serialized to executors. - let execution_plan = - SpillingHashJoinRule::default().optimize(execution_plan, config)?; let (new_plan, mut stages) = self.plan_query_stages_internal(job_id, execution_plan, config)?; stages.push(create_shuffle_writer_with_config( @@ -146,6 +138,18 @@ impl DefaultDistributedPlanner { // Apply broadcast-join promotion before recursing. let execution_plan = Self::maybe_promote_to_broadcast(execution_plan, config)?; + // Broadcast promotion has already had its chance to turn a small-build + // join into a `CollectLeft` broadcast. Any join still left as an + // `Inner`/`Partitioned` `HashJoinExec` (a large build side that was not + // broadcast) is now substituted with `SpillingHashJoinExec` when the + // feature flag is set. A promoted `CollectLeft` join is not + // `Partitioned`, so it is never eligible — broadcast wins for small + // build sides. + let execution_plan = crate::physical_optimizer::spilling_hash_join::maybe_substitute_spilling_hash_join( + execution_plan, + config, + )?; + // recurse down and replace children if execution_plan.children().is_empty() { return Ok((execution_plan, vec![])); @@ -1306,6 +1310,82 @@ order by .map(|s| format!("{}", displayable(s.as_ref()).indent(false))) .collect(); assert!(rendered.contains("SpillingHashJoinExec"), "{rendered}"); + // `SpillingHashJoinExec` ends in `HashJoinExec`, so a plain + // `contains("HashJoinExec")` would false-pass. Assert no *bare* + // `HashJoinExec` node survives the substitution. + assert!( + !rendered + .lines() + .any(|line| line.trim_start().starts_with("HashJoinExec")), + "{rendered}" + ); + + Ok(()) + } + + /// Broadcast-first ordering: a small-build-side `Inner`/`Partitioned` + /// `HashJoinExec` under the Ballista broadcast threshold, planned with BOTH + /// broadcast promotion AND spilling-hash-join substitution enabled, must be + /// promoted to a `CollectLeft` broadcast join and NOT substituted with + /// `SpillingHashJoinExec`. Broadcast promotion runs first and turns the join + /// `CollectLeft`, which is not `Partitioned` and therefore ineligible for + /// substitution — so broadcast wins the small-join case. + #[tokio::test] + async fn planner_broadcast_wins_over_spilling_for_small_build() + -> Result<(), BallistaError> { + use datafusion::config::ExtensionOptions; + use datafusion::physical_plan::joins::PartitionMode; + + // DF's own `hash_join_single_partition_threshold` is 0 in this fixture, + // so DataFusion leaves the join `Partitioned`; Ballista's broadcast + // threshold (10 MB) is what promotes the small side to `CollectLeft`. + let (ctx, mut options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; + // Turn on spilling substitution alongside broadcast promotion. + options + .extensions + .get_mut::() + .expect("BallistaConfig extension present") + .set("execution.spilling_hash_join.enabled", "true") + .unwrap(); + + let df = ctx + .sql("select count(*) from big join small on big.k = small.k") + .await?; + let plan = df.into_optimized_plan()?; + let plan = ctx.state().create_physical_plan(&plan).await?; + + let mut planner = DefaultDistributedPlanner::new(); + let job_uuid = Uuid::new_v4(); + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; + + let rendered: String = stages + .iter() + .map(|s| format!("{}", displayable(s.as_ref()).indent(false))) + .collect(); + // Broadcast won: no spilling substitution happened. + assert!( + !rendered.contains("SpillingHashJoinExec"), + "broadcast should win over spilling for a small build side\n{rendered}" + ); + + // And the surviving join is a `CollectLeft` broadcast join. + let mut found_broadcast_join = false; + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if let Some(hj) = node.downcast_ref::() { + assert_eq!(*hj.partition_mode(), PartitionMode::CollectLeft); + found_broadcast_join = true; + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + assert!( + found_broadcast_join, + "expected a CollectLeft broadcast HashJoinExec\n{rendered}" + ); Ok(()) } From d97b7a5f7dcbe81b4d13dab46da7a3272a210a68 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 12:44:31 -0600 Subject: [PATCH 13/15] feat(scheduler): apply spilling hash join rule in AQE planner --- ballista/scheduler/src/state/aqe/planner.rs | 16 ++ ballista/scheduler/src/state/aqe/test/mod.rs | 2 + .../src/state/aqe/test/spilling_hash_join.rs | 243 ++++++++++++++++++ 3 files changed, 261 insertions(+) create mode 100644 ballista/scheduler/src/state/aqe/test/spilling_hash_join.rs diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 67d303761a..806f56a134 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. use crate::physical_optimizer::filter_pushdown::FilterPushdown; +use crate::physical_optimizer::spilling_hash_join::SpillingHashJoinRule; use crate::state::aqe::adapter::BallistaAdapter; use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; @@ -370,6 +371,21 @@ impl AdaptivePlanner { // that would arise if the rule walked the entire residual // plan in `default_optimizers()`. let plan = CoalescePartitionsRule.optimize(plan, config)?; + // Broadcast-first ordering, mirroring the static planner: + // AQE's dynamic-join / broadcast-promotion machinery has + // already run (`DelayJoinSelectionRule` + + // `SelectJoinRule` in the optimizer chain), so any join + // still left as a plain `Inner`/`Partitioned` + // `HashJoinExec` is one AQE chose not to broadcast. Only + // those are substituted with `SpillingHashJoinExec` + // (flag-gated). A `DynamicJoinSelectionExec` stores its + // join parameters as fields and exposes only the join + // inputs as children, so a whole-tree `transform_up` + // cannot reach a wrapped join to clobber; and a + // `CollectLeft` broadcast join is not `Partitioned` and + // is therefore ineligible. + let plan = + SpillingHashJoinRule::default().optimize(plan, config)?; // adapt_to_ballista takes an job_id, we are passing a job_name. Need to transform to fix compiler. let job_id = self.job_name.clone().into(); BallistaAdapter::adapt_to_ballista(plan, &job_id, config) diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index a8640d3e30..894955a0a3 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -25,6 +25,8 @@ mod job_failure; mod join_selection; /// Tests if plan is going to be split to stages correctly mod plan_to_stages; +/// Covers spilling-hash-join substitution through the adaptive planner +mod spilling_hash_join; use ballista_core::config::BALLISTA_SHUFFLE_SORT_BASED_ENABLED; use ballista_core::extension::SessionConfigExt; diff --git a/ballista/scheduler/src/state/aqe/test/spilling_hash_join.rs b/ballista/scheduler/src/state/aqe/test/spilling_hash_join.rs new file mode 100644 index 0000000000..78ba96178f --- /dev/null +++ b/ballista/scheduler/src/state/aqe/test/spilling_hash_join.rs @@ -0,0 +1,243 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end coverage that the AQE (adaptive) planner substitutes eligible +//! `HashJoinExec` nodes with `SpillingHashJoinExec` when +//! `ballista.execution.spilling_hash_join.enabled` is set, using the SAME +//! broadcast-first ordering discipline as the static planner: the substitution +//! runs only after AQE's own dynamic-join / broadcast-promotion machinery has +//! resolved a join, so a plain leftover `Inner`/`Partitioned` `HashJoinExec` is +//! the only thing substituted. + +use crate::physical_optimizer::spilling_hash_join::SpillingHashJoinRule; +use crate::state::aqe::optimizer_rule::DelayJoinSelectionRule; +use crate::state::aqe::{ + planner::AdaptivePlanner, test::mock_partitions_with_statistics, +}; +use ballista_core::config::{BALLISTA_SPILLING_HASH_JOIN_ENABLED, BallistaConfig}; +use ballista_core::extension::SessionConfigExt; +use datafusion::{ + arrow::{ + array::{Int32Array, RecordBatch}, + datatypes::{DataType, Field, Schema}, + }, + common::config::ConfigOptions, + config::ExtensionOptions, + datasource::MemTable, + execution::{SessionStateBuilder, config::SessionConfig, context::SessionContext}, + physical_optimizer::PhysicalOptimizerRule, + physical_plan::{ExecutionPlan, displayable}, + physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}, +}; +use std::sync::Arc; + +fn make_table(schema: Arc) -> Arc { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0])), + Arc::new(Int32Array::from(vec![ + 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, + ])), + ], + ) + .unwrap(); + + Arc::new( + MemTable::try_new( + schema, + vec![ + vec![batch.clone()], + vec![batch.clone()], + vec![batch.clone()], + vec![batch], + ], + ) + .unwrap(), + ) +} + +/// A session that keeps a two-table equijoin as a plain `Inner`/`Partitioned` +/// `HashJoinExec` under AQE: `prefer_hash_join=true` picks hash over sort-merge, +/// and a broadcast threshold of `0` disables `CollectLeft` promotion so the join +/// is repartitioned rather than broadcast. `spilling` toggles the spilling flag. +fn make_ctx(spilling: bool) -> SessionContext { + let config = SessionConfig::new_with_ballista() + .set_bool("datafusion.optimizer.prefer_hash_join", true) + .set_u64( + "datafusion.optimizer.hash_join_single_partition_threshold", + 0, + ) + .set_u64( + "datafusion.optimizer.hash_join_single_partition_threshold_rows", + 0, + ) + .with_ballista_broadcast_join_threshold_bytes(0) + .set_str( + BALLISTA_SPILLING_HASH_JOIN_ENABLED, + if spilling { "true" } else { "false" }, + ) + .with_target_partitions(4) + .with_round_robin_repartition(false); + + let state = SessionStateBuilder::new_with_default_features() + .with_config(config) + .build(); + SessionContext::new_with_state(state) +} + +fn register_2tables(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + ])); + ctx.register_table("t1", make_table(Arc::clone(&schema))) + .unwrap(); + ctx.register_table("t2", make_table(schema)).unwrap(); +} + +fn plan_string(plan: &dyn ExecutionPlan) -> String { + format!("{}", displayable(plan).indent(false)) +} + +/// Whether the rendered plan contains a bare `HashJoinExec` node. A +/// `SpillingHashJoinExec` line ends with `HashJoinExec`, so a naive +/// `contains("HashJoinExec")` would false-positive; this matches only lines +/// whose (trimmed) node name starts with `HashJoinExec`. +fn has_bare_hash_join(s: &str) -> bool { + s.lines() + .any(|line| line.trim_start().starts_with("HashJoinExec")) +} + +/// Drives the AQE planner to the point where the two input shuffle stages have +/// finished and the join stage is runnable, then returns the rendered plan of +/// the runnable join stage handed to the executors. +async fn resolved_join_stage_string(spilling: bool) -> String { + let ctx = make_ctx(spilling); + register_2tables(&ctx); + + // `SELECT *` keeps every join-output column in natural order, so the + // physical planner folds no projection into the join — leaving a + // projection-less `HashJoinExec` that matches the strict eligibility shape. + let lp = ctx + .sql("SELECT * FROM t1 JOIN t2 ON t1.id = t2.id") + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + + let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()) + .await + .unwrap(); + + // First round: the two input scan stages become runnable. + let (stages, _) = planner.actionable_stages().unwrap(); + let stages = stages.unwrap(); + assert_eq!(2, stages.len(), "expected the two input scan stages"); + + // Finalise them so the dynamic join resolves to a Partitioned hash join. + planner + .finalise_stage_internal(0, mock_partitions_with_statistics()) + .unwrap(); + planner + .finalise_stage_internal(1, mock_partitions_with_statistics()) + .unwrap(); + + // Second round: the join stage is now runnable; its adapted plan is where + // the spilling substitution (if enabled) must appear. + let (stages, _) = planner.actionable_stages().unwrap(); + let stages = stages.unwrap(); + assert_eq!(1, stages.len(), "expected the single join stage"); + + plan_string(stages.first().unwrap().plan.as_ref()) +} + +/// Flag ON: a plain `Inner`/`Partitioned` `HashJoinExec` that AQE leaves +/// un-promoted must be substituted with `SpillingHashJoinExec` in the runnable +/// join stage. +#[tokio::test] +async fn aqe_planner_substitutes_when_enabled() { + let s = resolved_join_stage_string(true).await; + assert!(s.contains("SpillingHashJoinExec"), "{s}"); + assert!(!has_bare_hash_join(&s), "{s}"); +} + +/// Flag OFF: the same join is left as a bare `HashJoinExec`; no substitution. +#[tokio::test] +async fn aqe_planner_no_substitution_when_disabled() { + let s = resolved_join_stage_string(false).await; + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); +} + +/// Guards the exact clobbering risk called out for the AQE path: a join that is +/// still wrapped in `DynamicJoinSelectionExec` (i.e. AQE has NOT yet resolved +/// it) must never be substituted, even with the flag on. `DynamicJoinSelectionExec` +/// stores the join parameters as fields and exposes only the join *inputs* as +/// its children — it does not wrap a `HashJoinExec` — so a whole-tree +/// `transform_up` cannot reach an inner join to clobber. This proves the wrapper +/// and its (nonexistent) inner join both survive the rule. +#[tokio::test] +async fn dynamic_join_wrapper_not_substituted() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + ])); + let ctx = SessionContext::new(); + ctx.register_table("t1", make_table(Arc::clone(&schema))) + .unwrap(); + ctx.register_table("t2", make_table(schema)).unwrap(); + + let state = SessionStateBuilder::new() + .with_physical_optimizer_rules(vec![]) + .build(); + let lp = ctx + .sql("SELECT t1.id, t2.val FROM t1 JOIN t2 ON t1.id = t2.id") + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + let plan = DefaultPhysicalPlanner::default() + .create_physical_plan(&lp, &state) + .await + .unwrap(); + + // Wrap the join in a DynamicJoinSelectionExec (adaptive_join defaults on). + let dynamic = DelayJoinSelectionRule::default() + .optimize(plan, &ConfigOptions::default()) + .unwrap(); + let before = plan_string(dynamic.as_ref()); + assert!(before.contains("DynamicJoinSelectionExec"), "{before}"); + + // Apply the spilling rule with the flag ON over the whole tree. + let mut bc = BallistaConfig::default(); + bc.set("execution.spilling_hash_join.enabled", "true") + .unwrap(); + let mut cfg = ConfigOptions::default(); + cfg.extensions.insert(bc); + + let after = SpillingHashJoinRule::default() + .optimize(dynamic, &cfg) + .unwrap(); + let s = plan_string(after.as_ref()); + + // The dynamic wrapper survives and nothing was substituted: no HashJoinExec + // is exposed for the rule to reach, so no SpillingHashJoinExec is created. + assert!(s.contains("DynamicJoinSelectionExec"), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + assert!(!has_bare_hash_join(&s), "{s}"); +} From 9ea757cccd75e3221743ce644225b9d38f07f033 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 13:03:26 -0600 Subject: [PATCH 14/15] fix(spilling-hash-join): guard eligibility and construction, trim public api Close four correctness/API gaps found in whole-branch review: - Reject substitution when the source HashJoinExec uses NullEquality::NullEqualsNull (e.g. IS NOT DISTINCT FROM); the replacement operator hard-implements NullEqualsNothing semantics, so substituting here would silently drop null-key matches with no error. - Reject non-Partitioned partition_mode in try_new; execute() always runs Partitioned semantics, so a constructed CollectLeft instance would silently mis-execute. - Reject num_sub_partitions == 0 in try_new, before it can reach the build-side hash bucketing and panic on division by zero. - Narrow ProbeTable/assemble_output/RowPartitioner/PartitionedBatch out of the crate's public API; nothing outside the module reaches them through the mod-level re-export, only stream.rs's existing super:: paths. --- .../spilling_hash_join/exec.rs | 54 ++++++++++++++++++- .../execution_plans/spilling_hash_join/mod.rs | 10 +++- .../physical_optimizer/spilling_hash_join.rs | 49 ++++++++++++++++- 3 files changed, 108 insertions(+), 5 deletions(-) diff --git a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs index eda3670440..ae12c633f9 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/exec.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -27,11 +27,15 @@ // - join_type is always JoinType::Inner // - filter is always None // - projection is always None +// - partition_mode must be PartitionMode::Partitioned (execute() always +// runs Partitioned semantics; CollectLeft is rejected) +// - num_sub_partitions must be >= 1 (0 would divide by zero when bucketing +// the build side) use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{JoinType, NullEquality, Result, internal_err}; +use datafusion::common::{JoinType, NullEquality, Result, internal_err, not_impl_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::PhysicalExprRef; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; @@ -70,6 +74,12 @@ impl SpillingHashJoinExec { /// is constructed purely to borrow its output `schema()` and /// `properties()` (partitioning/equivalence info) — it is never /// executed. + /// + /// Returns an error if `partition_mode` is not `PartitionMode::Partitioned` + /// (v1 `execute()` always runs Partitioned semantics — `CollectLeft` is + /// accepted by the proto codec but not yet implemented here) or if + /// `num_sub_partitions` is `0` (the build-side hash bucketing divides by + /// this count and would panic). pub fn try_new( left: Arc, right: Arc, @@ -77,6 +87,17 @@ impl SpillingHashJoinExec { partition_mode: PartitionMode, num_sub_partitions: usize, ) -> Result { + if partition_mode != PartitionMode::Partitioned { + return not_impl_err!( + "SpillingHashJoinExec only supports PartitionMode::Partitioned in v1, got {partition_mode:?}" + ); + } + if num_sub_partitions == 0 { + return internal_err!( + "SpillingHashJoinExec requires num_sub_partitions >= 1, got 0" + ); + } + let hj = HashJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -274,6 +295,37 @@ mod tests { (left, right, on) } + #[test] + fn rejects_collect_left_partition_mode() { + let (left, right, on) = two_col_inputs(); + let err = SpillingHashJoinExec::try_new( + left, + right, + on, + PartitionMode::CollectLeft, + 16, + ) + .expect_err("CollectLeft is not implemented in v1 and must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("Partitioned"), + "error should name the unsupported mode, got: {msg}" + ); + } + + #[test] + fn rejects_zero_sub_partitions() { + let (left, right, on) = two_col_inputs(); + let err = + SpillingHashJoinExec::try_new(left, right, on, PartitionMode::Partitioned, 0) + .expect_err("num_sub_partitions == 0 must be rejected, not panic later"); + let msg = err.to_string(); + assert!( + msg.contains("num_sub_partitions"), + "error should name the bad argument, got: {msg}" + ); + } + #[test] fn scaffold_schema_and_name() { let (left, right, on) = two_col_inputs(); diff --git a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs index 8cd8534fc9..aa38b34981 100644 --- a/ballista/core/src/execution_plans/spilling_hash_join/mod.rs +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -21,6 +21,12 @@ mod partitioner; mod spill; mod stream; +// `SpillingHashJoinExec` is the only public surface of this module; it is +// re-exported further up in `execution_plans::mod`. pub use exec::SpillingHashJoinExec; -pub use hash_table::{ProbeTable, assemble_output}; -pub use partitioner::{PartitionedBatch, RowPartitioner}; + +// `hash_table` and `partitioner` are private submodules. Their `pub` items +// (`ProbeTable`, `assemble_output`, `RowPartitioner`, `PartitionedBatch`) are +// internal helpers consumed only by `stream.rs` via `super::hash_table::` / +// `super::partitioner::` — they are not re-exported here, so they are not +// part of `ballista_core`'s public API. diff --git a/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs index 25b7d87da2..4dc479ed6f 100644 --- a/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs +++ b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs @@ -27,6 +27,10 @@ //! - `*partition_mode() == PartitionMode::Partitioned` //! - `projection` is `None` (no output projection folded into the join) //! - `filter()` is `None` (no residual `JoinFilter`) +//! - `null_equality() == NullEquality::NullEqualsNothing` (the replacement +//! operator hard-implements "null keys never match"; substituting a join +//! with `NullEquality::NullEqualsNull` — e.g. `IS NOT DISTINCT FROM` — +//! would silently drop null-key matches) //! //! `SpillingHashJoinExec` itself only supports this exact shape (see its //! constructor), so the predicate here is not just a performance heuristic — @@ -48,6 +52,7 @@ use std::sync::Arc; use ballista_core::config::BallistaConfig; use ballista_core::execution_plans::SpillingHashJoinExec; use datafusion::common::JoinType; +use datafusion::common::NullEquality; use datafusion::common::Result; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::config::ConfigOptions; @@ -144,12 +149,14 @@ pub(crate) fn maybe_substitute_spilling_hash_join( } /// Whether `hj` matches the strict v1 shape `SpillingHashJoinExec` can -/// represent: inner, partitioned, no projection, no residual filter. +/// represent: inner, partitioned, no projection, no residual filter, and +/// `NullEqualsNothing` null semantics. fn is_eligible(hj: &HashJoinExec) -> bool { *hj.join_type() == JoinType::Inner && *hj.partition_mode() == PartitionMode::Partitioned && hj.projection.is_none() && hj.filter().is_none() + && hj.null_equality() == NullEquality::NullEqualsNothing } #[cfg(test)] @@ -201,6 +208,20 @@ mod tests { join_type: JoinType, partition_mode: PartitionMode, filter: Option, + ) -> Arc { + hash_join_with_null_equality( + join_type, + partition_mode, + filter, + NullEquality::NullEqualsNothing, + ) + } + + fn hash_join_with_null_equality( + join_type: JoinType, + partition_mode: PartitionMode, + filter: Option, + null_equality: NullEquality, ) -> Arc { let (left, right, on) = two_inputs(); Arc::new( @@ -212,7 +233,7 @@ mod tests { &join_type, None, partition_mode, - NullEquality::NullEqualsNothing, + null_equality, false, ) .unwrap(), @@ -340,6 +361,30 @@ mod tests { assert!(!s.contains("SpillingHashJoinExec"), "{s}"); } + #[test] + fn leaves_null_equals_null_untouched() { + // Otherwise-eligible (Inner, Partitioned, no projection, no filter), + // but `NullEquality::NullEqualsNull` (e.g. `IS NOT DISTINCT FROM`). + // `SpillingHashJoinExec` hard-implements `NullEqualsNothing` + // semantics, so substituting here would silently drop null-key + // matches — the join must be left untouched. + let plan = hash_join_with_null_equality( + JoinType::Inner, + PartitionMode::Partitioned, + None, + NullEquality::NullEqualsNull, + ); + let cfg = config_with_spilling(true); + + let out = SpillingHashJoinRule::default() + .optimize(plan, &cfg) + .unwrap(); + let s = plan_string(&out); + + assert!(has_bare_hash_join(&s), "{s}"); + assert!(!s.contains("SpillingHashJoinExec"), "{s}"); + } + #[test] fn config_extension_absent_defaults_to_off() { // No `BallistaConfig` extension inserted at all: the rule must fall From 45f48b12eedfca40350a12c170982d9a04e57a21 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 13:29:02 -0600 Subject: [PATCH 15/15] fix(python): sync python/Cargo.lock with ballista-core smallvec dependency --- python/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/python/Cargo.lock b/python/Cargo.lock index a0544afc65..86b14c430c 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -523,6 +523,7 @@ dependencies = [ "rand 0.10.2", "rustc_version", "serde", + "smallvec", "tokio", "tokio-stream", "tonic",