You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/02 14:32:38 UTC
(datafusion) branch gh-readonly-queue/main/pr-25658-a138ce4eee144977bbb0c2bdbc6d320d567b7940 created (now c922f8811a)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25658-a138ce4eee144977bbb0c2bdbc6d320d567b7940
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at c922f8811a perf: stop redoing per-chunk setup on the hash join probe path (#25658)
This branch includes the following new commits:
new c922f8811a perf: stop redoing per-chunk setup on the hash join probe path (#25658)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: stop redoing per-chunk setup on the hash join probe path (#25658)
Posted by gi...@apache.org.
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-25658-a138ce4eee144977bbb0c2bdbc6d320d567b7940
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit c922f8811a89bcf3f1cce875c8037e55e06bc053
Author: dustin <dw...@users.noreply.github.com>
AuthorDate: Fri Oct 2 14:32:15 2026 +0000
perf: stop redoing per-chunk setup on the hash join probe path (#25658)
## Which issue does this PR close?
- Closes #12131. The fused index compare that issue asked for landed in
#23843; this picks up the per-chunk allocations it left behind and
removes the stale note that still pointed at the old code.
## Rationale for this change
Three pieces of per-chunk work on the hash join probe path repeat work
whose inputs do not change within a probe batch. The multi-column key
comparator was rebuilt for every output chunk, and building it scans the
entire build side for negative zero when a key is a float. The
integer-map path cloned both index buffers per chunk. Every output batch
deep-cloned the output schema, metadata included, to wrap it in a new
`Arc`.
## What changes are included in this PR?
- The key comparator is built once per probe batch, cached on the
stream, and cleared when the next batch arrives or the stream stops.
- The integer-map path takes and returns its index buffers instead of
cloning them.
- `build_batch_from_indices` and `build_batch_empty_build_side` accept a
`SchemaRef` and clone the `Arc`.
Two-column join on `(l_orderkey, CAST(l_quantity AS DOUBLE))` over TPC-H
SF1 with a 60K row build side and 17M output rows, median of five: 150
ms before, 108 ms after; `HashJoinExec` `elapsed_compute` 147 ms to 104
ms. Low-fanout float joins, the `hj` benchmark queries 23 to 25, and
TPC-H q7 and q9 stay within run-to-run noise.
## What is the testing strategy for this PR?
New tests in `datafusion/physical-plan/src/joins/hash_join`:
`join_inner_multi_key_float_zero_and_nulls_across_chunks`,
`join_inner_multi_key_across_probe_batches` (fails if the per-batch
comparator reset is removed), and `join_inner_array_map_across_chunks`.
Existing join tests and the `join` sqllogictest files pass unchanged.
## Are there any user-facing changes?
No.
---
.../physical-plan/src/joins/hash_join/exec.rs | 368 ++++++++++++++++++++-
.../physical-plan/src/joins/hash_join/stream.rs | 46 ++-
.../physical-plan/src/joins/symmetric_hash_join.rs | 3 +-
datafusion/physical-plan/src/joins/utils.rs | 58 +++-
4 files changed, 450 insertions(+), 25 deletions(-)
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 9af40a9dfe..bd8103db0f 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -3488,8 +3488,8 @@ mod tests {
use arrow::array::{
Array, ArrayRef, AsArray, BinaryViewArray, Date32Array, DictionaryArray,
- Int32Array, Int64Array, StringArray, StringViewArray, StructArray, UInt32Array,
- UInt64Array,
+ Float64Array, Int32Array, Int64Array, StringArray, StringViewArray, StructArray,
+ UInt32Array, UInt64Array,
};
use arrow::buffer::NullBuffer;
use arrow::datatypes::{DataType, Field, Int32Type};
@@ -4166,6 +4166,221 @@ mod tests {
Ok(())
}
+ /// Two-column keys (Int32, Float64) take the general comparator path.
+ /// Small batch sizes split each probe batch into several lookup chunks,
+ /// so the key comparison must stay correct across chunk boundaries:
+ /// `-0.0` still equals `0.0`. NULL keys are also present on both sides;
+ /// under `NullEqualsNothing` those rows are skipped by the lookup before
+ /// any key comparison, so they must not appear in the output.
+ #[rstest]
+ #[tokio::test]
+ async fn join_inner_multi_key_float_zero_and_nulls_across_chunks(
+ #[values(8192, 3, 2, 1)] batch_size: usize,
+ ) -> Result<()> {
+ let task_ctx = prepare_task_ctx(batch_size, false);
+
+ let schema = |suffix: &str| {
+ Arc::new(Schema::new(vec![
+ Field::new(format!("k{suffix}"), DataType::Int32, true),
+ Field::new(format!("f{suffix}"), DataType::Float64, true),
+ Field::new(format!("v{suffix}"), DataType::Int32, true),
+ ]))
+ };
+ let table = |suffix: &str,
+ k: Vec<Option<i32>>,
+ f: Vec<Option<f64>>,
+ v: Vec<i32>|
+ -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = schema(suffix);
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(k)),
+ Arc::new(Float64Array::from(f)),
+ Arc::new(Int32Array::from(v)),
+ ],
+ )?;
+ TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)
+ .map(|exec| exec as _)
+ };
+
+ let left = table(
+ "1",
+ vec![Some(1), Some(1), Some(1), Some(2), Some(1)],
+ vec![Some(-0.0), Some(0.0), None, Some(1.5), Some(-0.0)],
+ vec![10, 11, 12, 13, 14],
+ )?;
+ let right = table(
+ "2",
+ vec![Some(1), Some(1), Some(2), Some(1)],
+ vec![Some(0.0), None, Some(1.5), Some(-0.0)],
+ vec![100, 101, 102, 103],
+ )?;
+
+ let on = vec![
+ (
+ Arc::new(Column::new_with_schema("k1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("k2", &right.schema())?) as _,
+ ),
+ (
+ Arc::new(Column::new_with_schema("f1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("f2", &right.schema())?) as _,
+ ),
+ ];
+
+ let (_, batches, _) = join_collect(
+ left,
+ right,
+ on,
+ &JoinType::Inner,
+ NullEquality::NullEqualsNothing,
+ task_ctx,
+ )
+ .await?;
+
+ allow_duplicates! {
+ assert_snapshot!(batches_to_sort_string(&batches), @r"
+ +----+------+----+----+------+-----+
+ | k1 | f1 | v1 | k2 | f2 | v2 |
+ +----+------+----+----+------+-----+
+ | 1 | -0.0 | 10 | 1 | -0.0 | 103 |
+ | 1 | -0.0 | 10 | 1 | 0.0 | 100 |
+ | 1 | -0.0 | 14 | 1 | -0.0 | 103 |
+ | 1 | -0.0 | 14 | 1 | 0.0 | 100 |
+ | 1 | 0.0 | 11 | 1 | -0.0 | 103 |
+ | 1 | 0.0 | 11 | 1 | 0.0 | 100 |
+ | 2 | 1.5 | 13 | 2 | 1.5 | 102 |
+ +----+------+----+----+------+-----+
+ ");
+ }
+
+ Ok(())
+ }
+
+ /// Two probe batches carry different multi-column keys at the same row
+ /// positions. Each batch has to be compared against its own key arrays;
+ /// a comparator left over from the first batch would drop the second
+ /// batch's match.
+ #[rstest]
+ #[tokio::test]
+ async fn join_inner_multi_key_across_probe_batches(
+ #[values(8192, 1)] batch_size: usize,
+ ) -> Result<()> {
+ let task_ctx = prepare_task_ctx(batch_size, false);
+ let left = build_table_two_cols(
+ ("a1", &vec![Some(1), Some(1)]),
+ ("b1", &vec![Some(10), Some(20)]),
+ );
+
+ let right_schema = Arc::new(Schema::new(vec![
+ Field::new("a2", DataType::Int32, true),
+ Field::new("b2", DataType::Int32, true),
+ ]));
+ let probe_batch = |b: i32| {
+ RecordBatch::try_new(
+ Arc::clone(&right_schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1])),
+ Arc::new(Int32Array::from(vec![b])),
+ ],
+ )
+ };
+ let right = TestMemoryExec::try_new_exec(
+ &[vec![probe_batch(10)?, probe_batch(20)?]],
+ Arc::clone(&right_schema),
+ None,
+ )?;
+
+ let on = vec![
+ (
+ Arc::new(Column::new_with_schema("a1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("a2", &right.schema())?) as _,
+ ),
+ (
+ Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+ ),
+ ];
+
+ let (_, batches, _) = join_collect(
+ left,
+ right,
+ on,
+ &JoinType::Inner,
+ NullEquality::NullEqualsNothing,
+ task_ctx,
+ )
+ .await?;
+
+ allow_duplicates! {
+ assert_snapshot!(batches_to_sort_string(&batches), @r"
+ +----+----+----+----+
+ | a1 | b1 | a2 | b2 |
+ +----+----+----+----+
+ | 1 | 10 | 1 | 10 |
+ | 1 | 20 | 1 | 20 |
+ +----+----+----+----+
+ ");
+ }
+
+ Ok(())
+ }
+
+ /// An integer key joined through the perfect hash (array map) path, with
+ /// more candidate pairs than the batch size so each probe batch is
+ /// emitted over several chunks.
+ #[rstest]
+ #[tokio::test]
+ async fn join_inner_array_map_across_chunks(
+ #[values(8192, 3, 2, 1)] batch_size: usize,
+ ) -> Result<()> {
+ let task_ctx = prepare_task_ctx(batch_size, true);
+ let left = build_table(
+ ("a1", &vec![1, 2, 3, 4, 5, 6]),
+ ("b1", &vec![1, 1, 1, 2, 2, 3]),
+ ("c1", &vec![7, 8, 9, 10, 11, 12]),
+ );
+ let right = build_table(
+ ("a2", &vec![10, 20, 30, 40]),
+ ("b1", &vec![1, 2, 3, 4]),
+ ("c2", &vec![70, 80, 90, 100]),
+ );
+ let on = vec![(
+ Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("b1", &right.schema())?) as _,
+ )];
+
+ let (_, batches, metrics) = join_collect(
+ left,
+ right,
+ on,
+ &JoinType::Inner,
+ NullEquality::NullEqualsNothing,
+ task_ctx,
+ )
+ .await?;
+
+ allow_duplicates! {
+ assert_snapshot!(batches_to_sort_string(&batches), @r"
+ +----+----+----+----+----+----+
+ | a1 | b1 | c1 | a2 | b1 | c2 |
+ +----+----+----+----+----+----+
+ | 1 | 1 | 7 | 10 | 1 | 70 |
+ | 2 | 1 | 8 | 10 | 1 | 70 |
+ | 3 | 1 | 9 | 10 | 1 | 70 |
+ | 4 | 2 | 10 | 20 | 2 | 80 |
+ | 5 | 2 | 11 | 20 | 2 | 80 |
+ | 6 | 3 | 12 | 30 | 3 | 90 |
+ +----+----+----+----+----+----+
+ ");
+ }
+
+ assert_join_metrics!(metrics, 6);
+ assert_phj_used(&metrics, true);
+
+ Ok(())
+ }
+
#[tokio::test]
async fn join_inner_one_no_shared_column_names() -> Result<()> {
let task_ctx = Arc::new(TaskContext::default());
@@ -6535,6 +6750,7 @@ mod tests {
(0, None),
&mut probe_indices_buffer,
&mut build_indices_buffer,
+ &mut None,
)?;
let left_ids: UInt64Array = vec![0, 1].into();
@@ -6597,6 +6813,7 @@ mod tests {
(0, None),
&mut probe_indices_buffer,
&mut build_indices_buffer,
+ &mut None,
)?;
// We still expect to match rows 0 and 1 on both sides
@@ -10061,6 +10278,153 @@ mod tests {
Ok(())
}
+ /// Correlated null-aware `LeftAnti` with a two-column correlation scope,
+ /// so scope lookups compare keys with the general multi-column comparator.
+ ///
+ /// Each NULL-marking direction yields several scope matches for a single
+ /// probe batch, which small batch sizes split across many lookup chunks.
+ #[apply(hash_join_exec_configs)]
+ #[tokio::test]
+ async fn test_null_aware_left_anti_multi_column_scope(
+ batch_size: usize,
+ ) -> Result<()> {
+ let task_ctx = prepare_task_ctx(batch_size, false);
+
+ fn table(
+ id: Vec<Option<i32>>,
+ g1: Vec<Option<i32>>,
+ g2: Vec<Option<i32>>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("g1", DataType::Int32, true),
+ Field::new("g2", DataType::Int32, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(id)),
+ Arc::new(Int32Array::from(g1)),
+ Arc::new(Int32Array::from(g2)),
+ ],
+ )
+ .unwrap();
+ TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap()
+ }
+
+ let left = table(
+ vec![
+ Some(1),
+ Some(2),
+ None,
+ None,
+ Some(6),
+ Some(8),
+ Some(9),
+ Some(10),
+ None,
+ ],
+ vec![
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(2),
+ Some(3),
+ ],
+ vec![
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(2),
+ Some(2),
+ Some(2),
+ Some(1),
+ Some(3),
+ ],
+ );
+ let right = table(
+ vec![
+ Some(2),
+ Some(3),
+ Some(4),
+ Some(5),
+ None,
+ None,
+ Some(7),
+ Some(11),
+ ],
+ vec![
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(2),
+ ],
+ vec![
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(1),
+ Some(2),
+ Some(2),
+ Some(2),
+ Some(1),
+ ],
+ );
+
+ let on = ["id", "g1", "g2"]
+ .into_iter()
+ .map(|name| {
+ Ok((
+ Arc::new(Column::new_with_schema(name, &left.schema())?) as _,
+ Arc::new(Column::new_with_schema(name, &right.schema())?) as _,
+ ))
+ })
+ .collect::<Result<JoinOn>>()?;
+
+ let join = HashJoinExec::try_new(
+ left,
+ right,
+ on,
+ None,
+ &JoinType::LeftAnti,
+ None,
+ PartitionMode::CollectLeft,
+ NullEquality::NullEqualsNothing,
+ true,
+ )?;
+
+ let stream = join.execute(0, task_ctx)?;
+ let batches = common::collect(stream).await?;
+
+ // Scope (1, 1) holds 2, 3, 4 and 5: `(1, 1, 1)` is kept, `(2, 1, 1)`
+ // matches, and both NULL-valued build rows there are UNKNOWN. Scope
+ // (1, 2) holds two NULLs, so all of its build rows are UNKNOWN. Scope
+ // (2, 1) holds only 11, so `(10, 2, 1)` is kept, and `(NULL, 3, 3)`
+ // has an empty scope.
+ allow_duplicates! {
+ assert_snapshot!(batches_to_sort_string(&batches), @r"
+ +----+----+----+
+ | id | g1 | g2 |
+ +----+----+----+
+ | | 3 | 3 |
+ | 1 | 1 | 1 |
+ | 10 | 2 | 1 |
+ +----+----+----+
+ ");
+ }
+
+ Ok(())
+ }
+
#[test]
fn test_lr_is_preserved() {
assert_eq!(lr_is_preserved(JoinType::Inner), (true, true));
diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs
index 03387ea574..3068bcb638 100644
--- a/datafusion/physical-plan/src/joins/hash_join/stream.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs
@@ -32,7 +32,9 @@ use crate::joins::hash_join::probe_completion::ProbeSideSummary;
use crate::joins::hash_join::shared_bounds::{
PartitionBounds, PartitionBuildData, SharedBuildAccumulator,
};
-use crate::joins::utils::{OnceFut, equal_rows_arr, matchable_join_keys};
+use crate::joins::utils::{
+ JoinKeyComparator, OnceFut, equal_rows_arr, matchable_join_keys,
+};
use crate::stream::EmptyRecordBatchStream;
use crate::{
RecordBatchStream, SendableRecordBatchStream, handle_state,
@@ -128,7 +130,7 @@ impl BuildSide {
/// [`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)]
+#[derive(Debug)]
pub(super) enum HashJoinStreamState {
/// Initial state for HashJoinStream indicating that build-side data not collected yet
WaitBuildSide,
@@ -173,7 +175,7 @@ impl HashJoinStreamState {
}
/// Container for HashJoinStreamState::ProcessProbeBatch related data
-#[derive(Debug, Clone)]
+#[derive(Debug)]
pub(super) struct ProcessProbeBatchState {
/// Current probe-side batch
batch: RecordBatch,
@@ -192,6 +194,9 @@ pub(super) struct ProcessProbeBatchState {
/// `probe_hit_rate` and `avg_fanout` count a probe row whose matches span
/// several chunks only once.
matched_probe_idx: Option<u32>,
+ /// Key comparator for this batch, built on first use and reused by every
+ /// chunk
+ key_comparator: Option<JoinKeyComparator>,
}
impl ProcessProbeBatchState {
@@ -475,6 +480,7 @@ pub(super) fn lookup_join_hashmap(
offset: MapOffset,
probe_indices_buffer: &mut Vec<u32>,
build_indices_buffer: &mut Vec<u64>,
+ key_comparator: &mut Option<JoinKeyComparator>,
) -> Result<(UInt64Array, UInt32Array, Option<MapOffset>)> {
let next_offset = build_hashmap.get_matched_indices_with_limit_offset(
hashes_buffer,
@@ -490,14 +496,13 @@ pub(super) fn lookup_join_hashmap(
let probe_indices_unfiltered: UInt32Array =
std::mem::take(probe_indices_buffer).into();
- // TODO: optimize equal_rows_arr to avoid allocation of intermediate arrays
- // https://github.com/apache/datafusion/issues/12131
let (build_indices, probe_indices) = equal_rows_arr(
&build_indices_unfiltered,
&probe_indices_unfiltered,
build_side_values,
probe_side_values,
null_equality,
+ key_comparator,
)?;
// Reclaim buffers
@@ -808,6 +813,7 @@ impl HashJoinStream {
offset: (0, None),
joined_probe_idx: None,
matched_probe_idx: None,
+ key_comparator: None,
});
}
Some(Err(err)) => return Poll::Ready(Err(err)),
@@ -880,6 +886,11 @@ impl HashJoinStream {
return Ok(StatefulStreamResult::Continue);
}
+ // Array map lookups move their index buffers into the arrays below.
+ // Keep a handle on those buffers so they can be reused as scratch
+ // space once this chunk's output batch is built.
+ let mut array_map_buffers = None;
+
// get the matched by join keys indices
let (left_indices, right_indices, next_offset) = match build_side.left_data.map()
{
@@ -894,6 +905,7 @@ impl HashJoinStream {
state.offset,
&mut self.probe_indices_buffer,
&mut self.build_indices_buffer,
+ &mut state.key_comparator,
)?,
Map::ArrayMap(array_map) => {
let next_offset = array_map.get_matched_indices_with_limit_offset(
@@ -903,11 +915,15 @@ impl HashJoinStream {
&mut self.probe_indices_buffer,
&mut self.build_indices_buffer,
)?;
- (
- UInt64Array::from(self.build_indices_buffer.clone()),
- UInt32Array::from(self.probe_indices_buffer.clone()),
- next_offset,
- )
+ let build_indices: UInt64Array =
+ std::mem::take(&mut self.build_indices_buffer).into();
+ let probe_indices: UInt32Array =
+ std::mem::take(&mut self.probe_indices_buffer).into();
+ array_map_buffers = Some((
+ build_indices.values().clone(),
+ probe_indices.values().clone(),
+ ));
+ (build_indices, probe_indices, next_offset)
}
};
@@ -1008,6 +1024,14 @@ impl HashJoinStream {
None,
)?;
+ // Reclaim the array map scratch buffers now that no index array
+ // refers to them.
+ drop((left_indices, right_indices));
+ if let Some((build_buffer, probe_buffer)) = array_map_buffers {
+ self.build_indices_buffer = build_buffer.into();
+ self.probe_indices_buffer = probe_buffer.into();
+ }
+
let push_status = self.output_buffer.push_batch(batch)?;
timer.done();
@@ -1556,6 +1580,7 @@ fn for_each_scope_match(
mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>,
) -> Result<()> {
let mut offset = (0, None);
+ let mut key_comparator = None;
loop {
let (build_indices, probe_indices, next_offset) = lookup_join_hashmap(
scope_map,
@@ -1568,6 +1593,7 @@ fn for_each_scope_match(
offset,
probe_indices_buffer,
build_indices_buffer,
+ &mut key_comparator,
)?;
if !build_indices.is_empty() {
diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
index 99a12796c6..d7398c2a90 100644
--- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
+++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
@@ -1232,7 +1232,7 @@ pub(crate) fn build_side_determined_results(
let empty_probe_batch = RecordBatch::new_empty(probe_schema);
// Build the final result from the indices of build and probe sides:
build_batch_from_indices(
- output_schema.as_ref(),
+ output_schema,
&build_hash_joiner.input_buffer,
&empty_probe_batch,
&build_indices,
@@ -1447,6 +1447,7 @@ fn lookup_join_hashmap(
&build_join_values,
&keys_values,
null_equality,
+ &mut None,
)?;
Ok((build_indices, probe_indices))
diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs
index 16f61d7409..1dfabe3d54 100644
--- a/datafusion/physical-plan/src/joins/utils.rs
+++ b/datafusion/physical-plan/src/joins/utils.rs
@@ -55,7 +55,7 @@ use arrow::array::{
use arrow::buffer::{BooleanBuffer, NullBuffer};
use arrow::compute::{self, take};
use arrow::datatypes::{
- ArrowNativeType, Field, Schema, SchemaBuilder, UInt32Type, UInt64Type,
+ ArrowNativeType, Field, Schema, SchemaBuilder, SchemaRef, UInt32Type, UInt64Type,
};
use arrow_ord::ord::{DynComparator, make_comparator};
use arrow_schema::{DataType, SortOptions, TimeUnit};
@@ -1288,10 +1288,10 @@ pub(crate) fn apply_join_filter_to_indices(
/// Creates a [RecordBatch] with zero columns but the given row count.
/// Used when a join has an empty projection (e.g. `SELECT count(1) ...`).
-fn new_empty_schema_batch(schema: &Schema, row_count: usize) -> Result<RecordBatch> {
+fn new_empty_schema_batch(schema: &SchemaRef, row_count: usize) -> Result<RecordBatch> {
let options = RecordBatchOptions::new().with_row_count(Some(row_count));
Ok(RecordBatch::try_new_with_options(
- Arc::new(schema.clone()),
+ Arc::clone(schema),
vec![],
&options,
)?)
@@ -1301,7 +1301,7 @@ fn new_empty_schema_batch(schema: &Schema, row_count: usize) -> Result<RecordBat
/// The resulting batch has [Schema] `schema`.
#[expect(clippy::too_many_arguments)]
pub(crate) fn build_batch_from_indices(
- schema: &Schema,
+ schema: &SchemaRef,
build_input_buffer: &RecordBatch,
probe_batch: &RecordBatch,
build_indices: &UInt64Array,
@@ -1358,7 +1358,7 @@ pub(crate) fn build_batch_from_indices(
columns.push(array);
}
- Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?)
+ Ok(RecordBatch::try_new(Arc::clone(schema), columns)?)
}
/// Builds the nullable mark column for a null-aware `LeftMark` join.
@@ -1426,7 +1426,7 @@ pub(crate) fn build_null_aware_left_mark_column(
/// rows or because none of its rows has a matchable (non-NULL) join key.
/// The resulting batch has [Schema] `schema`.
pub(crate) fn build_batch_empty_build_side(
- schema: &Schema,
+ schema: &SchemaRef,
build_batch: &RecordBatch,
probe_batch: &RecordBatch,
column_indices: &[ColumnIndex],
@@ -1434,7 +1434,7 @@ pub(crate) fn build_batch_empty_build_side(
) -> Result<RecordBatch> {
if join_type.empty_build_side_produces_empty_result() {
// These join types only return data if the left side is not empty.
- return Ok(RecordBatch::new_empty(Arc::new(schema.clone())));
+ return Ok(RecordBatch::new_empty(Arc::clone(schema)));
}
// The remaining joins return right-side rows and nulls for the left side.
@@ -1460,7 +1460,7 @@ pub(crate) fn build_batch_empty_build_side(
})
.collect();
- Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?)
+ Ok(RecordBatch::try_new(Arc::clone(schema), columns)?)
}
/// The input is the matched indices for left and right and
@@ -2286,12 +2286,21 @@ pub(crate) fn matchable_join_keys(
}
}
+/// Keeps only the candidate pairs whose join keys are equal.
+///
+/// `comparator` caches the general-path [`JoinKeyComparator`] between calls
+/// that share the same `left_arrays` and `right_arrays`, so its setup cost is
+/// paid once rather than per call. Pass an empty slot whenever either side's
+/// arrays change. Nothing checks this: a slot kept across different arrays
+/// silently returns wrong matches. An empty slot (`&mut None`) is always
+/// correct and only rebuilds the comparator on each call.
pub(super) fn equal_rows_arr(
indices_left: &UInt64Array,
indices_right: &UInt32Array,
left_arrays: &[ArrayRef],
right_arrays: &[ArrayRef],
null_equality: NullEquality,
+ comparator: &mut Option<JoinKeyComparator>,
) -> Result<(UInt64Array, UInt32Array)> {
if indices_left.len() != indices_right.len() {
return Err(internal_datafusion_err!(
@@ -2333,9 +2342,18 @@ pub(super) fn equal_rows_arr(
return Ok(res);
}
- let sort_options = vec![SortOptions::default(); left_arrays.len()];
- let comparator =
- JoinKeyComparator::new(left_arrays, right_arrays, &sort_options, null_equality)?;
+ let comparator = match comparator {
+ Some(comparator) => comparator,
+ None => {
+ let sort_options = vec![SortOptions::default(); left_arrays.len()];
+ comparator.insert(JoinKeyComparator::new(
+ left_arrays,
+ right_arrays,
+ &sort_options,
+ null_equality,
+ )?)
+ }
+ };
let mut left_filtered = Vec::with_capacity(indices_left.len());
let mut right_filtered = Vec::with_capacity(indices_right.len());
@@ -2462,6 +2480,12 @@ pub struct JoinKeyComparator {
rest: Vec<DynComparator>,
}
+impl Debug for JoinKeyComparator {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("JoinKeyComparator").finish_non_exhaustive()
+ }
+}
+
impl JoinKeyComparator {
/// Build comparators for each join key column pair.
pub fn new(
@@ -4625,7 +4649,7 @@ mod tests {
// When the output schema has no fields (empty projection pushed into
// the join), build_batch_empty_build_side should return a RecordBatch
// with the correct row count but no columns.
- let empty_schema = Schema::empty();
+ let empty_schema = Arc::new(Schema::empty());
let build_batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
@@ -4798,6 +4822,7 @@ mod tests {
&[left_a, left_b],
&[right_a, right_b],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
@@ -4816,6 +4841,7 @@ mod tests {
&[],
&[],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
@@ -4838,6 +4864,7 @@ mod tests {
&[Arc::clone(&left)],
&[Arc::clone(&right)],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0, 2]));
@@ -4849,6 +4876,7 @@ mod tests {
&[left],
&[right],
NullEquality::NullEqualsNull,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0, 1, 2, 3]));
@@ -4881,6 +4909,7 @@ mod tests {
&[Arc::clone(&left)],
&[Arc::clone(&right)],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0]));
@@ -4893,6 +4922,7 @@ mod tests {
&[left],
&[right],
NullEquality::NullEqualsNull,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0, 1]));
@@ -4911,6 +4941,7 @@ mod tests {
&[left],
&[right],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0]));
@@ -5020,6 +5051,7 @@ mod tests {
&[left],
&[right],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0]));
@@ -5037,6 +5069,7 @@ mod tests {
&[Arc::clone(&left)],
&[Arc::clone(&right)],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap_err();
assert!(
@@ -5050,6 +5083,7 @@ mod tests {
&[left, Arc::new(Int32Array::from(vec![3, 4]))],
&[right],
NullEquality::NullEqualsNothing,
+ &mut None,
)
.unwrap_err();
assert!(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org