Skip to content

perf: hash join probe loop has iterator overhead and per-row Growable extends #7098

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
QLiangong wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf: hash join probe loop has iterator overhead and per-row Growable extends #7098

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
72bf235
probe_inner
QLiangong Jun 10, 2026
69ae766
probe_inner_fix
QLiangong Jun 12, 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
171 changes: 124 additions & 47 deletions src/daft-local-execution/src/join/inner_join.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,73 +7,150 @@ use daft_recordbatch::{GrowableRecordBatch, ProbeState};

use crate::join::hash_join::HashJoinParams;

/// Threshold below which the per-row `GrowableRecordBatch::extend` path is faster
/// than bucketing + vectorized `RecordBatch::take`. Picked conservatively;
/// should be replaced with a value derived from a Rust-level micro-benchmark.
const TAKE_BATCH_MIN_MATCHES: usize = 1024;

/// Probe-side fan-out (matches per probe row) at or above which the vectorized
/// `take` path is preferred over the per-row Growable path. `take` has fixed
/// overhead from per-table `UInt64Array` construction + `RecordBatch::concat`;
/// it only pays off once each probe row produces enough matches.
const TAKE_BATCH_MIN_FANOUT: f64 = 3.0;

fn build_final_table(
left_table: daft_recordbatch::RecordBatch,
right_table: daft_recordbatch::RecordBatch,
params: &HashJoinParams,
) -> DaftResult<daft_recordbatch::RecordBatch> {
let common_join_keys: Vec<String> = params.common_join_cols.iter().cloned().collect();
let left_non_join_columns: Vec<String> = params
.left_schema
.field_names()
.filter(|c| !params.common_join_cols.contains(*c))
.map(ToString::to_string)
.collect();
let right_non_join_columns: Vec<String> = params
.right_schema
.field_names()
.filter(|c| !params.common_join_cols.contains(*c))
.map(ToString::to_string)
.collect();

let join_keys_table = daft_recordbatch::get_columns_by_name(&left_table, &common_join_keys)?;
let left_non_join_columns =
daft_recordbatch::get_columns_by_name(&left_table, &left_non_join_columns)?;
let right_non_join_columns =
daft_recordbatch::get_columns_by_name(&right_table, &right_non_join_columns)?;
join_keys_table
.union(&left_non_join_columns)?
.union(&right_non_join_columns)
}

pub(crate) fn probe_inner(
input: &MicroPartition,
probe_state: &ProbeState,
params: &HashJoinParams,
) -> DaftResult<MicroPartition> {
let build_side_tables = probe_state.get_record_batches().iter().collect::<Vec<_>>();
const DEFAULT_GROWABLE_SIZE: usize = 20;

let input_tables = input.record_batches();
let result_tables = input_tables
.iter()
.map(|input_table| {
let mut build_side_growable =
GrowableRecordBatch::new(&build_side_tables, false, DEFAULT_GROWABLE_SIZE)?;
let mut probe_side_idxs = Vec::new();

let join_keys = input_table.eval_expression_list(&params.probe_on)?;
let idx_iter = probe_state.probe_indices(join_keys)?;
for (probe_row_idx, inner_iter) in idx_iter.enumerate() {

// Materialize all matches into a flat Vec so we can (1) inspect the
// total count to pick a build-side strategy and (2) bucket by
// build_rb_idx for the take-based path. The probe-side hot loop
// is a tight `Vec::push`, which the compiler lowers to the same
// code as a `flat_map(...).collect()` form.
let mut matches: Vec<(u32, u64, u64)> = Vec::with_capacity(input_table.len());

@srilman srilman Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than putting everything into matches and then splitting by build_rb_idx, why not just do the fanout here? That way we can cut the memory usage by half

for (probe_idx, inner_iter) in idx_iter.enumerate() {
if let Some(inner_iter) = inner_iter {
for (build_rb_idx, build_row_idx) in inner_iter {
build_side_growable.extend(
build_rb_idx as usize,
build_row_idx as usize,
1,
);
probe_side_idxs.push(probe_row_idx as u64);
for (rb_idx, row_idx) in inner_iter {
matches.push((probe_idx as u32, rb_idx as u64, row_idx));

@greptile-apps greptile-apps Bot Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 probe_idx silently truncates on large input tables

idx_iter.enumerate() yields a usize counter cast to u32 before being stored in matches. On a 64-bit platform, any input table with more than u32::MAX (~4 billion) rows will silently produce wrong probe-side indices. Storing probe_idx as u64 would be consistent with row_idx and avoids the overflow.

}
}
}
let matches_len = matches.len();

let build_side_table = build_side_growable.build()?;
let probe_side_table = {
let indices_arr = UInt64Array::from_vec("", probe_side_idxs);
input_table.take(&indices_arr)?
};
// Fan-out aware path selection.
// - Below `TAKE_BATCH_MIN_MATCHES`: small total work, the per-row
// `GrowableRecordBatch::extend` path is faster.
// - Below `TAKE_BATCH_MIN_FANOUT`: `take + concat` overhead exceeds
// the per-row dispatch cost.
// - Otherwise: bucket by build table and use vectorized take.
let probe_rows = input_table.len().max(1);
let fanout = matches_len as f64 / probe_rows as f64;
let use_batch =
matches_len >= TAKE_BATCH_MIN_MATCHES && fanout >= TAKE_BATCH_MIN_FANOUT;

let (left_table, right_table) = if use_batch {
let build_tables = probe_state.get_record_batches();
let mut probe_indices = Vec::with_capacity(matches_len);
let mut per_table_indices = vec![Vec::new(); build_tables.len()];

for (probe_idx, rb_idx, row_idx) in matches {
probe_indices.push(probe_idx as u64);
per_table_indices[rb_idx as usize].push(row_idx);
}

let (left_table, right_table) = if params.build_on_left {
(build_side_table, probe_side_table)
let mut build_side_tables = Vec::new();
for (rb_idx, indices) in per_table_indices.into_iter().enumerate() {
if !indices.is_empty() {
let idx_arr = UInt64Array::from_vec("", indices);
let taken = build_tables[rb_idx].take(&idx_arr)?;
build_side_tables.push(taken);
}
}

let build_side_table = if build_side_tables.is_empty() {
daft_recordbatch::RecordBatch::empty(Some(params.output_schema.clone()))

@greptile-apps greptile-apps Bot Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Wrong schema used for the empty build-side fallback

RecordBatch::empty(Some(params.output_schema.clone())) creates a table whose schema is the full output schema, but this table is used as the build-side fragment before build_final_table tries to extract common_join_keys and left_non_join_columns from it. The empty table should be created with the build-side schema: left_schema when params.build_on_left is true, right_schema otherwise.

} else if build_side_tables.len() == 1 {
build_side_tables.pop().unwrap()
} else {
daft_recordbatch::RecordBatch::concat(&build_side_tables)?
};
Comment on lines +90 to +136

@greptile-apps greptile-apps Bot Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P0 Build/probe row alignment broken for multi-table build sides

In the batch path, probe_indices is filled in the original match order (interleaved across rb_idx values), but the build-side rows are concatenated in table-index order (rb_idx=0 first, then rb_idx=1, etc.). When the build side spans more than one RecordBatch, the two sides no longer correspond row-for-row and the join emits silently wrong results.

Concrete failure: suppose matches = [(probe=5, rb=1, row=10), (probe=3, rb=0, row=7), (probe=8, rb=1, row=2)]. probe_indices = [5, 3, 8]; the concatenated build side is [table[0][7], table[1][10], table[1][2]]. The final join pairs probe[5] with table[0][7] instead of table[1][10], and probe[3] with table[1][10] instead of table[0][7].

The fix is to sort matches by rb_idx before splitting into probe_indices and per_table_indices, so that probe_indices is accumulated in the same grouped order that the build-side concat produces.


let probe_indices_arr = UInt64Array::from_vec("", probe_indices);
let probe_side_table = input_table.take(&probe_indices_arr)?;

if params.build_on_left {
(build_side_table, probe_side_table)
} else {
(probe_side_table, build_side_table)
}
} else {
(probe_side_table, build_side_table)
let build_side_tables =
probe_state.get_record_batches().iter().collect::<Vec<_>>();
const DEFAULT_GROWABLE_SIZE: usize = 20;

let mut build_side_growable = GrowableRecordBatch::new(
&build_side_tables,
false,
matches_len.max(DEFAULT_GROWABLE_SIZE),
)?;
let mut probe_side_idxs = Vec::with_capacity(matches_len);

for (probe_idx, rb_idx, row_idx) in matches {
build_side_growable.extend(rb_idx as usize, row_idx as usize, 1);
probe_side_idxs.push(probe_idx as u64);
}

let build_side_table = build_side_growable.build()?;
let probe_side_table = {
let indices_arr = UInt64Array::from_vec("", probe_side_idxs);
input_table.take(&indices_arr)?
};

if params.build_on_left {
(build_side_table, probe_side_table)
} else {
(probe_side_table, build_side_table)
}
};

let common_join_keys: Vec<String> = params.common_join_cols.iter().cloned().collect();
let left_non_join_columns: Vec<String> = params
.left_schema
.field_names()
.filter(|c| !params.common_join_cols.contains(*c))
.map(ToString::to_string)
.collect();
let right_non_join_columns: Vec<String> = params
.right_schema
.field_names()
.filter(|c| !params.common_join_cols.contains(*c))
.map(ToString::to_string)
.collect();

let join_keys_table =
daft_recordbatch::get_columns_by_name(&left_table, &common_join_keys)?;
let left_non_join_columns =
daft_recordbatch::get_columns_by_name(&left_table, &left_non_join_columns)?;
let right_non_join_columns =
daft_recordbatch::get_columns_by_name(&right_table, &right_non_join_columns)?;
let final_table = join_keys_table
.union(&left_non_join_columns)?
.union(&right_non_join_columns)?;
Ok(final_table)
build_final_table(left_table, right_table, params)
})
.collect::<DaftResult<Vec<_>>>()?;

Expand Down
Loading