Skip to content

refactor join-key equality filtering #23843

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
shehab-ali wants to merge 1 commit into
base: main
Choose a base branch
from
+186 −73
Open

refactor join-key equality filtering #23843

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
31 changes: 31 additions & 0 deletions benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
name Q24
group hj

init sql_benchmarks/hj/init/set_config_no_stats.sql

load sql_benchmarks/hj/init/load.sql

assert I
SELECT count(*) > 0 FROM lineitem
----
true

expect_plan HashJoinExec

run
-- Q24: single-hot-bucket long string-key inner join.
-- Build rows all share one long string key, so each matching probe row fans
-- out to the whole build side. count(*) focuses the benchmark on hash match
-- and equality filtering without buffering joined rows.
-- Thresholds zeroed to force Partitioned mode (simulates absent row-count stats).
SELECT count(*)
FROM (
SELECT 'single_hot_bucket_string_join_key' as k
FROM supplier
WHERE s_suppkey <= 3000
) s
JOIN (
SELECT 'single_hot_bucket_string_join_key' as k
FROM lineitem
WHERE l_orderkey % 3000 = 0
) l ON s.k = l.k;
33 changes: 33 additions & 0 deletions benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
name Q25
group hj

init sql_benchmarks/hj/init/set_config_no_stats.sql

load sql_benchmarks/hj/init/load.sql

assert I
SELECT count(*) > 0 FROM lineitem
----
true

expect_plan HashJoinExec

run
-- Q25: skewed high-fanout multi-column string-key inner join.
-- This tracks candidate-pair filtering for composite keys: the first key is
-- skewed and the second long string key must also be checked before emitting
-- each match. count(*) isolates the match path.
-- Thresholds zeroed to force Partitioned mode (simulates absent row-count stats).
SELECT count(*)
FROM (
SELECT CAST((s_suppkey % 256) + 1 AS INT) as k1,
'multi_column_high_fanout_key' as k2
FROM supplier
WHERE s_suppkey <= 20000
) s
JOIN (
SELECT CAST(1 AS INT) as k1,
'multi_column_high_fanout_key' as k2
FROM lineitem
WHERE l_orderkey % 250 = 0
) l ON s.k1 = l.k1 AND s.k2 = l.k2;
46 changes: 46 additions & 0 deletions benchmarks/src/hj.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -472,6 +472,52 @@ const HASH_QUERIES: &[HashJoinQuery] = &[
probe_size: "2.3M_long_keys_count",
isolate_partitioned_join: true,
},
// Q24: single-hot-bucket long string-key inner join.
// Build rows all share one long string key, so each matching probe row fans
// out to the whole build side. The output is counted to focus on the hash
// match/equality path without buffering the joined rows.
HashJoinQuery {
sql: r###"SELECT count(*)
FROM (
SELECT 'single_hot_bucket_string_join_key' as k
FROM supplier
WHERE s_suppkey <= 3000
) s
JOIN (
SELECT 'single_hot_bucket_string_join_key' as k
FROM lineitem
WHERE l_orderkey % 3000 = 0
) l ON s.k = l.k"###,
density: 1.0,
prob_hit: 1.0,
build_size: "3K_(single_hot_bucket)",
probe_size: "20K_long_keys_count",
isolate_partitioned_join: true,
},
// Q25: skewed high-fanout multi-column string-key inner join.
// This tracks the same candidate-pair filtering path for composite join
// keys, where the first key is skewed and the second long string key must
// also be checked before emitting each match.
HashJoinQuery {
sql: r###"SELECT count(*)
FROM (
SELECT CAST((s_suppkey % 256) + 1 AS INT) as k1,
'multi_column_high_fanout_key' as k2
FROM supplier
WHERE s_suppkey <= 20000
) s
JOIN (
SELECT CAST(1 AS INT) as k1,
'multi_column_high_fanout_key' as k2
FROM lineitem
WHERE l_orderkey % 250 = 0
) l ON s.k1 = l.k1 AND s.k2 = l.k2"###,
density: 1.0,
prob_hit: 1.0,
build_size: "20K_(fanout~78_multi_key)",
probe_size: "240K_multi_key_count",
isolate_partitioned_join: true,
},
];

impl RunOpt {
Expand Down
149 changes: 76 additions & 73 deletions datafusion/physical-plan/src/joins/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ pub use crate::joins::{JoinOn, JoinOnRef};
use arrow::array::{
Array, ArrowPrimitiveType, BooleanBufferBuilder, NativeAdapter, PrimitiveArray,
RecordBatch, RecordBatchOptions, UInt32Array, UInt32Builder, UInt64Array,
builder::UInt64Builder, downcast_array, make_array, new_null_array,
builder::UInt64Builder, downcast_array, new_null_array,
};
use arrow::array::{
ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array,
Expand All @@ -54,14 +54,12 @@ use arrow::array::{
TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array,
};
use arrow::buffer::{BooleanBuffer, NullBuffer};
use arrow::compute::kernels::cmp::eq;
use arrow::compute::{self, FilterBuilder, and, take};
use arrow::compute::{self, take};
use arrow::datatypes::{
ArrowNativeType, Field, Schema, SchemaBuilder, UInt32Type, UInt64Type,
};
use arrow_ord::cmp::not_distinct;
use arrow_ord::ord::{DynComparator, make_comparator};
use arrow_schema::{ArrowError, DataType, SortOptions, TimeUnit};
use arrow_schema::{DataType, SortOptions, TimeUnit};
use datafusion_common::cast::as_boolean_array;
use datafusion_common::hash_utils::RandomState;
use datafusion_common::hash_utils::create_hashes;
Expand All @@ -71,7 +69,6 @@ use datafusion_common::{
DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult,
internal_datafusion_err, not_impl_err, plan_err,
};
use datafusion_expr::Operator;
use datafusion_expr::interval_arithmetic::Interval;
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::utils::collect_columns;
Expand All @@ -80,7 +77,6 @@ use datafusion_physical_expr::{
add_offset_to_physical_sort_exprs,
};

use datafusion_physical_expr_common::datum::compare_op_for_nested;
use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
use futures::future::{BoxFuture, Shared};
use futures::{FutureExt, ready};
Expand Down Expand Up @@ -2196,78 +2192,30 @@ pub(super) fn equal_rows_arr(
right_arrays: &[ArrayRef],
null_equality: NullEquality,
) -> Result<(UInt64Array, UInt32Array)> {
let mut iter = left_arrays.iter().zip(right_arrays.iter());

let Some((first_left, first_right)) = iter.next() else {
if left_arrays.is_empty() {
return Ok((Vec::<u64>::new().into(), Vec::<u32>::new().into()));
};

let arr_left = take(first_left.as_ref(), indices_left, None)?;
let arr_right = take(first_right.as_ref(), indices_right, None)?;

let mut equal: BooleanArray = eq_dyn_null(&arr_left, &arr_right, null_equality)?;
}

// Use map and try_fold to iterate over the remaining pairs of arrays.
// In each iteration, take is used on the pair of arrays and their equality is determined.
// The results are then folded (combined) using the and function to get a final equality result.
equal = iter
.map(|(left, right)| {
let arr_left = take(left.as_ref(), indices_left, None)?;
let arr_right = take(right.as_ref(), indices_right, None)?;
eq_dyn_null(arr_left.as_ref(), arr_right.as_ref(), null_equality)
})
.try_fold(equal, |acc, equal2| and(&acc, &equal2?))?;
let sort_options = vec![SortOptions::default(); left_arrays.len()];
let comparator =
JoinKeyComparator::new(left_arrays, right_arrays, &sort_options, null_equality)?;

let filter_builder = FilterBuilder::new(&equal).optimize().build();
let mut left_filtered = Vec::with_capacity(indices_left.len());
let mut right_filtered = Vec::with_capacity(indices_right.len());

let left_filtered = filter_builder.filter(indices_left)?;
let right_filtered = filter_builder.filter(indices_right)?;
for (left, right) in indices_left.values().iter().zip(indices_right.values()) {
let left_idx = usize::try_from(*left).map_err(|_| {
internal_datafusion_err!("Join index {left} can not be represented as usize")
})?;
let right_idx = *right as usize;

Ok((
downcast_array(left_filtered.as_ref()),
downcast_array(right_filtered.as_ref()),
))
}

// version of eq_dyn supporting equality on null arrays
fn eq_dyn_null(
left: &dyn Array,
right: &dyn Array,
null_equality: NullEquality,
) -> Result<BooleanArray, ArrowError> {
// Nested datatypes cannot use the underlying not_distinct/eq function and must use a special
// implementation
// <https://github.com/apache/datafusion/issues/10749>
if left.data_type().is_nested() {
let op = match null_equality {
NullEquality::NullEqualsNothing => Operator::Eq,
NullEquality::NullEqualsNull => Operator::IsNotDistinctFrom,
};
return Ok(compare_op_for_nested(op, &left, &right)?);
}
// Arrow's `eq` / `not_distinct` use IEEE 754 totalOrder semantics for
// floats, so `-0.0` and `+0.0` would compare unequal. Normalize float
// operands first; non-float types dispatch directly to avoid the
// `make_array(to_data())` round-trip.
if !matches!(
left.data_type(),
DataType::Float16 | DataType::Float32 | DataType::Float64
) {
return match null_equality {
NullEquality::NullEqualsNothing => eq(&left, &right),
NullEquality::NullEqualsNull => not_distinct(&left, &right),
};
}
let left_arr: ArrayRef = make_array(left.to_data());
let right_arr: ArrayRef = make_array(right.to_data());
let left_norm = normalize_float_zero(&left_arr);
let right_norm = normalize_float_zero(&right_arr);
let left = left_norm.as_ref();
let right = right_norm.as_ref();
match null_equality {
NullEquality::NullEqualsNothing => eq(&left, &right),
NullEquality::NullEqualsNull => not_distinct(&left, &right),
if comparator.is_equal(left_idx, right_idx) {
left_filtered.push(*left);
right_filtered.push(*right);
}
}

Ok((left_filtered.into(), right_filtered.into()))
}

/// Pre-built comparator for join key columns that eliminates per-row type
Expand Down Expand Up @@ -4553,6 +4501,61 @@ mod tests {
assert_eq!(cmp_nl.compare(1, 1), Ordering::Less);
}

#[test]
fn test_equal_rows_arr_filters_candidate_pairs() {
let left_a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 2, 3]));
let left_b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
let right_a: ArrayRef = Arc::new(Int32Array::from(vec![2, 2, 3, 4]));
let right_b: ArrayRef = Arc::new(StringArray::from(vec!["b", "d", "d", "a"]));

let left_indices = UInt64Array::from(vec![0, 1, 2, 3]);
let right_indices = UInt32Array::from(vec![0, 0, 1, 2]);

let (left_filtered, right_filtered) = equal_rows_arr(
&left_indices,
&right_indices,
&[left_a, left_b],
&[right_a, right_b],
NullEquality::NullEqualsNothing,
)
.unwrap();

assert_eq!(left_filtered, UInt64Array::from(vec![1, 3]));
assert_eq!(right_filtered, UInt32Array::from(vec![0, 2]));
}

#[test]
fn test_equal_rows_arr_respects_null_equality() {
let left: ArrayRef =
Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None]));
let right: ArrayRef =
Arc::new(Int32Array::from(vec![None, Some(1), Some(2), None]));
let left_indices = UInt64Array::from(vec![0, 1, 2, 3]);
let right_indices = UInt32Array::from(vec![1, 0, 2, 3]);

let (left_filtered, right_filtered) = equal_rows_arr(
&left_indices,
&right_indices,
&[Arc::clone(&left)],
&[Arc::clone(&right)],
NullEquality::NullEqualsNothing,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0, 2]));
assert_eq!(right_filtered, UInt32Array::from(vec![1, 2]));

let (left_filtered, right_filtered) = equal_rows_arr(
&left_indices,
&right_indices,
&[left],
&[right],
NullEquality::NullEqualsNull,
)
.unwrap();
assert_eq!(left_filtered, UInt64Array::from(vec![0, 1, 2, 3]));
assert_eq!(right_filtered, UInt32Array::from(vec![1, 0, 2, 3]));
}

#[test]
fn test_max_distinct_count_preserves_precision_when_not_capped() {
assert_eq!(
Expand Down
Loading