Skip to content

perf: replace PiecewiseMergeJoin's visited bitmap with a suffix watermark (~1.8x faster Left join) #24579

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jayzhan211 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf: replace PiecewiseMergeJoin's visited bitmap with a suffix watermark (~1.8x faster Left join) #24579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9461da5
perf: replace PiecewiseMergeJoin's visited bitmap with a suffix water…
jayzhan211 Aug 22, 2026
cbdcc72
refactor: simplify comments in join utility functions
jayzhan211 Aug 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 57 additions & 21 deletions datafusion/core/tests/fuzz_cases/join_fuzz.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<dyn ExecutionPlan>,
right: Arc<dyn ExecutionPlan>,
op: Operator,
Expand Down Expand Up @@ -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<dyn ExecutionPlan>,
task_ctx: Arc<TaskContext>,
) -> Vec<i32> {
) -> Vec<(Option<i32>, Option<i32>)> {
let streams = (0..plan.output_partitioning().partition_count())
.map(|partition| plan.execute(partition, Arc::clone(&task_ctx)).unwrap())
.collect::<Vec<_>>();
Expand All @@ -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::<Int32Array>()
.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::<Int32Array>()
.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);
Expand All @@ -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,
Expand All @@ -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),
Expand Down
64 changes: 43 additions & 21 deletions datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<ArrayRef> = self
Expand Down Expand Up @@ -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 {
Expand All @@ -389,6 +396,7 @@ impl BatchProcessState {
found: false,
continue_process: true,
processed_null_count: false,
batch_min_marked: usize::MAX,
}
}

Expand All @@ -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
Expand Down Expand Up @@ -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)?;
Expand All @@ -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
Expand Down Expand Up @@ -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<RecordBatch> {
// 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
Expand Down
Loading