diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 4fc17f1ebd7cc..cd02bf16b9797 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -20,7 +20,7 @@ use std::time::SystemTime; use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj}; -use arrow::array::{ArrayRef, BinaryArray, Int32Array}; +use arrow::array::{Array, ArrayRef, BinaryArray, Int32Array}; use arrow::compute::SortOptions; use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; @@ -1423,7 +1423,7 @@ fn pwmj_parts_exec( MemorySourceConfig::try_new_exec(&partitions, pwmj_kv_schema(), None).unwrap() } -fn pwmj_existence_plan( +fn pwmj_plan( left: Arc, right: Arc, op: Operator, @@ -1483,14 +1483,17 @@ fn pwmj_nlj_oracle_plan( ) } -/// Executes every output partition concurrently and returns the surviving left `id`s, sorted. +/// Executes every output partition concurrently and returns the join's rows as +/// `(left id, right id)` pairs, sorted so partition interleaving does not affect the +/// comparison. `None` means the join filled that side with NULLs, or -- for the existence +/// joins, whose output carries the left side only -- that the side is absent entirely. /// /// Concurrent rather than one partition at a time: the partitions share the watermark and race /// to be the one that runs the final pass, which is the part a sequential drain cannot reach. -async fn pwmj_collect_ids( +async fn pwmj_collect_id_pairs( plan: Arc, task_ctx: Arc, -) -> Vec { +) -> Vec<(Option, Option)> { let streams = (0..plan.output_partitioning().partition_count()) .map(|partition| plan.execute(partition, Arc::clone(&task_ctx)).unwrap()) .collect::>(); @@ -1500,30 +1503,63 @@ async fn pwmj_collect_ids( })) .await; - let mut ids = Vec::new(); + let id_column = |batch: &RecordBatch, col: usize| { + batch + .column(col) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }; + + let mut pairs = Vec::new(); for batches in per_partition { for batch in batches.unwrap().unwrap() { - let col = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - ids.extend((0..col.len()).map(|i| col.value(i))); + let left = id_column(&batch, 0); + // Left is (id, k), so the right side's `id` follows it -- when there is one. + let right = (batch.num_columns() > 2).then(|| id_column(&batch, 2)); + for row in 0..batch.num_rows() { + pairs.push(( + left.is_valid(row).then(|| left.value(row)), + right + .as_ref() + .filter(|r| r.is_valid(row)) + .map(|r| r.value(row)), + )); + } } } - ids.sort_unstable(); - ids + pairs.sort_unstable(); + pairs } +/// Differential test for every join type `PiecewiseMergeJoin` supports, against a +/// `NestedLoopJoin` oracle. +/// +/// `Left`/`Full` are the ones with teeth: their unmatched buffered rows are derived from the +/// shared `min_marked` watermark rather than materialized per row, and that encoding is only +/// valid because every match marks a *suffix* of the buffered side. The dimensions the +/// cheaper tests do not reach are the ones that matter here -- `pwmj.slt` and the unit tests +/// both run the streamed side at a single partition and the default batch size, so neither +/// covers several partitions racing to run the final pass, nor the mid-scan resume path a +/// small batch size forces. #[tokio::test(flavor = "multi_thread")] -async fn fuzz_pwmj_existence_matches_nested_loop() { - // A small batch size splits the final-pass output across several coalesced batches even - // for these tiny inputs, covering that boundary too. +async fn fuzz_pwmj_matches_nested_loop() { + // A small batch size splits output across several coalesced batches even for these tiny + // inputs, covering that boundary too. let task_ctx = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(3)), ); let ops = [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq]; + let join_types = [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + ]; for seed in 0..60u64 { let mut rng = StdRng::seed_from_u64(seed); @@ -1545,9 +1581,9 @@ async fn fuzz_pwmj_existence_matches_nested_loop() { let right_keys = gen_keys(right_len, &mut rng); for op in ops { - for join_type in [JoinType::LeftSemi, JoinType::LeftAnti] { - let got = pwmj_collect_ids( - pwmj_existence_plan( + for join_type in join_types { + let got = pwmj_collect_id_pairs( + pwmj_plan( pwmj_single_exec(&left_ids, &left_keys), pwmj_parts_exec(&right_ids, &right_keys, nparts), op, @@ -1556,7 +1592,7 @@ async fn fuzz_pwmj_existence_matches_nested_loop() { Arc::clone(&task_ctx), ) .await; - let want = pwmj_collect_ids( + let want = pwmj_collect_id_pairs( pwmj_nlj_oracle_plan( pwmj_single_exec(&left_ids, &left_keys), pwmj_single_exec(&right_ids, &right_keys), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index bc69ae5140831..ac1385223f0fa 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -31,14 +31,15 @@ use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_expr::{JoinType, Operator}; use datafusion_physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt}; +use std::sync::atomic::Ordering as AtomicOrdering; use std::{cmp::Ordering, task::ready}; use std::{sync::Arc, task::Poll}; use crate::handle_state; use crate::joins::piecewise_merge_join::exec::{BufferedSide, BufferedSideReadyState}; use crate::joins::piecewise_merge_join::utils::need_produce_result_in_final; +use crate::joins::utils::JoinKeyComparator; use crate::joins::utils::{BuildProbeJoinMetrics, StatefulStreamResult}; -use crate::joins::utils::{JoinKeyComparator, get_final_indices_from_shared_bitmap}; use crate::stream::EmptyRecordBatchStream; pub(super) enum PiecewiseMergeJoinStreamState { @@ -328,17 +329,19 @@ impl ClassicPWMJStream { return Ok(StatefulStreamResult::Continue); } - let buffered_data = - Arc::clone(&self.buffered_side.try_as_ready().unwrap().buffered_data); - - let (buffered_indices, _streamed_indices) = get_final_indices_from_shared_bitmap( - &buffered_data.visited_indices_bitmap, - self.join_type, - true, - ); - - let new_buffered_batch = - take_record_batch(buffered_data.batch(), &buffered_indices)?; + let buffered_data = Arc::clone(&self.buffered_side.try_as_ready()?.buffered_data); + let buffered_batch = buffered_data.batch(); + + // Every match marks the suffix `[k, buffered_len)`, so the buffered rows that were + // never matched are exactly the complementary prefix `[0, min_marked)` -- which + // includes the null-keyed rows, since nulls sort first and the scan starts past + // them. That makes the final pass a zero-copy slice instead of building an index + // array and running `take` over it. + let min_marked = buffered_data + .min_marked + .load(AtomicOrdering::SeqCst) + .min(buffered_batch.num_rows()); + let new_buffered_batch = buffered_batch.slice(0, min_marked); let mut buffered_columns = new_buffered_batch.columns().to_vec(); let streamed_columns: Vec = self @@ -377,6 +380,10 @@ struct BatchProcessState { continue_process: bool, // Skip nulls processed_null_count: bool, + // Smallest buffered index marked while scanning the current stream batch, or + // `usize::MAX` if nothing has been marked yet. Because `buffer_idx` only moves forward + // within a batch, this lets all but the batch's first match skip the shared atomic. + batch_min_marked: usize, } impl BatchProcessState { @@ -389,6 +396,7 @@ impl BatchProcessState { found: false, continue_process: true, processed_null_count: false, + batch_min_marked: usize::MAX, } } @@ -399,6 +407,7 @@ impl BatchProcessState { self.found = false; self.continue_process = true; self.processed_null_count = false; + self.batch_min_marked = usize::MAX; } // `None` guarantees the coalescer holds no pending rows, so the caller @@ -475,13 +484,14 @@ fn resolve_classic_join( batch_process_state.found = true; let count = buffered_len - buffer_idx; - let batch = build_matched_indices_and_set_buffered_bitmap( + let batch = build_matched_indices_and_mark_buffered( (buffer_idx, count), (row_idx, count), buffered_side, stream_batch, join_type, join_schema, + &mut batch_process_state.batch_min_marked, )?; batch_process_state.output_batches.push_batch(batch)?; @@ -503,13 +513,14 @@ fn resolve_classic_join( if matches!(compare, Ordering::Equal | Ordering::Less) { batch_process_state.found = true; let count = buffered_len - buffer_idx; - let batch = build_matched_indices_and_set_buffered_bitmap( + let batch = build_matched_indices_and_mark_buffered( (buffer_idx, count), (row_idx, count), buffered_side, stream_batch, join_type, join_schema, + &mut batch_process_state.batch_min_marked, )?; // Flush batch and update pointers if we have a completed batch @@ -570,20 +581,31 @@ fn resolve_classic_join( // // The two ranges are: buffered_range: (start index, count) and streamed_range: (start index, count) due // to batch.slice(start, count). -fn build_matched_indices_and_set_buffered_bitmap( +fn build_matched_indices_and_mark_buffered( buffered_range: (usize, usize), streamed_range: (usize, usize), buffered_side: &mut BufferedSideReadyState, stream_batch: &SortedStreamBatch, join_type: JoinType, join_schema: &SchemaRef, + batch_min_marked: &mut usize, ) -> Result { - // Mark the buffered indices as visited - if need_produce_result_in_final(join_type) { - let mut bitmap = buffered_side.buffered_data.visited_indices_bitmap.lock(); - for i in buffered_range.0..buffered_range.0 + buffered_range.1 { - bitmap.set_bit(i, true); - } + // Mark the matched buffered rows. `buffered_range` is always the suffix + // `[start, buffered_len)` -- a match emits every buffered row from the first match on -- + // so the union of everything marked is `[min over matches, buffered_len)` and lowering a + // single watermark records it exactly. That replaces a mutex plus one `set_bit` per + // matched row, which was `O(buffered_len)` work for *every* matched streamed row. + // + // `buffer_idx` is monotone non-decreasing across a stream batch, so only the batch's + // first match can lower the watermark; `batch_min_marked` keeps the atomic off the hot + // path for all the others. It survives the early returns that hand back a completed + // output batch mid-scan, and `reset()` clears it for the next stream batch. + if need_produce_result_in_final(join_type) && buffered_range.0 < *batch_min_marked { + *batch_min_marked = buffered_range.0; + buffered_side + .buffered_data + .min_marked + .fetch_min(buffered_range.0, AtomicOrdering::SeqCst); } let new_buffered_batch = buffered_side diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 5183d3aa0feb7..498bf9fd3dac2 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -17,9 +17,8 @@ use arrow::array::Array; use arrow::{ - array::{ArrayRef, BooleanBufferBuilder, RecordBatch}, + array::{ArrayRef, RecordBatch}, compute::concat_batches, - util::bit_util, }; use arrow_schema::{SchemaRef, SortOptions}; use datafusion_common::not_impl_err; @@ -37,7 +36,6 @@ use datafusion_physical_expr::{ }; use datafusion_physical_expr_common::physical_expr::fmt_sql; use futures::TryStreamExt; -use parking_lot::Mutex; use std::fmt::Formatter; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -49,8 +47,7 @@ use crate::joins::piecewise_merge_join::classic_join::{ }; use crate::joins::piecewise_merge_join::existence_join::ExistencePWMJStream; use crate::joins::piecewise_merge_join::utils::{ - build_visited_indices_map, is_existence_join, is_right_existence_join, - is_supported_existence_join, + is_existence_join, is_right_existence_join, is_supported_existence_join, }; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; @@ -60,10 +57,7 @@ use crate::{ }; use crate::{ ExecutionPlan, PlanProperties, - joins::{ - SharedBitmapBuilder, - utils::{BuildProbeJoinMetrics, OnceAsync, OnceFut, build_join_schema}, - }, + joins::utils::{BuildProbeJoinMetrics, OnceAsync, OnceFut, build_join_schema}, metrics::ExecutionPlanMetricsSet, spill::get_record_batch_memory_size, }; @@ -631,7 +625,6 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { Arc::clone(&on_buffered), metrics.clone(), reservation, - build_visited_indices_map(self.join_type), streamed_partitions, )) })?; @@ -716,22 +709,19 @@ async fn build_buffered_data( on_buffered: PhysicalExprRef, metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, - build_map: bool, remaining_partitions: usize, ) -> Result { let schema = buffered.schema(); // Combine batches and record number of rows - let initial = (Vec::new(), 0, metrics, reservation); - let (batches, num_rows, metrics, reservation) = buffered + let initial = (Vec::new(), metrics, reservation); + let (batches, metrics, reservation) = buffered .try_fold(initial, |mut acc, batch| async { let batch_size = get_record_batch_memory_size(&batch); - acc.3.try_grow(batch_size)?; - acc.2.build_mem_used.add(batch_size); - acc.2.build_input_batches.add(1); - acc.2.build_input_rows.add(batch.num_rows()); - // Update row count - acc.1 += batch.num_rows(); + acc.2.try_grow(batch_size)?; + acc.1.build_mem_used.add(batch_size); + acc.1.build_input_batches.add(1); + acc.1.build_input_rows.add(batch.num_rows()); // Push batch to output acc.0.push(batch); Ok(acc) @@ -752,23 +742,9 @@ async fn build_buffered_data( reservation.try_grow(size_estimation)?; metrics.build_mem_used.add(size_estimation); - // Created visited indices bitmap only if the join type requires it - let visited_indices_bitmap = if build_map { - let bitmap_size = bit_util::ceil(single_batch.num_rows(), 8); - reservation.try_grow(bitmap_size)?; - metrics.build_mem_used.add(bitmap_size); - - let mut bitmap_buffer = BooleanBufferBuilder::new(single_batch.num_rows()); - bitmap_buffer.append_n(num_rows, false); - bitmap_buffer - } else { - BooleanBufferBuilder::new(0) - }; - let buffered_data = BufferedSideData::new( single_batch, buffered_values, - Mutex::new(visited_indices_bitmap), remaining_partitions, reservation, ); @@ -779,13 +755,20 @@ async fn build_buffered_data( pub(super) struct BufferedSideData { pub(super) batch: RecordBatch, values: ArrayRef, - pub(super) visited_indices_bitmap: SharedBitmapBuilder, pub(super) remaining_partitions: AtomicUsize, - /// Existence joins only: the start of the matched suffix of the buffered side, or - /// `usize::MAX` before the first match. `[existence_min_marked, len)` *is* the matched - /// set -- no bitmap is allocated. Shared so each partition benefits from what the - /// others have marked; it only ever decreases, so a stale read is safe. - pub(super) existence_min_marked: AtomicUsize, + /// The start of the matched suffix of the buffered side, or `usize::MAX` before the + /// first match. `[min_marked, len)` *is* the matched set and `[0, min_marked)` the + /// unmatched one -- no bitmap is allocated. + /// + /// Both stream kinds only ever mark a suffix, which is what makes one index enough: + /// - `ExistencePWMJStream` marks `[k, len)` for the first buffered row `k` matching + /// a streamed batch's extreme key. + /// - `ClassicPWMJStream` emits `buffered[k..] x streamed_row` on each match, so the + /// rows it marks are exactly that same suffix. + /// + /// Shared so each partition benefits from what the others have marked; it only ever + /// decreases, so a stale read is safe. + pub(super) min_marked: AtomicUsize, _reservation: MemoryReservation, } @@ -793,16 +776,14 @@ impl BufferedSideData { pub(super) fn new( batch: RecordBatch, values: ArrayRef, - visited_indices_bitmap: SharedBitmapBuilder, remaining_partitions: usize, reservation: MemoryReservation, ) -> Self { Self { batch, values, - visited_indices_bitmap, remaining_partitions: AtomicUsize::new(remaining_partitions), - existence_min_marked: AtomicUsize::new(usize::MAX), + min_marked: AtomicUsize::new(usize::MAX), _reservation: reservation, } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs index b2c5212999f3b..887367b22eeb2 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs @@ -74,8 +74,8 @@ //! //! Marking only ever covers a suffix, and each mark lowers the watermark to its own start, //! so the matched set is always exactly `[min_marked, buffered_len)`. A bitmap would be a -//! less compact encoding of that one index, so none is allocated (see -//! `build_visited_indices_map`). +//! less compact encoding of that one index, so none is allocated. `ClassicPWMJStream` marks +//! the same way, which is why the watermark lives in `BufferedSideData` rather than here. //! //! Once every streamed partition has been consumed, the last one to finish slices the //! buffered batch: `LeftSemi` takes `[min_marked, len)`, `LeftAnti` the complementary @@ -256,9 +256,7 @@ impl ExistencePWMJStream { /// lets every other partition stop too. fn nothing_left_to_mark(&self) -> Result { let buffered_data = &self.buffered_side.try_as_ready()?.buffered_data; - let min_marked = buffered_data - .existence_min_marked - .load(AtomicOrdering::SeqCst); + let min_marked = buffered_data.min_marked.load(AtomicOrdering::SeqCst); let buffered_values = buffered_data.values(); Ok(min_marked.min(buffered_values.len()) <= buffered_values.null_count()) @@ -308,7 +306,7 @@ impl ExistencePWMJStream { // watermark: that bounds the comparisons this batch performs, not just the // bits it writes. let scan_limit = buffered_data - .existence_min_marked + .min_marked .load(AtomicOrdering::SeqCst) .min(buffered_len); @@ -366,7 +364,7 @@ impl ExistencePWMJStream { if buffer_idx < scan_limit { // Everything from `buffer_idx` on matches, so lowering the // watermark to it records the match: the marked set is exactly - // `[existence_min_marked, buffered_len)` and needs no bitmap. + // `[min_marked, buffered_len)` and needs no bitmap. // // INVARIANT: sound only because the buffered side and each // streamed batch are sorted the same way for this operator @@ -377,7 +375,7 @@ impl ExistencePWMJStream { // order, which is why the watermark takes a `min` rather than just // decreasing. buffered_data - .existence_min_marked + .min_marked .fetch_min(buffer_idx, AtomicOrdering::SeqCst); } } @@ -402,7 +400,7 @@ impl ExistencePWMJStream { // `k`, so the union is `[k, len)`. The result is therefore a slice, with no // index array to materialize and no `take`. let min_marked = buffered_data - .existence_min_marked + .min_marked .load(AtomicOrdering::SeqCst) .min(buffered_len); diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs index 5093be0ca19be..11ba328c01062 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs @@ -51,21 +51,3 @@ pub(super) fn is_supported_existence_join(join_type: JoinType) -> bool { pub(super) fn need_produce_result_in_final(join_type: JoinType) -> bool { matches!(join_type, JoinType::Full | JoinType::Left) } - -// Returns boolean for whether or not we need to build the buffered side -// bitmap for marking matched rows on the buffered side. -// -// `LeftSemi`/`LeftAnti` are absent on purpose: `ExistencePWMJStream` only ever marks a -// contiguous suffix of the buffered side, so it tracks the boundary as a single index -// (`BufferedSideData::existence_min_marked`) and needs no bitmap. -pub(super) fn build_visited_indices_map(join_type: JoinType) -> bool { - matches!( - join_type, - JoinType::Full - | JoinType::Left - | JoinType::RightAnti - | JoinType::RightSemi - | JoinType::LeftMark - | JoinType::RightMark - ) -}