From 653e123caa0f529f967d1c13563c1b6717620a79 Mon Sep 17 00:00:00 2001 From: Shehab Ali Date: Thu, 23 Jul 2026 10:56:44 -0400 Subject: [PATCH] refactor join-key equality filtering --- .../hj/benchmarks/q24.benchmark | 31 ++++ .../hj/benchmarks/q25.benchmark | 33 ++++ benchmarks/src/hj.rs | 46 ++++++ datafusion/physical-plan/src/joins/utils.rs | 149 +++++++++--------- 4 files changed, 186 insertions(+), 73 deletions(-) create mode 100644 benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark create mode 100644 benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark new file mode 100644 index 0000000000000..2ea60f0f87009 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark @@ -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; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark new file mode 100644 index 0000000000000..b29d6b959a853 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark @@ -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; diff --git a/benchmarks/src/hj.rs b/benchmarks/src/hj.rs index 7d33bc3aa9e50..9a91e2714ac86 100644 --- a/benchmarks/src/hj.rs +++ b/benchmarks/src/hj.rs @@ -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 { diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 654d873ae1b0e..7ab5a70cdebeb 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -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, @@ -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; @@ -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; @@ -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}; @@ -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::::new().into(), Vec::::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 { - // Nested datatypes cannot use the underlying not_distinct/eq function and must use a special - // implementation - // - 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 @@ -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!(