From 23d7f1c1339240a75dd2963eaa6be40d735df462 Mon Sep 17 00:00:00 2001 From: Shiv Bhatia Date: Tue, 25 Aug 2026 17:40:42 +0100 Subject: [PATCH 1/3] Check for overflows in left dataset concat in hash join --- .../physical-plan/src/joins/hash_join/exec.rs | 235 +++++++++++++++++- .../src/joins/hash_join/stream.rs | 33 ++- datafusion/physical-plan/src/joins/utils.rs | 89 +++++-- 3 files changed, 334 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index df10a8a5fcad5..be2a6347b716b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -69,8 +69,8 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{ArrayRef, BooleanBufferBuilder}; -use arrow::compute::concat_batches; +use arrow::array::{Array, ArrayRef, AsArray, BooleanBufferBuilder}; +use arrow::compute::{cast, concat_batches}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use arrow::util::bit_util; @@ -107,6 +107,111 @@ pub(crate) const HASH_JOIN_SEED: SeededRandomState = const ARRAY_MAP_CREATED_COUNT_METRIC_NAME: &str = "array_map_created_count"; +/// Largest payload addressable by the signed 32-bit offsets used by +/// [`DataType::Utf8`] and [`DataType::Binary`]. +const MAX_32BIT_OFFSET: usize = i32::MAX as usize; + +/// Returns the number of value bytes from `array` that `concat` would append. +/// +/// Sliced byte arrays can retain a larger backing buffer, so use their first +/// and last logical offsets rather than the backing buffer's allocation size. +fn byte_array_value_len(array: &dyn Array) -> Option { + fn offset_span(offsets: &[i32]) -> usize { + let first = offsets.first().copied().unwrap_or_default(); + let last = offsets.last().copied().unwrap_or(first); + (last - first) as usize + } + + match array.data_type() { + DataType::Utf8 => Some(offset_span(array.as_string::().value_offsets())), + DataType::Binary => Some(offset_span(array.as_binary::().value_offsets())), + _ => None, + } +} + +/// Chooses the physical schema used for the consolidated hash-join build batch. +/// +/// A regular `Utf8` / `Binary` array has one value buffer addressed by signed +/// 32-bit offsets. Concatenating build batches whose combined value buffers +/// exceed [`MAX_32BIT_OFFSET`] therefore fails even though every input batch is +/// valid. View arrays retain the input value buffers and concatenate only their +/// fixed-size views, preserving direct indexing without the single-buffer +/// limit. +fn build_storage_schema( + schema: &SchemaRef, + batches: &[&RecordBatch], + max_32bit_offset: usize, +) -> SchemaRef { + let fields: Vec<_> = schema + .fields() + .iter() + .enumerate() + .map(|(column_idx, field)| { + let target_type = match field.data_type() { + DataType::Utf8 | DataType::Binary => { + let exceeds_limit = batches + .iter() + .filter_map(|batch| { + byte_array_value_len(batch.column(column_idx)) + }) + .try_fold(0usize, |total, len| total.checked_add(len)) + .is_none_or(|total| total > max_32bit_offset); + + if exceeds_limit { + match field.data_type() { + DataType::Utf8 => DataType::Utf8View, + DataType::Binary => DataType::BinaryView, + _ => unreachable!(), + } + } else { + field.data_type().clone() + } + } + _ => field.data_type().clone(), + }; + Arc::new(field.as_ref().clone().with_data_type(target_type)) + }) + .collect(); + + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +/// Consolidates build batches, converting only overflowing 32-bit byte arrays +/// to their view equivalents. +fn concat_build_batches<'a>( + schema: &SchemaRef, + batches: impl IntoIterator, + max_32bit_offset: usize, +) -> Result { + let batches = batches.into_iter().collect::>(); + let storage_schema = build_storage_schema(schema, &batches, max_32bit_offset); + + if storage_schema == *schema { + return Ok(concat_batches(schema, batches)?); + } + + let converted = batches + .into_iter() + .map(|batch| { + let columns = batch + .columns() + .iter() + .zip(storage_schema.fields()) + .map(|(array, field)| { + if array.data_type() == field.data_type() { + Ok(Arc::clone(array)) + } else { + Ok(cast(array, field.data_type())?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(Arc::clone(&storage_schema), columns)?) + }) + .collect::>>()?; + + Ok(concat_batches(&storage_schema, &converted)?) +} + #[expect(clippy::too_many_arguments)] fn try_create_array_map( bounds: &Option, @@ -182,7 +287,7 @@ fn try_create_array_map( let mem_size = ArrayMap::estimate_memory_size(min_val, max_val, num_row); reservation.try_grow(mem_size)?; - let batch = concat_batches(schema, batches)?; + let batch = concat_build_batches(schema, batches, MAX_32BIT_OFFSET)?; let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; @@ -2733,7 +2838,8 @@ async fn collect_left_input( } // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_batches(&schema, batches_iter.clone())?; + let batch = + concat_build_batches(&schema, batches_iter.clone(), MAX_32BIT_OFFSET)?; let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; @@ -2858,8 +2964,8 @@ mod tests { }; use arrow::array::{ - Array, ArrayRef, Date32Array, DictionaryArray, Int32Array, Int64Array, - StructArray, UInt32Array, UInt64Array, + Array, ArrayRef, BinaryArray, Date32Array, DictionaryArray, Int32Array, + Int64Array, StringArray, StructArray, UInt32Array, UInt64Array, }; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field}; @@ -2882,6 +2988,87 @@ mod tests { use rstest::*; use rstest_reuse::*; + #[test] + fn concat_build_batches_uses_views_only_for_overflowing_columns() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("string", DataType::Utf8, false), + Field::new("binary", DataType::Binary, false), + ])); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["abc"])) as ArrayRef, + Arc::new(BinaryArray::from_iter_values([b"a".as_ref()])) as ArrayRef, + ], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["def"])) as ArrayRef, + Arc::new(BinaryArray::from_iter_values([b"b".as_ref()])) as ArrayRef, + ], + )?; + + let result = concat_build_batches(&schema, [&batch1, &batch2], 5)?; + + assert_eq!(result.column(0).data_type(), &DataType::Utf8View); + assert_eq!(result.column(1).data_type(), &DataType::Binary); + assert_eq!(result.num_rows(), 2); + Ok(()) + } + + #[test] + fn concat_build_batches_keeps_32_bit_offsets_at_the_limit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "string", + DataType::Utf8, + false, + )])); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["abc"]))], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["def"]))], + )?; + + let result = concat_build_batches(&schema, [&batch1, &batch2], 6)?; + + assert_eq!(result.column(0).data_type(), &DataType::Utf8); + Ok(()) + } + + #[test] + fn concat_build_batches_uses_binary_views_for_overflowing_binary() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "binary", + DataType::Binary, + false, + )])); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BinaryArray::from_iter_values([b"abc".as_ref()]))], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BinaryArray::from_iter_values([b"def".as_ref()]))], + )?; + + let result = concat_build_batches(&schema, [&batch1, &batch2], 5)?; + + assert_eq!(result.column(0).data_type(), &DataType::BinaryView); + Ok(()) + } + + #[test] + fn byte_array_value_len_uses_logical_slice_offsets() { + let array = StringArray::from(vec!["discarded backing value", "x"]); + let sliced = array.slice(1, 1); + + assert_eq!(byte_array_value_len(&sliced), Some(1)); + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, @@ -5537,6 +5724,42 @@ mod tests { Ok(()) } + #[test] + fn lookup_join_hashmap_compares_view_build_keys_with_utf8_probe_keys() -> Result<()> { + let build_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let probe_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let build_view_keys = cast(&build_keys, &DataType::Utf8View)?; + let random_state = RandomState::with_seed(0); + let mut build_hashes = vec![0; build_keys.len()]; + create_hashes([&build_keys], &random_state, &mut build_hashes)?; + + let mut table = HashTable::with_capacity(2); + table.insert_unique(build_hashes[0], (build_hashes[0], 1u32), |(hash, _)| *hash); + table.insert_unique(build_hashes[1], (build_hashes[1], 2u32), |(hash, _)| *hash); + let join_hash_map = JoinHashMapU32::new(table, vec![0, 0]); + + let mut probe_hashes = vec![0; probe_keys.len()]; + create_hashes([&probe_keys], &random_state, &mut probe_hashes)?; + let mut probe_indices_buffer = Vec::new(); + let mut build_indices_buffer = Vec::new(); + let (build_indices, probe_indices, _) = lookup_join_hashmap( + &join_hash_map, + &[build_view_keys], + &[probe_keys], + NullEquality::NullEqualsNothing, + &probe_hashes, + None, + 8192, + (0, None), + &mut probe_indices_buffer, + &mut build_indices_buffer, + )?; + + assert_eq!(build_indices, UInt64Array::from(vec![0, 1])); + assert_eq!(probe_indices, UInt32Array::from(vec![0, 1])); + Ok(()) + } + #[tokio::test] async fn join_with_duplicated_column_names() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 63a94e1987d91..c13da59a9fb34 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -49,6 +49,7 @@ use crate::{ use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; use arrow::buffer::NullBuffer; +use arrow::compute::cast; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ @@ -419,13 +420,43 @@ pub(super) fn lookup_join_hashmap( let probe_indices_unfiltered: UInt32Array = std::mem::take(probe_indices_buffer).into(); + // The consolidated build batch may use byte-view arrays internally to + // avoid overflowing 32-bit offsets. Hashing is representation independent, + // but collision checks require matching physical array types, so adapt the + // bounded probe-side key arrays to the build representation. + let probe_side_values = build_side_values + .iter() + .zip(probe_side_values) + .map( + |(build, probe)| match (build.data_type(), probe.data_type()) { + ( + arrow::datatypes::DataType::Utf8View, + arrow::datatypes::DataType::Utf8, + ) + | ( + arrow::datatypes::DataType::Utf8View, + arrow::datatypes::DataType::LargeUtf8, + ) + | ( + arrow::datatypes::DataType::BinaryView, + arrow::datatypes::DataType::Binary, + ) + | ( + arrow::datatypes::DataType::BinaryView, + arrow::datatypes::DataType::LargeBinary, + ) => Ok(cast(probe, build.data_type())?), + _ => Ok(Arc::clone(probe)), + }, + ) + .collect::>>()?; + // 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, + &probe_side_values, null_equality, )?; diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index fe19123cb20c5..65e85eff89f71 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1355,7 +1355,8 @@ pub(crate) fn build_batch_from_indices( // 2. based on the pick, `take` items from the different RecordBatches let mut columns: Vec> = Vec::with_capacity(schema.fields().len()); - for column_index in column_indices { + for (output_idx, column_index) in column_indices.iter().enumerate() { + let expected_type = schema.field(output_idx).data_type(); let array = if column_index.side == JoinSide::None { // For mark joins, the mark column is a true if the indices is not null, otherwise it will be false Arc::new(compute::is_not_null(probe_indices)?) @@ -1366,7 +1367,7 @@ pub(crate) fn build_batch_from_indices( // Therefore, it's possible we are empty but need to populate an n-length null array, // where n is the length of the index array. assert_eq!(build_indices.null_count(), build_indices.len()); - new_null_array(array.data_type(), build_indices.len()) + new_null_array(expected_type, build_indices.len()) } else { take(array.as_ref(), build_indices, None)? } @@ -1380,6 +1381,12 @@ pub(crate) fn build_batch_from_indices( } }; + let array = if array.data_type() == expected_type { + array + } else { + compute::cast(&array, expected_type)? + }; + columns.push(array); } Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) @@ -1391,7 +1398,7 @@ pub(crate) fn build_batch_from_indices( /// The resulting batch has [Schema] `schema`. pub(crate) fn build_batch_empty_build_side( schema: &Schema, - build_batch: &RecordBatch, + _build_batch: &RecordBatch, probe_batch: &RecordBatch, column_indices: &[ColumnIndex], join_type: JoinType, @@ -1409,20 +1416,30 @@ pub(crate) fn build_batch_empty_build_side( let columns = column_indices .iter() - .map(|column_index| match column_index.side { - // left -> null array - JoinSide::Left => new_null_array( - build_batch.column(column_index.index).data_type(), - num_rows, - ), - // right -> respective right array - JoinSide::Right => Arc::clone(probe_batch.column(column_index.index)), - // right mark -> unset boolean array as there are no matches on the left side - JoinSide::None => { - Arc::new(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None)) - } + .enumerate() + .map(|(output_idx, column_index)| -> Result { + Ok(match column_index.side { + // left -> null array + JoinSide::Left => { + new_null_array(schema.field(output_idx).data_type(), num_rows) + } + // right -> respective right array + JoinSide::Right => { + let array = probe_batch.column(column_index.index); + let expected_type = schema.field(output_idx).data_type(); + if array.data_type() == expected_type { + Arc::clone(array) + } else { + compute::cast(array, expected_type)? + } + } + // right mark -> unset boolean array as there are no matches on the left side + JoinSide::None => { + Arc::new(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None)) + } + }) }) - .collect(); + .collect::>>()?; Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) } @@ -2551,6 +2568,7 @@ mod tests { use super::*; + use arrow::array::AsArray; use arrow::datatypes::{DataType, Fields}; use arrow::error::{ArrowError, Result as ArrowResult}; use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; @@ -2563,6 +2581,45 @@ mod tests { assert_eq!(array.values().as_ref(), expected); } + #[test] + fn build_batch_from_indices_restores_declared_string_type() -> Result<()> { + let output_schema = + Schema::new(vec![Field::new("build_string", DataType::Utf8, false)]); + let internal_schema = Arc::new(Schema::new(vec![Field::new( + "build_string", + DataType::Utf8View, + false, + )])); + let build_batch = RecordBatch::try_new( + internal_schema, + vec![Arc::new(StringViewArray::from(vec!["a", "b"]))], + )?; + let probe_batch = RecordBatch::new_empty(Arc::new(Schema::empty())); + let build_indices = UInt64Array::from(vec![1, 0]); + let probe_indices = UInt32Array::from(vec![0, 0]); + let column_indices = vec![ColumnIndex { + index: 0, + side: JoinSide::Left, + }]; + + let result = build_batch_from_indices( + &output_schema, + &build_batch, + &probe_batch, + &build_indices, + &probe_indices, + &column_indices, + JoinSide::Left, + JoinType::Inner, + )?; + + assert_eq!(result.column(0).data_type(), &DataType::Utf8); + let values = result.column(0).as_string::(); + assert_eq!(values.value(0), "b"); + assert_eq!(values.value(1), "a"); + Ok(()) + } + #[test] fn get_anti_indices_returns_unmatched_range_indices() { let input = UInt32Array::from(vec![3, 5, 5]); From 548bf6b4cebafe6c5b7927c339b706404ee4b1f5 Mon Sep 17 00:00:00 2001 From: Shiv Bhatia Date: Tue, 25 Aug 2026 17:56:49 +0100 Subject: [PATCH 2/3] Fix test --- .../physical-plan/src/joins/hash_join/exec.rs | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index be2a6347b716b..18220fd462061 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -5726,17 +5726,29 @@ mod tests { #[test] fn lookup_join_hashmap_compares_view_build_keys_with_utf8_probe_keys() -> Result<()> { + let build_schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Utf8, false)])); let build_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); let probe_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let build_batch = RecordBatch::try_new( + Arc::clone(&build_schema), + vec![Arc::clone(&build_keys)], + )?; let build_view_keys = cast(&build_keys, &DataType::Utf8View)?; let random_state = RandomState::with_seed(0); let mut build_hashes = vec![0; build_keys.len()]; - create_hashes([&build_keys], &random_state, &mut build_hashes)?; - - let mut table = HashTable::with_capacity(2); - table.insert_unique(build_hashes[0], (build_hashes[0], 1u32), |(hash, _)| *hash); - table.insert_unique(build_hashes[1], (build_hashes[1], 2u32), |(hash, _)| *hash); - let join_hash_map = JoinHashMapU32::new(table, vec![0, 0]); + let mut join_hash_map = JoinHashMapU32::with_capacity(build_keys.len()); + update_hash( + &[Arc::new(Column::new("key", 0))], + &build_batch, + &mut join_hash_map, + 0, + &random_state, + &mut build_hashes, + 0, + true, + NullEquality::NullEqualsNothing, + )?; let mut probe_hashes = vec![0; probe_keys.len()]; create_hashes([&probe_keys], &random_state, &mut probe_hashes)?; From daa5229c6b930947bc4093d5dd31de331ec84e91 Mon Sep 17 00:00:00 2001 From: Shiv Bhatia Date: Wed, 26 Aug 2026 12:54:32 +0100 Subject: [PATCH 3/3] Always compute build key expressions while reading batches --- .../physical-plan/src/joins/hash_join/exec.rs | 296 ++++++++++++++---- .../src/joins/hash_join/stream.rs | 64 ++-- datafusion/physical-plan/src/joins/utils.rs | 39 ++- 3 files changed, 307 insertions(+), 92 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 18220fd462061..b22e7526ded6a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -42,7 +42,7 @@ use crate::joins::hash_join::stream::{ use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64}; use crate::joins::utils::{ OnceAsync, OnceFut, asymmetric_join_output_partitioning, reorder_output_after_swap, - swap_join_projection, update_hash, + swap_join_projection, update_hash_from_values, }; use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; use crate::metrics::{Count, MetricBuilder, MetricCategory}; @@ -69,8 +69,8 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{Array, ArrayRef, AsArray, BooleanBufferBuilder}; -use arrow::compute::{cast, concat_batches}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanBufferBuilder, new_empty_array}; +use arrow::compute::{cast, concat, concat_batches}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use arrow::util::bit_util; @@ -212,11 +212,81 @@ fn concat_build_batches<'a>( Ok(concat_batches(&storage_schema, &converted)?) } +/// Concatenates the already-evaluated build-side join keys in the same row +/// order as `build_batch`. +/// +/// Direct column keys can reuse the corresponding consolidated build column. +/// More complex expressions are never re-evaluated: their per-batch results +/// are concatenated directly, using byte views if their combined 32-bit value +/// buffers would overflow. +fn concat_build_key_values<'a>( + on_left: &[PhysicalExprRef], + key_schema: &SchemaRef, + key_batches: impl IntoIterator, + build_batch: &RecordBatch, + max_32bit_offset: usize, +) -> Result> { + let key_batches = key_batches.into_iter().collect::>(); + + on_left + .iter() + .enumerate() + .map(|(key_idx, expr)| { + if let Some(column) = expr.downcast_ref::() { + return Ok(Arc::clone(build_batch.column(column.index()))); + } + + let data_type = key_schema.field(key_idx).data_type(); + let arrays = key_batches + .iter() + .map(|batch| batch.column(key_idx).as_ref()) + .collect::>(); + + if arrays.is_empty() { + return Ok(new_empty_array(data_type)); + } + + let target_type = match data_type { + DataType::Utf8 | DataType::Binary => { + let exceeds_limit = arrays + .iter() + .filter_map(|array| byte_array_value_len(*array)) + .try_fold(0usize, |total, len| total.checked_add(len)) + .is_none_or(|total| total > max_32bit_offset); + + match (data_type, exceeds_limit) { + (DataType::Utf8, true) => DataType::Utf8View, + (DataType::Binary, true) => DataType::BinaryView, + _ => data_type.clone(), + } + } + _ => data_type.clone(), + }; + + if &target_type == data_type { + return Ok(concat(&arrays)?); + } + + let converted = arrays + .into_iter() + .map(|array| cast(array, &target_type)) + .collect::, _>>()?; + let converted = converted + .iter() + .map(|array| array.as_ref()) + .collect::>(); + Ok(concat(&converted)?) + }) + .collect() +} + #[expect(clippy::too_many_arguments)] fn try_create_array_map( bounds: &Option, schema: &SchemaRef, batches: &[RecordBatch], + key_schema: &SchemaRef, + key_batches: &[RecordBatch], on_left: &[PhysicalExprRef], reservation: &mut MemoryReservation, perfect_hash_join_small_build_threshold: usize, @@ -228,9 +298,8 @@ fn try_create_array_map( } if null_equality == NullEquality::NullEqualsNull { - for batch in batches.iter() { - let arrays = evaluate_expressions_to_arrays(on_left, batch)?; - if arrays[0].null_count() > 0 { + for key_batch in key_batches { + if key_batch.column(0).null_count() > 0 { return Ok(None); } } @@ -288,7 +357,13 @@ fn try_create_array_map( reservation.try_grow(mem_size)?; let batch = concat_build_batches(schema, batches, MAX_32BIT_OFFSET)?; - let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; + let left_values = concat_build_key_values( + on_left, + key_schema, + key_batches, + &batch, + MAX_32BIT_OFFSET, + )?; let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; @@ -2550,8 +2625,6 @@ fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { /// based on the actual data ranges can be pushed down to the probe side to /// eliminate unnecessary data early. struct CollectLeftAccumulator { - /// The physical expression to evaluate for each batch - expr: Arc, /// Accumulator for tracking the minimum value across all batches min: MinAccumulator, /// Accumulator for tracking the maximum value across all batches @@ -2583,26 +2656,24 @@ impl CollectLeftAccumulator { // Min/Max can operate on dictionary data but expect to be initialized with the underlying value type .map(|dt| dictionary_value_type(&dt))?; Ok(Self { - expr, min: MinAccumulator::try_new(&data_type)?, max: MaxAccumulator::try_new(&data_type)?, }) } - /// Updates the accumulators with values from a new batch. + /// Updates the accumulators with values from one evaluated key array. /// - /// Evaluates the expression on the batch and updates both min and max - /// accumulators with the resulting values. + /// Updates both min and max accumulators with an already-evaluated join + /// key array. /// /// # Arguments - /// * `batch` - The record batch to process + /// * `array` - The evaluated join key values to process /// /// # Returns - /// Ok(()) if the update succeeds, or an error if expression evaluation fails - fn update_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let array = self.expr.evaluate(batch)?.into_array(batch.num_rows())?; - self.min.update_batch(std::slice::from_ref(&array))?; - self.max.update_batch(std::slice::from_ref(&array))?; + /// Ok(()) if the update succeeds + fn update_array(&mut self, array: &ArrayRef) -> Result<()> { + self.min.update_batch(std::slice::from_ref(array))?; + self.max.update_batch(std::slice::from_ref(array))?; Ok(()) } @@ -2623,6 +2694,9 @@ impl CollectLeftAccumulator { /// State for collecting the build-side data during hash join struct BuildSideState { batches: Vec, + /// Join-key expressions evaluated once for each corresponding input batch. + key_batches: Vec, + key_schema: SchemaRef, num_rows: usize, metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, @@ -2642,8 +2716,23 @@ impl BuildSideState { schema: &SchemaRef, should_compute_dynamic_filters: bool, ) -> Result { + let key_schema = Arc::new(Schema::new( + on_left + .iter() + .enumerate() + .map(|(idx, expr)| { + Ok(arrow_schema::Field::new( + format!("join_key_{idx}"), + expr.data_type(schema)?, + true, + )) + }) + .collect::>>()?, + )); Ok(Self { batches: Vec::new(), + key_batches: Vec::new(), + key_schema, num_rows: 0, metrics, reservation, @@ -2730,33 +2819,52 @@ async fn collect_left_input( )?; let state = left_stream - .try_fold(initial, |mut state, batch| async move { - // Update accumulators if computing bounds - if let Some(ref mut accumulators) = state.bounds_accumulators { - for accumulator in accumulators { - accumulator.update_batch(&batch)?; + .try_fold(initial, |mut state, batch| { + let keys_values = evaluate_expressions_to_arrays(&on_left, &batch); + async move { + // Join-key expressions may not accept a different physical string + // representation, so evaluate them once against the original batch + // and retain their results for hashing and collision checks. + let keys_values = keys_values?; + + // Update accumulators if computing bounds + if let Some(ref mut accumulators) = state.bounds_accumulators { + for (accumulator, array) in accumulators.iter_mut().zip(&keys_values) + { + accumulator.update_array(array)?; + } } - } - // Decide if we spill or not - let batch_size = state.memory_counter.count_batch(&batch); - // Reserve memory for incoming batch - state.reservation.try_grow(batch_size)?; - // Update metrics - state.metrics.build_mem_used.add(batch_size); - state.metrics.build_input_batches.add(1); - state.metrics.build_input_rows.add(batch.num_rows()); - // Update row count - state.num_rows += batch.num_rows(); - // Push batch to output - state.batches.push(batch); - Ok(state) + let key_batch = + RecordBatch::try_new(Arc::clone(&state.key_schema), keys_values)?; + + // Decide if we spill or not + let batch_size = state.memory_counter.count_batch(&batch); + let key_batch_size = state.memory_counter.count_batch(&key_batch); + let Some(batch_size) = batch_size.checked_add(key_batch_size) else { + return internal_err!("Build-side batch memory size overflow"); + }; + // Reserve memory for incoming batch + state.reservation.try_grow(batch_size)?; + // Update metrics + state.metrics.build_mem_used.add(batch_size); + state.metrics.build_input_batches.add(1); + state.metrics.build_input_rows.add(batch.num_rows()); + // Update row count + state.num_rows += batch.num_rows(); + // Push batch to output + state.batches.push(batch); + state.key_batches.push(key_batch); + Ok(state) + } }) .await?; // Extract fields from state let BuildSideState { batches, + key_batches, + key_schema, num_rows, metrics, mut reservation, @@ -2781,6 +2889,8 @@ async fn collect_left_input( &bounds, &schema, &batches, + &key_schema, + &key_batches, &on_left, &mut reservation, config.execution.perfect_hash_join_small_build_threshold, @@ -2818,14 +2928,15 @@ async fn collect_left_input( let mut offset = 0; let batches_iter = batches.iter().rev(); + let key_batches_iter = key_batches.iter().rev(); // Updating hashmap starting from the last batch - for batch in batches_iter.clone() { + for key_batch in key_batches_iter.clone() { hashes_buffer.clear(); - hashes_buffer.resize(batch.num_rows(), 0); - update_hash( - &on_left, - batch, + hashes_buffer.resize(key_batch.num_rows(), 0); + update_hash_from_values( + key_batch.columns(), + key_batch.num_rows(), &mut *hashmap, offset, &random_state, @@ -2834,14 +2945,20 @@ async fn collect_left_input( true, null_equality, )?; - offset += batch.num_rows(); + offset += key_batch.num_rows(); } // Merge all batches into a single batch, so we can directly index into the arrays let batch = concat_build_batches(&schema, batches_iter.clone(), MAX_32BIT_OFFSET)?; - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + let left_values = concat_build_key_values( + &on_left, + &key_schema, + key_batches_iter, + &batch, + MAX_32BIT_OFFSET, + )?; (Map::HashMap(hashmap), batch, left_values) }; @@ -2955,7 +3072,7 @@ mod tests { use crate::coalesce_partitions::CoalescePartitionsExec; use crate::execution_plan::Boundedness; use crate::filter::FilterExecBuilder; - use crate::joins::hash_join::stream::lookup_join_hashmap; + use crate::joins::hash_join::stream::{adapt_probe_side_values, lookup_join_hashmap}; use crate::test::{TestMemoryExec, assert_join_metrics}; use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ @@ -3061,6 +3178,72 @@ mod tests { Ok(()) } + #[test] + fn concat_build_key_values_reuses_direct_build_column() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "string", + DataType::Utf8, + false, + )])); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["abc"]))], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(vec!["def"]))], + )?; + let build_batch = concat_build_batches(&schema, [&batch1, &batch2], 5)?; + let on_left: Vec = vec![Arc::new(Column::new("string", 0))]; + + let values = concat_build_key_values( + &on_left, + &schema, + [&batch1, &batch2], + &build_batch, + 5, + )?; + + assert!(Arc::ptr_eq(&values[0], build_batch.column(0))); + assert_eq!(values[0].data_type(), &DataType::Utf8View); + Ok(()) + } + + #[test] + fn concat_build_key_values_uses_retained_expression_results() -> Result<()> { + let key_schema = Arc::new(Schema::new(vec![Field::new( + "computed_key", + DataType::Utf8, + true, + )])); + let key_batch1 = RecordBatch::try_new( + Arc::clone(&key_schema), + vec![Arc::new(StringArray::from(vec!["abc"]))], + )?; + let key_batch2 = RecordBatch::try_new( + Arc::clone(&key_schema), + vec![Arc::new(StringArray::from(vec!["def"]))], + )?; + let build_batch = RecordBatch::new_empty(Arc::new(Schema::empty())); + let on_left: Vec = vec![Arc::new(Literal::new( + ScalarValue::Utf8(Some("not re-evaluated".to_string())), + ))]; + + let values = concat_build_key_values( + &on_left, + &key_schema, + [&key_batch1, &key_batch2], + &build_batch, + 5, + )?; + + assert_eq!(values[0].data_type(), &DataType::Utf8View); + let values = values[0].as_string_view(); + assert_eq!(values.value(0), "abc"); + assert_eq!(values.value(1), "def"); + Ok(()) + } + #[test] fn byte_array_value_len_uses_logical_slice_offsets() { let array = StringArray::from(vec!["discarded backing value", "x"]); @@ -5725,22 +5908,16 @@ mod tests { } #[test] - fn lookup_join_hashmap_compares_view_build_keys_with_utf8_probe_keys() -> Result<()> { - let build_schema = - Arc::new(Schema::new(vec![Field::new("key", DataType::Utf8, false)])); + fn probe_values_are_adapted_once_for_view_build_keys() -> Result<()> { let build_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); let probe_keys: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); - let build_batch = RecordBatch::try_new( - Arc::clone(&build_schema), - vec![Arc::clone(&build_keys)], - )?; let build_view_keys = cast(&build_keys, &DataType::Utf8View)?; let random_state = RandomState::with_seed(0); let mut build_hashes = vec![0; build_keys.len()]; let mut join_hash_map = JoinHashMapU32::with_capacity(build_keys.len()); - update_hash( - &[Arc::new(Column::new("key", 0))], - &build_batch, + update_hash_from_values( + &[build_keys], + 2, &mut join_hash_map, 0, &random_state, @@ -5752,12 +5929,17 @@ mod tests { let mut probe_hashes = vec![0; probe_keys.len()]; create_hashes([&probe_keys], &random_state, &mut probe_hashes)?; + let probe_view_keys = adapt_probe_side_values( + std::slice::from_ref(&build_view_keys), + vec![probe_keys], + )?; + assert_eq!(probe_view_keys[0].data_type(), &DataType::Utf8View); let mut probe_indices_buffer = Vec::new(); let mut build_indices_buffer = Vec::new(); let (build_indices, probe_indices, _) = lookup_join_hashmap( &join_hash_map, &[build_view_keys], - &[probe_keys], + &probe_view_keys, NullEquality::NullEqualsNothing, &probe_hashes, None, diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index c13da59a9fb34..3fa20099239d0 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -393,6 +393,28 @@ impl RecordBatchStream for HashJoinStream { /// Build indices: 4, 5, 6, 6 /// Probe indices: 3, 3, 4, 5 /// ``` +pub(super) fn adapt_probe_side_values( + build_side_values: &[ArrayRef], + mut probe_side_values: Vec, +) -> Result> { + for (build, probe) in build_side_values.iter().zip(&mut probe_side_values) { + if matches!( + (build.data_type(), probe.data_type()), + ( + arrow::datatypes::DataType::Utf8View, + arrow::datatypes::DataType::Utf8 | arrow::datatypes::DataType::LargeUtf8, + ) | ( + arrow::datatypes::DataType::BinaryView, + arrow::datatypes::DataType::Binary + | arrow::datatypes::DataType::LargeBinary, + ) + ) { + *probe = cast(probe, build.data_type())?; + } + } + Ok(probe_side_values) +} + #[expect(clippy::too_many_arguments)] pub(super) fn lookup_join_hashmap( build_hashmap: &dyn JoinHashMapType, @@ -420,43 +442,13 @@ pub(super) fn lookup_join_hashmap( let probe_indices_unfiltered: UInt32Array = std::mem::take(probe_indices_buffer).into(); - // The consolidated build batch may use byte-view arrays internally to - // avoid overflowing 32-bit offsets. Hashing is representation independent, - // but collision checks require matching physical array types, so adapt the - // bounded probe-side key arrays to the build representation. - let probe_side_values = build_side_values - .iter() - .zip(probe_side_values) - .map( - |(build, probe)| match (build.data_type(), probe.data_type()) { - ( - arrow::datatypes::DataType::Utf8View, - arrow::datatypes::DataType::Utf8, - ) - | ( - arrow::datatypes::DataType::Utf8View, - arrow::datatypes::DataType::LargeUtf8, - ) - | ( - arrow::datatypes::DataType::BinaryView, - arrow::datatypes::DataType::Binary, - ) - | ( - arrow::datatypes::DataType::BinaryView, - arrow::datatypes::DataType::LargeBinary, - ) => Ok(cast(probe, build.data_type())?), - _ => Ok(Arc::clone(probe)), - }, - ) - .collect::>>()?; - // 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, + probe_side_values, null_equality, )?; @@ -747,13 +739,21 @@ impl HashJoinStream { None }; + // Collision checks require matching physical key types. Adapt + // once per probe batch and reuse the result if high fanout + // requires multiple hash-map lookup calls for this batch. + let values = adapt_probe_side_values( + self.build_side.try_as_ready()?.left_data.values(), + keys_values, + )?; + self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); self.state = HashJoinStreamState::ProcessProbeBatch(ProcessProbeBatchState { batch, - values: keys_values, + values, valid_keys, offset: (0, None), joined_probe_idx: None, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 65e85eff89f71..3648bd15cbb00 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -2155,14 +2155,47 @@ pub fn update_hash( // evaluate the keys let keys_values = evaluate_expressions_to_arrays(on, batch)?; + update_hash_from_values( + &keys_values, + batch.num_rows(), + hash_map, + offset, + random_state, + hashes_buffer, + deleted_offset, + fifo_hashmap, + null_equality, + ) +} + +/// Updates `hash_map` from already-evaluated join key arrays. +/// +/// This is equivalent to [`update_hash`] without evaluating physical +/// expressions, allowing callers that retain key arrays to avoid evaluating +/// arbitrary expressions more than once. +#[expect(clippy::too_many_arguments)] +pub(crate) fn update_hash_from_values( + keys_values: &[ArrayRef], + num_rows: usize, + hash_map: &mut dyn JoinHashMapType, + offset: usize, + random_state: &RandomState, + hashes_buffer: &mut [u64], + deleted_offset: usize, + fifo_hashmap: bool, + null_equality: NullEquality, +) -> Result<()> { + assert_eq!(hashes_buffer.len(), num_rows); + assert!(keys_values.iter().all(|array| array.len() == num_rows)); + // calculate the hash values - let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?; + let hash_values = create_hashes(keys_values, random_state, hashes_buffer)?; // For usual JoinHashmap, the implementation is void. - hash_map.extend_zero(batch.num_rows()); + hash_map.extend_zero(num_rows); // Unmatchable NULL-key rows are filtered out below. - let valid_keys = matchable_join_keys(&keys_values, null_equality); + let valid_keys = matchable_join_keys(keys_values, null_equality); // Updating JoinHashMap from hash values iterator let hash_values_iter = hash_values