diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index d71af206c78d5..5081d4a8ab1e6 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -552,6 +552,32 @@ config_namespace! { /// (i.e., if there's no `DiskManager` configured). pub sort_spill_reservation_bytes: usize, default = 10 * 1024 * 1024 + /// When enabled, a hash join whose build side exceeds its memory budget + /// partitions both sides to disk (grace hash join) and joins partition + /// pairs within the budget instead of failing with a resources-exhausted + /// error. The in-memory fast path is unchanged; spilling engages only at + /// the moment the build-side memory reservation first fails. + /// + /// Requires a configured `DiskManager` (spilling stays disabled otherwise). + pub enable_hash_join_spill: bool, default = false + + /// Number of disk partitions a spilling hash join scatters each side + /// into. 0 (default) sizes automatically from the observed build size + /// and the memory budget. + pub hash_join_spill_partition_count: usize, default = 0 + + /// Reserved memory headroom for a spilling hash join's scatter scratch + /// space and per-partition write buffers (analogous to + /// `sort_spill_reservation_bytes`). + pub hash_join_spill_headroom_bytes: usize, default = 32 * 1024 * 1024 + + /// Maximum number of recursive repartition passes a spilling hash join + /// applies to a partition whose build side still exceeds the memory + /// budget (key skew). When the limit is reached the join falls back to + /// a chunked build where supported, or reports a clean + /// resources-exhausted error. + pub hash_join_spill_max_recursion_depth: usize, default = 2 + /// When sorting, below what size should data be concatenated /// and sorted in a single RecordBatch rather than sorted in /// batches and merged. diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 25b320f985507..49509cbda95e8 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -37,6 +37,10 @@ use crate::joins::hash_join::inlist_builder::build_struct_inlist_values; use crate::joins::hash_join::shared_bounds::{ ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator, }; +use crate::joins::hash_join::spill::{ + HashJoinSpillContext, SharedSpilledBuild, SideScatter, SpilledBuild, + compute_partition_count, scatter_random_state, +}; use crate::joins::hash_join::stream::{ BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, }; @@ -46,6 +50,7 @@ use crate::joins::utils::{ swap_join_projection, update_hash, }; use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; +use crate::metrics::SpillMetrics; use crate::metrics::{Count, MetricBuilder}; use crate::projection::{ EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, @@ -53,6 +58,7 @@ use crate::projection::{ }; use crate::repartition::REPARTITION_RANDOM_STATE; use crate::spill::get_record_batch_memory_size; +use crate::spill::spill_manager::SpillManager; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, @@ -74,8 +80,8 @@ use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::utils::memory::estimate_memory_size; use datafusion_common::{ - JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, - plan_err, project_schema, + DataFusionError, JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, + internal_datafusion_err, internal_err, plan_err, project_schema, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -194,11 +200,13 @@ pub(super) struct JoinLeftData { batch: RecordBatch, /// The build side on expressions values values: Vec, - /// Shared bitmap builder for visited left indices - visited_indices_bitmap: SharedBitmapBuilder, + /// Shared bitmap builder for visited left indices. Arc'd so a spilled + /// CollectLeft build can share one bitmap per disk partition across the + /// independently-built per-probe-partition copies of the hash table. + visited_indices_bitmap: Arc, /// Counter of running probe-threads, potentially - /// able to update `visited_indices_bitmap` - probe_threads_counter: AtomicUsize, + /// able to update `visited_indices_bitmap` (Arc'd for the same reason) + probe_threads_counter: Arc, /// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it. /// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage. /// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption. @@ -236,7 +244,7 @@ impl JoinLeftData { /// returns a reference to the visited indices bitmap pub(super) fn visited_indices_bitmap(&self) -> &SharedBitmapBuilder { - &self.visited_indices_bitmap + self.visited_indices_bitmap.as_ref() } /// returns a reference to the InList values for filter pushdown @@ -251,6 +259,21 @@ impl JoinLeftData { } } +/// Result of the build phase: either the classic fully in-memory build, or +/// (when `enable_hash_join_spill` is on and the memory budget was exceeded) +/// a build side that has been scattered into disk partitions to be joined +/// pair-by-pair within the budget. +/// +/// The `Spilled` payload sits behind a `Mutex>` because `OnceFut` +/// shares its output as an `Arc`, while the spill-join driver needs to take +/// ownership; in `Partitioned` mode there is exactly one consumer per build. +pub(super) enum BuildPhaseOutput { + InMemory(Arc), + Spilled(Mutex>), + /// CollectLeft: the scattered build is shared by all probe partitions. + SpilledShared(Arc), +} + /// Helps to build [`HashJoinExec`]. /// /// Builder can be created from an existing [`HashJoinExec`] using [`From::from`]. @@ -734,7 +757,7 @@ pub struct HashJoinExec { /// /// Each output stream waits on the `OnceAsync` to signal the completion of /// the hash table creation. - left_fut: Arc>, + left_fut: Arc>, /// Shared the `SeededRandomState` for the hashing algorithm (seeds preserved for serialization) random_state: SeededRandomState, /// Partitioning mode to use @@ -1295,12 +1318,83 @@ impl ExecutionPlan for HashJoinExec { let array_map_created_count = MetricBuilder::new(&self.metrics) .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, partition); + // When the spill fallback is enabled, the build-side reservation is + // spillable: on overflow the operator scatters the build side to disk + // instead of failing, so spilling pools may treat it accordingly. + let spill_flag = context + .session_config() + .options() + .execution + .enable_hash_join_spill; + + // Spill applies to Partitioned-mode joins over a working DiskManager. + // (CollectLeft's shared build is a follow-up; null-aware anti joins + // need global probe-side knowledge and are excluded.) + // + // Every join type is supported: disk partitions are key-disjoint, so + // each pair's in-memory join sees all rows for its keys — unmatched + // detection (visited bitmaps, probe-side null padding) is complete + // within a pair. + // + // A spilled join emits probe rows partition-by-partition, so it CANNOT + // preserve probe-side ordering. When this operator advertises that it + // maintains the right input's order and that input is actually + // ordered, spilling stays disabled rather than risk downstream + // operators consuming out-of-order rows. + let would_break_output_order = Self::maintains_input_order(self.join_type)[1] + && self.right.output_ordering().is_some(); + let spill_context = if spill_flag + && matches!( + self.mode, + PartitionMode::Partitioned | PartitionMode::CollectLeft + ) + && !self.null_aware + && !would_break_output_order + && context.runtime_env().disk_manager.tmp_files_enabled() + { + let options = context.session_config().options(); + let exec_options = &options.execution; + let build_spill = SpillManager::new( + context.runtime_env(), + SpillMetrics::new(&self.metrics, partition), + self.left.schema(), + ) + .with_compression_type(exec_options.spill_compression); + let probe_spill = SpillManager::new( + context.runtime_env(), + SpillMetrics::new(&self.metrics, partition), + self.right.schema(), + ) + .with_compression_type(exec_options.spill_compression); + Some(Arc::new(HashJoinSpillContext { + pool: Arc::clone(context.memory_pool()), + build_spill: Arc::new(build_spill), + probe_spill: Arc::new(probe_spill), + spill_engaged: MetricBuilder::new(&self.metrics) + .counter("join_spill_engaged", partition), + repartition_passes: MetricBuilder::new(&self.metrics) + .counter("join_spill_repartition_passes", partition), + fallback_chunks: MetricBuilder::new(&self.metrics) + .counter("join_spill_fallback_chunks", partition), + partition_count_override: exec_options.hash_join_spill_partition_count, + headroom_bytes: exec_options.hash_join_spill_headroom_bytes, + max_recursion_depth: exec_options.hash_join_spill_max_recursion_depth, + max_spill_file_size: exec_options.max_spill_file_size_bytes, + partition, + on_left: on_left.clone(), + config: Arc::clone(context.session_config().options()), + })) + } else { + None + }; + let left_fut = match self.mode { PartitionMode::CollectLeft => self.left_fut.try_once(|| { let left_stream = self.left.execute(0, Arc::clone(&context))?; - let reservation = - MemoryConsumer::new("HashJoinInput").register(context.memory_pool()); + let reservation = MemoryConsumer::new("HashJoinInput") + .with_can_spill(spill_context.is_some()) + .register(context.memory_pool()); Ok(collect_left_input( self.random_state.random_state().clone(), @@ -1314,6 +1408,10 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + spill_context.clone(), + // CollectLeft: the scattered build is shared by all + // probe partitions. + true, )) })?, PartitionMode::Partitioned => { @@ -1321,6 +1419,7 @@ impl ExecutionPlan for HashJoinExec { let reservation = MemoryConsumer::new(format!("HashJoinInput[{partition}]")) + .with_can_spill(spill_context.is_some()) .register(context.memory_pool()); OnceFut::new(collect_left_input( @@ -1335,6 +1434,8 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + spill_context.clone(), + false, )) } PartitionMode::Auto => { @@ -1413,6 +1514,7 @@ impl ExecutionPlan for HashJoinExec { self.mode, self.null_aware, self.fetch, + spill_context, ))) } @@ -1759,10 +1861,27 @@ impl CollectLeftAccumulator { /// State for collecting the build-side data during hash join struct BuildSideState { batches: Vec, + /// Reserved memory size recorded per buffered batch, so scattering can + /// shrink the reservation by exactly what was grown. + batch_sizes: Vec, num_rows: usize, metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, bounds_accumulators: Option>, + /// Bytes of the hash-table estimate already reserved (spill mode only: + /// the estimate is reserved incrementally so overflow surfaces during + /// batch collection, where the scatter fallback can engage). + table_reserved: usize, + /// How many spilled-partition tables can be resident at once (1 for + /// Partitioned mode; the probe partition count for a shared CollectLeft + /// build). Sizes the scatter's partition count. + concurrent_table_consumers: usize, + /// Present when the spill fallback is enabled and applicable. + spill: Option>, + /// Present once the build has switched into scatter mode. + scatter: Option, + /// Headroom reservation acquired at spill engage; handed to the driver. + headroom: Option, } impl BuildSideState { @@ -1776,6 +1895,7 @@ impl BuildSideState { ) -> Result { Ok(Self { batches: Vec::new(), + batch_sizes: Vec::new(), num_rows: 0, metrics, reservation, @@ -1787,8 +1907,91 @@ impl BuildSideState { .collect::>>() }) .transpose()?, + table_reserved: 0, + concurrent_table_consumers: 1, + spill: None, + scatter: None, + headroom: None, }) } + + /// Switch into scatter mode after the build reservation first failed: + /// acquire the scatter headroom, then re-scatter every already-buffered + /// batch to disk, releasing its memory. + fn engage_spill(&mut self) -> Result<()> { + let ctx = self + .spill + .clone() + .ok_or_else(|| internal_datafusion_err!("engage_spill without context"))?; + ctx.spill_engaged.add(1); + log::info!( + "hash join spill engaged: partition={} buffered_bytes={} rows={}", + ctx.partition, + self.batch_sizes.iter().sum::(), + self.num_rows + ); + + let buffered: usize = self.batch_sizes.iter().sum(); + let partition_count = compute_partition_count( + buffered, + &ctx.pool.memory_limit(), + ctx.partition_count_override, + self.concurrent_table_consumers, + ); + let mut scatter = SideScatter::new( + partition_count, + ctx.on_left.clone(), + scatter_random_state(0), + &ctx.build_spill, + ctx.max_spill_file_size, + ctx.headroom_bytes / 2, + "hash_join_build_spill", + ); + + // Acquire scratch headroom; if even that fails, scatter buffered + // batches one at a time (each frees more than the transient scratch + // uses) and retry. With nothing left to free, give up cleanly. + let headroom = + MemoryConsumer::new(format!("HashJoinSpillHeadroom[{}]", ctx.partition)) + .with_can_spill(true) + .register(&ctx.pool); + loop { + match headroom.try_grow(ctx.headroom_bytes) { + Ok(()) => break, + Err(e) => { + if self.batches.is_empty() { + return Err(e); + } + let batch = self.batches.remove(0); + let size = self.batch_sizes.remove(0); + scatter.scatter_batch(&batch)?; + drop(batch); + self.reservation.shrink(size); + } + } + } + + // Re-scatter the remaining buffered batches oldest-first. + for (batch, size) in self.batches.drain(..).zip(self.batch_sizes.drain(..)) { + scatter.scatter_batch(&batch)?; + drop(batch); + self.reservation.shrink(size); + } + // No table will be built for this build side anymore. + self.reservation.shrink(self.table_reserved); + self.table_reserved = 0; + + self.scatter = Some(scatter); + self.headroom = Some(headroom); + Ok(()) + } + + fn scatter_batch(&mut self, batch: &RecordBatch) -> Result<()> { + self.scatter + .as_mut() + .ok_or_else(|| internal_datafusion_err!("scatter_batch without scatter"))? + .scatter_batch(batch) + } } fn should_collect_min_max_for_perfect_hash( @@ -1845,41 +2048,75 @@ async fn collect_left_input( config: Arc, null_equality: NullEquality, array_map_created_count: Count, -) -> Result { + spill: Option>, + shared_build: bool, +) -> Result { let schema = left_stream.schema(); let should_collect_min_max_for_phj = should_collect_min_max_for_perfect_hash(&on_left, &schema)?; - let initial = BuildSideState::try_new( + let mut initial = BuildSideState::try_new( metrics, reservation, on_left.clone(), &schema, should_compute_dynamic_filters || should_collect_min_max_for_phj, )?; + initial.spill = spill; + if shared_build { + initial.concurrent_table_consumers = probe_threads_count; + } let state = left_stream .try_fold(initial, |mut state, batch| async move { - // Update accumulators if computing bounds + // Update accumulators if computing bounds. These keep running in + // scatter mode too, so dynamic-filter bounds stay exact under + // spill. if let Some(ref mut accumulators) = state.bounds_accumulators { for accumulator in accumulators { accumulator.update_batch(&batch)?; } } - // Decide if we spill or not - let batch_size = get_record_batch_memory_size(&batch); - // Reserve memory for incoming batch - state.reservation.try_grow(batch_size)?; - // Update metrics - state.metrics.build_mem_used.add(batch_size); state.metrics.build_input_batches.add(1); state.metrics.build_input_rows.add(batch.num_rows()); - // Update row count state.num_rows += batch.num_rows(); - // Push batch to output - state.batches.push(batch); + + if state.scatter.is_some() { + // Scatter mode: batches go to disk partitions, not memory. + state.scatter_batch(&batch)?; + return Ok(state); + } + + let batch_size = get_record_batch_memory_size(&batch); + // In spill mode, also reserve the hash table's estimated growth + // for the rows in this batch, so "data fits but the table will + // not" is detected here — where the scatter fallback can engage — + // rather than at table construction time. + let table_delta = if state.spill.is_some() { + hash_table_estimate(state.num_rows)?.saturating_sub(state.table_reserved) + } else { + 0 + }; + // Reserve memory for incoming batch + match state.reservation.try_grow(batch_size + table_delta) { + Ok(()) => { + state.metrics.build_mem_used.add(batch_size); + state.table_reserved += table_delta; + state.batches.push(batch); + state.batch_sizes.push(batch_size); + } + Err(DataFusionError::ResourcesExhausted(_)) if state.spill.is_some() => { + // The build side no longer fits: switch into scatter + // mode instead of failing. Already-buffered batches are + // re-scattered (releasing their memory), then the + // current batch follows. + state.engage_spill()?; + state.scatter_batch(&batch)?; + } + Err(e) => return Err(e), + } Ok(state) }) .await?; @@ -1887,14 +2124,20 @@ async fn collect_left_input( // Extract fields from state let BuildSideState { batches, + batch_sizes: _, num_rows, metrics, - mut reservation, + reservation, bounds_accumulators, + table_reserved, + concurrent_table_consumers: _, + spill: _, + scatter, + headroom, } = state; // Compute bounds - let mut bounds = match bounds_accumulators { + let bounds = match bounds_accumulators { Some(accumulators) if num_rows > 0 => { let bounds = accumulators .into_iter() @@ -1905,12 +2148,113 @@ async fn collect_left_input( _ => None, }; + // Overflowed: the build side lives in disk partitions now. The buffered + // batches were all scattered, so the build reservation is back to zero; + // the headroom reservation transfers to the spill-join driver. + if let Some(scatter) = scatter { + let headroom = headroom.ok_or_else(|| { + internal_datafusion_err!("scatter engaged without headroom reservation") + })?; + debug_assert_eq!(reservation.size(), 0); + let partitions = scatter.finish()?; + if shared_build { + return Ok(BuildPhaseOutput::SpilledShared(Arc::new( + SharedSpilledBuild::new( + partitions, + 0, + bounds, + probe_threads_count, + headroom, + ), + ))); + } + let partition_count = partitions.len(); + return Ok(BuildPhaseOutput::Spilled(Mutex::new(Some(SpilledBuild { + partitions, + partition_count, + level: 0, + bounds, + headroom, + })))); + } + + let data = build_left_data( + &batches, + num_rows, + &schema, + &on_left, + &random_state, + reservation, + &metrics, + with_visited_indices_bitmap, + probe_threads_count, + bounds, + should_collect_min_max_for_phj && !should_compute_dynamic_filters, + &config, + null_equality, + &array_map_created_count, + table_reserved, + None, + )?; + + Ok(BuildPhaseOutput::InMemory(Arc::new(data))) +} + +/// Builds the in-memory join structures (hash table or perfect-hash +/// [`ArrayMap`], concatenated build batch, join-key arrays, visited bitmap, +/// membership pushdown) from an already-collected set of build batches. +/// +/// Extracted from [`collect_left_input`] so the same code serves both the +/// classic fully in-memory build and per-partition builds when the join +/// spills (grace hash join). All memory is accounted against `reservation`, +/// which is moved into the returned [`JoinLeftData`]. +/// +/// Estimated memory footprint of the join hash table for `num_rows` build +/// rows — the same formula [`build_left_data`] uses, exposed so spill-mode +/// loaders can pre-reserve it incrementally while batches stream in. +pub(super) fn hash_table_estimate(num_rows: usize) -> Result { + if num_rows > u32::MAX as usize { + estimate_memory_size::<(u64, u64)>(num_rows, size_of::()) + } else { + estimate_memory_size::<(u32, u64)>(num_rows, size_of::()) + } +} + +/// `mask_bounds` clears the computed bounds from the result when they were +/// only collected for the perfect-hash decision and dynamic-filter pushdown +/// is disabled. +/// +/// `prereserved_table_bytes` is how much of the hash table's estimated size +/// the caller already holds in `reservation` (spill-mode loaders reserve the +/// estimate incrementally, via [`hash_table_estimate`], so overflow surfaces +/// while batches stream rather than at table construction). The fast path +/// passes 0 and this function grows the full estimate itself, exactly as +/// before. +#[expect(clippy::too_many_arguments)] +pub(super) fn build_left_data( + batches: &[RecordBatch], + num_rows: usize, + schema: &SchemaRef, + on_left: &[PhysicalExprRef], + random_state: &RandomState, + mut reservation: MemoryReservation, + metrics: &BuildProbeJoinMetrics, + with_visited_indices_bitmap: bool, + probe_threads_count: usize, + mut bounds: Option, + mask_bounds: bool, + config: &ConfigOptions, + null_equality: NullEquality, + array_map_created_count: &Count, + prereserved_table_bytes: usize, + shared_state: Option<(Arc, Arc)>, +) -> Result { let (join_hash_map, batch, left_values) = if let Some((array_map, batch, left_value)) = try_create_array_map( &bounds, - &schema, - &batches, - &on_left, + schema, + batches, + on_left, &mut reservation, config.execution.perfect_hash_join_small_build_threshold, config.execution.perfect_hash_join_min_key_density, @@ -1918,11 +2262,15 @@ async fn collect_left_input( )? { array_map_created_count.add(1); metrics.build_mem_used.add(array_map.size()); + // The perfect-hash map replaces the hash table the caller + // pre-reserved for; release that estimate. + reservation.shrink(prereserved_table_bytes); (Map::ArrayMap(array_map), batch, left_value) } else { // Estimation of memory size, required for hashtable, prior to allocation. // Final result can be verified using `RawTable.allocation_info()` + // (Same formula as `hash_table_estimate`; keep them in sync.) let fixed_size_u32 = size_of::(); let fixed_size_u64 = size_of::(); @@ -1932,13 +2280,17 @@ async fn collect_left_input( let mut hashmap: Box = if num_rows > u32::MAX as usize { let estimated_hashtable_size = estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)?; - reservation.try_grow(estimated_hashtable_size)?; + reservation.try_grow( + estimated_hashtable_size.saturating_sub(prereserved_table_bytes), + )?; metrics.build_mem_used.add(estimated_hashtable_size); Box::new(JoinHashMapU64::with_capacity(num_rows)) } else { let estimated_hashtable_size = estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)?; - reservation.try_grow(estimated_hashtable_size)?; + reservation.try_grow( + estimated_hashtable_size.saturating_sub(prereserved_table_bytes), + )?; metrics.build_mem_used.add(estimated_hashtable_size); Box::new(JoinHashMapU32::with_capacity(num_rows)) }; @@ -1953,11 +2305,11 @@ async fn collect_left_input( hashes_buffer.clear(); hashes_buffer.resize(batch.num_rows(), 0); update_hash( - &on_left, + on_left, batch, &mut *hashmap, offset, - &random_state, + random_state, &mut hashes_buffer, 0, true, @@ -1966,24 +2318,47 @@ async fn collect_left_input( } // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_batches(&schema, batches_iter.clone())?; + let batch = concat_batches(schema, batches_iter.clone())?; - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; (Map::HashMap(hashmap), batch, left_values) }; - // Reserve additional memory for visited indices bitmap and create shared builder - let visited_indices_bitmap = if with_visited_indices_bitmap { - let bitmap_size = bit_util::ceil(batch.num_rows(), 8); - reservation.try_grow(bitmap_size)?; - metrics.build_mem_used.add(bitmap_size); - - let mut bitmap_buffer = BooleanBufferBuilder::new(batch.num_rows()); - bitmap_buffer.append_n(num_rows, false); - bitmap_buffer - } else { - BooleanBufferBuilder::new(0) + // Reserve additional memory for visited indices bitmap and create shared + // builder. An injected shared state (spilled CollectLeft) is initialized + // by whichever probe partition builds this disk partition first. + let (visited_indices_bitmap, probe_threads_counter) = match shared_state { + Some((bitmap, counter)) => { + if with_visited_indices_bitmap { + let mut guard = bitmap.lock(); + if guard.len() < num_rows { + debug_assert_eq!(guard.len(), 0); + let bitmap_size = bit_util::ceil(batch.num_rows(), 8); + reservation.try_grow(bitmap_size)?; + metrics.build_mem_used.add(bitmap_size); + guard.append_n(num_rows, false); + } + } + (bitmap, counter) + } + None => { + let bitmap_buffer = if with_visited_indices_bitmap { + let bitmap_size = bit_util::ceil(batch.num_rows(), 8); + reservation.try_grow(bitmap_size)?; + metrics.build_mem_used.add(bitmap_size); + + let mut bitmap_buffer = BooleanBufferBuilder::new(batch.num_rows()); + bitmap_buffer.append_n(num_rows, false); + bitmap_buffer + } else { + BooleanBufferBuilder::new(0) + }; + ( + Arc::new(Mutex::new(bitmap_buffer)), + Arc::new(AtomicUsize::new(probe_threads_count)), + ) + } }; let map = Arc::new(join_hash_map); @@ -2014,7 +2389,7 @@ async fn collect_left_input( } }; - if should_collect_min_max_for_phj && !should_compute_dynamic_filters { + if mask_bounds { bounds = None; } @@ -2022,8 +2397,8 @@ async fn collect_left_input( map, batch, values: left_values, - visited_indices_bitmap: Mutex::new(visited_indices_bitmap), - probe_threads_counter: AtomicUsize::new(probe_threads_count), + visited_indices_bitmap, + probe_threads_counter, _reservation: reservation, bounds, membership, diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index b915802ea4015..1308229298c42 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -24,4 +24,5 @@ mod exec; mod inlist_builder; mod partitioned_hash_eval; mod shared_bounds; +mod spill; mod stream; diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index f32dc7fa80268..78e5fdad13b80 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -136,6 +136,9 @@ fn create_membership_predicate( )) as Arc)), // Empty partition - should not create a filter for this PushdownStrategy::Empty => Ok(None), + // Spilled partition: has data but no in-memory membership structure. + // No membership predicate; bounds (handled by the caller) still apply. + PushdownStrategy::BoundsOnly => Ok(None), } } @@ -237,6 +240,10 @@ pub(crate) enum PushdownStrategy { Map(Arc), /// There was no data in this partition, do not build a dynamic filter for it Empty, + /// The partition HAS data but no in-memory membership structure exists + /// (its build side spilled to disk). Only bounds may be pushed down; + /// unlike `Empty`, this partition's probe rows must NOT be pruned. + BoundsOnly, } /// Build-side data reported by a single partition diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs new file mode 100644 index 0000000000000..4e2d9e90c4a68 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -0,0 +1,2800 @@ +// 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. + +//! Disk-spilling fallback for [`super::HashJoinExec`] (grace hash join). +//! +//! The fast path is untouched: the operator buffers the build side in memory +//! exactly as before. Only when the build-side [`MemoryReservation`] first +//! fails does the build switch into *scatter mode*: already-buffered batches +//! and the rest of the build input are hash-partitioned into `K` disk +//! partitions (releasing their memory), the probe side is scattered with the +//! same hash, and the join then runs partition-pair by partition-pair, each +//! pair building an ordinary in-memory hash table within the memory budget. +//! +//! Memory ownership rules (every reservation has exactly one owner): +//! - `HashJoinInput[p]` (the existing build reservation): grows per buffered +//! batch on the fast path, shrinks batch-by-batch as buffers are scattered, +//! reaches zero once scatter mode is fully engaged. +//! - `HashJoinSpillHeadroom[p]`: a fixed reservation covering scatter scratch +//! and per-partition write buffers; held until the spill join completes. +//! - `HashJoinSpillPartition[p.k]`: per partition-pair; covers the loaded +//! build batches plus the hash table built from them (moved into the +//! pair's [`JoinLeftData`]); dropped when the pair finishes. +//! +//! The scatter hash uses seeds distinct from both `RepartitionExec`'s +//! `(0,0,0,0)` routing seeds and the join hash map's `HASH_JOIN_SEED` +//! (`'J','O','I','N'`), so disk partitions do not correlate with either the +//! upstream partitioning or the per-pair hash tables. + +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::task::{Context, Poll}; + +use crate::hash_utils::create_hashes; +use crate::joins::PartitionMode; +use crate::joins::SharedBitmapBuilder; +use crate::joins::hash_join::exec::{ + BuildPhaseOutput, build_left_data, hash_table_estimate, +}; +use crate::joins::hash_join::shared_bounds::PartitionBounds; +use crate::joins::hash_join::stream::{ + BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, +}; +use crate::joins::utils::{ + BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceFut, need_produce_result_in_final, +}; +use crate::metrics::{Count, ExecutionPlanMetricsSet}; +use crate::spill::get_record_batch_memory_size; +use crate::spill::in_progress_spill_file::InProgressSpillFile; +use crate::spill::spill_manager::SpillManager; +use crate::stream::RecordBatchStreamAdapter; +use crate::{RecordBatchStream, SendableRecordBatchStream}; + +use arrow::array::{Array, BooleanBufferBuilder, StringViewArray, UInt32Array}; +use arrow::compute::{concat_batches, take_record_batch}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{DataFusionError, JoinType, NullEquality, Result, internal_err}; +use datafusion_execution::disk_manager::RefCountedTempFile; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; + +use ahash::RandomState; +use futures::future::BoxFuture; +use futures::{Stream, StreamExt, ready}; +use parking_lot::Mutex; + +/// Per-partition write-buffer flush threshold: coalesce scattered slices into +/// IPC batches of roughly this size before writing. +const WRITE_BUFFER_FLUSH_BYTES: usize = 1024 * 1024; + +/// Fallback per-partition target size when the memory pool cannot report a +/// finite limit. +const DEFAULT_PARTITION_TARGET_BYTES: usize = 128 * 1024 * 1024; + +/// SplitMix64 finalizer: a strong 64-bit mixer for seed derivation. +fn splitmix64(seed: u64) -> u64 { + let mut z = seed.wrapping_add(0x9E3779B97F4A7C15); + z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB); + z ^ (z >> 31) +} + +/// Hash seeds for scattering rows into disk partitions. +/// +/// Must stay independent from `REPARTITION_RANDOM_STATE` (0,0,0,0) and +/// `HASH_JOIN_SEED` ('J','O','I','N'): rows arrive pre-bucketed by the former +/// and are hashed by the latter inside each pair's hash table. +/// +/// All four seeds are derived from the level through a strong mixer: ahash +/// folds some seed words in ADDITIVELY, so states differing only by a small +/// constant in one word produce hashes shifted by a constant — under `% K` +/// that maps an entire parent partition into a single child, making +/// recursive repartitioning useless. (Caught by +/// `recursive_scatter_levels_are_decorrelated`.) +pub(super) fn scatter_random_state(level: usize) -> RandomState { + let base = 0x53_50_4C_00_C0_FF_EE_00u64 ^ (level as u64); + RandomState::with_seeds( + splitmix64(base), + splitmix64(base ^ 0xA5A5_A5A5_A5A5_A5A5), + splitmix64(base.rotate_left(17)), + splitmix64(base.rotate_left(43) ^ 0x5A5A_5A5A_5A5A_5A5A), + ) +} + +/// Choose the number of disk partitions when spilling engages. +/// +/// `buffered_bytes` is what fit in memory before the reservation failed; the +/// 4x factor assumes overflow happened early in the stream. Recursion +/// corrects underestimates, so precision is not important. +/// `concurrent_consumers` is how many partition tables can be resident at +/// once: 1 for Partitioned mode (one pair at a time), N for a shared +/// CollectLeft build whose N probe partitions walk the k's out of sync — +/// the per-partition target shrinks accordingly. +pub(super) fn compute_partition_count( + buffered_bytes: usize, + pool_limit: &MemoryLimit, + override_count: usize, + concurrent_consumers: usize, +) -> usize { + if override_count > 0 { + return override_count; + } + let base_target = match pool_limit { + MemoryLimit::Finite(limit) => { + (*limit / 16).clamp(32 * 1024 * 1024, 256 * 1024 * 1024) + } + _ => DEFAULT_PARTITION_TARGET_BYTES, + }; + let target = base_target / concurrent_consumers.max(1); + buffered_bytes + .saturating_mul(4) + .div_ceil(target.max(1)) + .clamp(16, 64) +} + +/// Everything a hash join partition needs to engage and drive spilling. +/// Built once per output partition in `HashJoinExec::execute` when the +/// feature is enabled and applicable; shared by the build phase and the +/// spill-join driver. +pub(super) struct HashJoinSpillContext { + pub(super) pool: Arc, + /// Spill manager for build-side partitions (build schema). + pub(super) build_spill: Arc, + /// Spill manager for probe-side partitions (probe schema). + pub(super) probe_spill: Arc, + /// Incremented once when a build side switches into scatter mode. + pub(super) spill_engaged: Count, + /// Incremented per recursive repartition pass on an oversized pair. + pub(super) repartition_passes: Count, + /// Incremented per chunk executed by the chunked build fallback. + pub(super) fallback_chunks: Count, + pub(super) partition_count_override: usize, + pub(super) headroom_bytes: usize, + pub(super) max_recursion_depth: usize, + pub(super) max_spill_file_size: usize, + /// Output partition index (consumer naming / diagnostics). + pub(super) partition: usize, + /// Build-side join key expressions (needed to rebuild per-pair tables). + pub(super) on_left: Vec, + pub(super) config: Arc, +} + +/// One side's scattered partition: the rotated spill files plus totals. +pub(super) struct SpilledSide { + pub(super) files: Vec, + pub(super) bytes: usize, + pub(super) rows: usize, +} + +/// The build phase's output when it overflowed and scattered to disk. +pub(super) struct SpilledBuild { + pub(super) partitions: Vec, + pub(super) partition_count: usize, + /// Recursion level the scatter used (level 0 at first engage). + pub(super) level: usize, + /// Dynamic-filter bounds computed before/during the scatter (bounds + /// accumulators keep running, so these stay exact under spill). + pub(super) bounds: Option, + /// Headroom reservation kept alive for the probe scatter and the + /// partition-pair loop; dropped when the spill join completes. + pub(super) headroom: MemoryReservation, +} + +/// A spilled build shared by all `CollectLeft` probe partitions. +/// +/// The single build was scattered once into `K` disk partitions; each of the +/// `probe_threads` output partitions scatters ITS OWN probe stream with the +/// same seed and then walks k = 0..K, loading its OWN copy of k's hash +/// table (memory: at most one table per probe partition at a time, which is +/// what the auto partition-count sizing assumes). What IS shared per k is +/// the visited bitmap and the probe-completion counter injected into each +/// copy's [`JoinLeftData`]: matches from every probe partition mark one +/// bitmap, and the existing last-finisher logic emits unmatched build rows +/// exactly once. Row order is deterministic across copies (same files, same +/// order), so bitmap indices agree. +/// +/// [`JoinLeftData`]: super::exec::JoinLeftData +pub(super) struct SharedSpilledBuild { + pub(super) partitions: Vec, + pub(super) partition_count: usize, + pub(super) level: usize, + /// Bounds for dynamic-filter reporting (exact — accumulators keep + /// running during the scatter). + pub(super) bounds: Option, + /// Number of probe partitions that will walk the shared build. + pub(super) probe_threads: usize, + /// Per-k shared visited bitmap + probe-completion counter. + shared_state: Vec<(Arc, Arc)>, + /// Scatter headroom, held until the whole shared build is dropped. + _headroom: MemoryReservation, +} + +/// One shared build partition: file handles are `Arc`d so every output +/// partition can read them. +pub(super) struct SharedSpilledSide { + pub(super) files: Arc>, +} + +impl SharedSpilledBuild { + pub(super) fn new( + partitions: Vec, + level: usize, + bounds: Option, + probe_threads: usize, + headroom: MemoryReservation, + ) -> Self { + let partition_count = partitions.len(); + let partitions = partitions + .into_iter() + .map(|side| SharedSpilledSide { + files: Arc::new(side.files), + }) + .collect(); + Self { + partitions, + partition_count, + level, + bounds, + probe_threads, + shared_state: (0..partition_count) + .map(|_| { + ( + Arc::new(Mutex::new(BooleanBufferBuilder::new(0))), + Arc::new(AtomicUsize::new(probe_threads)), + ) + }) + .collect(), + _headroom: headroom, + } + } +} + +/// A single partition's spill writer: buffers scattered slices and appends +/// them as coalesced IPC batches, rotating files at `max_file_size` bytes. +struct PartitionSpillWriter { + spill_manager: Arc, + description: String, + max_file_size: usize, + current: Option, + current_bytes: usize, + files: Vec, + buffered: Vec, + buffered_bytes: usize, + bytes: usize, + rows: usize, +} + +impl PartitionSpillWriter { + fn new( + spill_manager: Arc, + description: String, + max_file_size: usize, + ) -> Self { + Self { + spill_manager, + description, + max_file_size, + current: None, + current_bytes: 0, + files: Vec::new(), + buffered: Vec::new(), + buffered_bytes: 0, + bytes: 0, + rows: 0, + } + } + + fn append(&mut self, batch: RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + self.rows += batch.num_rows(); + self.buffered_bytes += get_record_batch_memory_size(&batch); + self.buffered.push(batch); + if self.buffered_bytes >= WRITE_BUFFER_FLUSH_BYTES { + self.flush()?; + } + Ok(()) + } + + /// Write the buffered slices out as one coalesced IPC batch. + fn flush(&mut self) -> Result<()> { + if self.buffered.is_empty() { + return Ok(()); + } + let schema = Arc::clone(self.spill_manager.schema()); + let batch = concat_batches(&schema, self.buffered.iter())?; + // Compact StringView columns: `take`d views reference the source + // batch's full data buffers, which the IPC writer would otherwise + // serialize wholesale into every partition. + let batch = compact_string_views(batch)?; + self.buffered.clear(); + self.buffered_bytes = 0; + + if self.current.is_none() { + self.current = Some( + self.spill_manager + .create_in_progress_file(&self.description)?, + ); + self.current_bytes = 0; + } + if let Some(writer) = self.current.as_mut() { + writer.append_batch(&batch)?; + } + let written = get_record_batch_memory_size(&batch); + self.current_bytes += written; + self.bytes += written; + + if self.current_bytes >= self.max_file_size { + self.rotate_file()?; + } + Ok(()) + } + + /// Close the in-progress file (if any) and add it to the finished set. + fn rotate_file(&mut self) -> Result<()> { + if let Some(mut writer) = self.current.take() + && let Some(file) = writer.finish()? + { + self.files.push(file); + } + Ok(()) + } + + fn finish(mut self) -> Result { + self.flush()?; + self.rotate_file()?; + Ok(SpilledSide { + files: self.files, + bytes: self.bytes, + rows: self.rows, + }) + } +} + +/// Rebuild any StringView columns so they own only the data they reference. +fn compact_string_views(batch: RecordBatch) -> Result { + let mut mutated = false; + let mut columns = Vec::with_capacity(batch.num_columns()); + for array in batch.columns() { + if let Some(string_view) = array.as_any().downcast_ref::() { + columns.push(Arc::new(string_view.gc()) as _); + mutated = true; + } else { + columns.push(Arc::clone(array)); + } + } + if mutated { + Ok(RecordBatch::try_new(batch.schema(), columns)?) + } else { + Ok(batch) + } +} + +/// Scatters one side's batches into `K` disk partitions by join-key hash. +pub(super) struct SideScatter { + on_exprs: Vec, + random_state: RandomState, + writers: Vec, + /// Flush all write buffers when their sum exceeds this. + buffer_cap: usize, + hashes_buffer: Vec, +} + +impl SideScatter { + pub(super) fn new( + partition_count: usize, + on_exprs: Vec, + random_state: RandomState, + spill_manager: &Arc, + max_file_size: usize, + buffer_cap: usize, + description: &str, + ) -> Self { + let writers = (0..partition_count) + .map(|i| { + PartitionSpillWriter::new( + Arc::clone(spill_manager), + format!("{description}[{i}]"), + max_file_size, + ) + }) + .collect(); + Self { + on_exprs, + random_state, + writers, + buffer_cap, + hashes_buffer: Vec::new(), + } + } + + pub(super) fn scatter_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + let keys = evaluate_expressions_to_arrays(&self.on_exprs, batch)?; + self.hashes_buffer.clear(); + self.hashes_buffer.resize(num_rows, 0); + create_hashes(&keys, &self.random_state, &mut self.hashes_buffer)?; + + let partition_count = self.writers.len() as u64; + let mut indices: Vec> = vec![Vec::new(); self.writers.len()]; + for (row, hash) in self.hashes_buffer.iter().enumerate() { + indices[(hash % partition_count) as usize].push(row as u32); + } + + for (partition, rows) in indices.into_iter().enumerate() { + if rows.is_empty() { + continue; + } + let take_indices = UInt32Array::from(rows); + let slice = take_record_batch(batch, &take_indices)?; + self.writers[partition].append(slice)?; + } + + let total_buffered: usize = self.writers.iter().map(|w| w.buffered_bytes).sum(); + if total_buffered > self.buffer_cap { + for writer in &mut self.writers { + writer.flush()?; + } + } + Ok(()) + } + + pub(super) fn finish(self) -> Result> { + self.writers.into_iter().map(|w| w.finish()).collect() + } +} + +/// Lazily chains a partition's spill files into one record batch stream. +struct ChainedSpillReader { + schema: SchemaRef, + files: VecDeque, + spill_manager: Arc, + current: Option, +} + +impl Stream for ChainedSpillReader { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + if let Some(current) = &mut self.current { + match ready!(current.poll_next_unpin(cx)) { + Some(item) => return Poll::Ready(Some(item)), + None => self.current = None, + } + } + match self.files.pop_front() { + Some(file) => { + let this = &mut *self; + match this.spill_manager.read_spill_as_stream(file, None) { + Ok(stream) => this.current = Some(stream), + Err(e) => return Poll::Ready(Some(Err(e))), + } + } + None => return Poll::Ready(None), + } + } + } +} + +impl RecordBatchStream for ChainedSpillReader { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +fn spilled_side_stream( + files: Vec, + spill_manager: Arc, +) -> SendableRecordBatchStream { + Box::pin(ChainedSpillReader { + schema: Arc::clone(spill_manager.schema()), + files: files.into(), + spill_manager, + current: None, + }) +} + +/// An immediately-empty record batch stream (used to take ownership of the +/// outer probe stream when handing it to the spill driver). +pub(super) fn empty_record_batch_stream(schema: SchemaRef) -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::empty(), + )) +} + +/// Parameters copied from the outer [`HashJoinStream`] needed to construct +/// the per-partition-pair inner join streams. +pub(super) struct InnerJoinSpec { + pub(super) partition: usize, + pub(super) schema: SchemaRef, + pub(super) on_right: Vec, + pub(super) filter: Option, + pub(super) join_type: JoinType, + /// The join hash map's random state (`HASH_JOIN_SEED`). + pub(super) random_state: RandomState, + pub(super) column_indices: Vec, + pub(super) null_equality: NullEquality, + pub(super) batch_size: usize, +} + +/// Fan-out when an oversized pair is recursively repartitioned. +const CHILD_FANOUT: usize = 8; + +/// One partition pair waiting to be joined. +struct PairWork { + build: SpilledSide, + probe: SpilledSide, + /// Recursion level this pair's data was scattered at. + level: usize, + /// False once a repartition pass failed to shrink the data + /// (duplicate-key mass): further passes cannot help, go straight to the + /// chunked fallback. + allow_recursion: bool, +} + +/// Per-pair state carried through the chunked fallback: the remaining build +/// stream and the probe files that get re-read once per chunk. +struct ChunkedState { + pair_id: usize, + level: usize, + /// None once the build partition is fully consumed. + build_stream: Option, + /// A batch that did not fit in the previous chunk; leads the next one. + pending: Option, + probe_files: Vec, + chunk_index: usize, +} + +impl ChunkedState { + fn exhausted(&self) -> bool { + self.build_stream.is_none() && self.pending.is_none() + } +} + +enum DriverPhase { + /// Scatter the entire probe stream into `K` disk partitions. + ScatterProbe { + probe: SendableRecordBatchStream, + scatter: SideScatter, + }, + /// Pick the next partition pair to join. + NextPair, + /// Stream the pair's build partition into memory under its reservation. + LoadBuild { + pair_id: usize, + level: usize, + allow_recursion: bool, + probe: SpilledSide, + stream: SendableRecordBatchStream, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + /// Bytes of the hash-table estimate reserved so far (grows with the + /// row count so table overflow surfaces during the load). + table_reserved: usize, + parent_build_bytes: usize, + }, + /// Run the pair's in-memory join and forward its output. + RunPair { + inner: SendableRecordBatchStream, + }, + /// The pair's build side did not fit: re-scatter the remaining build + /// rows (already-loaded ones went in first) at the next seed level. + RepartitionBuild { + level: usize, + probe: SpilledSide, + stream: SendableRecordBatchStream, + scatter: SideScatter, + parent_build_bytes: usize, + }, + /// Re-scatter the oversized pair's probe partition with the same child + /// seed, then enqueue the child pairs. + RepartitionProbe { + children_build: Vec, + child_level: usize, + allow_recursion: bool, + stream: SendableRecordBatchStream, + scatter: SideScatter, + }, + /// Chunked fallback: load the next slice of build rows that fits. + ChunkedLoad { + state: ChunkedState, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + table_reserved: usize, + }, + /// Chunked fallback: run one chunk's join over a full probe re-read. + ChunkedRun { + inner: SendableRecordBatchStream, + state: ChunkedState, + }, + Done, +} + +/// Drives the spilled join: probe scatter, then a sequential loop over +/// partition pairs, each joined by an ordinary in-memory [`HashJoinStream`]. +/// +/// A pair whose build side exceeds the budget is recursively repartitioned +/// with a fresh seed (up to `hash_join_spill_max_recursion_depth` passes); +/// when repartitioning cannot shrink it (duplicate-key mass), build-side +/// emission join types degrade to a chunked build over full probe re-reads, +/// and the rest surface a descriptive resources-exhausted error. +pub(super) struct SpillJoinDriver { + ctx: Arc, + spec: InnerJoinSpec, + phase: DriverPhase, + /// Build partitions from the initial scatter, consumed when the probe + /// scatter finishes and the pair queue is built. + initial_build: Vec, + base_level: usize, + queue: VecDeque, + /// Monotonic pair counter for reservation naming across recursion. + pair_seq: usize, + /// Outer metrics: probe input rows/batches are recorded during scatter. + outer_metrics: BuildProbeJoinMetrics, + /// Private metrics for inner per-pair streams (not exposed; prevents + /// double counting on the operator's metrics set). + inner_metrics_set: ExecutionPlanMetricsSet, + /// Held for the driver's lifetime; released on drop. + _headroom: MemoryReservation, +} + +impl SpillJoinDriver { + pub(super) fn new( + spilled: SpilledBuild, + probe: SendableRecordBatchStream, + ctx: Arc, + spec: InnerJoinSpec, + outer_metrics: BuildProbeJoinMetrics, + ) -> Self { + let SpilledBuild { + partitions, + partition_count, + level, + bounds: _, + headroom, + } = spilled; + + let scatter = SideScatter::new( + partition_count, + spec.on_right.clone(), + scatter_random_state(level), + &ctx.probe_spill, + ctx.max_spill_file_size, + ctx.headroom_bytes / 2, + "hash_join_probe_spill", + ); + + Self { + spec, + phase: DriverPhase::ScatterProbe { probe, scatter }, + initial_build: partitions, + base_level: level, + queue: VecDeque::new(), + pair_seq: 0, + outer_metrics, + inner_metrics_set: ExecutionPlanMetricsSet::new(), + _headroom: headroom, + ctx, + } + } + + /// Whether an empty side lets the pair be skipped entirely, i.e. the + /// join type provably produces no rows for it. Pairs that emit unmatched + /// rows from a non-empty side must still run: the inner stream handles + /// an empty probe (unmatched build emission) and an empty build + /// (probe-side null padding) natively. + fn can_skip_pair(&self, build_rows: usize, probe_rows: usize) -> bool { + use JoinType::*; + match self.spec.join_type { + // Emit only matches: either side empty means no output. + Inner | LeftSemi | RightSemi => build_rows == 0 || probe_rows == 0, + // Emit (un)matched build rows: only an empty build side is silent. + Left | LeftAnti | LeftMark => build_rows == 0, + // Emit (un)matched probe rows: only an empty probe side is silent. + Right | RightAnti | RightMark => probe_rows == 0, + // Emits unmatched rows from both sides. + Full => build_rows == 0 && probe_rows == 0, + } + } + + /// Chunking the build side is correct only for join types whose output + /// is decided per build row (each build row lives in exactly one chunk). + /// Types that emit unmatched PROBE rows would need cross-chunk match + /// tracking, which chunked execution cannot provide. + fn chunked_fallback_supported(&self) -> bool { + matches!( + self.spec.join_type, + JoinType::Inner + | JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::LeftMark + ) + } + + pub(super) fn poll_next( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + loop { + match &self.phase { + DriverPhase::ScatterProbe { .. } => { + ready!(self.poll_scatter_probe(cx))?; + } + DriverPhase::NextPair => self.select_next_pair()?, + DriverPhase::LoadBuild { .. } => { + ready!(self.poll_load_build(cx))?; + } + DriverPhase::RunPair { .. } => match ready!(self.poll_run_pair(cx)) { + Some(item) => return Poll::Ready(Some(item)), + None => self.phase = DriverPhase::NextPair, + }, + DriverPhase::RepartitionBuild { .. } => { + ready!(self.poll_repartition_build(cx))?; + } + DriverPhase::RepartitionProbe { .. } => { + ready!(self.poll_repartition_probe(cx))?; + } + DriverPhase::ChunkedLoad { .. } => { + ready!(self.poll_chunked_load(cx))?; + } + DriverPhase::ChunkedRun { .. } => { + match ready!(self.poll_chunked_run(cx)) { + Some(item) => return Poll::Ready(Some(item)), + None => self.finish_chunk()?, + } + } + DriverPhase::Done => return Poll::Ready(None), + } + } + } + + /// One step of probe scatter. Returns `Ready(Ok(()))` after a state + /// change; the caller loops. + fn poll_scatter_probe(&mut self, cx: &mut Context<'_>) -> Poll> { + let polled = { + let DriverPhase::ScatterProbe { probe, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + ready!(probe.poll_next_unpin(cx)) + }; + match polled { + Some(batch) => { + let batch = batch?; + self.outer_metrics.input_batches.add(1); + self.outer_metrics.input_rows.add(batch.num_rows()); + let DriverPhase::ScatterProbe { scatter, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + scatter.scatter_batch(&batch)?; + } + None => { + let DriverPhase::ScatterProbe { scatter, .. } = + std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + let probe_parts = scatter.finish()?; + let level = self.base_level; + self.queue = std::mem::take(&mut self.initial_build) + .into_iter() + .zip(probe_parts) + .map(|(build, probe)| PairWork { + build, + probe, + level, + allow_recursion: true, + }) + .collect(); + } + } + Poll::Ready(Ok(())) + } + + /// Pop the next pair; skip pairs this join type provably emits nothing + /// for. + fn select_next_pair(&mut self) -> Result<()> { + loop { + let Some(work) = self.queue.pop_front() else { + self.phase = DriverPhase::Done; + return Ok(()); + }; + + if self.can_skip_pair(work.build.rows, work.probe.rows) { + // Dropping the sides deletes their spill files. + continue; + } + + let pair_id = self.pair_seq; + self.pair_seq += 1; + + let reservation = MemoryConsumer::new(format!( + "HashJoinSpillPartition[{}.{}]", + self.ctx.partition, pair_id + )) + .with_can_spill(true) + .register(&self.ctx.pool); + + let parent_build_bytes = work.build.bytes; + let stream = + spilled_side_stream(work.build.files, Arc::clone(&self.ctx.build_spill)); + self.phase = DriverPhase::LoadBuild { + pair_id, + level: work.level, + allow_recursion: work.allow_recursion, + probe: work.probe, + stream, + batches: Vec::new(), + num_rows: 0, + reservation, + table_reserved: 0, + parent_build_bytes, + }; + return Ok(()); + } + } + + /// One step of loading the current pair's build partition. + fn poll_load_build(&mut self, cx: &mut Context<'_>) -> Poll> { + let polled = { + let DriverPhase::LoadBuild { stream, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected LoadBuild phase")); + }; + ready!(stream.poll_next_unpin(cx)) + }; + match polled { + Some(batch) => { + let batch = batch?; + let size = get_record_batch_memory_size(&batch); + let grow_result = { + let DriverPhase::LoadBuild { + reservation, + num_rows, + table_reserved, + .. + } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected LoadBuild phase")); + }; + // Also reserve the hash table's estimated growth so the + // pair overflows here (where recursion/fallback applies) + // rather than at table construction. + let table_delta = hash_table_estimate(*num_rows + batch.num_rows())? + .saturating_sub(*table_reserved); + reservation + .try_grow(size + table_delta) + .map(|()| table_delta) + }; + match grow_result { + Ok(table_delta) => { + let DriverPhase::LoadBuild { + batches, + num_rows, + table_reserved, + .. + } = &mut self.phase + else { + return Poll::Ready(internal_err!( + "expected LoadBuild phase" + )); + }; + *num_rows += batch.num_rows(); + *table_reserved += table_delta; + batches.push(batch); + } + Err(e) => { + return Poll::Ready(self.handle_oversized_build(batch, e)); + } + } + } + None => { + let DriverPhase::LoadBuild { + pair_id, + probe, + batches, + num_rows, + reservation, + table_reserved, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return Poll::Ready(internal_err!("expected LoadBuild phase")); + }; + let inner = self.build_inner_join( + pair_id, + batches, + num_rows, + reservation, + table_reserved, + probe.files, + )?; + self.phase = DriverPhase::RunPair { inner }; + } + } + Poll::Ready(Ok(())) + } + + /// The current pair's build partition exceeded the budget mid-load. + /// Choose: recursive repartition, chunked fallback, or a clean error. + fn handle_oversized_build( + &mut self, + failed_batch: RecordBatch, + source: DataFusionError, + ) -> Result<()> { + let (level, allow_recursion) = { + let DriverPhase::LoadBuild { + level, + allow_recursion, + .. + } = &self.phase + else { + return internal_err!("expected LoadBuild phase"); + }; + (*level, *allow_recursion) + }; + + if allow_recursion && level < self.ctx.max_recursion_depth { + self.begin_repartition(failed_batch) + } else if self.chunked_fallback_supported() { + self.begin_chunked(failed_batch) + } else { + let join_type = self.spec.join_type; + Err(source.context(format!( + "spilled hash join partition remains oversized after {level} \ + repartition pass(es) (extreme key skew), and join type \ + {join_type:?} decides probe-side rows from matches against \ + the WHOLE build side, which the chunked build fallback \ + cannot track across chunks; increase the memory budget to \ + run this join" + ))) + } + } + + /// Re-scatter the oversized pair's build side at the next seed level: + /// already-loaded batches first (releasing their memory), then the rest + /// of the pair's build stream. + fn begin_repartition(&mut self, failed_batch: RecordBatch) -> Result<()> { + let DriverPhase::LoadBuild { + level, + probe, + stream, + batches, + reservation, + parent_build_bytes, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return internal_err!("expected LoadBuild phase"); + }; + + self.ctx.repartition_passes.add(1); + log::info!( + "hash join spill: repartitioning oversized pair (partition={} level={} parent_bytes={})", + self.ctx.partition, + level, + parent_build_bytes + ); + let child_level = level + 1; + let mut scatter = SideScatter::new( + CHILD_FANOUT, + self.ctx.on_left.clone(), + scatter_random_state(child_level), + &self.ctx.build_spill, + self.ctx.max_spill_file_size, + self.ctx.headroom_bytes / 2, + "hash_join_build_respill", + ); + + for batch in batches { + scatter.scatter_batch(&batch)?; + } + // The freed bytes exactly cover the scattered batches — this + // reservation held nothing else. + drop(reservation); + scatter.scatter_batch(&failed_batch)?; + drop(failed_batch); + + self.phase = DriverPhase::RepartitionBuild { + level, + probe, + stream, + scatter, + parent_build_bytes, + }; + Ok(()) + } + + fn poll_repartition_build(&mut self, cx: &mut Context<'_>) -> Poll> { + let polled = { + let DriverPhase::RepartitionBuild { stream, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected RepartitionBuild phase")); + }; + ready!(stream.poll_next_unpin(cx)) + }; + match polled { + Some(batch) => { + let batch = batch?; + let DriverPhase::RepartitionBuild { scatter, .. } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected RepartitionBuild phase")); + }; + scatter.scatter_batch(&batch)?; + } + None => { + let DriverPhase::RepartitionBuild { + level, + probe, + scatter, + parent_build_bytes, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return Poll::Ready(internal_err!("expected RepartitionBuild phase")); + }; + let children_build = scatter.finish()?; + + // No-shrink shortcut: when the largest child still holds + // most of the parent, the mass sits on duplicate keys and + // more passes cannot split it — the children go straight to + // the fallback if they do not fit. + let max_child = children_build.iter().map(|c| c.bytes).max(); + let allow_recursion = max_child + .is_none_or(|m| m.saturating_mul(4) <= parent_build_bytes * 3); + + let child_level = level + 1; + let probe_scatter = SideScatter::new( + CHILD_FANOUT, + self.spec.on_right.clone(), + scatter_random_state(child_level), + &self.ctx.probe_spill, + self.ctx.max_spill_file_size, + self.ctx.headroom_bytes / 2, + "hash_join_probe_respill", + ); + let probe_stream = + spilled_side_stream(probe.files, Arc::clone(&self.ctx.probe_spill)); + self.phase = DriverPhase::RepartitionProbe { + children_build, + child_level, + allow_recursion, + stream: probe_stream, + scatter: probe_scatter, + }; + } + } + Poll::Ready(Ok(())) + } + + fn poll_repartition_probe(&mut self, cx: &mut Context<'_>) -> Poll> { + let polled = { + let DriverPhase::RepartitionProbe { stream, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected RepartitionProbe phase")); + }; + ready!(stream.poll_next_unpin(cx)) + }; + match polled { + Some(batch) => { + let batch = batch?; + let DriverPhase::RepartitionProbe { scatter, .. } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected RepartitionProbe phase")); + }; + scatter.scatter_batch(&batch)?; + } + None => { + let DriverPhase::RepartitionProbe { + children_build, + child_level, + allow_recursion, + scatter, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return Poll::Ready(internal_err!("expected RepartitionProbe phase")); + }; + let children_probe = scatter.finish()?; + // Children go to the FRONT so their disk space is reclaimed + // before other top-level pairs run. + for (build, probe) in children_build.into_iter().zip(children_probe).rev() + { + self.queue.push_front(PairWork { + build, + probe, + level: child_level, + allow_recursion, + }); + } + } + } + Poll::Ready(Ok(())) + } + + /// Enter the chunked fallback: the already-loaded batches (which fit) + /// become chunk 0; the failed batch leads chunk 1. + fn begin_chunked(&mut self, failed_batch: RecordBatch) -> Result<()> { + let DriverPhase::LoadBuild { + pair_id, + level, + probe, + stream, + batches, + num_rows, + reservation, + table_reserved, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return internal_err!("expected LoadBuild phase"); + }; + + if batches.is_empty() { + // Not even one batch fits: nothing to chunk. + return internal_err!( + "spilled hash join partition {pair_id}: a single build batch \ + exceeds the memory budget" + ); + } + + let state = ChunkedState { + pair_id, + level, + build_stream: Some(stream), + pending: Some(failed_batch), + probe_files: probe.files, + chunk_index: 0, + }; + let inner = self.start_chunk_join( + &state, + batches, + num_rows, + reservation, + table_reserved, + )?; + self.phase = DriverPhase::ChunkedRun { inner, state }; + Ok(()) + } + + /// One step of loading the next chunk of build rows. + fn poll_chunked_load(&mut self, cx: &mut Context<'_>) -> Poll> { + // A pending batch from the previous chunk leads this one. + let pending_step = { + let DriverPhase::ChunkedLoad { + state, + batches, + num_rows, + reservation, + table_reserved, + } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected ChunkedLoad phase")); + }; + if let Some(batch) = state.pending.take() { + let size = get_record_batch_memory_size(&batch); + let table_delta = hash_table_estimate(*num_rows + batch.num_rows())? + .saturating_sub(*table_reserved); + match reservation.try_grow(size + table_delta) { + Ok(()) => { + *num_rows += batch.num_rows(); + *table_reserved += table_delta; + batches.push(batch); + true + } + Err(e) => { + if batches.is_empty() { + let pair_id = state.pair_id; + return Poll::Ready(Err(e.context(format!( + "spilled hash join partition {pair_id}: a single \ + build batch exceeds the memory budget" + )))); + } + // Chunk full before the pending batch: run with what + // we have, keep the batch pending. + state.pending = Some(batch); + return Poll::Ready(self.start_chunk_run()); + } + } + } else { + false + } + }; + if pending_step { + return Poll::Ready(Ok(())); + } + + // Poll the build stream (it may already be exhausted). + let polled = { + let DriverPhase::ChunkedLoad { state, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected ChunkedLoad phase")); + }; + match &mut state.build_stream { + Some(stream) => ready!(stream.poll_next_unpin(cx)), + None => None, + } + }; + match polled { + Some(batch) => { + let batch = batch?; + let size = get_record_batch_memory_size(&batch); + let DriverPhase::ChunkedLoad { + state, + batches, + num_rows, + reservation, + table_reserved, + } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected ChunkedLoad phase")); + }; + let table_delta = hash_table_estimate(*num_rows + batch.num_rows())? + .saturating_sub(*table_reserved); + match reservation.try_grow(size + table_delta) { + Ok(()) => { + *num_rows += batch.num_rows(); + *table_reserved += table_delta; + batches.push(batch); + } + Err(e) => { + if batches.is_empty() { + let pair_id = state.pair_id; + return Poll::Ready(Err(e.context(format!( + "spilled hash join partition {pair_id}: a single \ + build batch exceeds the memory budget" + )))); + } + state.pending = Some(batch); + return Poll::Ready(self.start_chunk_run()); + } + } + } + None => { + { + let DriverPhase::ChunkedLoad { state, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected ChunkedLoad phase")); + }; + state.build_stream = None; + } + return Poll::Ready(self.start_chunk_run()); + } + } + Poll::Ready(Ok(())) + } + + /// Transition ChunkedLoad -> ChunkedRun with the accumulated chunk. + fn start_chunk_run(&mut self) -> Result<()> { + let DriverPhase::ChunkedLoad { + state, + batches, + num_rows, + reservation, + table_reserved, + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return internal_err!("expected ChunkedLoad phase"); + }; + if batches.is_empty() { + // Final poll found no more rows: the pair is done. + debug_assert!(state.exhausted()); + self.phase = DriverPhase::NextPair; + return Ok(()); + } + let inner = self.start_chunk_join( + &state, + batches, + num_rows, + reservation, + table_reserved, + )?; + self.phase = DriverPhase::ChunkedRun { inner, state }; + Ok(()) + } + + fn poll_run_pair( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + let DriverPhase::RunPair { inner } = &mut self.phase else { + return Poll::Ready(Some(internal_err!("expected RunPair phase"))); + }; + inner.poll_next_unpin(cx) + } + + fn poll_chunked_run( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + let DriverPhase::ChunkedRun { inner, .. } = &mut self.phase else { + return Poll::Ready(Some(internal_err!("expected ChunkedRun phase"))); + }; + inner.poll_next_unpin(cx) + } + + /// A chunk's join finished: move to the next chunk or the next pair. + fn finish_chunk(&mut self) -> Result<()> { + let DriverPhase::ChunkedRun { state, .. } = + std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return internal_err!("expected ChunkedRun phase"); + }; + if state.exhausted() { + // Dropping the state deletes the probe partition's files. + return Ok(()); + } + let mut state = state; + state.chunk_index += 1; + let reservation = MemoryConsumer::new(format!( + "HashJoinSpillChunk[{}.{}.{}]", + self.ctx.partition, state.pair_id, state.chunk_index + )) + .with_can_spill(true) + .register(&self.ctx.pool); + self.phase = DriverPhase::ChunkedLoad { + state, + batches: Vec::new(), + num_rows: 0, + reservation, + table_reserved: 0, + }; + Ok(()) + } + + /// Build one chunk's `JoinLeftData` and join it against a full re-read + /// of the pair's probe partition. + fn start_chunk_join( + &mut self, + state: &ChunkedState, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + table_reserved: usize, + ) -> Result { + self.ctx.fallback_chunks.add(1); + log::info!( + "hash join spill: chunked fallback chunk (partition={} pair={} chunk={} rows={})", + self.ctx.partition, + state.pair_id, + state.chunk_index, + num_rows + ); + let _ = state.level; // recorded for diagnostics via consumer names + self.build_inner_join( + state.pair_id, + batches, + num_rows, + reservation, + table_reserved, + state.probe_files.clone(), + ) + } + + fn build_inner_join( + &mut self, + pair_id: usize, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + table_reserved: usize, + probe_files: Vec, + ) -> Result { + let inner_metrics = BuildProbeJoinMetrics::new(pair_id, &self.inner_metrics_set); + let array_map_count = Count::new(); + let need_bitmap = need_produce_result_in_final(self.spec.join_type); + + let left_data = build_left_data( + &batches, + num_rows, + self.ctx.build_spill.schema(), + &self.ctx.on_left, + &self.spec.random_state, + reservation, + &inner_metrics, + need_bitmap, + 1, + None, + false, + &self.ctx.config, + self.spec.null_equality, + &array_map_count, + table_reserved, + None, + )?; + drop(batches); + + let probe_stream = + spilled_side_stream(probe_files, Arc::clone(&self.ctx.probe_spill)); + + let build_output = BuildPhaseOutput::InMemory(Arc::new(left_data)); + let left_fut = OnceFut::new(std::future::ready(Ok(build_output))); + + let inner = HashJoinStream::new( + self.spec.partition, + Arc::clone(&self.spec.schema), + self.spec.on_right.clone(), + self.spec.filter.clone(), + self.spec.join_type, + probe_stream, + self.spec.random_state.clone(), + inner_metrics, + self.spec.column_indices.clone(), + self.spec.null_equality, + HashJoinStreamState::WaitBuildSide, + BuildSide::Initial(BuildSideInitialState { left_fut }), + self.spec.batch_size, + vec![], + false, + None, + PartitionMode::Partitioned, + false, + None, + None, + ); + + Ok(Box::pin(inner)) + } +} + +enum SharedDriverPhase { + /// Scatter this output partition's OWN probe stream into `K` partitions. + ScatterProbe { + probe: SendableRecordBatchStream, + scatter: SideScatter, + }, + /// Move to the next shared build partition. + NextK, + /// Run k's join: shared build (lazily built once across partitions) + /// against this partition's probe files for k. + RunK { + inner: SendableRecordBatchStream, + }, + Done, +} + +/// Drives one output partition's share of a spilled `CollectLeft` join. +/// +/// Unlike the Partitioned-mode driver, every k MUST run even when this +/// partition's probe side is empty: the shared `JoinLeftData` counts probe +/// completions, and skipping would leave the counter high so unmatched +/// build rows would never be emitted. Oversized shared partitions surface a +/// clean error (no recursion/chunking for the shared build in v1 — the +/// shared visited bitmap cannot be split across chunks). +pub(super) struct SharedSpillJoinDriver { + ctx: Arc, + spec: InnerJoinSpec, + shared: Arc, + phase: SharedDriverPhase, + probe_parts: Vec>, + next_k: usize, + /// Outer metrics: probe input rows/batches are recorded during scatter. + outer_metrics: BuildProbeJoinMetrics, + /// Private metrics for inner per-k streams. + inner_metrics_set: ExecutionPlanMetricsSet, +} + +impl SharedSpillJoinDriver { + pub(super) fn new( + shared: Arc, + probe: SendableRecordBatchStream, + ctx: Arc, + spec: InnerJoinSpec, + outer_metrics: BuildProbeJoinMetrics, + ) -> Self { + let scatter = SideScatter::new( + shared.partition_count, + spec.on_right.clone(), + scatter_random_state(shared.level), + &ctx.probe_spill, + ctx.max_spill_file_size, + ctx.headroom_bytes / 2, + "hash_join_probe_spill", + ); + Self { + spec, + phase: SharedDriverPhase::ScatterProbe { probe, scatter }, + probe_parts: Vec::new(), + next_k: 0, + outer_metrics, + inner_metrics_set: ExecutionPlanMetricsSet::new(), + ctx, + shared, + } + } + + pub(super) fn poll_next( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + loop { + match &self.phase { + SharedDriverPhase::ScatterProbe { .. } => { + ready!(self.poll_scatter_probe(cx))?; + } + SharedDriverPhase::NextK => self.start_next_k()?, + SharedDriverPhase::RunK { .. } => { + let polled = { + let SharedDriverPhase::RunK { inner } = &mut self.phase else { + return Poll::Ready(Some(internal_err!( + "expected RunK phase" + ))); + }; + ready!(inner.poll_next_unpin(cx)) + }; + match polled { + Some(item) => return Poll::Ready(Some(item)), + None => { + self.phase = SharedDriverPhase::NextK; + } + } + } + SharedDriverPhase::Done => return Poll::Ready(None), + } + } + } + + fn poll_scatter_probe(&mut self, cx: &mut Context<'_>) -> Poll> { + let polled = { + let SharedDriverPhase::ScatterProbe { probe, .. } = &mut self.phase else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + ready!(probe.poll_next_unpin(cx)) + }; + match polled { + Some(batch) => { + let batch = batch?; + self.outer_metrics.input_batches.add(1); + self.outer_metrics.input_rows.add(batch.num_rows()); + let SharedDriverPhase::ScatterProbe { scatter, .. } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + scatter.scatter_batch(&batch)?; + } + None => { + let SharedDriverPhase::ScatterProbe { scatter, .. } = + std::mem::replace(&mut self.phase, SharedDriverPhase::NextK) + else { + return Poll::Ready(internal_err!("expected ScatterProbe phase")); + }; + self.probe_parts = scatter.finish()?.into_iter().map(Some).collect(); + } + } + Poll::Ready(Ok(())) + } + + fn start_next_k(&mut self) -> Result<()> { + let k = self.next_k; + if k >= self.shared.partition_count { + self.phase = SharedDriverPhase::Done; + return Ok(()); + } + self.next_k += 1; + + let probe = self.probe_parts[k].take().ok_or_else(|| { + DataFusionError::Internal(format!( + "shared spill probe partition {k} already consumed" + )) + })?; + + let (bitmap, counter) = &self.shared.shared_state[k]; + let left_fut = OnceFut::new(shared_build_loader( + Arc::clone(&self.ctx), + Arc::clone(&self.shared.partitions[k].files), + k, + self.spec.random_state.clone(), + self.spec.join_type, + self.spec.null_equality, + self.shared.probe_threads, + (Arc::clone(bitmap), Arc::clone(counter)), + )); + + let inner_metrics = BuildProbeJoinMetrics::new(k, &self.inner_metrics_set); + let probe_stream = + spilled_side_stream(probe.files, Arc::clone(&self.ctx.probe_spill)); + let inner = HashJoinStream::new( + self.spec.partition, + Arc::clone(&self.spec.schema), + self.spec.on_right.clone(), + self.spec.filter.clone(), + self.spec.join_type, + probe_stream, + self.spec.random_state.clone(), + inner_metrics, + self.spec.column_indices.clone(), + self.spec.null_equality, + HashJoinStreamState::WaitBuildSide, + BuildSide::Initial(BuildSideInitialState { left_fut }), + self.spec.batch_size, + vec![], + false, + None, + PartitionMode::CollectLeft, + false, + None, + None, + ); + self.phase = SharedDriverPhase::RunK { + inner: Box::pin(inner), + }; + Ok(()) + } +} + +/// Load shared build partition `k` into memory and construct its +/// `JoinLeftData` with `probe_threads_count = N` so the existing +/// last-finisher machinery handles shared unmatched-row emission. +#[expect(clippy::too_many_arguments)] +fn shared_build_loader( + ctx: Arc, + files: Arc>, + k: usize, + random_state: RandomState, + join_type: JoinType, + null_equality: NullEquality, + probe_threads: usize, + shared_state: (Arc, Arc), +) -> BoxFuture<'static, Result> { + Box::pin(async move { + let reservation = + MemoryConsumer::new(format!("HashJoinSpillShared[{}.{k}]", ctx.partition)) + .with_can_spill(true) + .register(&ctx.pool); + + let mut stream = + spilled_side_stream(files.as_ref().clone(), Arc::clone(&ctx.build_spill)); + let mut batches = Vec::new(); + let mut num_rows = 0usize; + let mut table_reserved = 0usize; + while let Some(batch) = stream.next().await { + let batch = batch?; + let size = get_record_batch_memory_size(&batch); + let table_delta = hash_table_estimate(num_rows + batch.num_rows())? + .saturating_sub(table_reserved); + reservation.try_grow(size + table_delta).map_err(|e| { + e.context(format!( + "shared spilled hash join build partition {k} does not fit \ + in the memory budget (CollectLeft spill does not \ + repartition recursively); increase the memory budget" + )) + })?; + table_reserved += table_delta; + num_rows += batch.num_rows(); + batches.push(batch); + } + + let metrics_set = ExecutionPlanMetricsSet::new(); + let inner_metrics = BuildProbeJoinMetrics::new(k, &metrics_set); + let array_map_count = Count::new(); + let need_bitmap = need_produce_result_in_final(join_type); + let data = build_left_data( + &batches, + num_rows, + ctx.build_spill.schema(), + &ctx.on_left, + &random_state, + reservation, + &inner_metrics, + need_bitmap, + probe_threads, + None, + false, + &ctx.config, + null_equality, + &array_map_count, + table_reserved, + Some(shared_state), + )?; + Ok(BuildPhaseOutput::InMemory(Arc::new(data))) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::joins::HashJoinExec; + use crate::joins::PartitionMode; + use crate::metrics::SpillMetrics; + use crate::repartition::REPARTITION_RANDOM_STATE; + use crate::test::TestMemoryExec; + use crate::{ExecutionPlan, common}; + + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::util::pretty::pretty_format_batches; + use datafusion_common::JoinType; + use datafusion_common::NullEquality; + use datafusion_execution::TaskContext; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_physical_expr::expressions::Column; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])) + } + + fn make_batch(start: i32, len: usize) -> RecordBatch { + let keys = Int32Array::from_iter_values(start..start + len as i32); + let vals = Int32Array::from_iter_values((0..len as i32).map(|v| v * 10)); + RecordBatch::try_new(test_schema(), vec![Arc::new(keys), Arc::new(vals)]).unwrap() + } + + fn key_exprs() -> Vec { + vec![Arc::new(Column::new("k", 0)) as _] + } + + fn test_spill_manager(env: Arc) -> Arc { + Arc::new(SpillManager::new( + env, + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + test_schema(), + )) + } + + async fn read_all( + side: SpilledSide, + manager: &Arc, + ) -> Vec { + let stream = spilled_side_stream(side.files, Arc::clone(manager)); + common::collect(stream).await.unwrap() + } + + fn sorted_rows(batches: &[RecordBatch]) -> Vec { + let non_empty: Vec = batches + .iter() + .filter(|b| b.num_rows() > 0) + .cloned() + .collect(); + if non_empty.is_empty() { + return vec![]; + } + let formatted = pretty_format_batches(&non_empty).unwrap().to_string(); + let mut rows: Vec = formatted + .lines() + .filter(|l| l.starts_with('|') && !l.contains(" k ")) + .map(|s| s.to_string()) + .collect(); + rows.sort(); + rows + } + + /// Cheap multiset fingerprint of a result set: (row count, wrapping sum + /// of row hashes). Equal row multisets yield equal fingerprints; the + /// converse holds with overwhelming probability, which is plenty for + /// comparing a spilled run against its in-memory reference. + fn result_fingerprint(batches: &[RecordBatch]) -> (usize, u64) { + let state = RandomState::with_seeds(1, 2, 3, 4); + let mut count = 0usize; + let mut sum = 0u64; + for batch in batches.iter().filter(|b| b.num_rows() > 0) { + let mut hashes = vec![0u64; batch.num_rows()]; + create_hashes(batch.columns(), &state, &mut hashes).unwrap(); + count += batch.num_rows(); + sum = hashes.iter().fold(sum, |acc, h| acc.wrapping_add(*h)); + } + (count, sum) + } + + #[tokio::test] + async fn scatter_round_trip_preserves_all_rows() { + let env = Arc::new(RuntimeEnv::default()); + let manager = test_spill_manager(Arc::clone(&env)); + let mut scatter = SideScatter::new( + 8, + key_exprs(), + scatter_random_state(0), + &manager, + 128 * 1024 * 1024, + 1024 * 1024, + "test_scatter", + ); + + let inputs: Vec = + (0..10).map(|i| make_batch(i * 1000, 1000)).collect(); + for batch in &inputs { + scatter.scatter_batch(batch).unwrap(); + } + let sides = scatter.finish().unwrap(); + assert_eq!(sides.len(), 8); + let total_rows: usize = sides.iter().map(|s| s.rows).sum(); + assert_eq!(total_rows, 10_000); + + let mut all = Vec::new(); + for side in sides { + all.extend(read_all(side, &manager).await); + } + assert_eq!(sorted_rows(&all), sorted_rows(&inputs)); + } + + /// Rows that all land in ONE bucket of the repartition hash (the shape + /// Partitioned-mode inputs actually have) must still spread across the + /// scatter's disk partitions — i.e. the seeds are independent. + #[tokio::test] + async fn scatter_spreads_rows_prebucketed_by_repartition_hash() { + const K: u64 = 8; + + // Select keys whose REPARTITION hash lands in bucket 0. + let candidates = Int32Array::from_iter_values(0..200_000); + let mut hashes = vec![0u64; candidates.len()]; + let candidate_arrays: Vec = + vec![Arc::new(candidates.clone())]; + create_hashes( + &candidate_arrays, + REPARTITION_RANDOM_STATE.random_state(), + &mut hashes, + ) + .unwrap(); + let selected: Vec = hashes + .iter() + .enumerate() + .filter(|(_, h)| *h % K == 0) + .map(|(i, _)| candidates.value(i)) + .collect(); + assert!(selected.len() > 10_000, "want a meaningful sample"); + + let keys = Int32Array::from(selected.clone()); + let vals = Int32Array::from(vec![0; selected.len()]); + let batch = + RecordBatch::try_new(test_schema(), vec![Arc::new(keys), Arc::new(vals)]) + .unwrap(); + + let env = Arc::new(RuntimeEnv::default()); + let manager = test_spill_manager(Arc::clone(&env)); + let mut scatter = SideScatter::new( + K as usize, + key_exprs(), + scatter_random_state(0), + &manager, + 128 * 1024 * 1024, + 1024 * 1024, + "test_seed_independence", + ); + scatter.scatter_batch(&batch).unwrap(); + let sides = scatter.finish().unwrap(); + + let non_empty = sides.iter().filter(|s| s.rows > 0).count(); + assert_eq!(non_empty, K as usize, "scatter must use every partition"); + let max = sides.iter().map(|s| s.rows).max().unwrap(); + let min = sides.iter().map(|s| s.rows).min().unwrap(); + assert!( + max < min * 3, + "scatter should be roughly uniform, got min={min} max={max}" + ); + } + + #[tokio::test] + async fn partition_writer_rotates_and_cleans_up_files() { + let env = Arc::new(RuntimeEnv::default()); + let manager = test_spill_manager(Arc::clone(&env)); + // Tiny rotation threshold: every flushed buffer closes its file. + let mut writer = PartitionSpillWriter::new( + Arc::clone(&manager), + "test_rotation".to_string(), + 1, + ); + for i in 0..4 { + writer.append(make_batch(i * 100_000, 100_000)).unwrap(); + } + let side = writer.finish().unwrap(); + assert!( + side.files.len() >= 2, + "expected file rotation, got {} file(s)", + side.files.len() + ); + assert_eq!(side.rows, 400_000); + + let paths: Vec = + side.files.iter().map(|f| f.path().to_path_buf()).collect(); + for p in &paths { + assert!(p.exists()); + } + drop(side); + for p in &paths { + assert!(!p.exists(), "spill file must be deleted on drop: {p:?}"); + } + } + + // ---- end-to-end HashJoinExec tests ---- + + fn partitioned_join( + build_rows: usize, + join_type: JoinType, + ) -> (Arc, RecordBatch, RecordBatch) { + // Feed the build side in ~100k-row batches so the overflow happens + // mid-stream (buffered batches get re-scattered, the rest stream). + let left_batches: Vec = (0..build_rows.div_ceil(100_000)) + .map(|i| { + make_batch(i as i32 * 100_000, 100_000.min(build_rows - i * 100_000)) + }) + .collect(); + let left_batch = left_batches[0].clone(); + let right_batch = make_batch(0, 4096); + let left = + TestMemoryExec::try_new_exec(&[left_batches], left_batch.schema(), None) + .unwrap(); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()]], + right_batch.schema(), + None, + ) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("k", &left_batch.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("k", &right_batch.schema()).unwrap()) as _, + )]; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + (Arc::new(join), left_batch, right_batch) + } + + fn spill_task_ctx(memory_limit: usize) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .build_arc() + .unwrap(); + let mut session_config = SessionConfig::default().with_batch_size(4096); + { + let exec = &mut session_config.options_mut().execution; + exec.enable_hash_join_spill = true; + exec.hash_join_spill_headroom_bytes = 256 * 1024; + exec.hash_join_spill_partition_count = 16; + } + Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ) + } + + #[tokio::test] + async fn forced_spill_inner_join_matches_in_memory() { + // ~2M rows * 8B ≈ 16MB build side vs an 8MB pool: must spill. With + // K=16, each pair holds ~1MB of data plus its hash table — well + // within the budget. + let (join, _, _) = partitioned_join(2_000_000, JoinType::Inner); + + let spill_ctx = spill_task_ctx(8 * 1024 * 1024); + let stream = join.execute(0, Arc::clone(&spill_ctx)).unwrap(); + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert_eq!( + metrics + .sum_by_name("join_spill_engaged") + .map(|m| m.as_usize()), + Some(1), + "join must have engaged spilling" + ); + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "spill files must have been written" + ); + + // Reference: same join, no limit, feature off. + let (reference_join, _, _) = partitioned_join(2_000_000, JoinType::Inner); + let reference_ctx = Arc::new(TaskContext::default()); + let reference = + common::collect(reference_join.execute(0, reference_ctx).unwrap()) + .await + .unwrap(); + + let spilled_rows: usize = spilled_result.iter().map(|b| b.num_rows()).sum(); + let reference_rows: usize = reference.iter().map(|b| b.num_rows()).sum(); + assert_eq!(spilled_rows, reference_rows); + assert_eq!(sorted_rows(&spilled_result), sorted_rows(&reference)); + + // All join memory must be released. + assert_eq!(spill_ctx.memory_pool().reserved(), 0); + } + + #[tokio::test] + async fn fast_path_untouched_when_memory_sufficient() { + let (join, _, _) = partitioned_join(100_000, JoinType::Inner); + // Plenty of memory: flag on, but no spill may occur. + let ctx = spill_task_ctx(512 * 1024 * 1024); + let result = common::collect(join.execute(0, ctx).unwrap()) + .await + .unwrap(); + + let metrics = join.metrics().unwrap(); + assert_eq!( + metrics + .sum_by_name("join_spill_engaged") + .map(|m| m.as_usize()), + Some(0) + ); + assert_eq!(metrics.spill_count(), Some(0)); + let rows: usize = result.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4096); + } + + #[tokio::test] + async fn pool_released_when_stream_dropped_mid_spill_join() { + let (join, _, _) = partitioned_join(2_000_000, JoinType::Inner); + let ctx = spill_task_ctx(8 * 1024 * 1024); + let mut stream = join.execute(0, Arc::clone(&ctx)).unwrap(); + + // Pull a single batch, then drop the stream mid-flight. + let first = stream.next().await; + assert!(matches!(first, Some(Ok(_))), "expected at least one batch"); + drop(stream); + + assert_eq!( + ctx.memory_pool().reserved(), + 0, + "dropping the stream must release all join memory" + ); + } + + #[tokio::test] + async fn disk_cap_exhaustion_is_a_clean_error() { + let (join, _, _) = partitioned_join(500_000, JoinType::Inner); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(2 * 1024 * 1024, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .with_max_temp_directory_size(64 * 1024), + ) + .build_arc() + .unwrap(); + let mut session_config = SessionConfig::default().with_batch_size(4096); + { + let exec = &mut session_config.options_mut().execution; + exec.enable_hash_join_spill = true; + exec.hash_join_spill_headroom_bytes = 256 * 1024; + exec.hash_join_spill_partition_count = 16; + } + let ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + let err = common::collect(join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("disk") || err.contains("Resources exhausted"), + "expected a clean disk/resources error, got: {err}" + ); + assert_eq!(ctx.memory_pool().reserved(), 0); + } + + #[tokio::test] + async fn spill_disabled_keeps_todays_clean_error() { + let (join, _, _) = partitioned_join(500_000, JoinType::Inner); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(2 * 1024 * 1024, 1.0) + .build_arc() + .unwrap(); + let ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::default().with_batch_size(4096)) + .with_runtime(runtime), + ); + let err = common::collect(join.execute(0, ctx).unwrap()) + .await + .unwrap_err() + .to_string(); + assert!(err.contains("Resources exhausted"), "got: {err}"); + } + /// Build a Partitioned-mode join from explicit batch lists. + fn join_from_batches( + left_batches: Vec, + right_batches: Vec, + join_type: JoinType, + null_equality: NullEquality, + ) -> Arc { + let left_schema = left_batches[0].schema(); + let right_schema = right_batches[0].schema(); + let left = + TestMemoryExec::try_new_exec(&[left_batches], Arc::clone(&left_schema), None) + .unwrap(); + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + ) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("k", &left_schema).unwrap()) as _, + Arc::new(Column::new_with_schema("k", &right_schema).unwrap()) as _, + )]; + Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::Partitioned, + null_equality, + false, + ) + .unwrap(), + ) + } + + fn chunked_build_batches(rows: usize) -> Vec { + (0..rows.div_ceil(100_000)) + .map(|i| make_batch(i as i32 * 100_000, 100_000.min(rows - i * 100_000))) + .collect() + } + + /// Every join type must produce identical results spilled vs in-memory. + /// The probe range half-overlaps the build range so matches, unmatched + /// build rows, and unmatched probe rows all occur. + #[tokio::test] + async fn forced_spill_matches_in_memory_for_all_join_types() { + use JoinType::*; + let join_types = [ + Inner, Left, Right, Full, LeftSemi, LeftAnti, RightSemi, RightAnti, LeftMark, + RightMark, + ]; + for join_type in join_types { + let build = chunked_build_batches(1_000_000); + // 8192 probe rows: first half matches build keys, second half + // (>= 1M) has no build match. + let probe = vec![make_batch(996_000, 8192)]; + + let spilled_join = join_from_batches( + build.clone(), + probe.clone(), + join_type, + NullEquality::NullEqualsNothing, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let spilled = + common::collect(spilled_join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap_or_else(|e| panic!("{join_type:?} spilled run failed: {e}")); + let metrics = spilled_join.metrics().unwrap(); + assert_eq!( + metrics + .sum_by_name("join_spill_engaged") + .map(|m| m.as_usize()), + Some(1), + "{join_type:?} must have engaged spilling" + ); + + let reference_join = join_from_batches( + build, + probe, + join_type, + NullEquality::NullEqualsNothing, + ); + let reference = common::collect( + reference_join + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!( + result_fingerprint(&spilled), + result_fingerprint(&reference), + "{join_type:?} results differ between spilled and in-memory" + ); + assert_eq!(ctx.memory_pool().reserved(), 0, "{join_type:?} leaked"); + } + } + + fn nullable_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("v", DataType::Int32, false), + ])) + } + + /// Like `make_batch` but every 100th key is NULL. All NULL keys hash to + /// one disk partition, and under NullEqualsNull they cross-join — the + /// density keeps that product (and the partition skew, which is the + /// chunked fallback's job) small. + fn make_nullable_batch(start: i32, len: usize) -> RecordBatch { + let keys = Int32Array::from_iter( + (start..start + len as i32).map(|i| (i % 100 != 0).then_some(i)), + ); + let vals = Int32Array::from_iter_values((0..len as i32).map(|v| v * 10)); + RecordBatch::try_new(nullable_schema(), vec![Arc::new(keys), Arc::new(vals)]) + .unwrap() + } + + /// NULL join keys must behave identically spilled vs in-memory under + /// both null-equality semantics (all NULLs scatter to one partition, so + /// NullEqualsNull matching stays complete). + #[tokio::test] + async fn forced_spill_null_equality_matrix() { + use JoinType::*; + for join_type in [Inner, Left, Full] { + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let build: Vec = (0..10) + .map(|i| make_nullable_batch(i * 100_000, 100_000)) + .collect(); + // Half the probe keys match build keys, half exceed them. + let probe = vec![make_nullable_batch(999_500, 1024)]; + + let spilled_join = join_from_batches( + build.clone(), + probe.clone(), + join_type, + null_equality, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let spilled = + common::collect(spilled_join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap_or_else(|e| { + panic!( + "{join_type:?}/{null_equality:?} spilled run failed: {e}" + ) + }); + let metrics = spilled_join.metrics().unwrap(); + assert_eq!( + metrics + .sum_by_name("join_spill_engaged") + .map(|m| m.as_usize()), + Some(1), + "{join_type:?}/{null_equality:?} must have engaged spilling" + ); + + let reference_join = + join_from_batches(build, probe, join_type, null_equality); + let reference = common::collect( + reference_join + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!( + result_fingerprint(&spilled), + result_fingerprint(&reference), + "{join_type:?}/{null_equality:?} spilled vs in-memory mismatch" + ); + } + } + } + /// Extreme duplicate-key skew: every build row has the same key, so + /// repartitioning cannot shrink the partition and the chunked fallback + /// must carry the join. Inner join: matches only. + #[tokio::test] + async fn extreme_skew_falls_back_to_chunked_inner() { + // 600k rows, all key=42 (~4.8MB) vs a 4MB pool. + let key = Int32Array::from(vec![42; 100_000]); + let build: Vec = (0..6) + .map(|i| { + let vals = Int32Array::from_iter_values(i * 100_000..(i + 1) * 100_000); + RecordBatch::try_new( + test_schema(), + vec![Arc::new(key.clone()), Arc::new(vals)], + ) + .unwrap() + }) + .collect(); + // 4 probe rows match key 42; 4 miss. + let probe = vec![ + RecordBatch::try_new( + test_schema(), + vec![ + Arc::new(Int32Array::from(vec![42, 42, 42, 42, 1, 2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7])), + ], + ) + .unwrap(), + ]; + + let spilled_join = join_from_batches( + build.clone(), + probe.clone(), + JoinType::Inner, + NullEquality::NullEqualsNothing, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let spilled = common::collect(spilled_join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap(); + + let metrics = spilled_join.metrics().unwrap(); + assert!( + metrics + .sum_by_name("join_spill_repartition_passes") + .map(|m| m.as_usize()) + .unwrap_or(0) + >= 1, + "skewed pair should attempt at least one repartition pass" + ); + assert!( + metrics + .sum_by_name("join_spill_fallback_chunks") + .map(|m| m.as_usize()) + .unwrap_or(0) + >= 2, + "skewed pair must run multiple fallback chunks" + ); + + let reference_join = join_from_batches( + build, + probe, + JoinType::Inner, + NullEquality::NullEqualsNothing, + ); + let reference = common::collect( + reference_join + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(result_fingerprint(&spilled), result_fingerprint(&reference)); + assert_eq!(ctx.memory_pool().reserved(), 0); + } + + /// Chunked fallback with a build-side-emission type: LeftAnti emits + /// every build row exactly once across all chunks. + #[tokio::test] + async fn extreme_skew_falls_back_to_chunked_left_anti() { + let key = Int32Array::from(vec![42; 100_000]); + let build: Vec = (0..6) + .map(|i| { + let vals = Int32Array::from_iter_values(i * 100_000..(i + 1) * 100_000); + RecordBatch::try_new( + test_schema(), + vec![Arc::new(key.clone()), Arc::new(vals)], + ) + .unwrap() + }) + .collect(); + // No probe row matches key 42: all 600k build rows are anti-matches. + let probe = vec![make_batch(100, 64)]; + + let spilled_join = join_from_batches( + build.clone(), + probe.clone(), + JoinType::LeftAnti, + NullEquality::NullEqualsNothing, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let spilled = common::collect(spilled_join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap(); + let spilled_rows: usize = spilled.iter().map(|b| b.num_rows()).sum(); + assert_eq!(spilled_rows, 600_000); + + let reference_join = join_from_batches( + build, + probe, + JoinType::LeftAnti, + NullEquality::NullEqualsNothing, + ); + let reference = common::collect( + reference_join + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(result_fingerprint(&spilled), result_fingerprint(&reference)); + } + + /// Probe-side-emission types cannot use the chunked fallback: extreme + /// skew must surface a descriptive clean error, not wrong results. + #[tokio::test] + async fn extreme_skew_unsupported_type_is_a_clean_error() { + let key = Int32Array::from(vec![42; 100_000]); + let build: Vec = (0..6) + .map(|i| { + let vals = Int32Array::from_iter_values(i * 100_000..(i + 1) * 100_000); + RecordBatch::try_new( + test_schema(), + vec![Arc::new(key.clone()), Arc::new(vals)], + ) + .unwrap() + }) + .collect(); + let probe = vec![make_batch(0, 1024)]; + + let join = join_from_batches( + build, + probe, + JoinType::RightSemi, + NullEquality::NullEqualsNothing, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let err = common::collect(join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("extreme key skew") && err.contains("RightSemi"), + "expected a descriptive skew error, got: {err}" + ); + assert_eq!(ctx.memory_pool().reserved(), 0); + } + + /// Moderate skew (too few initial partitions) is fixed by one recursion + /// pass — no fallback chunks needed. + #[tokio::test] + async fn recursion_splits_oversized_uniform_partitions() { + let build = chunked_build_batches(1_000_000); + let probe = vec![make_batch(996_000, 8192)]; + + let spilled_join = join_from_batches( + build.clone(), + probe.clone(), + JoinType::Inner, + NullEquality::NullEqualsNothing, + ); + // Force only 2 initial partitions: each ~4MB against a 4MB pool, so + // both need one repartition pass into 8 children each. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(4 * 1024 * 1024, 1.0) + .build_arc() + .unwrap(); + let mut session_config = SessionConfig::default().with_batch_size(4096); + { + let exec = &mut session_config.options_mut().execution; + exec.enable_hash_join_spill = true; + exec.hash_join_spill_headroom_bytes = 256 * 1024; + exec.hash_join_spill_partition_count = 2; + } + let ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + let spilled = common::collect(spilled_join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap(); + + let metrics = spilled_join.metrics().unwrap(); + assert!( + metrics + .sum_by_name("join_spill_repartition_passes") + .map(|m| m.as_usize()) + .unwrap_or(0) + >= 1, + "expected at least one repartition pass" + ); + assert_eq!( + metrics + .sum_by_name("join_spill_fallback_chunks") + .map(|m| m.as_usize()), + Some(0), + "uniform keys must not need the chunked fallback" + ); + + let reference_join = join_from_batches( + build, + probe, + JoinType::Inner, + NullEquality::NullEqualsNothing, + ); + let reference = common::collect( + reference_join + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(result_fingerprint(&spilled), result_fingerprint(&reference)); + assert_eq!(ctx.memory_pool().reserved(), 0); + } + /// Re-scattering one level's partition at the next level must spread it + /// across all children. (Catches additive seed correlation: with weakly + /// derived per-level seeds, `hash_l1 = hash_l0 + const`, so a parent + /// partition maps into a single child under `% K` and recursion never + /// shrinks anything.) + #[tokio::test] + async fn recursive_scatter_levels_are_decorrelated() { + const K: u64 = 8; + let candidates = Int32Array::from_iter_values(0..200_000); + let candidate_arrays: Vec = + vec![Arc::new(candidates.clone())]; + for parent_level in 0..3usize { + let mut hashes = vec![0u64; candidates.len()]; + create_hashes( + &candidate_arrays, + &scatter_random_state(parent_level), + &mut hashes, + ) + .unwrap(); + // Keys in the parent scatter's bucket 0. + let selected: Vec = hashes + .iter() + .enumerate() + .filter(|(_, h)| *h % K == 0) + .map(|(i, _)| candidates.value(i)) + .collect(); + let keys = Int32Array::from(selected.clone()); + let vals = Int32Array::from(vec![0; selected.len()]); + let batch = + RecordBatch::try_new(test_schema(), vec![Arc::new(keys), Arc::new(vals)]) + .unwrap(); + + let env = Arc::new(RuntimeEnv::default()); + let manager = test_spill_manager(Arc::clone(&env)); + let mut scatter = SideScatter::new( + K as usize, + key_exprs(), + scatter_random_state(parent_level + 1), + &manager, + 128 * 1024 * 1024, + 1024 * 1024, + "test_level_decorrelation", + ); + scatter.scatter_batch(&batch).unwrap(); + let sides = scatter.finish().unwrap(); + let non_empty = sides.iter().filter(|s| s.rows > 0).count(); + assert_eq!( + non_empty, + K as usize, + "level {parent_level}->{}: child scatter must use every \ + partition", + parent_level + 1 + ); + let max = sides.iter().map(|s| s.rows).max().unwrap(); + let min = sides.iter().map(|s| s.rows).min().unwrap(); + assert!( + max < min * 3, + "level {parent_level}->{}: got min={min} max={max}", + parent_level + 1 + ); + } + } + /// Build a CollectLeft-mode join: one shared build, two probe partitions. + fn collect_left_join( + left_batches: Vec, + right_parts: Vec>, + join_type: JoinType, + ) -> Arc { + let left_schema = left_batches[0].schema(); + let right_schema = right_parts[0][0].schema(); + let left = + TestMemoryExec::try_new_exec(&[left_batches], Arc::clone(&left_schema), None) + .unwrap(); + let right = + TestMemoryExec::try_new_exec(&right_parts, Arc::clone(&right_schema), None) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("k", &left_schema).unwrap()) as _, + Arc::new(Column::new_with_schema("k", &right_schema).unwrap()) as _, + )]; + Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) + } + + /// CollectLeft spill: the shared scattered build must produce identical + /// results to the in-memory join across all probe partitions — including + /// exactly-once unmatched build emission (Left/Full/LeftAnti), which + /// exercises the shared per-k bitmap and last-finisher accounting. + #[tokio::test] + async fn collect_left_forced_spill_matches_in_memory() { + use JoinType::*; + for join_type in [Inner, Left, LeftAnti, Full] { + let build = chunked_build_batches(2_000_000); + // Partition 0 half-overlaps the build key range; partition 1 is + // fully matched. + let right_parts = + vec![vec![make_batch(1_996_000, 8192)], vec![make_batch(0, 4096)]]; + + let spilled_join = + collect_left_join(build.clone(), right_parts.clone(), join_type); + // Two probe partitions walk the shared k's out of sync, so the + // pool must admit two resident per-k tables at once: K=32 keeps + // each table ~2.7MB against the 8MB pool. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(8 * 1024 * 1024, 1.0) + .build_arc() + .unwrap(); + let mut session_config = SessionConfig::default().with_batch_size(4096); + { + let exec = &mut session_config.options_mut().execution; + exec.enable_hash_join_spill = true; + exec.hash_join_spill_headroom_bytes = 256 * 1024; + exec.hash_join_spill_partition_count = 32; + } + let ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + // Poll both output partitions concurrently (as real plans do): + // out-of-sync consumption is what the shared per-k accounting + // must survive. + let s0 = spilled_join.execute(0, Arc::clone(&ctx)).unwrap(); + let s1 = spilled_join.execute(1, Arc::clone(&ctx)).unwrap(); + let (r0, r1) = futures::join!(common::collect(s0), common::collect(s1)); + let mut spilled = + r0.unwrap_or_else(|e| panic!("{join_type:?} partition 0 failed: {e}")); + spilled.extend( + r1.unwrap_or_else(|e| panic!("{join_type:?} partition 1 failed: {e}")), + ); + + let metrics = spilled_join.metrics().unwrap(); + assert_eq!( + metrics + .sum_by_name("join_spill_engaged") + .map(|m| m.as_usize()), + Some(1), + "{join_type:?} must have engaged spilling" + ); + + let reference_join = collect_left_join(build, right_parts, join_type); + let ref_ctx = Arc::new(TaskContext::default()); + let (f0, f1) = futures::join!( + common::collect(reference_join.execute(0, Arc::clone(&ref_ctx)).unwrap()), + common::collect(reference_join.execute(1, ref_ctx).unwrap()) + ); + let mut reference = f0.unwrap(); + reference.extend(f1.unwrap()); + + assert_eq!( + result_fingerprint(&spilled), + result_fingerprint(&reference), + "{join_type:?} CollectLeft spilled vs in-memory mismatch" + ); + // The exec node's left_fut cache retains the shared build (and + // its scatter headroom) until the plan drops — the same lifetime + // the in-memory CollectLeft build has today. + drop(spilled_join); + assert_eq!( + ctx.memory_pool().reserved(), + 0, + "{join_type:?} leaked memory" + ); + } + } + /// Dropping the stream mid-way through the chunked fallback must release + /// every reservation (chunk data, table estimate, headroom, scratch). + #[tokio::test] + async fn pool_released_when_dropped_mid_chunked_fallback() { + let key = Int32Array::from(vec![42; 100_000]); + let build: Vec = (0..6) + .map(|i| { + let vals = Int32Array::from_iter_values(i * 100_000..(i + 1) * 100_000); + RecordBatch::try_new( + test_schema(), + vec![Arc::new(key.clone()), Arc::new(vals)], + ) + .unwrap() + }) + .collect(); + let probe = vec![ + RecordBatch::try_new( + test_schema(), + vec![ + Arc::new(Int32Array::from(vec![42, 42, 42, 42])), + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + ], + ) + .unwrap(), + ]; + + let join = join_from_batches( + build, + probe, + JoinType::Inner, + NullEquality::NullEqualsNothing, + ); + let ctx = spill_task_ctx(4 * 1024 * 1024); + let mut stream = join.execute(0, Arc::clone(&ctx)).unwrap(); + // Pull a couple of batches (deep inside the chunked fallback given + // the all-equal keys), then drop mid-flight. + let first = stream.next().await; + assert!(matches!(first, Some(Ok(_))), "expected output before drop"); + let _ = stream.next().await; + drop(stream); + + assert_eq!( + ctx.memory_pool().reserved(), + 0, + "mid-chunk drop must release all join memory" + ); + } + + /// A probe-side error arriving while the spilled probe is being + /// scattered must propagate as-is and release all memory. + #[tokio::test] + async fn probe_error_during_spill_scatter_propagates() { + use crate::test::exec::MockExec; + use datafusion_common::exec_err; + + let build = chunked_build_batches(1_000_000); + let schema = test_schema(); + let good = make_batch(0, 4096); + let right = Arc::new(MockExec::new( + vec![Ok(good), exec_err!("bad probe data")], + Arc::clone(&schema), + )); + + let left = TestMemoryExec::try_new_exec(&[build], schema, None).unwrap(); + let on = vec![( + Arc::new(Column::new("k", 0)) as _, + Arc::new(Column::new("k", 0)) as _, + )]; + let join = Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + + let ctx = spill_task_ctx(4 * 1024 * 1024); + let err = common::collect(join.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .unwrap_err() + .to_string(); + assert!(err.contains("bad probe data"), "got: {err}"); + assert_eq!(ctx.memory_pool().reserved(), 0); + } +} diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index b31982ea3b7b4..f567480131e69 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -28,9 +28,13 @@ use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::joins::Map; use crate::joins::MapOffset; use crate::joins::PartitionMode; -use crate::joins::hash_join::exec::JoinLeftData; +use crate::joins::hash_join::exec::{BuildPhaseOutput, JoinLeftData}; use crate::joins::hash_join::shared_bounds::{ - PartitionBounds, PartitionBuildData, SharedBuildAccumulator, + PartitionBounds, PartitionBuildData, PushdownStrategy, SharedBuildAccumulator, +}; +use crate::joins::hash_join::spill::{ + HashJoinSpillContext, InnerJoinSpec, SharedSpillJoinDriver, SpillJoinDriver, + empty_record_batch_stream, }; use crate::joins::utils::{ OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, @@ -69,7 +73,7 @@ pub(super) enum BuildSide { /// Container for BuildSide::Initial related data pub(super) struct BuildSideInitialState { /// Future for building hash table from build-side input - pub(super) left_fut: OnceFut, + pub(super) left_fut: OnceFut, } /// Container for BuildSide::Ready related data @@ -133,6 +137,10 @@ pub(super) enum HashJoinStreamState { ProcessProbeBatch(ProcessProbeBatchState), /// Indicates that probe-side has been fully processed ExhaustedProbeSide, + /// The build side spilled to disk partitions; `spill_driver` scatters the + /// probe side and joins partition pairs, and this stream forwards its + /// output. + SpillJoin, /// Indicates that HashJoinStream execution is completed Completed, } @@ -226,6 +234,29 @@ pub(super) struct HashJoinStream { output_buffer: LimitedBatchCoalescer, /// Whether this is a null-aware anti join null_aware: bool, + /// Spill support (None when the feature is disabled or inapplicable) + spill_ctx: Option>, + /// Drives the partitioned disk join after the build side spilled + spill_driver: Option, +} + +/// The two spill-join drivers: per-partition (Partitioned mode) and +/// shared-build (CollectLeft mode). +pub(super) enum SpillDriver { + Owned(Box), + Shared(Box), +} + +impl SpillDriver { + fn poll_next( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>> { + match self { + SpillDriver::Owned(driver) => driver.poll_next(cx), + SpillDriver::Shared(driver) => driver.poll_next(cx), + } + } } impl RecordBatchStream for HashJoinStream { @@ -375,6 +406,7 @@ impl HashJoinStream { mode: PartitionMode, null_aware: bool, fetch: Option, + spill_ctx: Option>, ) -> Self { // Create output buffer with coalescing and optional fetch limit. let output_buffer = @@ -403,6 +435,8 @@ impl HashJoinStream { mode, output_buffer, null_aware, + spill_ctx, + spill_driver: None, } } @@ -442,6 +476,9 @@ impl HashJoinStream { HashJoinStreamState::ExhaustedProbeSide => { handle_state!(self.process_unmatched_build_batch()) } + HashJoinStreamState::SpillJoin => { + handle_state!(ready!(self.process_spill_join(cx))) + } HashJoinStreamState::Completed if !self.output_buffer.is_empty() => { // Flush any remaining buffered data self.output_buffer.finish()?; @@ -469,10 +506,139 @@ impl HashJoinStream { if let Some(ref mut fut) = self.build_waiter { ready!(fut.get_shared(cx))?; } - self.state = HashJoinStreamState::FetchProbeBatch; + self.state = if self.spill_driver.is_some() { + HashJoinStreamState::SpillJoin + } else { + HashJoinStreamState::FetchProbeBatch + }; + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Report bounds-only pushdown for a spilled build (so peer partitions + /// neither deadlock nor prune this partition's probe rows) and route the + /// state machine into the spill join. + fn report_spill_bounds(&mut self, bounds: Option) -> Result<()> { + if let Some(ref build_accumulator) = self.build_accumulator { + let build_accumulator = Arc::clone(build_accumulator); + let bounds = bounds.unwrap_or_else(|| PartitionBounds::new(vec![])); + let build_data = match self.mode { + PartitionMode::Partitioned => PartitionBuildData::Partitioned { + partition_id: self.partition, + pushdown: PushdownStrategy::BoundsOnly, + bounds, + }, + PartitionMode::CollectLeft => PartitionBuildData::CollectLeft { + pushdown: PushdownStrategy::BoundsOnly, + bounds, + }, + PartitionMode::Auto => { + return internal_err!( + "PartitionMode::Auto should not be present at execution time" + ); + } + }; + self.build_waiter = Some(OnceFut::new(async move { + build_accumulator.report_build_data(build_data).await + })); + self.state = HashJoinStreamState::WaitPartitionBoundsReport; + } else { + self.state = HashJoinStreamState::SpillJoin; + } + Ok(()) + } + + /// Take ownership of the probe stream and snapshot the parameters the + /// spill drivers need to construct inner per-pair join streams. + fn take_probe_and_spec(&mut self) -> (SendableRecordBatchStream, InnerJoinSpec) { + let probe_schema = self.right.schema(); + let probe = + std::mem::replace(&mut self.right, empty_record_batch_stream(probe_schema)); + let spec = InnerJoinSpec { + partition: self.partition, + schema: Arc::clone(&self.schema), + on_right: self.on_right.clone(), + filter: self.filter.clone(), + join_type: self.join_type, + random_state: self.random_state.clone(), + column_indices: self.column_indices.clone(), + null_equality: self.null_equality, + batch_size: self.batch_size, + }; + (probe, spec) + } + + /// Hand off to the Partitioned-mode spill-join driver after this + /// partition's build side scattered to disk. + fn begin_spill_join( + &mut self, + spilled: crate::joins::hash_join::spill::SpilledBuild, + ) -> Poll>>> { + let ctx = self.spill_ctx.clone().ok_or_else(|| { + internal_datafusion_err!("spilled build without a spill context") + })?; + self.report_spill_bounds(spilled.bounds.clone())?; + let (probe, spec) = self.take_probe_and_spec(); + self.spill_driver = Some(SpillDriver::Owned(Box::new(SpillJoinDriver::new( + spilled, + probe, + ctx, + spec, + self.join_metrics.clone(), + )))); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Hand off to the CollectLeft shared-build spill driver: this output + /// partition scatters its own probe stream and walks the shared build + /// partitions. + fn begin_shared_spill_join( + &mut self, + shared: Arc, + ) -> Poll>>> { + let ctx = self.spill_ctx.clone().ok_or_else(|| { + internal_datafusion_err!("spilled build without a spill context") + })?; + self.report_spill_bounds(shared.bounds.clone())?; + let (probe, spec) = self.take_probe_and_spec(); + self.spill_driver = + Some(SpillDriver::Shared(Box::new(SharedSpillJoinDriver::new( + shared, + probe, + ctx, + spec, + self.join_metrics.clone(), + )))); Poll::Ready(Ok(StatefulStreamResult::Continue)) } + /// Forward the spill-join driver's output through the output buffer. + fn process_spill_join( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>>> { + let driver = self.spill_driver.as_mut().ok_or_else(|| { + internal_datafusion_err!("SpillJoin state without a driver") + })?; + match ready!(driver.poll_next(cx)) { + Some(Ok(batch)) => { + let push_status = self.output_buffer.push_batch(batch)?; + if push_status == PushBatchStatus::LimitReached { + self.spill_driver = None; + self.output_buffer.finish()?; + self.state = HashJoinStreamState::Completed; + } + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + Some(Err(e)) => Poll::Ready(Err(e)), + None => { + self.spill_driver = None; + self.output_buffer.finish()?; + self.state = HashJoinStreamState::Completed; + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + } + } + /// Collects build-side data by polling `OnceFut` future from initialized build-side /// /// Updates build-side to `Ready`, and state to `FetchProbeSide` @@ -482,7 +648,7 @@ impl HashJoinStream { ) -> Poll>>> { let build_timer = self.join_metrics.build_time.timer(); // build hash table from left (build) side, if not yet done - let left_data = ready!( + let build_output = ready!( self.build_side .try_as_initial_mut()? .left_fut @@ -490,6 +656,20 @@ impl HashJoinStream { )?; build_timer.done(); + let left_data = match build_output.as_ref() { + BuildPhaseOutput::InMemory(data) => Arc::clone(data), + BuildPhaseOutput::Spilled(slot) => { + let spilled = slot.lock().take().ok_or_else(|| { + internal_datafusion_err!("spilled hash join build already consumed") + })?; + return self.begin_spill_join(spilled); + } + BuildPhaseOutput::SpilledShared(shared) => { + let shared = Arc::clone(shared); + return self.begin_shared_spill_join(shared); + } + }; + // Note: For null-aware anti join, we need to check the probe side (right) for NULLs, // not the build side (left). The probe-side NULL check happens during process_probe_batch. // The probe_side_has_null flag will be set there if any probe batch contains NULL. diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b61ceecb24fc0..1ae93eafd07a6 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -218,8 +218,12 @@ datafusion.execution.batch_size 8192 datafusion.execution.coalesce_batches true datafusion.execution.collect_statistics true datafusion.execution.enable_ansi_mode false +datafusion.execution.enable_hash_join_spill false datafusion.execution.enable_recursive_ctes true datafusion.execution.enforce_batch_size_in_joins false +datafusion.execution.hash_join_spill_headroom_bytes 33554432 +datafusion.execution.hash_join_spill_max_recursion_depth 2 +datafusion.execution.hash_join_spill_partition_count 0 datafusion.execution.keep_partition_by_columns false datafusion.execution.listing_table_factory_infer_partitions true datafusion.execution.listing_table_ignore_subdirectory true @@ -356,8 +360,12 @@ datafusion.execution.batch_size 8192 Default batch size while creating new batch datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true. datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. +datafusion.execution.enable_hash_join_spill false When enabled, a hash join whose build side exceeds its memory budget partitions both sides to disk (grace hash join) and joins partition pairs within the budget instead of failing with a resources-exhausted error. The in-memory fast path is unchanged; spilling engages only at the moment the build-side memory reservation first fails. Requires a configured `DiskManager` (spilling stays disabled otherwise). datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. +datafusion.execution.hash_join_spill_headroom_bytes 33554432 Reserved memory headroom for a spilling hash join's scatter scratch space and per-partition write buffers (analogous to `sort_spill_reservation_bytes`). +datafusion.execution.hash_join_spill_max_recursion_depth 2 Maximum number of recursive repartition passes a spilling hash join applies to a partition whose build side still exceeds the memory budget (key skew). When the limit is reached the join falls back to a chunked build where supported, or reports a clean resources-exhausted error. +datafusion.execution.hash_join_spill_partition_count 0 Number of disk partitions a spilling hash join scatters each side into. 0 (default) sizes automatically from the observed build size and the memory budget. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches datafusion.execution.listing_table_factory_infer_partitions true Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). datafusion.execution.listing_table_ignore_subdirectory true Should sub directories be ignored when scanning directories for data files. Defaults to true (ignores subdirectories), consistent with Hive. Note that this setting does not affect reading partitioned tables (e.g. `/table/year=2021/month=01/data.parquet`). diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 11a1a8a2d6831..3fa6dee811923 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -116,6 +116,10 @@ The following configuration settings are available: | datafusion.execution.skip_physical_aggregate_schema_check | false | When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. | | datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. | | datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). | +| datafusion.execution.enable_hash_join_spill | false | When enabled, a hash join whose build side exceeds its memory budget partitions both sides to disk (grace hash join) and joins partition pairs within the budget instead of failing with a resources-exhausted error. The in-memory fast path is unchanged; spilling engages only at the moment the build-side memory reservation first fails. Requires a configured `DiskManager` (spilling stays disabled otherwise). | +| datafusion.execution.hash_join_spill_partition_count | 0 | Number of disk partitions a spilling hash join scatters each side into. 0 (default) sizes automatically from the observed build size and the memory budget. | +| datafusion.execution.hash_join_spill_headroom_bytes | 33554432 | Reserved memory headroom for a spilling hash join's scatter scratch space and per-partition write buffers (analogous to `sort_spill_reservation_bytes`). | +| datafusion.execution.hash_join_spill_max_recursion_depth | 2 | Maximum number of recursive repartition passes a spilling hash join applies to a partition whose build side still exceeds the memory budget (key skew). When the limit is reached the join falls back to a chunked build where supported, or reports a clean resources-exhausted error. | | datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. | | datafusion.execution.max_spill_file_size_bytes | 134217728 | Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB | | datafusion.execution.meta_fetch_concurrency | 32 | Number of files to read in parallel when inferring schema and statistics |