This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24976-9f21155c48faac3f1658853ad6a55bd314a322ed in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit cc61eeaa5eb775258e1d520dfbb6622c436aaf42 Author: Jay Zhan <[email protected]> AuthorDate: Wed Sep 9 14:07:17 2026 +0000 perf: emit unmatched build rows in batch_size chunks in HashJoinExec (#25028) ## Which issue does this PR close? - Related to #24768 (bounded-memory hash join) and #18942 (hash join operator optimizations); no dedicated issue. ## Rationale for this change For join types that emit build-side rows after the probe side is exhausted (`Left`, `Full`, `LeftAnti`, `LeftSemi`, `LeftMark`), `HashJoinExec` computed the final indices over the *whole* build side and materialized them as **one** `RecordBatch` before handing it to the output coalescer. That batch is not bounded by `batch_size`, and it is not covered by the memory reservation. Worse, `LimitedBatchCoalescer` configures arrow's `BatchCoalescer` with `biggest_coalesce_batch_size = batch_size / 2`, which passes any larger batch through untouched. So a `LEFT ANTI` join over a 10M-row build side with few matches emitted a single ~10M-row batch downstream, regardless of `datafusion.execution.batch_size`. `NestedLoopJoinExec` already emits its unmatched build rows in `batch_size` chunks; this PR brings `HashJoinExec` in line. ## What changes are included in this PR? - `HashJoinStream` gets a new state, `EmitUnmatchedBuildRows`, entered from `ExhaustedProbeSide` by the last probe partition. It holds a `BooleanBuffer` snapshot of the visited bitmap (taken once, after every partition reported completion, so the lock is not held while emitting) and a cursor. - `next_final_indices_chunk` scans the snapshot from the cursor and returns at most `batch_size` final indices per call (`LeftMark` emits every row, so its chunks are plain ranges). Null-aware `LeftAnti`/`LeftMark` post-processing and `fetch` handling are unchanged and now run per chunk. - `input_batches` is still bumped once for the final phase and `input_rows` once per chunk, so metric values are identical to before. - `get_final_indices_from_bit_map` / `get_final_indices_from_shared_bitmap` in `joins/utils.rs` had no other callers and are removed. - New benchmark cases in `hash_join_semi_anti.rs` with a 1M-row build side and a 100K-row probe side (`left_semi_build1m_h10`, `left_anti_build1m_h10`, `left_build1m_h10`). Benchmark (this branch vs. `main`, Apple Silicon, `cargo bench --bench hash_join_semi_anti -- build1m`): | case | main | this PR | change | |---|---|---|---| | left_semi_build1m_h10 | 2.05 ms | 1.97 ms | -4% | | left_anti_build1m_h10 | 6.21 ms | 4.75 ms | -24% | | left_build1m_h10 | 7.37 ms | 6.79 ms | -8% | Peak RSS of the `left_anti_build1m_h10` bench binary: ~369 MB on `main` vs ~160 MB on this branch (`/usr/bin/time -l`, 1M build rows, 900K unmatched). ## Are these changes tested? - New `join_emits_final_build_rows_in_batch_size_chunks` test (Left/Full/LeftAnti/LeftSemi/LeftMark × batch sizes 1/7/8192 × perfect-hash-join on/off) asserts the output rows and that every output batch respects `batch_size`. On `main` 14 of its 30 cases fail the batch-size assertion. - New `join_fetch_stops_final_build_rows_mid_chunk` test checks a `fetch` that is reached in the middle of the final rows. - Existing hash join unit tests (493), join sqllogictests (`joins`, `join_limit_pushdown`, `join_disable_repartition_joins`, `subquery`), and the core join fuzz tests pass. ## Are there any user-facing changes? No result changes. Output batches of the affected join types are now bounded by `batch_size` instead of arriving as one batch holding every unmatched build row. --- .../physical-plan/benches/hash_join_semi_anti.rs | 59 +++++++ .../physical-plan/src/joins/hash_join/exec.rs | 162 ++++++++++++++++++- .../physical-plan/src/joins/hash_join/stream.rs | 174 ++++++++++++++++++--- datafusion/physical-plan/src/joins/utils.rs | 56 ------- 4 files changed, 374 insertions(+), 77 deletions(-) diff --git a/datafusion/physical-plan/benches/hash_join_semi_anti.rs b/datafusion/physical-plan/benches/hash_join_semi_anti.rs index 1e11da36be..40ba41272c 100644 --- a/datafusion/physical-plan/benches/hash_join_semi_anti.rs +++ b/datafusion/physical-plan/benches/hash_join_semi_anti.rs @@ -380,6 +380,65 @@ fn bench_hash_join_semi_anti(c: &mut Criterion) { }); } + // ========================================================================= + // Build-side output benchmarks (LeftSemi / LeftAnti / Left) + // ========================================================================= + // + // These join types emit build-side rows only after the probe side is + // exhausted, so they exercise the final "unmatched build rows" emission. + // Build side: 1M rows, Probe side: 100K rows, so most build rows stay + // unmatched and the final emission dominates. + let large_build_rows = 1_000_000; + let small_probe_rows = 100_000; + + // LeftSemi - 10% of build rows matched -> 100K output rows + { + let left_batches = build_batches(large_build_rows, large_build_rows, 0, &s); + let right_batches = build_batches(small_probe_rows, small_probe_rows, 0, &s); + group.bench_function( + BenchmarkId::new("left_semi_build1m_h10", small_probe_rows), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::LeftSemi, &rt) + }) + }, + ); + } + + // LeftAnti - 10% of build rows matched -> 900K output rows + { + let left_batches = build_batches(large_build_rows, large_build_rows, 0, &s); + let right_batches = build_batches(small_probe_rows, small_probe_rows, 0, &s); + group.bench_function( + BenchmarkId::new("left_anti_build1m_h10", small_probe_rows), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::LeftAnti, &rt) + }) + }, + ); + } + + // Left - 10% of build rows matched -> 100K matched + 900K unmatched output rows + { + let left_batches = build_batches(large_build_rows, large_build_rows, 0, &s); + let right_batches = build_batches(small_probe_rows, small_probe_rows, 0, &s); + group.bench_function( + BenchmarkId::new("left_build1m_h10", small_probe_rows), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::Left, &rt) + }) + }, + ); + } + group.finish(); } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 5ead46ed63..01dec6570b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3041,11 +3041,11 @@ mod tests { }; use arrow::array::{ - Array, ArrayRef, Date32Array, DictionaryArray, Int32Array, Int64Array, + Array, ArrayRef, AsArray, Date32Array, DictionaryArray, Int32Array, Int64Array, StructArray, UInt32Array, UInt64Array, }; use arrow::buffer::NullBuffer; - use arrow::datatypes::{DataType, Field}; + use arrow::datatypes::{DataType, Field, Int32Type}; use datafusion_common::hash_utils::create_hashes; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{ @@ -4127,6 +4127,164 @@ mod tests { return Ok(()); } + /// Build side keyed `0..num_build_rows`, probe side holding only + /// `matched_keys`, joined on the key column. + fn final_build_rows_inputs( + num_build_rows: i32, + matched_keys: &[i32], + ) -> (Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>, JoinOn) { + let build_keys: Vec<i32> = (0..num_build_rows).collect(); + let probe_keys = matched_keys.to_vec(); + let left = build_table( + ("a1", &build_keys), + ("b1", &build_keys), + ("c1", &build_keys), + ); + let right = build_table( + ("a2", &probe_keys), + ("b2", &probe_keys), + ("c2", &probe_keys), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + (left, right, on) + } + + /// Returns the sorted `(a1, matched)` pairs of the joined output, where + /// `matched` is read from the mark column for `LeftMark`, from the + /// presence of probe-side data for `Left`/`Full`, and is `None` for + /// existence joins whose output carries no match information. + fn build_rows_with_match_flag(batches: &[RecordBatch]) -> Vec<(i32, Option<bool>)> { + let mut rows = vec![]; + for batch in batches { + let a1 = batch + .column_by_name("a1") + .unwrap() + .as_primitive::<Int32Type>(); + let matched: Vec<Option<bool>> = match batch.column_by_name("mark") { + Some(mark) => mark.as_boolean().iter().collect(), + None => match batch.column_by_name("a2") { + Some(a2) => (0..a2.len()).map(|i| Some(a2.is_valid(i))).collect(), + None => vec![None; batch.num_rows()], + }, + }; + rows.extend(a1.values().iter().copied().zip(matched)); + } + rows.sort_unstable(); + rows + } + + /// The `(a1, matched)` pairs a join of [`final_build_rows_inputs`] must + /// produce, in the format of [`build_rows_with_match_flag`]. + fn expected_final_build_rows( + join_type: JoinType, + num_build_rows: i32, + matched_keys: &[i32], + ) -> Vec<(i32, Option<bool>)> { + (0..num_build_rows) + .filter_map(|key| { + let matched = matched_keys.contains(&key); + match join_type { + JoinType::LeftSemi => matched.then_some((key, None)), + JoinType::LeftAnti => (!matched).then_some((key, None)), + _ => Some((key, Some(matched))), + } + }) + .collect() + } + + /// The final build-side rows (unmatched rows, or matched rows for + /// `LeftSemi`) must be emitted in chunks bounded by `batch_size` instead + /// of one batch over the whole build side. + #[rstest] + #[tokio::test] + async fn join_emits_final_build_rows_in_batch_size_chunks( + #[values( + JoinType::Left, + JoinType::Full, + JoinType::LeftAnti, + JoinType::LeftSemi, + JoinType::LeftMark + )] + join_type: JoinType, + #[values(1, 7, 8192)] batch_size: usize, + #[values(true, false)] use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let num_build_rows = 100; + let matched_keys = [3, 50, 97]; + let (left, right, on) = final_build_rows_inputs(num_build_rows, &matched_keys); + + let (_, batches, metrics) = join_collect( + left, + right, + on, + &join_type, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + let expected = + expected_final_build_rows(join_type, num_build_rows, &matched_keys); + assert_eq!(build_rows_with_match_flag(&batches), expected); + + for batch in &batches { + assert!( + batch.num_rows() <= batch_size, + "{join_type} join emitted a batch of {} rows with batch_size {batch_size}", + batch.num_rows() + ); + } + assert!( + batches.len() >= expected.len().div_ceil(batch_size), + "{join_type} join emitted {} batches for {} rows with batch_size {batch_size}", + batches.len(), + expected.len() + ); + + Ok(()) + } + + /// A `fetch` limit that is reached in the middle of the final build-side + /// rows stops the emission at exactly `fetch` rows. + #[rstest] + #[tokio::test] + async fn join_fetch_stops_final_build_rows_mid_chunk( + #[values(JoinType::Left, JoinType::LeftAnti, JoinType::LeftMark)] + join_type: JoinType, + ) -> Result<()> { + let batch_size = 7; + let fetch = 20; + let num_build_rows = 100; + let matched_keys = [3, 50, 97]; + let task_ctx = prepare_task_ctx(batch_size, false); + let (left, right, on) = final_build_rows_inputs(num_build_rows, &matched_keys); + + let join = HashJoinExecBuilder::new(left, right, on, join_type) + .with_partition_mode(PartitionMode::CollectLeft) + .with_fetch(Some(fetch)) + .build()?; + let batches = common::collect(join.execute(0, task_ctx)?).await?; + + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, fetch); + for batch in &batches { + assert!(batch.num_rows() <= batch_size); + } + // Every emitted row is a genuine join result row. + let expected = + expected_final_build_rows(join_type, num_build_rows, &matched_keys); + for row in build_rows_with_match_flag(&batches) { + assert!(expected.contains(&row), "unexpected output row {row:?}"); + } + + Ok(()) + } + #[apply(hash_join_exec_configs)] #[tokio::test] async fn join_full_multi_batch( diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index ff08828e94..2c1ad94460 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -32,9 +32,7 @@ use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode}; use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, SharedBuildAccumulator, }; -use crate::joins::utils::{ - OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, matchable_join_keys, -}; +use crate::joins::utils::{OnceFut, equal_rows_arr, matchable_join_keys}; use crate::stream::EmptyRecordBatchStream; use crate::{ RecordBatchStream, SendableRecordBatchStream, handle_state, @@ -49,7 +47,7 @@ use crate::{ }; use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; -use arrow::buffer::NullBuffer; +use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ @@ -119,11 +117,18 @@ impl BuildSide { /// WaitBuildSide /// │ /// ▼ -/// ┌─► FetchProbeBatch ───► ExhaustedProbeSide ───► Completed -/// │ │ -/// │ ▼ -/// └─ ProcessProbeBatch +/// ┌─► FetchProbeBatch ───► ExhaustedProbeSide ──────────► Completed +/// │ │ │ ▲ +/// │ ▼ ▼ │ +/// └─ ProcessProbeBatch ┌─► EmitUnmatchedBuildRows ───────┘ +/// └──────────┘ /// ``` +/// +/// `ExhaustedProbeSide` moves to `EmitUnmatchedBuildRows` only for join types +/// that emit build-side rows after the probe side is exhausted (see +/// [`need_produce_result_in_final`]), and only in the partition that finished +/// probing last. That state re-enters itself once per emitted chunk of at most +/// `batch_size` rows. #[derive(Debug, Clone)] pub(super) enum HashJoinStreamState { /// Initial state for HashJoinStream indicating that build-side data not collected yet @@ -136,6 +141,10 @@ pub(super) enum HashJoinStreamState { ProcessProbeBatch(ProcessProbeBatchState), /// Indicates that probe-side has been fully processed ExhaustedProbeSide, + /// Indicates that the probe side has been fully processed by every + /// partition, and this stream is emitting the final build-side rows + /// (unmatched rows, or matched rows for `LeftSemi`) in chunks + EmitUnmatchedBuildRows(EmitUnmatchedBuildRowsState), /// Indicates that HashJoinStream execution is completed Completed, } @@ -149,6 +158,19 @@ impl HashJoinStreamState { _ => internal_err!("Expected hash join stream in ProcessProbeBatch state"), } } + + /// Tries to extract EmitUnmatchedBuildRowsState from HashJoinStreamState enum. + /// Returns an error if state is not EmitUnmatchedBuildRows. + fn try_as_emit_unmatched_build_rows_mut( + &mut self, + ) -> Result<&mut EmitUnmatchedBuildRowsState> { + match self { + HashJoinStreamState::EmitUnmatchedBuildRows(state) => Ok(state), + _ => { + internal_err!("Expected hash join stream in EmitUnmatchedBuildRows state") + } + } + } } /// Container for HashJoinStreamState::ProcessProbeBatch related data @@ -177,6 +199,19 @@ impl ProcessProbeBatchState { } } +/// Container for HashJoinStreamState::EmitUnmatchedBuildRows related data +#[derive(Debug, Clone)] +pub(super) struct EmitUnmatchedBuildRowsState { + /// Snapshot of the build-side visited bitmap, one bit per build row. + /// + /// Taken once every probe partition has reported completion, so no + /// further updates to the shared bitmap are possible and the snapshot + /// can be scanned without holding its lock. + visited: BooleanBuffer, + /// Index of the next build row to examine + cursor: usize, +} + /// Lifecycle of this partition's build-data report to the shared coordinator. /// /// `Scheduled` means the reporting `OnceFut` has been constructed but is lazy: @@ -638,7 +673,10 @@ impl HashJoinStream { handle_state!(self.process_probe_batch()) } HashJoinStreamState::ExhaustedProbeSide => { - handle_state!(self.process_unmatched_build_batch()) + handle_state!(self.prepare_unmatched_build_rows()) + } + HashJoinStreamState::EmitUnmatchedBuildRows(_) => { + handle_state!(self.emit_unmatched_build_rows()) } HashJoinStreamState::Completed if !self.output_buffer.is_empty() => { // Flush any remaining buffered data @@ -971,10 +1009,14 @@ impl HashJoinStream { Ok(StatefulStreamResult::Continue) } - /// Processes unmatched build-side rows for certain join types and produces output batch + /// Decides whether this stream emits the final build-side rows (unmatched + /// rows, or matched rows for `LeftSemi`) once the probe side is exhausted. /// - /// Updates state to `Completed` - fn process_unmatched_build_batch( + /// Only the last partition to finish probing emits them, since the shared + /// visited bitmap is complete only at that point. Updates state to + /// `EmitUnmatchedBuildRows` when there are rows to emit, and to `Completed` + /// otherwise. + fn prepare_unmatched_build_rows( &mut self, ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { let timer = self.join_metrics.join_time.timer(); @@ -1006,11 +1048,56 @@ impl HashJoinStream { return Ok(StatefulStreamResult::Continue); } + // Every probe partition has finished, so the bitmap is final: snapshot + // it once and release the lock for the whole emission phase. + let visited = build_side + .left_data + .visited_indices_bitmap() + .lock() + .finish_cloned(); + + // The final build rows count as one logical input batch; the rows are + // added chunk by chunk in `emit_unmatched_build_rows`. + self.join_metrics.input_batches.add(1); + + timer.done(); + self.state = + HashJoinStreamState::EmitUnmatchedBuildRows(EmitUnmatchedBuildRowsState { + visited, + cursor: 0, + }); + + Ok(StatefulStreamResult::Continue) + } + + /// Emits the next chunk of at most `batch_size` final build-side rows and + /// stays in `EmitUnmatchedBuildRows` until every build row has been + /// examined, then updates state to `Completed`. + /// + /// Emitting in chunks instead of one batch over the whole build side keeps + /// the output within `batch_size` and bounds the memory materialized at + /// once, which otherwise grows with the number of unmatched build rows + /// (e.g. a highly selective `LeftAnti` join over a large build side). + fn emit_unmatched_build_rows( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + let timer = self.join_metrics.join_time.timer(); + + let state = self.state.try_as_emit_unmatched_build_rows_mut()?; + if state.cursor >= state.visited.len() { + timer.done(); + self.state = HashJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); + } + + let build_side = self.build_side.try_as_ready()?; + // use the global left bitmap to produce the left indices and right indices - let (left_side, right_side) = get_final_indices_from_shared_bitmap( - build_side.left_data.visited_indices_bitmap(), + let (left_side, right_side) = next_final_indices_chunk( + &state.visited, + &mut state.cursor, self.join_type, - true, + self.batch_size, ); // Null-aware joins post-process the build rows under SQL three-valued @@ -1035,10 +1122,9 @@ impl HashJoinStream { _ => (left_side, right_side, None), }; - self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(left_side.len()); - // Push final unmatched indices to output buffer + // Push this chunk of final indices to output buffer if !left_side.is_empty() { let empty_right_batch = RecordBatch::new_empty(self.right.schema()); let batch = build_batch_from_indices( @@ -1054,19 +1140,69 @@ impl HashJoinStream { )?; let push_status = self.output_buffer.push_batch(batch)?; - // If limit reached, finish the coalescer + // If limit reached, finish the coalescer and stop emitting if push_status == PushBatchStatus::LimitReached { self.output_buffer.finish()?; + self.state = HashJoinStreamState::Completed; } } timer.done(); - self.state = HashJoinStreamState::Completed; Ok(StatefulStreamResult::Continue) } } +/// Returns the next chunk of final build-side indices for join types that +/// produce build rows once the probe side is exhausted, starting at `cursor` +/// and holding at most `batch_size` rows. Advances `cursor` past the build +/// rows examined; `cursor == visited.len()` means every row has been examined. +/// +/// The build indices are always valid. The probe indices are NULL for every +/// row (`Left`, `LeftAnti`, `Full`, `LeftSemi`), except for `LeftMark`, where +/// every build row is emitted and a NULL probe index marks an unmatched row. +/// +/// For example, with `visited = [true, false, true, true, false]`: +/// - `Left`: build `[1, 4]`, probe `[null, null]` +/// - `LeftSemi`: build `[0, 2, 3]`, probe `[null, null, null]` +/// - `LeftMark`: build `[0, 1, 2, 3, 4]`, probe `[0, null, 0, 0, null]` +fn next_final_indices_chunk( + visited: &BooleanBuffer, + cursor: &mut usize, + join_type: JoinType, + batch_size: usize, +) -> (UInt64Array, UInt32Array) { + let num_rows = visited.len(); + let start = *cursor; + + if join_type == JoinType::LeftMark { + // Every build row is emitted, so a chunk is a plain range of rows. + let end = (start + batch_size).min(num_rows); + let build_indices = (start as u64..end as u64).collect::<UInt64Array>(); + let probe_indices = (start..end) + .map(|idx| visited.value(idx).then_some(0)) + .collect::<UInt32Array>(); + *cursor = end; + return (build_indices, probe_indices); + } + + // `LeftSemi` emits the matched build rows; `Left`, `LeftAnti` and `Full` + // emit the unmatched ones. + let emit_visited = join_type == JoinType::LeftSemi; + let mut build_indices = Vec::with_capacity(batch_size.min(num_rows - start)); + let mut idx = start; + while idx < num_rows && build_indices.len() < batch_size { + if visited.value(idx) == emit_visited { + build_indices.push(idx as u64); + } + idx += 1; + } + *cursor = idx; + + let probe_indices = UInt32Array::new_null(build_indices.len()); + (UInt64Array::from(build_indices), probe_indices) +} + /// Applies the pre-lookup bookkeeping of a null-aware join to a probe batch /// and returns `true` if the batch cannot contribute any output row. /// diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 04ff9e30bf..4e01588e5a 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -26,7 +26,6 @@ use std::ops::Range; use std::sync::Arc; use std::task::{Context, Poll}; -use crate::joins::SharedBitmapBuilder; use crate::metrics::{ self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, @@ -1217,61 +1216,6 @@ pub(crate) fn is_existence_join(join_type: JoinType) -> bool { ) } -pub(crate) fn get_final_indices_from_shared_bitmap( - shared_bitmap: &SharedBitmapBuilder, - join_type: JoinType, - piecewise: bool, -) -> (UInt64Array, UInt32Array) { - let bitmap = shared_bitmap.lock(); - get_final_indices_from_bit_map(&bitmap, join_type, piecewise) -} - -/// In the end of join execution, need to use bit map of the matched -/// indices to generate the final left and right indices. -/// -/// For example: -/// -/// 1. left_bit_map: `[true, false, true, true, false]` -/// 2. join_type: `Left` -/// -/// The result is: `([1,4], [null, null])` -pub(crate) fn get_final_indices_from_bit_map( - left_bit_map: &BooleanBufferBuilder, - join_type: JoinType, - // We add a flag for whether this is being passed from the `PiecewiseMergeJoin` - // because the bitmap can be for left + right `JoinType`s - piecewise: bool, -) -> (UInt64Array, UInt32Array) { - let left_size = left_bit_map.len(); - if join_type == JoinType::LeftMark || (join_type == JoinType::RightMark && piecewise) - { - let left_indices = (0..left_size as u64).collect::<UInt64Array>(); - let right_indices = (0..left_size) - .map(|idx| left_bit_map.get_bit(idx).then_some(0)) - .collect::<UInt32Array>(); - return (left_indices, right_indices); - } - let left_indices = if join_type == JoinType::LeftSemi - || (join_type == JoinType::RightSemi && piecewise) - { - (0..left_size) - .filter_map(|idx| (left_bit_map.get_bit(idx)).then_some(idx as u64)) - .collect::<UInt64Array>() - } else { - // just for `Left`, `LeftAnti` and `Full` join - // `LeftAnti`, `Left` and `Full` will produce the unmatched left row finally - (0..left_size) - .filter_map(|idx| (!left_bit_map.get_bit(idx)).then_some(idx as u64)) - .collect::<UInt64Array>() - }; - // right_indices - // all the element in the right side is None - let mut builder = UInt32Builder::with_capacity(left_indices.len()); - builder.append_nulls(left_indices.len()); - let right_indices = builder.finish(); - (left_indices, right_indices) -} - #[expect(clippy::too_many_arguments)] pub(crate) fn apply_join_filter_to_indices( build_input_buffer: &RecordBatch, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
