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/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/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()); + } } 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..ae12c633f9 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/exec.rs @@ -0,0 +1,738 @@ +// 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. `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): +// - 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, 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, +}; + +use super::stream::execute_join; + +/// 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` 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, + 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. + /// + /// 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, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + 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), + 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, + )?)) + } + + /// 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 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, Arc::clone(&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, + context, + &self.metrics, + partition, + )) + } + + 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::array::RecordBatch; + 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 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(); + 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}"); + } + + // --- 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); + } + + // --- 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) + } + + /// 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. + 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. 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/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 new file mode 100644 index 0000000000..aa38b34981 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/mod.rs @@ -0,0 +1,32 @@ +// 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; +mod hash_table; +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; + +// `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/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..c7e825b498 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/partitioner.rs @@ -0,0 +1,223 @@ +// 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; + +/// 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. +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::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(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); + } + } +} 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 new file mode 100644 index 0000000000..e224e62776 --- /dev/null +++ b/ballista/core/src/execution_plans/spilling_hash_join/stream.rs @@ -0,0 +1,832 @@ +// 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 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 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, 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 +// 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; + +use datafusion::arrow::array::{ArrayRef, RecordBatch}; +use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{Result, exec_err}; +use datafusion::execution::TaskContext; +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}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{Stream, StreamExt, TryStreamExt, stream}; + +use super::hash_table::{ProbeTable, assemble_output}; +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)] +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, + /// 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 { + 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), + repartition_count: MetricBuilder::new(metrics) + .counter("repartition_count", partition), + } + } + + /// Records one spilled batch. + fn record(&self, batch: &RecordBatch) { + 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`. +/// +/// `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, + right: SendableRecordBatchStream, + 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); + + 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)) +} + +/// 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 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); + + // 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, + &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, + right_schema, + output_schema, + num_sub: num_sub_partitions, + pool_limit, + disk_manager: Arc::clone(&runtime.disk_manager), + 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` 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, + 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_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; + } + } + } + } + } + + // 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`, 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`. +struct JoinState { + phase: Phase, + right: SendableRecordBatchStream, + 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>, + 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, +} + +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(()) + } + + /// 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(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; 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; + } + } + continue; + } + + // Load the next spilled bucket. + let Some(bucket) = self.drain_buckets.pop_front() else { + return Ok(false); + }; + + // 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 + }; + 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); + } + } + } + } +} + +/// 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, + )?); + } + } + 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 + }; + 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 +/// 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)), + }, + } + } + }) +} 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..5424a4d568 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,96 @@ 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. + // + // `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; + 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 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(&left_schema), None) + .expect("left MemorySourceConfig"); + let right_source = + MemorySourceConfig::try_new(&partitions, Arc::clone(&right_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. 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("l_k2", 1)), + Arc::new(Column::new("r_k2", 1)), + ), + ( + Arc::new(Column::new("l_k1", 0)), + Arc::new(Column::new("r_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); + } } 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..4dc479ed6f --- /dev/null +++ b/ballista/scheduler/src/physical_optimizer/spilling_hash_join.rs @@ -0,0 +1,404 @@ +// 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`) +//! - `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 — +//! 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::NullEquality; +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); + } + + plan.transform_up(|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)) + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "SpillingHashJoinRule" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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, 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)] +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 { + 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( + HashJoinExec::try_new( + left, + right, + on, + filter, + &join_type, + None, + partition_mode, + null_equality, + 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 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 + // 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}"); + } +} diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 0128540c5b..4c37f745b5 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -138,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![])); @@ -744,10 +756,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 +1224,200 @@ 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}"); + // `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(()) + } + + /// 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 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}"); +} 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",