From fe69472606e9cc33c4f3be1d1c8577a6af6e57e0 Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 14:30:07 +0300 Subject: [PATCH 1/8] Hash join spill M1: neutral build-phase refactor + config plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Groundwork for the spillable hash join (no behavior change): - Extract build_left_data() from collect_left_input() — the hash-table/ ArrayMap/concat/bitmap/membership construction — so the same code will serve both the classic in-memory build and per-partition builds when the join spills. - collect_left_input() now returns BuildPhaseOutput (single InMemory variant for now; Spilled arrives with the scatter phase). - Register the build-side reservation with_can_spill(true) when the new execution.enable_hash_join_spill flag is on. - New config knobs (default off/auto): enable_hash_join_spill, hash_join_spill_partition_count, hash_join_spill_headroom_bytes, hash_join_spill_max_recursion_depth; config docs + information_schema expectations regenerated. Verified: full physical-plan suite (1268 passed; spill_compression zstd feature quirk pre-exists on clean base), joins suite 926 passed, full sqllogictest suite green modulo the pre-existing local encrypted_parquet quirk, clippy -D warnings clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- datafusion/common/src/config.rs | 26 +++++ .../physical-plan/src/joins/hash_join/exec.rs | 96 ++++++++++++++++--- .../src/joins/hash_join/stream.rs | 10 +- .../test_files/information_schema.slt | 8 ++ docs/source/user-guide/configs.md | 4 + 5 files changed, 127 insertions(+), 17 deletions(-) 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..f844d5ad961e4 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -251,6 +251,14 @@ 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. +pub(super) enum BuildPhaseOutput { + InMemory(Arc), +} + /// Helps to build [`HashJoinExec`]. /// /// Builder can be created from an existing [`HashJoinExec`] using [`From::from`]. @@ -734,7 +742,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 +1303,22 @@ 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 can_spill = context + .session_config() + .options() + .execution + .enable_hash_join_spill; + 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(can_spill) + .register(context.memory_pool()); Ok(collect_left_input( self.random_state.random_state().clone(), @@ -1321,6 +1339,7 @@ impl ExecutionPlan for HashJoinExec { let reservation = MemoryConsumer::new(format!("HashJoinInput[{partition}]")) + .with_can_spill(can_spill) .register(context.memory_pool()); OnceFut::new(collect_left_input( @@ -1845,7 +1864,7 @@ async fn collect_left_input( config: Arc, null_equality: NullEquality, array_map_created_count: Count, -) -> Result { +) -> Result { let schema = left_stream.schema(); let should_collect_min_max_for_phj = @@ -1889,12 +1908,12 @@ async fn collect_left_input( batches, num_rows, metrics, - mut reservation, + reservation, bounds_accumulators, } = 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 +1924,61 @@ async fn collect_left_input( _ => None, }; + 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, + )?; + + 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`]. +/// +/// `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. +#[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, +) -> 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, @@ -1953,11 +2021,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,9 +2034,9 @@ 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) }; @@ -2014,7 +2082,7 @@ async fn collect_left_input( } }; - if should_collect_min_max_for_phj && !should_compute_dynamic_filters { + if mask_bounds { bounds = None; } diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index b31982ea3b7b4..748b65e317bb2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -28,7 +28,7 @@ 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, }; @@ -69,7 +69,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 @@ -482,7 +482,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 +490,10 @@ impl HashJoinStream { )?; build_timer.done(); + let left_data = match build_output.as_ref() { + BuildPhaseOutput::InMemory(data) => Arc::clone(data), + }; + // 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 | From d5badbbedcf88985bf98db57117a49de4a1b46f9 Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 15:08:34 +0300 Subject: [PATCH 2/8] Hash join spill M2: Partitioned-mode inner-join grace spill path When enable_hash_join_spill is on and the build-side reservation first fails, the build switches into scatter mode instead of erroring: buffered batches and the rest of the build input are hash-partitioned into K disk partitions (new spill.rs), the probe side is scattered with the same seeds, and the join runs partition-pair by pair, each pair building an ordinary in-memory hash table under its own reservation. The in-memory fast path is unchanged. - spill.rs: SideScatter / PartitionSpillWriter (coalesced IPC writes, file rotation, StringView compaction), chained spill readers, SpillJoinDriver state machine, auto partition-count sizing. - exec.rs: engage-on-overflow inside the build fold with exact reservation shrinks; BuildPhaseOutput::Spilled; per-partition spill context in execute(); spill disabled when it would break the advertised probe-side output ordering. - stream.rs: SpillJoin stream state; bounds-only dynamic-filter reporting for spilled partitions (prevents both the bounds-report deadlock and wrong pruning); output flows through the existing coalescer/fetch. - shared_bounds.rs: PushdownStrategy::BoundsOnly (unlike Empty, a spilled partition's probe rows must not be pruned by the CASE filter). M2 scope: PartitionMode::Partitioned + JoinType::Inner only (gate in execute()); oversized partitions surface a clean resources-exhausted error until recursion lands (M4). Tests (8 new): scatter round-trip, seed independence from repartition hash, file rotation + temp cleanup, forced-spill inner join equals in-memory results with spill metrics recorded, fast path untouched with ample memory, pool fully released after completion and after mid-stream drop, disk-cap exhaustion and flag-off both keep today's clean errors. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 216 ++- .../physical-plan/src/joins/hash_join/mod.rs | 1 + .../src/joins/hash_join/shared_bounds.rs | 7 + .../src/joins/hash_join/spill.rs | 1167 +++++++++++++++++ .../src/joins/hash_join/stream.rs | 125 +- 5 files changed, 1501 insertions(+), 15 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/hash_join/spill.rs diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f844d5ad961e4..dfd5d915a2aa7 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, 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}; @@ -255,8 +261,13 @@ impl JoinLeftData { /// (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>), } /// Helps to build [`HashJoinExec`]. @@ -1312,6 +1323,57 @@ impl ExecutionPlan for HashJoinExec { .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. Non-inner join + // types arrive with per-partition bitmap support.) + // + // 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 can_spill + && self.mode == PartitionMode::Partitioned + && !self.null_aware + && self.join_type == JoinType::Inner + && !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), + 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))?; @@ -1332,6 +1394,8 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + // CollectLeft spill (shared scattered build) is a follow-up. + None, )) })?, PartitionMode::Partitioned => { @@ -1354,6 +1418,7 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + spill_context.clone(), )) } PartitionMode::Auto => { @@ -1432,6 +1497,7 @@ impl ExecutionPlan for HashJoinExec { self.mode, self.null_aware, self.fetch, + spill_context, ))) } @@ -1778,10 +1844,19 @@ 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>, + /// 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 { @@ -1795,6 +1870,7 @@ impl BuildSideState { ) -> Result { Ok(Self { batches: Vec::new(), + batch_sizes: Vec::new(), num_rows: 0, metrics, reservation, @@ -1806,8 +1882,79 @@ impl BuildSideState { .collect::>>() }) .transpose()?, + 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); + + let buffered: usize = self.batch_sizes.iter().sum(); + let partition_count = compute_partition_count( + buffered, + &ctx.pool.memory_limit(), + ctx.partition_count_override, + ); + 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); + } + + 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( @@ -1864,41 +2011,61 @@ async fn collect_left_input( config: Arc, null_equality: NullEquality, array_map_created_count: Count, + spill: Option>, ) -> 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; 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); + // Reserve memory for incoming batch + match state.reservation.try_grow(batch_size) { + Ok(()) => { + state.metrics.build_mem_used.add(batch_size); + 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?; @@ -1906,10 +2073,14 @@ async fn collect_left_input( // Extract fields from state let BuildSideState { batches, + batch_sizes: _, num_rows, metrics, reservation, bounds_accumulators, + spill: _, + scatter, + headroom, } = state; // Compute bounds @@ -1924,6 +2095,25 @@ 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()?; + 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, 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..2e6b51cc73a65 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -0,0 +1,1167 @@ +// 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::task::{Context, Poll}; + +use crate::hash_utils::create_hashes; +use crate::joins::PartitionMode; +use crate::joins::hash_join::exec::{BuildPhaseOutput, build_left_data}; +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, 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::{Stream, StreamExt, ready}; + +/// 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; + +/// 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. The seed +/// varies per recursion level so a repartition pass re-randomizes bucket +/// assignment. +pub(super) fn scatter_random_state(level: usize) -> RandomState { + RandomState::with_seeds('S' as u64, 'P' as u64, 'L' as u64, 0xC0FFEE + level as u64) +} + +/// 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. +pub(super) fn compute_partition_count( + buffered_bytes: usize, + pool_limit: &MemoryLimit, + override_count: usize, +) -> usize { + if override_count > 0 { + return override_count; + } + let target = match pool_limit { + MemoryLimit::Finite(limit) => { + (*limit / 16).clamp(32 * 1024 * 1024, 256 * 1024 * 1024) + } + _ => DEFAULT_PARTITION_TARGET_BYTES, + }; + 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, + pub(super) partition_count_override: usize, + pub(super) headroom_bytes: usize, + #[expect(dead_code)] // used by recursive repartitioning (M4) + 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, + #[expect(dead_code)] // used by recursion sizing decisions (M4) + 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 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, +} + +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: usize, + stream: SendableRecordBatchStream, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + }, + /// Run the pair's in-memory join and forward its output. + RunPair { + inner: SendableRecordBatchStream, + }, + Done, +} + +/// Drives the spilled join: probe scatter, then a sequential loop over +/// partition pairs, each joined by an ordinary in-memory [`HashJoinStream`]. +pub(super) struct SpillJoinDriver { + ctx: Arc, + spec: InnerJoinSpec, + phase: DriverPhase, + build_parts: Vec>, + probe_parts: Vec>, + next_pair: usize, + level: 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 }, + build_parts: partitions.into_iter().map(Some).collect(), + probe_parts: Vec::new(), + next_pair: 0, + level, + outer_metrics, + inner_metrics_set: ExecutionPlanMetricsSet::new(), + _headroom: headroom, + ctx, + } + } + + /// Whether an empty side lets the pair be skipped entirely. Inner joins + /// produce nothing when either side is empty. (Join types that emit + /// unmatched rows arrive in M3 and must still run such pairs.) + fn can_skip_pair(&self, build_rows: usize, probe_rows: usize) -> bool { + match self.spec.join_type { + JoinType::Inner => build_rows == 0 || probe_rows == 0, + _ => false, + } + } + + 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::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")); + }; + self.probe_parts = scatter.finish()?.into_iter().map(Some).collect(); + } + } + Poll::Ready(Ok(())) + } + + /// Pick the next pair; skip pairs an inner join cannot produce rows for. + fn select_next_pair(&mut self) -> Result<()> { + loop { + let pair = self.next_pair; + if pair >= self.build_parts.len() { + self.phase = DriverPhase::Done; + return Ok(()); + } + self.next_pair += 1; + + let (Some(build), Some(probe)) = + (self.build_parts[pair].take(), self.probe_parts[pair].take()) + else { + return internal_err!( + "spilled hash join partition {pair} already consumed" + ); + }; + + if self.can_skip_pair(build.rows, probe.rows) { + // Dropping the sides deletes their spill files. + continue; + } + + let reservation = MemoryConsumer::new(format!( + "HashJoinSpillPartition[{}.{}]", + self.ctx.partition, pair + )) + .with_can_spill(true) + .register(&self.ctx.pool); + + let stream = + spilled_side_stream(build.files, Arc::clone(&self.ctx.build_spill)); + // Stash the probe side back; consumed when the pair's join starts. + self.probe_parts[pair] = Some(probe); + self.phase = DriverPhase::LoadBuild { + pair, + stream, + batches: Vec::new(), + num_rows: 0, + reservation, + }; + 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 DriverPhase::LoadBuild { + pair, + batches, + num_rows, + reservation, + .. + } = &mut self.phase + else { + return Poll::Ready(internal_err!("expected LoadBuild phase")); + }; + if let Err(e) = reservation.try_grow(size) { + // Recursive repartitioning arrives in M4; for now an + // oversized partition surfaces the clean + // resources-exhausted error. + return Poll::Ready(Err(oversize_partition_err( + e, *pair, self.level, + ))); + } + *num_rows += batch.num_rows(); + batches.push(batch); + } + None => { + let DriverPhase::LoadBuild { + pair, + batches, + num_rows, + reservation, + .. + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) + else { + return Poll::Ready(internal_err!("expected LoadBuild phase")); + }; + let inner = self.start_pair_join(pair, batches, num_rows, reservation)?; + self.phase = DriverPhase::RunPair { inner }; + } + } + Poll::Ready(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) + } + + /// Build the pair's `JoinLeftData` from the loaded batches and construct + /// the inner in-memory join stream over the pair's probe partition. + fn start_pair_join( + &mut self, + pair: usize, + batches: Vec, + num_rows: usize, + reservation: MemoryReservation, + ) -> Result { + let inner_metrics = BuildProbeJoinMetrics::new(pair, &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, + )?; + drop(batches); + + let probe = self.probe_parts[pair].take().ok_or_else(|| { + DataFusionError::Internal(format!( + "probe partition {pair} missing for spilled hash join" + )) + })?; + 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)) + } +} + +fn oversize_partition_err( + source: DataFusionError, + pair: usize, + level: usize, +) -> DataFusionError { + source.context(format!( + "spilled hash join partition {pair} (repartition level {level}) does not fit \ + in the memory budget" + )) +} + +#[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 + } + + #[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}"); + } +} diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 748b65e317bb2..b3d3dfe12ccd4 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -30,7 +30,10 @@ use crate::joins::MapOffset; use crate::joins::PartitionMode; 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, SpillJoinDriver, empty_record_batch_stream, }; use crate::joins::utils::{ OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, @@ -133,6 +136,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 +233,10 @@ 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>, } impl RecordBatchStream for HashJoinStream { @@ -375,6 +386,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 +415,8 @@ impl HashJoinStream { mode, output_buffer, null_aware, + spill_ctx, + spill_driver: None, } } @@ -442,6 +456,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 +486,108 @@ 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)) } + /// Hand off to the spill-join driver after the build side scattered to + /// disk: report bounds-only pushdown (so peer partitions neither deadlock + /// nor prune this partition's probe rows), take ownership of the probe + /// stream, and construct the driver. + 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") + })?; + + if let Some(ref build_accumulator) = self.build_accumulator { + let build_accumulator = Arc::clone(build_accumulator); + let bounds = spilled + .bounds + .clone() + .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 Poll::Ready(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; + } + + 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, + }; + self.spill_driver = Some(Box::new(SpillJoinDriver::new( + spilled, + 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` @@ -492,6 +607,12 @@ impl HashJoinStream { 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); + } }; // Note: For null-aware anti join, we need to check the probe side (right) for NULLs, From 94b9aaa9bc5897d31f7495d7416ea3a51ebc0aac Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 15:23:55 +0300 Subject: [PATCH 3/8] Hash join spill M3: all join types under the Partitioned spill path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the inner-join-only gate: disk partitions are key-disjoint, so each pair's in-memory join has complete knowledge for its keys — per-pair visited bitmaps (unmatched build emission) and probe-side null padding work unchanged. The driver's pair-skip logic is now a per-type matrix: pairs are skipped only when the join type provably emits nothing for an empty side; pairs that emit unmatched rows from the non-empty side still run (the inner stream handles the empty other side natively). The ordering guard from M2 (no spilling when the operator's advertised probe-side order maintenance is live) now carries the weight for Inner/Right/RightSemi/RightAnti/RightMark. Null-aware anti joins remain excluded. Tests: forced-spill matrix over all 10 join types vs in-memory reference (multiset-fingerprint comparison), and an Inner/Left/Full x NullEqualsNothing/NullEqualsNull matrix with NULL join keys (all NULLs scatter to one partition, so NullEqualsNull matching stays complete; NULL-heavy skew is deferred to the chunked fallback in M4). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 9 +- .../src/joins/hash_join/spill.rs | 218 +++++++++++++++++- 2 files changed, 219 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index dfd5d915a2aa7..f19c9fb1be763 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1325,8 +1325,12 @@ impl ExecutionPlan for HashJoinExec { // 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. Non-inner join - // types arrive with per-partition bitmap support.) + // 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 @@ -1338,7 +1342,6 @@ impl ExecutionPlan for HashJoinExec { let spill_context = if can_spill && self.mode == PartitionMode::Partitioned && !self.null_aware - && self.join_type == JoinType::Inner && !would_break_output_order && context.runtime_env().disk_manager.tmp_files_enabled() { diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index 2e6b51cc73a65..b8fbf60110000 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -526,13 +526,22 @@ impl SpillJoinDriver { } } - /// Whether an empty side lets the pair be skipped entirely. Inner joins - /// produce nothing when either side is empty. (Join types that emit - /// unmatched rows arrive in M3 and must still run such pairs.) + /// 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 { - JoinType::Inner => build_rows == 0 || probe_rows == 0, - _ => false, + // 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, } } @@ -848,6 +857,23 @@ mod tests { 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()); @@ -1164,4 +1190,186 @@ mod tests { .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" + ); + } + } + } } From 046474f24ef5e33ea154472bcdfdb4cac1b7ab1f Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 15:43:25 +0300 Subject: [PATCH 4/8] Hash join spill M4: recursive repartitioning + chunked skew fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A pair whose build partition exceeds the budget is now handled in three tiers instead of erroring: 1. Recursive repartition (up to hash_join_spill_max_recursion_depth passes): already-loaded batches and the rest of the pair's build stream re-scatter into 8 children at the next seed level, the pair's probe partition follows with the same seed, and the children are queued front-first so their disk space is reclaimed promptly. A no-shrink shortcut (largest child > 3/4 of the parent) marks children as duplicate-key mass and sends them straight to the fallback. 2. Chunked build fallback for build-side-emission types (Inner, Left, LeftSemi, LeftAnti, LeftMark): build rows load in budget-sized chunks, each joined against a full re-read of the pair's probe partition — correct because each build row lives in exactly one chunk. 3. Probe-side-emission types (Right family, Full) in this corner surface a descriptive resources-exhausted error (their outputs depend on matches against the whole build side, which chunks cannot track) — never wrong results. Two real bugs found by the new tests and fixed: - Seed derivation: per-level scatter seeds differed only by a small additive constant in one ahash seed word, and ahash folds that word in additively — hashes at consecutive levels differed by a constant, so under % K a whole parent partition landed in ONE child and recursion could never split anything. Seeds are now derived through a SplitMix64 mixer (regression test: recursive_scatter_levels_are_decorrelated). - Memory accounting: loaders reserved only batch bytes, so a chunk/pair could fill the budget with data and then fail on the hash table's estimate (~2x the data for narrow rows). Spill-mode loaders (build fold, pair load, chunk load) now pre-reserve the incremental table estimate per batch — overflow surfaces where the fallback logic can engage — and build_left_data takes prereserved_table_bytes so the accounting stays exact (fast path passes 0: unchanged). New metrics: join_spill_repartition_passes, join_spill_fallback_chunks. Tests (15 total now): extreme duplicate-key skew via chunked fallback (Inner + LeftAnti) vs in-memory reference, unsupported-type skew clean error, moderate-skew recursion with zero fallback chunks, per-level scatter decorrelation. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 57 +- .../src/joins/hash_join/spill.rs | 1025 +++++++++++++++-- 2 files changed, 1003 insertions(+), 79 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f19c9fb1be763..296c5be503a8c 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1365,6 +1365,10 @@ impl ExecutionPlan for HashJoinExec { 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, @@ -1854,6 +1858,10 @@ struct BuildSideState { 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, /// Present when the spill fallback is enabled and applicable. spill: Option>, /// Present once the build has switched into scatter mode. @@ -1885,6 +1893,7 @@ impl BuildSideState { .collect::>>() }) .transpose()?, + table_reserved: 0, spill: None, scatter: None, headroom: None, @@ -1946,6 +1955,9 @@ impl BuildSideState { 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); @@ -2052,10 +2064,20 @@ async fn collect_left_input( } 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) { + 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); } @@ -2081,6 +2103,7 @@ async fn collect_left_input( metrics, reservation, bounds_accumulators, + table_reserved, spill: _, scatter, headroom, @@ -2132,6 +2155,7 @@ async fn collect_left_input( &config, null_equality, &array_map_created_count, + table_reserved, )?; Ok(BuildPhaseOutput::InMemory(Arc::new(data))) @@ -2146,9 +2170,27 @@ async fn collect_left_input( /// 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], @@ -2165,6 +2207,7 @@ pub(super) fn build_left_data( config: &ConfigOptions, null_equality: NullEquality, array_map_created_count: &Count, + prereserved_table_bytes: usize, ) -> Result { let (join_hash_map, batch, left_values) = if let Some((array_map, batch, left_value)) = try_create_array_map( @@ -2179,11 +2222,15 @@ pub(super) fn build_left_data( )? { 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::(); @@ -2193,13 +2240,17 @@ pub(super) fn build_left_data( 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)) }; diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index b8fbf60110000..88e1a79cb01a2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -46,7 +46,9 @@ use std::task::{Context, Poll}; use crate::hash_utils::create_hashes; use crate::joins::PartitionMode; -use crate::joins::hash_join::exec::{BuildPhaseOutput, build_left_data}; +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, @@ -85,15 +87,34 @@ const WRITE_BUFFER_FLUSH_BYTES: usize = 1024 * 1024; /// 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. The seed -/// varies per recursion level so a repartition pass re-randomizes bucket -/// assignment. +/// 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 { - RandomState::with_seeds('S' as u64, 'P' as u64, 'L' as u64, 0xC0FFEE + level as u64) + 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. @@ -133,9 +154,12 @@ pub(super) struct HashJoinSpillContext { 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, - #[expect(dead_code)] // used by recursive repartitioning (M4) pub(super) max_recursion_depth: usize, pub(super) max_spill_file_size: usize, /// Output partition index (consumer naming / diagnostics). @@ -148,7 +172,6 @@ pub(super) struct HashJoinSpillContext { /// One side's scattered partition: the rotated spill files plus totals. pub(super) struct SpilledSide { pub(super) files: Vec, - #[expect(dead_code)] // used by recursion sizing decisions (M4) pub(super) bytes: usize, pub(super) rows: usize, } @@ -444,6 +467,40 @@ pub(super) struct InnerJoinSpec { 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 { @@ -454,29 +511,76 @@ enum DriverPhase { NextPair, /// Stream the pair's build partition into memory under its reservation. LoadBuild { - pair: usize, + 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_parts: Vec>, - probe_parts: Vec>, - next_pair: usize, - level: usize, + /// 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 @@ -515,10 +619,10 @@ impl SpillJoinDriver { Self { spec, phase: DriverPhase::ScatterProbe { probe, scatter }, - build_parts: partitions.into_iter().map(Some).collect(), - probe_parts: Vec::new(), - next_pair: 0, - level, + initial_build: partitions, + base_level: level, + queue: VecDeque::new(), + pair_seq: 0, outer_metrics, inner_metrics_set: ExecutionPlanMetricsSet::new(), _headroom: headroom, @@ -545,6 +649,21 @@ impl SpillJoinDriver { } } + /// 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<'_>, @@ -562,6 +681,21 @@ impl SpillJoinDriver { 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), } } @@ -592,52 +726,61 @@ impl SpillJoinDriver { else { return Poll::Ready(internal_err!("expected ScatterProbe phase")); }; - self.probe_parts = scatter.finish()?.into_iter().map(Some).collect(); + 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(())) } - /// Pick the next pair; skip pairs an inner join cannot produce rows for. + /// Pop the next pair; skip pairs this join type provably emits nothing + /// for. fn select_next_pair(&mut self) -> Result<()> { loop { - let pair = self.next_pair; - if pair >= self.build_parts.len() { + let Some(work) = self.queue.pop_front() else { self.phase = DriverPhase::Done; return Ok(()); - } - self.next_pair += 1; - - let (Some(build), Some(probe)) = - (self.build_parts[pair].take(), self.probe_parts[pair].take()) - else { - return internal_err!( - "spilled hash join partition {pair} already consumed" - ); }; - if self.can_skip_pair(build.rows, probe.rows) { + 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 + 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(build.files, Arc::clone(&self.ctx.build_spill)); - // Stash the probe side back; consumed when the pair's join starts. - self.probe_parts[pair] = Some(probe); + spilled_side_stream(work.build.files, Arc::clone(&self.ctx.build_spill)); self.phase = DriverPhase::LoadBuild { - pair, + 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(()); } @@ -655,45 +798,442 @@ impl SpillJoinDriver { 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, + pair_id, + probe, batches, num_rows, reservation, + table_reserved, .. - } = &mut self.phase + } = std::mem::replace(&mut self.phase, DriverPhase::NextPair) else { return Poll::Ready(internal_err!("expected LoadBuild phase")); }; - if let Err(e) = reservation.try_grow(size) { - // Recursive repartitioning arrives in M4; for now an - // oversized partition surfaces the clean - // resources-exhausted error. - return Poll::Ready(Err(oversize_partition_err( - e, *pair, self.level, - ))); - } - *num_rows += batch.num_rows(); - batches.push(batch); - } - None => { - let DriverPhase::LoadBuild { - pair, + 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); + 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 LoadBuild phase")); + return Poll::Ready(internal_err!("expected RepartitionBuild phase")); }; - let inner = self.start_pair_join(pair, batches, num_rows, reservation)?; - self.phase = DriverPhase::RunPair { inner }; + 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<'_>, @@ -704,16 +1244,77 @@ impl SpillJoinDriver { inner.poll_next_unpin(cx) } - /// Build the pair's `JoinLeftData` from the loaded batches and construct - /// the inner in-memory join stream over the pair's probe partition. - fn start_pair_join( + 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); + 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: usize, + pair_id: usize, batches: Vec, num_rows: usize, reservation: MemoryReservation, + table_reserved: usize, + probe_files: Vec, ) -> Result { - let inner_metrics = BuildProbeJoinMetrics::new(pair, &self.inner_metrics_set); + 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); @@ -732,16 +1333,12 @@ impl SpillJoinDriver { &self.ctx.config, self.spec.null_equality, &array_map_count, + table_reserved, )?; drop(batches); - let probe = self.probe_parts[pair].take().ok_or_else(|| { - DataFusionError::Internal(format!( - "probe partition {pair} missing for spilled hash join" - )) - })?; let probe_stream = - spilled_side_stream(probe.files, Arc::clone(&self.ctx.probe_spill)); + 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))); @@ -773,17 +1370,6 @@ impl SpillJoinDriver { } } -fn oversize_partition_err( - source: DataFusionError, - pair: usize, - level: usize, -) -> DataFusionError { - source.context(format!( - "spilled hash join partition {pair} (repartition level {level}) does not fit \ - in the memory budget" - )) -} - #[cfg(test)] mod tests { use super::*; @@ -1372,4 +1958,291 @@ mod tests { } } } + /// 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 + ); + } + } } From 0f70b0067147c4bc81764bddc7fac1d0a723750c Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 17:14:10 +0300 Subject: [PATCH 5/8] Hash join spill M5: CollectLeft shared-build spill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CollectLeft joins (single shared build, N probe partitions) can now spill: the build scatters once into K disk partitions, each output partition scatters ITS OWN probe stream with the same seed and walks k = 0..K, and every k must run in every partition (skipping would leave the shared completion counter high and suppress unmatched emission). Design decision recorded by the tests: per-k hash tables are NOT shared across probe partitions. Sharing them (first attempt, via per-k OnceAsync) meant a retained-table WINDOW spanning the progress gap between the fastest and slowest partition — up to K tables resident, observed directly as three concurrent multi-MB tables blowing a pool sized for two. Instead each partition builds its own copy of k's table from the shared files (bounded: at most one resident table per partition, matching the auto-K sizing, which now also divides its per-partition target by the consumer count). What IS shared per k is exactly what correctness needs: the visited bitmap and the probe-completion counter, injected into each copy's JoinLeftData — JoinLeftData's bitmap/counter are Arc'd now — so matches from all partitions mark one bitmap and the existing last-finisher logic emits unmatched build rows exactly once. Row order across copies is deterministic (same files, same order), so bitmap indices agree. Oversized shared partitions surface a clean descriptive error (no recursion for the shared build: a bitmap cannot be split across repartition levels). Null-aware anti joins remain excluded. Test: forced CollectLeft spill over Inner/Left/LeftAnti/Full with two concurrently-polled output partitions — fingerprint-identical to the in-memory join (exactly-once unmatched emission), spill engaged, pool drained after the plan drops (the exec node's left_fut cache retains the scatter headroom until then, as it retains the whole build today). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 103 +++- .../src/joins/hash_join/spill.rs | 452 +++++++++++++++++- .../src/joins/hash_join/stream.rs | 99 +++- 3 files changed, 607 insertions(+), 47 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 296c5be503a8c..1ce51ad25b725 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -38,8 +38,8 @@ use crate::joins::hash_join::shared_bounds::{ ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator, }; use crate::joins::hash_join::spill::{ - HashJoinSpillContext, SideScatter, SpilledBuild, compute_partition_count, - scatter_random_state, + HashJoinSpillContext, SharedSpilledBuild, SideScatter, SpilledBuild, + compute_partition_count, scatter_random_state, }; use crate::joins::hash_join::stream::{ BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, @@ -200,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. @@ -242,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 @@ -268,6 +270,8 @@ impl JoinLeftData { pub(super) enum BuildPhaseOutput { InMemory(Arc), Spilled(Mutex>), + /// CollectLeft: the scattered build is shared by all probe partitions. + SpilledShared(Arc), } /// Helps to build [`HashJoinExec`]. @@ -1340,7 +1344,10 @@ impl ExecutionPlan for HashJoinExec { let would_break_output_order = Self::maintains_input_order(self.join_type)[1] && self.right.output_ordering().is_some(); let spill_context = if can_spill - && self.mode == PartitionMode::Partitioned + && matches!( + self.mode, + PartitionMode::Partitioned | PartitionMode::CollectLeft + ) && !self.null_aware && !would_break_output_order && context.runtime_env().disk_manager.tmp_files_enabled() @@ -1401,8 +1408,10 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, - // CollectLeft spill (shared scattered build) is a follow-up. - None, + spill_context.clone(), + // CollectLeft: the scattered build is shared by all + // probe partitions. + true, )) })?, PartitionMode::Partitioned => { @@ -1426,6 +1435,7 @@ impl ExecutionPlan for HashJoinExec { self.null_equality, array_map_created_count, spill_context.clone(), + false, )) } PartitionMode::Auto => { @@ -1862,6 +1872,10 @@ struct BuildSideState { /// 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. @@ -1894,6 +1908,7 @@ impl BuildSideState { }) .transpose()?, table_reserved: 0, + concurrent_table_consumers: 1, spill: None, scatter: None, headroom: None, @@ -1915,6 +1930,7 @@ impl BuildSideState { buffered, &ctx.pool.memory_limit(), ctx.partition_count_override, + self.concurrent_table_consumers, ); let mut scatter = SideScatter::new( partition_count, @@ -2027,6 +2043,7 @@ async fn collect_left_input( null_equality: NullEquality, array_map_created_count: Count, spill: Option>, + shared_build: bool, ) -> Result { let schema = left_stream.schema(); @@ -2041,6 +2058,9 @@ async fn collect_left_input( 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 { @@ -2104,6 +2124,7 @@ async fn collect_left_input( reservation, bounds_accumulators, table_reserved, + concurrent_table_consumers: _, spill: _, scatter, headroom, @@ -2130,6 +2151,17 @@ async fn collect_left_input( })?; 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, @@ -2156,6 +2188,7 @@ async fn collect_left_input( null_equality, &array_map_created_count, table_reserved, + None, )?; Ok(BuildPhaseOutput::InMemory(Arc::new(data))) @@ -2208,6 +2241,7 @@ pub(super) fn build_left_data( 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( @@ -2285,17 +2319,40 @@ pub(super) fn build_left_data( (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); @@ -2334,8 +2391,8 @@ pub(super) fn build_left_data( 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/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index 88e1a79cb01a2..fb3d536b4b20d 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -42,10 +42,12 @@ 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, }; @@ -63,7 +65,7 @@ use crate::spill::spill_manager::SpillManager; use crate::stream::RecordBatchStreamAdapter; use crate::{RecordBatchStream, SendableRecordBatchStream}; -use arrow::array::{Array, StringViewArray, UInt32Array}; +use arrow::array::{Array, BooleanBufferBuilder, StringViewArray, UInt32Array}; use arrow::compute::{concat_batches, take_record_batch}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -77,7 +79,9 @@ 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. @@ -122,20 +126,26 @@ pub(super) fn scatter_random_state(level: usize) -> RandomState { /// `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 target = match pool_limit { + 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)) @@ -190,6 +200,75 @@ pub(super) struct SpilledBuild { 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 { @@ -1334,6 +1413,7 @@ impl SpillJoinDriver { self.spec.null_equality, &array_map_count, table_reserved, + None, )?; drop(batches); @@ -1370,6 +1450,259 @@ impl SpillJoinDriver { } } +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::*; @@ -2245,4 +2578,119 @@ mod tests { ); } } + /// 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" + ); + } + } } diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index b3d3dfe12ccd4..f567480131e69 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -33,7 +33,8 @@ use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, PushdownStrategy, SharedBuildAccumulator, }; use crate::joins::hash_join::spill::{ - HashJoinSpillContext, InnerJoinSpec, SpillJoinDriver, empty_record_batch_stream, + HashJoinSpillContext, InnerJoinSpec, SharedSpillJoinDriver, SpillJoinDriver, + empty_record_batch_stream, }; use crate::joins::utils::{ OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, @@ -236,7 +237,26 @@ pub(super) struct HashJoinStream { /// 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>, + 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 { @@ -494,24 +514,13 @@ impl HashJoinStream { Poll::Ready(Ok(StatefulStreamResult::Continue)) } - /// Hand off to the spill-join driver after the build side scattered to - /// disk: report bounds-only pushdown (so peer partitions neither deadlock - /// nor prune this partition's probe rows), take ownership of the probe - /// stream, and construct the driver. - 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") - })?; - + /// 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 = spilled - .bounds - .clone() - .unwrap_or_else(|| PartitionBounds::new(vec![])); + let bounds = bounds.unwrap_or_else(|| PartitionBounds::new(vec![])); let build_data = match self.mode { PartitionMode::Partitioned => PartitionBuildData::Partitioned { partition_id: self.partition, @@ -523,9 +532,9 @@ impl HashJoinStream { bounds, }, PartitionMode::Auto => { - return Poll::Ready(internal_err!( + return internal_err!( "PartitionMode::Auto should not be present at execution time" - )); + ); } }; self.build_waiter = Some(OnceFut::new(async move { @@ -535,7 +544,12 @@ impl HashJoinStream { } 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)); @@ -550,13 +564,50 @@ impl HashJoinStream { null_equality: self.null_equality, batch_size: self.batch_size, }; - self.spill_driver = Some(Box::new(SpillJoinDriver::new( + (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)) } @@ -613,6 +664,10 @@ impl HashJoinStream { })?; 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, From 52b12e050235eb614e3f2091b177a1a59698ffa7 Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 17:20:17 +0300 Subject: [PATCH 6/8] Hash join spill M6: lifecycle hardening tests + full-suite regression pass Two more cancel/error-safety tests: dropping the output stream deep inside the chunked skew fallback releases every reservation (chunk data, table estimate, headroom, scratch), and a probe-side error arriving mid-scatter propagates verbatim with the pool drained. Regression status at this point: 18 spill tests, full physical-plan suite 1286 passed (spill_compression zstd feature quirk pre-exists on clean base), full sqllogictest suite green except the pre-existing local encrypted_parquet environment quirk, clippy -D warnings and fmt clean. Performance parity gates (flag-on vs flag-off at SF1/SF10) run with the rustice integration, where a release build and the benchmark harness exist. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/joins/hash_join/spill.rs | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index fb3d536b4b20d..b17e25d2e9453 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -2693,4 +2693,95 @@ mod tests { ); } } + /// 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); + } } From fd92d106023da15213d62d6d45b6522f43462085 Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 20:18:28 +0300 Subject: [PATCH 7/8] Hash join spill: log engage/repartition/chunk progress at info Operational visibility for long spilled joins on remote deployments: one line when a build side switches to scatter mode (with buffered bytes/rows), one per recursive repartition pass, one per chunked- fallback chunk. Mirrors SortExec's spill logging. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 6 ++++++ .../physical-plan/src/joins/hash_join/spill.rs | 13 +++++++++++++ 2 files changed, 19 insertions(+) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1ce51ad25b725..959eb49bbaef3 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1924,6 +1924,12 @@ impl BuildSideState { .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( diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index b17e25d2e9453..4e2d9e90c4a68 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -999,6 +999,12 @@ impl SpillJoinDriver { }; 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, @@ -1373,6 +1379,13 @@ impl SpillJoinDriver { 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, From 8ac1843480b361f0129598be39635aed61f8b03c Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Wed, 8 Jul 2026 23:50:31 +0300 Subject: [PATCH 8/8] Hash join spill: mark reservations spillable only when spilling is actually available MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit with_can_spill() was driven by the raw config flag, but under FairSpillPool a spillable consumer is capped at pool/n_spillable — a join excluded from the spill fallback (null-aware anti join, order-preserving probe side, disabled DiskManager) got the tiny fair share with no fallback to honor it, failing queries that passed before the feature existed. Found by TPC-H Q16 at SF100: its NOT IN subquery plans as a null-aware anti join, which hit a ~155 MB share of a 16 GB pool and errored. The reservation is now marked spillable exactly when the spill context exists (the full gate), so non-degradable joins keep unspillable semantics. Co-Authored-By: Claude Opus 4.8 (1M context) --- datafusion/physical-plan/src/joins/hash_join/exec.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 959eb49bbaef3..49509cbda95e8 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1321,7 +1321,7 @@ impl ExecutionPlan for HashJoinExec { // 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 can_spill = context + let spill_flag = context .session_config() .options() .execution @@ -1343,7 +1343,7 @@ impl ExecutionPlan for HashJoinExec { // 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 can_spill + let spill_context = if spill_flag && matches!( self.mode, PartitionMode::Partitioned | PartitionMode::CollectLeft @@ -1393,7 +1393,7 @@ impl ExecutionPlan for HashJoinExec { let left_stream = self.left.execute(0, Arc::clone(&context))?; let reservation = MemoryConsumer::new("HashJoinInput") - .with_can_spill(can_spill) + .with_can_spill(spill_context.is_some()) .register(context.memory_pool()); Ok(collect_left_input( @@ -1419,7 +1419,7 @@ impl ExecutionPlan for HashJoinExec { let reservation = MemoryConsumer::new(format!("HashJoinInput[{partition}]")) - .with_can_spill(can_spill) + .with_can_spill(spill_context.is_some()) .register(context.memory_pool()); OnceFut::new(collect_left_input(