From 6f07b359229c644a9ba355ec419b2725f94a41fa Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 00:12:03 +0000 Subject: [PATCH 1/6] perf: retain batched build storage for large hash joins Keep small builds compact, retain large builds with batch-local keys and bounded coalescing, and gather only referenced sources. Cover computed and composite keys, encoded and nested payloads, perfect hashing, memory admission, and bounded final output. Add a reproducible build-layout benchmark suite. --- datafusion/physical-plan/Cargo.toml | 5 + .../benches/hash_join_batches.rs | 432 +++++ .../physical-plan/src/joins/array_map.rs | 172 ++ .../physical-plan/src/joins/hash_join/exec.rs | 149 +- .../src/joins/hash_join/exec/build_data.rs | 573 +++++++ .../physical-plan/src/joins/hash_join/mod.rs | 2 + .../src/joins/hash_join/multi_batch_tests.rs | 1419 +++++++++++++++++ .../src/joins/hash_join/stream.rs | 177 +- datafusion/physical-plan/src/joins/utils.rs | 7 + .../src/joins/utils/multi_batch.rs | 633 ++++++++ 10 files changed, 3492 insertions(+), 77 deletions(-) create mode 100644 datafusion/physical-plan/benches/hash_join_batches.rs create mode 100644 datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs create mode 100644 datafusion/physical-plan/src/joins/hash_join/multi_batch_tests.rs create mode 100644 datafusion/physical-plan/src/joins/utils/multi_batch.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 17e98a4764c23..2cbc268041329 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -182,6 +182,11 @@ harness = false name = "prepared_hash_join" required-features = ["test_utils"] +[[bench]] +harness = false +name = "hash_join_batches" +required-features = ["test_utils"] + [[bench]] harness = false name = "multi_group_by" diff --git a/datafusion/physical-plan/benches/hash_join_batches.rs b/datafusion/physical-plan/benches/hash_join_batches.rs new file mode 100644 index 0000000000000..bf4046d6aa135 --- /dev/null +++ b/datafusion/physical-plan/benches/hash_join_batches.rs @@ -0,0 +1,432 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Hash-join build storage across controlled input batch boundaries. +//! +//! SQL can prune payloads or change the input batch layout, so this benchmark +//! uses the public physical execution API. Input generation is excluded; fresh +//! plan construction, build, probe and output draining are included. Large cases +//! retain over 64 MiB of unique Arrow backing. Cases include tiny independent +//! batches, shared slices and oversized backing buffers. The small case measures +//! the compact-build path. Every case probes 4096 rows and checks the matched +//! build-row checksum before timing. Inner joins have 50% matches; the outer case +//! has 6.25%. Perfect-hash selection is enabled and disabled over identical inputs. +//! The separate correctness run reports peak reserved bytes, not process RSS; +//! timed runs use the normal memory pool without reservation instrumentation. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, ListArray, RecordBatch, StringArray, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; +use datafusion_common::{JoinType, Result}; +use datafusion_execution::memory_pool::{ + MemoryPool, PeakRecordingPool, UnboundedMemoryPool, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::{TaskContext, config::SessionConfig}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, col, lit}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::joins::{HashJoinExec, HashJoinExecBuilder, PartitionMode}; +use datafusion_physical_plan::test::TestMemoryExec; +use futures::TryStreamExt; +use tokio::runtime::Builder; + +const PROBE_ROWS: usize = 4096; +const INPUT_BATCH_ROWS: usize = 4096; +const COMPACT_BUILD_BYTES: usize = 64 * 1024 * 1024; + +#[derive(Clone, Copy)] +enum Keys { + Integer, + Computed, + Dictionary, +} + +#[derive(Clone, Copy)] +enum Payload { + Plain, + Dictionary, + List, +} + +#[derive(Clone, Copy)] +enum Layout { + Independent, + Tiny, + Sliced, + Overallocated, +} + +struct Workload { + build: Vec, + probe: RecordBatch, + keys: Keys, + join_type: JoinType, + expected_rows: usize, + expected_sum: i64, +} + +fn key_array(keys: &[i32], kind: Keys) -> ArrayRef { + if matches!(kind, Keys::Dictionary) { + Arc::new( + DictionaryArray::::try_new( + Int32Array::from_iter_values(0..keys.len() as i32), + Arc::new(StringArray::from_iter_values( + keys.iter().map(|key| format!("key-{key:08}")), + )), + ) + .unwrap(), + ) + } else { + Arc::new(Int32Array::from_iter_values(keys.iter().copied())) + } +} + +impl Workload { + fn new( + rows: usize, + width: usize, + keys: Keys, + payload: Payload, + layout: Layout, + low_match_outer: bool, + ) -> Self { + let value = "x".repeat(width); + let batch = |start, end| { + let ids = (start..end).map(|id| id as i32).collect::>(); + let key = key_array(&ids, keys); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + ids.iter().map(|_| value.as_str()), + )); + let values: ArrayRef = match payload { + Payload::Plain => values, + Payload::Dictionary => Arc::new( + DictionaryArray::::try_new( + Int32Array::from_iter_values(0..ids.len() as i32), + values, + ) + .unwrap(), + ), + Payload::List => Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, false)), + OffsetBuffer::new(ScalarBuffer::from( + (0..=ids.len() as i32).collect::>(), + )), + values, + None, + )), + }; + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("key", key.data_type().clone(), false), + Field::new("payload", values.data_type().clone(), false), + Field::new("id", DataType::Int32, false), + ])), + vec![key, values, Arc::new(Int32Array::from(ids))], + ) + .unwrap() + }; + let build = if matches!(layout, Layout::Sliced) { + let backing = batch(0, rows); + (0..rows) + .step_by(INPUT_BATCH_ROWS) + .map(|start| backing.slice(start, (rows - start).min(INPUT_BATCH_ROWS))) + .collect::>() + } else { + let batch_rows = if matches!(layout, Layout::Tiny) { + 64 + } else { + INPUT_BATCH_ROWS + }; + (0..rows) + .step_by(batch_rows) + .map(|start| { + let count = (rows - start).min(batch_rows); + let allocated = if matches!(layout, Layout::Overallocated) { + count * 2 + } else { + count + }; + batch(start, start + allocated).slice(0, count) + }) + .collect::>() + }; + let hit_every = if low_match_outer { 16 } else { 2 }; + let probe_ids = (0..PROBE_ROWS) + .map(|index| { + let candidate = (index * rows / PROBE_ROWS) % rows; + (candidate + usize::from(index % hit_every != 0) * rows) as i32 + }) + .collect::>(); + let probe_keys = key_array(&probe_ids, keys); + let probe = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "key", + probe_keys.data_type().clone(), + false, + )])), + vec![probe_keys], + ) + .unwrap(); + Self { + build, + probe, + keys, + join_type: if low_match_outer { + JoinType::Right + } else { + JoinType::Inner + }, + expected_rows: if low_match_outer { + PROBE_ROWS + } else { + PROBE_ROWS / hit_every + }, + expected_sum: probe_ids + .iter() + .step_by(hit_every) + .map(|&id| i64::from(id)) + .sum(), + } + } + + fn join(&self) -> Result { + let left_schema = self.build[0].schema(); + let right_schema = self.probe.schema(); + let mut left_key = col("key", &left_schema)?; + let mut right_key = col("key", &right_schema)?; + if matches!(self.keys, Keys::Computed) { + left_key = Arc::new(BinaryExpr::new(left_key, Operator::Plus, lit(7i32))); + right_key = Arc::new(BinaryExpr::new(right_key, Operator::Plus, lit(7i32))); + } + HashJoinExecBuilder::new( + TestMemoryExec::try_new_exec( + std::slice::from_ref(&self.build), + left_schema, + None, + )?, + TestMemoryExec::try_new_exec( + &[vec![self.probe.clone()]], + right_schema, + None, + )?, + vec![(left_key, right_key)], + self.join_type, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .with_projection(Some(vec![2, 1])) + .build() + } + + async fn run(&self, context: Arc) -> Result<(usize, i64)> { + let mut stream = self.join()?.execute(0, context)?; + let mut rows = 0; + let mut sum = 0; + while let Some(batch) = stream.try_next().await? { + rows += batch.num_rows(); + sum += batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(i64::from) + .sum::(); + black_box(batch.column(1)); + } + Ok((rows, sum)) + } +} + +fn benchmark(c: &mut Criterion) { + let runtime = Builder::new_current_thread().enable_all().build().unwrap(); + let mut group = c.benchmark_group("hash_join_batches"); + for (name, rows, width, keys, payload, layout, outer) in [ + ( + "small_plain", + 8192, + 16, + Keys::Integer, + Payload::Plain, + Layout::Independent, + false, + ), + ( + "large_plain", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Independent, + false, + ), + ( + "large_computed", + 65536, + 1024, + Keys::Computed, + Payload::Plain, + Layout::Independent, + false, + ), + ( + "large_dictionary_keys", + 65536, + 1024, + Keys::Dictionary, + Payload::Plain, + Layout::Independent, + false, + ), + ( + "large_dictionary_payload", + 65536, + 1024, + Keys::Integer, + Payload::Dictionary, + Layout::Independent, + false, + ), + ( + "large_list_payload", + 65536, + 1024, + Keys::Integer, + Payload::List, + Layout::Independent, + false, + ), + ( + "large_tiny_batches", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Tiny, + false, + ), + ( + "large_sliced", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Sliced, + false, + ), + ( + "large_overallocated", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Overallocated, + false, + ), + ( + "large_low_match_outer", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Independent, + true, + ), + ] { + let workload = Workload::new(rows, width, keys, payload, layout, outer); + let mut counter = RecordBatchMemoryCounter::new(); + for batch in &workload.build { + counter.count_batch(batch); + } + assert_eq!( + counter.memory_usage() > COMPACT_BUILD_BYTES, + name.starts_with("large") + ); + for perfect_hash in [false, true] { + let mut config = SessionConfig::default(); + config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = false; + config + .options_mut() + .execution + .perfect_hash_join_small_build_threshold = + if perfect_hash { usize::MAX } else { 0 }; + config + .options_mut() + .execution + .perfect_hash_join_min_key_density = + if perfect_hash { 0.0 } else { f64::INFINITY }; + let recording = Arc::new(PeakRecordingPool::new(Arc::new( + UnboundedMemoryPool::default(), + ))); + let recording_context = Arc::new( + TaskContext::default() + .with_session_config(config.clone()) + .with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool( + Arc::clone(&recording) as Arc + ) + .build_arc() + .unwrap(), + ), + ); + assert_eq!( + runtime + .block_on(workload.run(Arc::clone(&recording_context))) + .unwrap(), + (workload.expected_rows, workload.expected_sum), + ); + drop(recording_context); + assert_eq!(recording.reserved(), 0); + let mode = if perfect_hash { + "perfect_hash_on" + } else { + "perfect_hash_off" + }; + eprintln!( + "HASH_JOIN_RESERVATION name={name} mode={mode} input_unique_bytes={} peak_reserved_bytes={}", + counter.memory_usage(), + recording.peak_reserved(), + ); + let context = Arc::new(TaskContext::default().with_session_config(config)); + group.bench_function(BenchmarkId::new(name, mode), |b| { + b.iter(|| { + black_box( + runtime + .block_on(workload.run(Arc::clone(&context))) + .unwrap(), + ) + }) + }); + } + } + group.finish(); +} + +criterion_group!(benches, benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index 1a1cd8d32970b..24ab22bd01be8 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -202,6 +202,62 @@ impl ArrayMap { }) } + /// Builds a map over several key arrays, using their concatenated row order + /// without concatenating the arrays themselves. + pub(crate) fn try_new_batched( + arrays: &[ArrayRef], + min_val: u64, + max_val: u64, + ) -> Result { + let Some(first) = arrays.first() else { + return internal_err!("ArrayMap requires at least one build-side key array"); + }; + if arrays.len() == 1 { + return Self::try_new(first, min_val, max_val); + } + let Some(num_rows) = arrays + .iter() + .try_fold(0usize, |rows, array| rows.checked_add(array.len())) + .filter(|&rows| u32::try_from(rows).is_ok()) + else { + return internal_err!("ArrayMap build-side row count exceeds u32::MAX"); + }; + if arrays + .iter() + .any(|array| array.data_type() != first.data_type()) + { + return internal_err!( + "ArrayMap build-side key arrays must have the same type" + ); + } + let range = Self::calculate_range(min_val, max_val); + if range >= usize::MAX as u64 { + return internal_err!("ArrayMap key range is too large to be allocated."); + } + let mut data = vec![0; (range + 1) as usize]; + let mut next = vec![]; + let mut num_of_distinct_key = 0; + + downcast_supported_integer!( + first.data_type() => ( + fill_data_batched, + arrays, + num_rows, + min_val, + &mut data, + &mut next, + &mut num_of_distinct_key + ) + )?; + + Ok(Self { + data, + offset: min_val, + next, + num_of_distinct_key, + }) + } + fn fill_data( array: &ArrayRef, offset_val: u64, @@ -235,6 +291,46 @@ impl ArrayMap { Ok(()) } + fn fill_data_batched( + arrays: &[ArrayRef], + num_rows: usize, + offset_val: u64, + data: &mut [u32], + next: &mut Vec, + num_of_distinct_key: &mut usize, + ) -> Result<()> + where + T::Native: AsPrimitive, + { + // Reverse both dimensions so duplicate chains retain the original + // global row order, including duplicates in different batches. + let mut batch_offset = num_rows; + for array in arrays.iter().rev() { + let array = array.as_primitive::(); + batch_offset -= array.len(); + for (row, value) in array.iter().enumerate().rev() { + if let Some(value) = value { + let Some(index) = + Self::key_to_index(value.as_(), offset_val, data.len()) + else { + return internal_err!("failed build Array idx >= data.len()"); + }; + let global_row = batch_offset + row; + if data[index] != 0 { + if next.is_empty() { + *next = vec![0; num_rows]; + } + next[global_row] = data[index]; + } else { + *num_of_distinct_key += 1; + } + data[index] = global_row as u32 + 1; + } + } + } + Ok(()) + } + pub fn num_of_distinct_key(&self) -> usize { self.num_of_distinct_key } @@ -425,6 +521,82 @@ mod tests { use arrow::array::UInt64Array; use std::sync::Arc; + #[test] + fn test_array_map_batched_duplicate_order_and_nulls() -> Result<()> { + let first: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(99), + Some(-2), + None, + Some(2), + Some(-2), + ])); + let arrays = vec![ + first.slice(1, 4), + Arc::new(Int64Array::from(Vec::::new())) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(2), None, Some(-2), Some(0)])), + Arc::new(Int64Array::from(vec![None, None])), + ]; + let map = ArrayMap::try_new_batched(&arrays, -2_i64 as u64, 2)?; + assert_eq!(map.num_of_distinct_key(), 3); + let probe = [Arc::new(Int64Array::from(vec![ + None, + Some(-3), + Some(-2), + Some(0), + Some(2), + Some(3), + ])) as ArrayRef]; + + for limit in [1, 2, 4, 16] { + let mut next = Some((0, None)); + let mut probe_indices = vec![]; + let mut build_indices = vec![]; + let mut actual = vec![]; + while let Some(offset) = next { + next = map.get_matched_indices_with_limit_offset( + &probe, + limit, + offset, + &mut probe_indices, + &mut build_indices, + )?; + assert!(build_indices.len() <= limit); + actual.extend( + probe_indices + .iter() + .copied() + .zip(build_indices.iter().copied()), + ); + } + assert_eq!(actual, vec![(2, 0), (2, 3), (2, 6), (3, 7), (4, 2), (4, 4)]); + } + Ok(()) + } + + #[test] + fn test_array_map_batched_unique_and_all_null_keys() -> Result<()> { + let arrays = vec![ + Arc::new(Int32Array::from(vec![Some(-1), None])) as ArrayRef, + Arc::new(Int32Array::from(vec![None, Some(1)])), + ]; + let map = ArrayMap::try_new_batched(&arrays, -1_i64 as u64, 1)?; + assert!(map.next.is_empty()); + assert_eq!(map.num_of_distinct_key(), 2); + assert_eq!(map.get_value(-1_i64 as u64), Some(1)); + assert_eq!(map.get_value(1), Some(4)); + assert_eq!(map.get_value(0), None); + + let nulls = vec![ + Arc::new(Int32Array::from(vec![None, None])) as ArrayRef, + Arc::new(Int32Array::from(vec![None])), + ]; + let map = ArrayMap::try_new_batched(&nulls, 0, 0)?; + assert!(map.next.is_empty()); + assert_eq!(map.num_of_distinct_key(), 0); + assert_eq!(map.get_value(0), None); + Ok(()) + } + #[test] fn test_array_map_limit_offset_duplicate_elements() -> Result<()> { let build: ArrayRef = Arc::new(Int32Array::from(vec![1, 1, 2])); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1bffdb1ab2d7c..1b7e383b8582e 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -106,6 +106,9 @@ use super::partitioned_hash_eval::SeededRandomState; mod prepared; pub use prepared::PreparedHashJoinBuild; +mod build_data; +pub(super) use build_data::MultiBatchBuildData; +use build_data::{coalesce_build_batches, should_preserve_batches}; /// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions. pub(crate) const HASH_JOIN_SEED: SeededRandomState = @@ -294,12 +297,14 @@ impl NullAwareMode { /// Immutable build buffers and their durable reservation. struct JoinBuildData { - /// Hash table with row indices into `batch`, also shared with dynamic filters. + /// Logical build-row indices, also shared with dynamic filters. map: Arc, - /// The input rows for the build side. + /// The compact build, or the first retained batch for schema-only access. batch: RecordBatch, - /// Evaluated build-side key expressions. + /// Compact keys, or the first retained batch's keys (see `multi_batch`). values: Vec, + /// Large builds retain bounded batches instead of concatenating the payload. + multi_batch: Option, /// Bounds computed from the build side; absent for an empty partition. bounds: Option, /// IN-list values or a hash-table reference used for filter pushdown. @@ -347,7 +352,8 @@ impl JoinLeftData { self.null_value_build_rows.as_ref() } - /// returns a reference to the build side batch + /// Returns the compact build or a representative batch for schema-only access. + /// Row access must use `multi_batch` when present. pub(super) fn batch(&self) -> &RecordBatch { &self.build.batch } @@ -357,7 +363,16 @@ impl JoinLeftData { /// This is distinct from [`Self::has_matchable_build_rows`]: a build side /// can hold rows while its hash map is empty (see that method). pub(super) fn has_build_rows(&self) -> bool { - self.batch().num_rows() > 0 + self.num_rows() > 0 + } + + pub(super) fn num_rows(&self) -> usize { + self.multi_batch() + .map_or_else(|| self.batch().num_rows(), MultiBatchBuildData::num_rows) + } + + pub(super) fn multi_batch(&self) -> Option<&MultiBatchBuildData> { + self.build.multi_batch.as_ref() } /// Returns `true` if the build-side hash map has any matchable entries. @@ -369,7 +384,8 @@ impl JoinLeftData { !self.map().is_empty() } - /// returns a reference to the build side expressions values + /// Compact build keys, or the representative batch's keys. Multi-batch + /// comparison must use the batch-local keys in `multi_batch`. pub(super) fn values(&self) -> &[ArrayRef] { &self.build.values } @@ -3007,11 +3023,12 @@ fn concat_build_batches( /// bitmaps and correlation-scope maps are required /// /// # Memory Accounting -/// Build batches are added to `reservation` as they arrive. They are then copied -/// into a single batch, see [`concat_build_batches`]: the copy is reserved -/// before it is made, and the reservation is trimmed to what the single batch -/// retains once the input batches are dropped. Join key arrays that do not -/// share the buffers of that batch are reserved as well. +/// Build batches are added to `reservation` as they arrive, deduplicating shared +/// backing buffers. Small builds are compacted with [`concat_build_batches`]; +/// large ordinary joins retain batches and coalesce independent flat inputs in +/// bounded groups. Copies are reserved before allocation and their input charge +/// is released only after the inputs are dropped. Computed key buffers and the +/// multi-batch row directory are reserved for the lifetime of the build. /// /// # Dynamic Filter Coordination /// When `should_compute_dynamic_filters` is true, this function computes the min/max bounds @@ -3021,7 +3038,7 @@ fn concat_build_batches( /// before updating the filter exactly once. /// /// # Returns -/// `JoinLeftData` containing the hash map, consolidated batch, join key values, +/// `JoinLeftData` containing the hash map, retained batches, join key values, /// visited indices bitmap, and computed bounds (if requested). #[expect(clippy::too_many_arguments)] async fn collect_left_input( @@ -3093,6 +3110,25 @@ async fn collect_left_input( memory_counter: _, } = state; + // Prepared builds and null-aware joins retain their existing contiguous + // key/state contract. Ordinary joins have no key or payload type gate. + let preserve_batches = !prepared + && null_aware.is_none() + && num_rows > 0 + && batches.len() > 1 + && should_preserve_batches(&batches, input_bytes); + if preserve_batches { + batches = coalesce_build_batches( + &schema, + batches, + input_bytes, + &mut reservation, + &metrics, + )?; + max_batch_rows = batches.iter().map(RecordBatch::num_rows).max().unwrap_or(0); + } + let mut multi_batch = None; + // Admit concatenation copies while the original batches are retained. // Arrow keeps a single batch as an inexpensive slice. let copy_bytes = if prepared && batches.len() > 1 { @@ -3126,7 +3162,15 @@ async fn collect_left_input( config.execution.perfect_hash_join_min_key_density, null_equality, )? { - let batch = if prepared { + let batch = if preserve_batches { + multi_batch = Some(MultiBatchBuildData::try_new( + std::mem::take(&mut batches), + &on_left, + &reservation, + &metrics, + )?); + multi_batch.as_ref().unwrap().batches()[0].clone() + } else if prepared { concat_batches(&schema, batches.iter())? } else { concat_build_batches( @@ -3138,8 +3182,21 @@ async fn collect_left_input( &metrics, )? }; - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; - let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; + let left_values = if let Some(data) = &multi_batch { + data.values()[0].clone() + } else { + evaluate_expressions_to_arrays(&on_left, &batch)? + }; + let array_map = if let Some(data) = &multi_batch { + let keys = data + .values() + .iter() + .map(|values| Arc::clone(&values[0])) + .collect::>(); + ArrayMap::try_new_batched(&keys, min_val, max_val)? + } else { + ArrayMap::try_new(&left_values[0], min_val, max_val)? + }; array_map_created_count.add(1); metrics.build_mem_used.add(array_map.size()); @@ -3183,8 +3240,17 @@ async fn collect_left_input( offset += batch.num_rows(); } - // Merge all batches into a single batch, so we can directly index into the arrays - let batch = if prepared { + // Match the logical row order used while populating the hash table. + let batch = if preserve_batches { + batches.reverse(); + multi_batch = Some(MultiBatchBuildData::try_new( + std::mem::take(&mut batches), + &on_left, + &reservation, + &metrics, + )?); + multi_batch.as_ref().unwrap().batches()[0].clone() + } else if prepared { concat_batches(&schema, batches.iter().rev())? } else { concat_build_batches( @@ -3197,7 +3263,11 @@ async fn collect_left_input( )? }; - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + let left_values = if let Some(data) = &multi_batch { + data.values()[0].clone() + } else { + evaluate_expressions_to_arrays(&on_left, &batch)? + }; (Map::HashMap(hashmap), batch, left_values) }; @@ -3205,20 +3275,31 @@ async fn collect_left_input( // Join keys that are plain columns share the buffers of `batch`, any other // expression evaluates to new arrays that are kept for the whole join. let mut key_counter = RecordBatchMemoryCounter::new(); - key_counter.count_batch(&batch); - let keys_size = left_values - .iter() - .map(|values| key_counter.count_array(values.as_ref())) - .sum::(); + let keys_size = if let Some(data) = &multi_batch { + for batch in data.batches() { + key_counter.count_batch(batch); + } + data.values() + .iter() + .flatten() + .map(|values| key_counter.count_array(values.as_ref())) + .sum::() + } else { + key_counter.count_batch(&batch); + left_values + .iter() + .map(|values| key_counter.count_array(values.as_ref())) + .sum::() + }; reservation.try_grow(keys_size)?; metrics.build_mem_used.add(keys_size); let allocate_bitmap = || -> Result { - let bitmap_size = bit_util::ceil(batch.num_rows(), 8); + let bitmap_size = bit_util::ceil(num_rows, 8); reservation.try_grow(bitmap_size)?; metrics.build_mem_used.add(bitmap_size); - let mut bitmap = BooleanBufferBuilder::new(batch.num_rows()); + let mut bitmap = BooleanBufferBuilder::new(num_rows); bitmap.append_n(num_rows, false); Ok(bitmap) }; @@ -3315,6 +3396,21 @@ async fn collect_left_input( let membership = if num_rows == 0 { PushdownStrategy::Empty + } else if let Some(data) = &multi_batch { + if map.num_of_distinct_key() + <= config + .optimizer + .hash_join_inlist_pushdown_max_distinct_values + && let Some(inlist) = data.try_inlist_values( + config.optimizer.hash_join_inlist_pushdown_max_size, + &reservation, + &metrics, + )? + { + PushdownStrategy::InList(inlist) + } else { + PushdownStrategy::Map(Arc::clone(&map)) + } } else { // If the build side is small enough we can use IN list pushdown. // If it's too big we fall back to pushing down a reference to the hash table. @@ -3365,6 +3461,7 @@ async fn collect_left_input( map, batch, values: left_values, + multi_batch, bounds, membership, reservation, @@ -6381,6 +6478,7 @@ mod tests { let (l, r, _) = lookup_join_hashmap( &join_hash_map, &[left_keys_values], + None, &[right_keys_values], NullEquality::NullEqualsNothing, &hashes_buffer, @@ -6443,6 +6541,7 @@ mod tests { let (l, r, _) = lookup_join_hashmap( &join_hash_map, &[left_keys_values], + None, &[right_keys_values], NullEquality::NullEqualsNothing, &hashes_buffer, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs new file mode 100644 index 0000000000000..3446b5cd741e2 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs @@ -0,0 +1,573 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::*; +use arrow::array::{AsArray, GenericListArray, OffsetSizeTrait}; +use datafusion_common::utils::offset_span; + +pub(super) const MAX_COMPACT_BUILD_BYTES: usize = 64 * 1024 * 1024; +const TARGET_BUILD_BATCH_BYTES: usize = 8 * 1024 * 1024; +const TARGET_BUILD_BATCH_ROWS: usize = 8192; +const BUILD_ROW_DIRECTORY_STRIDE: usize = 1024; + +/// Keep the compact path for small logical inputs, even when their slices pin +/// much larger allocations. Unsupported estimates only disable this optimization. +pub(super) fn should_preserve_batches( + batches: &[RecordBatch], + input_bytes: usize, +) -> bool { + if input_bytes <= MAX_COMPACT_BUILD_BYTES { + return false; + } + let mut copy_bytes = 0usize; + for batch in batches { + let Ok(bytes) = estimate_batch_concat_allocation(batch) else { + return true; + }; + copy_bytes = copy_bytes.saturating_add(bytes); + if copy_bytes > MAX_COMPACT_BUILD_BYTES { + return true; + } + } + false +} + +fn estimate_batch_concat_allocation(batch: &RecordBatch) -> Result { + batch.columns().iter().try_fold(0usize, |bytes, array| { + Ok(bytes.saturating_add(estimate_compact_copy_size(array.as_ref())?)) + }) +} + +/// Estimate visible children for the layout decision, not memory admission. +/// `concat_build_batches` still reserves its existing conservative estimate. +/// Arrow's generic slice estimate includes unsliced List and Map children, +/// whereas concat copies only the child ranges covered by their offsets. +fn estimate_compact_copy_size(array: &dyn Array) -> Result { + let bytes = match array.data_type() { + DataType::List(_) => estimate_list_copy_size(array.as_list::())?, + DataType::LargeList(_) => estimate_list_copy_size(array.as_list::())?, + DataType::Map(_, _) => { + let map = array.as_map(); + let (start, len) = offset_span(map.offsets()); + let entries = map.entries().slice(start, len); + (map.len() + 1) + .saturating_mul(size_of::()) + .saturating_add(estimate_compact_copy_size(&entries)?) + } + DataType::Struct(_) => array.as_struct().columns().iter().try_fold( + 0usize, + |bytes, child| -> Result { + Ok(bytes.saturating_add(estimate_compact_copy_size(child.as_ref())?)) + }, + )?, + DataType::FixedSizeList(_, _) => { + estimate_compact_copy_size(array.as_fixed_size_list().values().as_ref())? + } + _ => return estimate_concat_allocation(array), + }; + Ok(bytes.saturating_add(array.nulls().map_or(0, |_| array.len().div_ceil(8)))) +} + +fn estimate_list_copy_size( + list: &GenericListArray, +) -> Result { + let (start, len) = offset_span(list.offsets()); + let values = list.values().slice(start, len); + Ok((list.len() + 1) + .saturating_mul(size_of::()) + .saturating_add(estimate_compact_copy_size(values.as_ref())?)) +} + +/// Batch-local keys and a sparse directory over logical hash-table row indices. +pub(in crate::joins::hash_join) struct MultiBatchBuildData { + batches: Vec, + values: Vec>, + batch_offsets: Vec, + row_directory: Vec, +} + +impl MultiBatchBuildData { + pub(super) fn try_new( + batches: Vec, + on_left: &[PhysicalExprRef], + reservation: &MemoryReservation, + metrics: &BuildProbeJoinMetrics, + ) -> Result { + let rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let directory_len = rows.div_ceil(BUILD_ROW_DIRECTORY_STRIDE); + let metadata_size = (batches.len() + 1 + directory_len) * size_of::() + + batches.capacity() * size_of::() + + batches.len() + * (size_of::>() + on_left.len() * size_of::()); + reservation.try_grow(metadata_size)?; + metrics.build_mem_used.add(metadata_size); + + let values = batches + .iter() + .map(|batch| evaluate_expressions_to_arrays(on_left, batch)) + .collect::>>()?; + let mut batch_offsets = Vec::with_capacity(batches.len() + 1); + batch_offsets.push(0); + for batch in &batches { + batch_offsets.push(batch_offsets.last().unwrap() + batch.num_rows()); + } + let mut row_directory = Vec::with_capacity(directory_len); + let mut batch_index = 0; + for row in (0..rows).step_by(BUILD_ROW_DIRECTORY_STRIDE) { + while batch_offsets[batch_index + 1] <= row { + batch_index += 1; + } + row_directory.push(batch_index); + } + Ok(Self { + batches, + values, + batch_offsets, + row_directory, + }) + } + + pub(in crate::joins::hash_join) fn batches(&self) -> &[RecordBatch] { + &self.batches + } + + pub(in crate::joins::hash_join) fn values(&self) -> &[Vec] { + &self.values + } + + pub(in crate::joins::hash_join) fn num_rows(&self) -> usize { + *self.batch_offsets.last().unwrap() + } + + pub(in crate::joins::hash_join) fn gather_indices( + &self, + indices: &UInt64Array, + ) -> Vec<(usize, usize)> { + indices + .iter() + .map(|index| { + index.map_or((0, 0), |row| { + let row = row as usize; + let mut batch = self.row_directory[row / BUILD_ROW_DIRECTORY_STRIDE]; + while self.batch_offsets[batch + 1] <= row { + batch += 1; + } + (batch + 1, row - self.batch_offsets[batch]) + }) + }) + .collect() + } + + /// Only small key arrays are copied for IN-list pushdown, never the payload. + pub(super) fn try_inlist_values( + &self, + max_size: usize, + reservation: &MemoryReservation, + metrics: &BuildProbeJoinMetrics, + ) -> Result> { + let max_size = max_size.min(MAX_COMPACT_BUILD_BYTES); + let num_keys = self.values[0].len(); + let mut copy_size = num_keys * (self.num_rows().div_ceil(8) + 64); + for array in self.values.iter().flatten() { + let Ok(bytes) = estimate_concat_allocation(array.as_ref()) else { + return Ok(None); + }; + copy_size = copy_size.saturating_add(bytes); + if copy_size > max_size { + return Ok(None); + } + } + if reservation.try_grow(copy_size).is_err() { + return Ok(None); + } + // Concatenation can fail for otherwise valid batch-local keys, e.g. + // when the union of dictionaries exceeds their key type's capacity. + // Membership pushdown is optional, so keep the hash predicate then. + let result = (|| -> Result> { + let keys = (0..num_keys) + .map(|key| { + let arrays = self + .values + .iter() + .map(|values| values[key].as_ref()) + .collect::>(); + Ok(arrow::compute::concat(&arrays)?) + }) + .collect::>>()?; + let Some(inlist) = build_struct_inlist_values(&keys)? else { + return Ok(None); + }; + let mut counter = RecordBatchMemoryCounter::new(); + for batch in &self.batches { + counter.count_batch(batch); + } + for values in self.values.iter().flatten() { + counter.count_array(values.as_ref()); + } + let retained = counter.count_array(inlist.as_ref()); + Ok((retained <= max_size && retained <= copy_size) + .then_some((inlist, retained))) + })() + .unwrap_or(None); + if let Some((inlist, retained)) = result { + reservation.shrink(copy_size - retained); + metrics.build_mem_used.add(retained); + Ok(Some(inlist)) + } else { + reservation.shrink(copy_size); + Ok(None) + } + } +} + +/// Coalesce independent flat inputs a bounded group at a time. Shared buffers +/// stay intact: replacing one slice must not release another slice's charge. +pub(super) fn coalesce_build_batches( + schema: &SchemaRef, + mut batches: Vec, + mut input_bytes: usize, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result> { + if batches.iter().any(|batch| batch.num_rows() == 0) { + batches.retain(|batch| batch.num_rows() != 0); + // Empty slices can pin whole allocations, but buffers shared with a + // remaining batch must keep their charge. + let mut counter = RecordBatchMemoryCounter::new(); + for batch in &batches { + counter.count_batch(batch); + } + let retained = counter.memory_usage(); + let released = input_bytes - retained; + reservation.shrink(released); + metrics.build_mem_used.sub(released); + input_bytes = retained; + } + let independent_bytes = batches + .iter() + .map(get_record_batch_memory_size) + .sum::(); + if independent_bytes != input_bytes + || schema.fields().iter().any(|field| { + matches!( + field.data_type(), + DataType::Dictionary(_, _) + | DataType::Union(_, _) + | DataType::RunEndEncoded(_, _) + ) || field.data_type().is_nested() + }) + { + return Ok(batches); + } + + let mut output = Vec::new(); + let mut pending = Vec::new(); + let mut pending_reserved_bytes = 0usize; + let mut pending_copy_bytes = 0usize; + let mut pending_rows = 0usize; + for batch in batches { + let reserved_bytes = get_record_batch_memory_size(&batch); + let copy_bytes = estimate_batch_concat_allocation(&batch).unwrap_or(usize::MAX); + let rows = batch.num_rows(); + if !pending.is_empty() + && (pending_copy_bytes.saturating_add(copy_bytes) > TARGET_BUILD_BATCH_BYTES + || pending_rows.saturating_add(rows) > TARGET_BUILD_BATCH_ROWS) + { + output.push(coalesce_build_group( + schema, + std::mem::take(&mut pending), + pending_reserved_bytes, + pending_copy_bytes, + reservation, + metrics, + )?); + pending_reserved_bytes = 0; + pending_copy_bytes = 0; + pending_rows = 0; + } + pending_reserved_bytes += reserved_bytes; + pending_copy_bytes = pending_copy_bytes.saturating_add(copy_bytes); + pending_rows += rows; + pending.push(batch); + } + if !pending.is_empty() { + output.push(coalesce_build_group( + schema, + pending, + pending_reserved_bytes, + pending_copy_bytes, + reservation, + metrics, + )?); + } + Ok(output) +} + +fn coalesce_build_group( + schema: &SchemaRef, + mut batches: Vec, + reserved_bytes: usize, + copy_bytes: usize, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result { + if batches.len() == 1 + && copy_bytes <= TARGET_BUILD_BATCH_BYTES + && reserved_bytes > copy_bytes.saturating_mul(2) + && !schema.fields().iter().any(|field| { + matches!(field.data_type(), DataType::Utf8View | DataType::BinaryView) + }) + { + // Arrow's single-input concat is zero-copy. A second, empty slice + // forces a bounded copy without allocating another input buffer. + batches.push(batches[0].slice(0, 0)); + } + concat_build_batches(schema, batches, false, reserved_bytes, reservation, metrics) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + FixedSizeListArray, Int64Array, LargeListArray, ListArray, MapArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow_schema::Field; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + + fn primitive_batch(rows: usize) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1; rows]))]) + .unwrap() + } + + fn input_bytes(batches: &[RecordBatch]) -> usize { + let mut counter = RecordBatchMemoryCounter::new(); + for batch in batches { + counter.count_batch(batch); + } + counter.memory_usage() + } + + #[test] + fn small_slices_of_large_shared_input_stay_compact() -> Result<()> { + let parent = primitive_batch((MAX_COMPACT_BUILD_BYTES + 8) / 8); + let batches = vec![parent.slice(0, 1), parent.slice(1, 1)]; + drop(parent); + let bytes = input_bytes(&batches); + assert!(bytes > MAX_COMPACT_BUILD_BYTES); + assert!(!should_preserve_batches(&batches, bytes)); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes + 128)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let compact = concat_build_batches( + &batches[0].schema(), + batches, + false, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(compact.num_rows(), 2); + assert_eq!(get_record_batch_memory_size(&compact), 16); + assert_eq!(reservation.size(), 16); + Ok(()) + } + + #[test] + fn large_logical_input_keeps_batches() { + let parent = primitive_batch((MAX_COMPACT_BUILD_BYTES + 8) / 8); + let batches = vec![parent.clone(), parent]; + assert!(should_preserve_batches(&batches, input_bytes(&batches))); + let small = vec![primitive_batch(2)]; + assert!(!should_preserve_batches(&small, input_bytes(&small))); + } + + #[test] + fn small_nested_slices_of_large_shared_input_stay_compact() -> Result<()> { + let len = (MAX_COMPACT_BUILD_BYTES + 8) / 8; + let values: ArrayRef = Arc::new(Int64Array::from(vec![1; len])); + let item = Arc::new(Field::new_list_field(DataType::Int64, false)); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::clone(&item), + OffsetBuffer::new(vec![0, 1, 2, len as i32].into()), + Arc::clone(&values), + None, + )); + let large_list: ArrayRef = Arc::new(LargeListArray::new( + item, + OffsetBuffer::new(vec![0, 1, 2, len as i64].into()), + Arc::clone(&values), + None, + )); + let entries = StructArray::new( + vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ] + .into(), + vec![Arc::clone(&values), values], + None, + ); + let map: ArrayRef = Arc::new(MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new(vec![0, 1, 2, len as i32].into()), + entries, + None, + false, + )); + let structure: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("list", list.data_type().clone(), false)].into(), + vec![Arc::clone(&list)], + None, + )); + let fixed_size_list: ArrayRef = Arc::new(FixedSizeListArray::new( + Arc::new(Field::new_list_field(list.data_type().clone(), false)), + 1, + Arc::clone(&list), + None, + )); + + for array in [list, large_list, map, structure, fixed_size_list] { + let schema = Arc::new(Schema::new(vec![Field::new( + "nested", + array.data_type().clone(), + false, + )])); + let parent = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; + let batches = vec![parent.slice(0, 1), parent.slice(1, 1)]; + drop(parent); + let bytes = input_bytes(&batches); + assert!(bytes > MAX_COMPACT_BUILD_BYTES); + assert!(!should_preserve_batches(&batches, bytes)); + + // Admission remains deliberately more conservative than the layout + // estimate, as on the original contiguous build path. + let admission_bytes = batches.iter().try_fold(0usize, |bytes, batch| { + Ok::<_, datafusion_common::DataFusionError>( + bytes + estimate_concat_allocation(batch.column(0).as_ref())?, + ) + })?; + assert!(admission_bytes > MAX_COMPACT_BUILD_BYTES); + let pool: Arc = + Arc::new(GreedyMemoryPool::new(bytes + admission_bytes + 1024)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let compact = concat_build_batches( + &schema, + batches, + false, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(compact.num_rows(), 2); + assert!(get_record_batch_memory_size(&compact) < 1024); + assert_eq!(reservation.size(), get_record_batch_memory_size(&compact)); + } + Ok(()) + } + + #[test] + fn coalesce_independent_slices_uses_copy_size() -> Result<()> { + let batches = (0..2) + .map(|_| primitive_batch(2 * 1024 * 1024).slice(0, 4096)) + .collect::>(); + let bytes = input_bytes(&batches); + assert_eq!(bytes, 32 * 1024 * 1024); + let pool: Arc = + Arc::new(GreedyMemoryPool::new(bytes + TARGET_BUILD_BATCH_BYTES)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let output = coalesce_build_batches( + &batches[0].schema(), + batches, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(output.len(), 1); + assert_eq!(output[0].num_rows(), 8192); + assert_eq!(input_bytes(&output), 8192 * 8); + assert_eq!(reservation.size(), input_bytes(&output)); + Ok(()) + } + + #[test] + fn coalesce_repackages_wasteful_singleton() -> Result<()> { + let batches = vec![primitive_batch(2 * 1024 * 1024).slice(0, 4096)]; + let bytes = input_bytes(&batches); + let pool: Arc = + Arc::new(GreedyMemoryPool::new(bytes + TARGET_BUILD_BATCH_BYTES)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let output = coalesce_build_batches( + &batches[0].schema(), + batches, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(output.len(), 1); + assert_eq!(input_bytes(&output), 4096 * 8); + assert_eq!(reservation.size(), input_bytes(&output)); + Ok(()) + } + + #[test] + fn coalesce_releases_buffers_retained_only_by_empty_batches() -> Result<()> { + let parent = primitive_batch(16); + let batches = vec![ + primitive_batch(2 * 1024 * 1024).slice(0, 0), + parent.slice(0, 8), + parent.slice(8, 8), + parent.slice(0, 0), + ]; + drop(parent); + let bytes = input_bytes(&batches); + assert_eq!(bytes, 16 * 1024 * 1024 + 128); + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let output = coalesce_build_batches( + &batches[0].schema(), + batches, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(output.len(), 2); + assert!(output.iter().all(|batch| batch.num_rows() == 8)); + assert_eq!(input_bytes(&output), 128); + assert_eq!(reservation.size(), 128); + assert_eq!(metrics.build_mem_used.value(), 128); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index 7b50435e18792..f566186ee1927 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -22,6 +22,8 @@ pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState mod exec; mod inlist_builder; +#[cfg(test)] +mod multi_batch_tests; mod partitioned_hash_eval; mod probe_completion; mod shared_bounds; diff --git a/datafusion/physical-plan/src/joins/hash_join/multi_batch_tests.rs b/datafusion/physical-plan/src/joins/hash_join/multi_batch_tests.rs new file mode 100644 index 0000000000000..f04b31d9f60e5 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/multi_batch_tests.rs @@ -0,0 +1,1419 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Public-plan regressions for large, multi-batch hash-join builds. + +use std::sync::{Arc, OnceLock}; + +use arrow::array::{ + Array, ArrayRef, DictionaryArray, FixedSizeListBuilder, Int8Array, Int32Array, + Int32Builder, ListArray, ListBuilder, MapBuilder, RecordBatch, StringArray, + StringBuilder, StructArray, new_null_array, +}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::{cast, concat}; +use arrow::datatypes::{DataType, Field, Int8Type, Int32Type, Schema, SchemaRef}; +use arrow::error::ArrowError; +use datafusion_common::utils::memory::get_record_batch_memory_size; +use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue}; +use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::{TaskContext, config::SessionConfig}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{ + BinaryExpr, IsNotNullExpr, cast as physical_cast, col, lit, +}; + +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::joins::{HashJoinExec, HashJoinExecBuilder, PartitionMode}; +use crate::test::TestMemoryExec; +use crate::{ExecutionPlan, common}; + +const COMPACT_BUILD_BYTES: usize = 64 * 1024 * 1024; +const PADDING_BYTES: usize = 65 * 1024 * 1024; + +fn string_with_backing(rows: usize, values: Buffer) -> ArrayRef { + let mut offsets = vec![values.len() as i32; rows + 1]; + offsets[0] = 0; + Arc::new(StringArray::new( + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + None, + )) +} + +/// Two batches exceed the compact-build threshold while sharing one allocation. +fn padding(rows: usize) -> ArrayRef { + static VALUES: OnceLock = OnceLock::new(); + let values = VALUES.get_or_init(|| Buffer::from_vec(vec![b'p'; PADDING_BYTES])); + string_with_backing(rows, values.clone()) +} + +fn assert_large_build(batches: &[RecordBatch]) { + assert!(batches.len() > 1); + assert!(batches.iter().all(|batch| batch.num_rows() > 0)); + assert!( + batches + .iter() + .map(get_record_batch_memory_size) + .sum::() + > COMPACT_BUILD_BYTES + ); +} + +fn padded_build(keys: Vec) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("build_key", keys[0].data_type().clone(), true), + Field::new("padding", DataType::Utf8, false), + ])); + let batches = keys + .into_iter() + .map(|key| { + let rows = key.len(); + RecordBatch::try_new(Arc::clone(&schema), vec![key, padding(rows)]) + .map_err(Into::into) + }) + .collect::>>()?; + assert_large_build(&batches); + Ok(batches) +} + +fn with_payload( + batches: Vec, + payload: ArrayRef, +) -> Result> { + let mut fields = batches[0].schema().fields().to_vec(); + fields.push(Arc::new(Field::new( + "payload", + payload.data_type().clone(), + true, + ))); + let schema = Arc::new(Schema::new(fields)); + let mut offset = 0; + let output = batches + .into_iter() + .map(|batch| { + let mut columns = batch.columns().to_vec(); + columns.push(payload.slice(offset, batch.num_rows())); + offset += batch.num_rows(); + RecordBatch::try_new(Arc::clone(&schema), columns).map_err(Into::into) + }) + .collect::>>()?; + assert_eq!(offset, payload.len()); + Ok(output) +} + +fn nested_payload(ids: &[i32]) -> ArrayRef { + let mut lists = ListBuilder::new(StringBuilder::new()); + for id in ids { + if id % 3 != 0 { + lists.values().append_value(format!("response{id}")); + lists.values().append_null(); + } + lists.append(id % 3 != 0); + } + let lists: ArrayRef = Arc::new(lists.finish()); + Arc::new(StructArray::new( + vec![Field::new("labels", lists.data_type().clone(), true)].into(), + vec![lists], + Some(NullBuffer::from( + ids.iter().map(|id| id % 3 != 1).collect::>(), + )), + )) +} + +fn scalar_rows(batches: &[RecordBatch]) -> Result>> { + batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(move |row| { + batch + .columns() + .iter() + .map(|array| ScalarValue::try_from_array(array.as_ref(), row)) + .collect() + }) + }) + .collect() +} + +fn probe(values: ArrayRef) -> Result<(SchemaRef, Vec)> { + let schema = Arc::new(Schema::new(vec![Field::new( + "probe_key", + values.data_type().clone(), + true, + )])); + let batches = if values.is_empty() { + vec![] + } else { + vec![RecordBatch::try_new(Arc::clone(&schema), vec![values])?] + }; + Ok((schema, batches)) +} + +fn join_builder( + build: &[RecordBatch], + probe: (SchemaRef, Vec), + join_type: JoinType, + projection: Vec, +) -> Result { + let (probe_schema, probe_batches) = probe; + let left = TestMemoryExec::try_new_exec(&[build.to_vec()], build[0].schema(), None)?; + let right = TestMemoryExec::try_new_exec(&[probe_batches], probe_schema, None)?; + join_builder_from_plans(left, right, join_type, projection) +} + +fn join_builder_from_plans( + left: Arc, + right: Arc, + join_type: JoinType, + projection: Vec, +) -> Result { + let on = vec![( + col("build_key", &left.schema())?, + col("probe_key", &right.schema())?, + )]; + Ok(HashJoinExecBuilder::new(left, right, on, join_type) + .with_projection(Some(projection)) + .with_partition_mode(PartitionMode::CollectLeft) + .with_null_equality(NullEquality::NullEqualsNothing)) +} + +fn task_context(batch_size: usize, use_perfect_hash: bool) -> Arc { + let mut config = SessionConfig::default().with_batch_size(batch_size); + config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = false; + config + .options_mut() + .execution + .perfect_hash_join_small_build_threshold = + if use_perfect_hash { usize::MAX } else { 0 }; + config + .options_mut() + .execution + .perfect_hash_join_min_key_density = + if use_perfect_hash { 0.0 } else { f64::INFINITY }; + Arc::new(TaskContext::default().with_session_config(config)) +} + +fn assert_perfect_hash(join: &HashJoinExec, expected: bool) { + let count = join + .metrics() + .and_then(|metrics| metrics.sum_by_name("array_map_created_count")) + .map_or(0, |metric| metric.as_usize()); + assert_eq!(count > 0, expected); +} + +async fn collect_join( + join: HashJoinExec, + batch_size: usize, + use_perfect_hash: bool, +) -> Result> { + let output = + common::collect(join.execute(0, task_context(batch_size, use_perfect_hash))?) + .await?; + assert_perfect_hash(&join, use_perfect_hash); + Ok(output) +} + +fn int_values(batches: &[RecordBatch]) -> Vec> { + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect() +} + +#[track_caller] +fn assert_batch_size(batches: &[RecordBatch], batch_size: usize) { + assert!( + batches.iter().all(|batch| batch.num_rows() <= batch_size), + "expected at most {batch_size} rows, got {:?}", + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + ); +} + +fn string_key_rows(batches: &[RecordBatch]) -> Result> { + let mut rows = vec![]; + for batch in batches { + let left = cast(batch.column(0), &DataType::Utf8)?; + let right = cast(batch.column(1), &DataType::Utf8)?; + let left = left.as_any().downcast_ref::().unwrap(); + let right = right.as_any().downcast_ref::().unwrap(); + assert_eq!(left.null_count(), 0); + assert_eq!(right.null_count(), 0); + for row in 0..batch.num_rows() { + rows.push((left.value(row).to_owned(), right.value(row).to_owned())); + } + } + Ok(rows) +} + +#[tokio::test] +async fn coerced_dictionary_and_plain_keys_join_across_batches() -> Result<()> { + let values: ArrayRef = Arc::new(StringArray::from(vec!["x", "y", "missing"])); + let dictionary = |keys: Vec| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&values), + )?)) + }; + + for dictionary_on_build in [true, false] { + let (build_keys, probe_keys): (Vec, ArrayRef) = if dictionary_on_build { + ( + vec![dictionary(vec![0])?, dictionary(vec![1])?], + Arc::new(StringArray::from(vec!["x", "x", "y", "missing"])), + ) + } else { + ( + vec![ + Arc::new(StringArray::from(vec!["x"])), + Arc::new(StringArray::from(vec!["y"])), + ], + dictionary(vec![0, 0, 1, 2])?, + ) + }; + let build = padded_build(build_keys)?; + let probe = probe(probe_keys)?; + let mut left_key = col("build_key", &build[0].schema())?; + let mut right_key = col("probe_key", &probe.0)?; + if dictionary_on_build { + left_key = physical_cast(left_key, &build[0].schema(), DataType::Utf8)?; + } else { + right_key = physical_cast(right_key, &probe.0, DataType::Utf8)?; + } + let join = join_builder(&build, probe, JoinType::Inner, vec![0, 2])? + .with_on(vec![(left_key, right_key)]) + .build()?; + let output = collect_join(join, 8192, false).await?; + let mut actual = string_key_rows(&output)?; + actual.sort_unstable(); + let expected = [("x", "x"), ("x", "x"), ("y", "y")] + .map(|(left, right)| (left.to_owned(), right.to_owned())); + assert_eq!(actual, expected); + } + Ok(()) +} + +#[tokio::test] +async fn computed_dictionary_keys_join_across_batches() -> Result<()> { + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let dictionary = |values: Vec<&str>| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from_iter_values(0..values.len() as i32), + Arc::new(StringArray::from(values)), + )?)) + }; + + for computed_on_build in [true, false] { + let build_keys: Vec = if computed_on_build { + vec![ + Arc::new(StringArray::from(vec!["x"])), + Arc::new(StringArray::from(vec!["y"])), + ] + } else { + vec![dictionary(vec!["x"])?, dictionary(vec!["y"])?] + }; + let probe_keys = if computed_on_build { + dictionary(vec!["y", "missing", "x", "y"])? + } else { + Arc::new(StringArray::from(vec!["y", "missing", "x", "y"])) as ArrayRef + }; + let build = padded_build(build_keys)?; + let build_schema = build[0].schema(); + let probe = probe(probe_keys)?; + let mut left_key = col("build_key", &build_schema)?; + let mut right_key = col("probe_key", &probe.0)?; + if computed_on_build { + left_key = physical_cast(left_key, &build_schema, dictionary_type.clone())?; + } else { + right_key = physical_cast(right_key, &probe.0, dictionary_type.clone())?; + } + let join = join_builder(&build, probe, JoinType::Inner, vec![0, 2])? + .with_on(vec![(left_key, right_key)]) + .build()?; + let output = collect_join(join, 8192, false).await?; + let expected = [("y", "y"), ("x", "x"), ("y", "y")] + .map(|(left, right)| (left.to_owned(), right.to_owned())); + assert_eq!(string_key_rows(&output)?, expected); + } + Ok(()) +} + +/// Find the small Arrow dictionary shape from the partial-concat regression. +/// Every trial references 128 dictionary slots, so the three-way merge uses +/// the same interner capacity. Keep only values that do not collide there. +/// Using Arrow itself avoids a new ahash dependency or platform-specific hashes. +fn partial_concat_dictionaries() -> Result<(Vec, String)> { + let mut values = Vec::::new(); + for candidate in 0..4096 { + values.push(format!("payload{candidate:06}")); + let dictionaries = (0..3) + .map(|_| { + let strings = (0..128) + .map(|index| values.get(index).unwrap_or(&values[0]).as_str()); + Ok(Arc::new(DictionaryArray::::try_new( + Int8Array::from_iter_values(0..=i8::MAX), + Arc::new(StringArray::from_iter_values(strings)), + )?) as ArrayRef) + }) + .collect::>>()?; + let arrays = dictionaries + .iter() + .map(|array| array.as_ref()) + .collect::>(); + let distinct = match concat(&arrays) { + Ok(array) => Some( + array + .as_any() + .downcast_ref::>() + .unwrap() + .values() + .len(), + ), + Err(ArrowError::DictionaryKeyOverflowError) => None, + Err(error) => return Err(error.into()), + }; + if distinct != Some(values.len()) { + values.pop(); + } else if values.len() == 128 { + if matches!( + concat(&arrays[..2]), + Err(ArrowError::DictionaryKeyOverflowError) + ) { + return Ok((dictionaries, values[0].clone())); + } + values.pop(); + } + } + panic!("could not construct the partial-concat dictionary regression"); +} + +#[tokio::test] +async fn dictionary_payloads_preserve_values_across_batches() -> Result<()> { + let (payloads, expected) = partial_concat_dictionaries()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("build_key", DataType::Int32, false), + Field::new("payload", payloads[0].data_type().clone(), false), + Field::new("padding", DataType::Utf8, false), + ])); + let build = payloads + .into_iter() + .enumerate() + .map(|(index, payload)| { + let bytes = if index == 2 { 65 * 1024 * 1024 } else { 0 }; + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values( + (index * 128..(index + 1) * 128).map(|key| key as i32), + )), + payload, + string_with_backing(128, Buffer::from_vec(vec![b'p'; bytes])), + ], + ) + .map_err(Into::into) + }) + .collect::>>()?; + assert_large_build(&build); + assert!( + build[..2] + .iter() + .map(get_record_batch_memory_size) + .sum::() + < 8 * 1024 * 1024 + ); + + // The original all-at-once concat succeeds, but coalescing the two small + // batches first overflows Int8 dictionary keys. Both hash implementations + // must preserve values without introducing that partial-concat failure. + for use_perfect_hash in [false, true] { + let probe = probe(Arc::new(Int32Array::from(vec![0, 128])))?; + let join = join_builder(&build, probe, JoinType::Inner, vec![1])?.build()?; + let output = collect_join(join, 8192, use_perfect_hash).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 2); + for batch in output { + let decoded = cast(batch.column(0), &DataType::Utf8)?; + let decoded = decoded.as_any().downcast_ref::().unwrap(); + assert!(decoded.iter().all(|value| value == Some(expected.as_str()))); + } + } + Ok(()) +} + +#[tokio::test] +async fn hidden_null_utf8_payload_does_not_overflow_on_fanout() -> Result<()> { + const MATCHES: usize = 66; + const BATCH_SIZE: usize = 8192; + + // Repeating the hidden span for every match would overflow Utf8 offsets. + let payload = StringArray::new( + OffsetBuffer::new(ScalarBuffer::from(vec![0, PADDING_BYTES as i32])), + Buffer::from_vec(vec![b'p'; PADDING_BYTES]), + Some(NullBuffer::new_null(1)), + ); + assert!(payload.value_length(0) as usize * MATCHES > i32::MAX as usize); + let payload: ArrayRef = Arc::new(payload); + let schema = Arc::new(Schema::new(vec![ + Field::new("build_key", DataType::Int32, false), + Field::new("payload", DataType::Utf8, true), + ])); + let build = [7, 8] + .into_iter() + .map(|key| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![key])), Arc::clone(&payload)], + ) + .map_err(Into::into) + }) + .collect::>>()?; + assert_large_build(&build); + let expected = (0..MATCHES) + .map(|index| 7 + (index % 2) as i32) + .collect::>(); + + for use_perfect_hash in [false, true] { + let probe = probe(Arc::new(Int32Array::from(expected.clone())))?; + let join = + join_builder(&build, probe, JoinType::Inner, vec![0, 1, 2])?.build()?; + let output = collect_join(join, BATCH_SIZE, use_perfect_hash).await?; + assert_batch_size(&output, BATCH_SIZE); + let mut actual = vec![]; + for batch in &output { + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let payload = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let probe_keys = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(keys.null_count(), 0); + assert_eq!(payload.null_count(), batch.num_rows()); + assert_eq!(probe_keys, keys); + actual.extend(keys.values().iter().copied()); + } + assert_eq!(actual, expected); + } + Ok(()) +} + +#[tokio::test] +async fn nested_payloads_survive_join_filters_and_outer_rows() -> Result<()> { + let mut maps = MapBuilder::new( + None, + StringBuilder::new(), + ListBuilder::new(StringBuilder::new()), + ); + for id in 0..4 { + if id == 0 || id == 3 { + maps.keys().append_value(format!("detector{id}")); + if id == 0 { + maps.values().values().append_value("email"); + maps.values().values().append_null(); + } + maps.values().append(id == 0); + } + maps.append(id != 2)?; + } + let payload: ArrayRef = Arc::new(maps.finish()); + let build_keys: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 2, 3])); + let build = with_payload( + padded_build(vec![build_keys.slice(0, 2), build_keys.slice(2, 2)])?, + Arc::clone(&payload), + )?; + let probe_keys: ArrayRef = Arc::new(Int32Array::from(vec![2, 2, 5])); + let response = nested_payload(&[0, 1, 2]); + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("probe_key", DataType::Int32, false), + Field::new("response", response.data_type().clone(), true), + ])); + let probe_batch = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![Arc::clone(&probe_keys), Arc::clone(&response)], + )?; + let filter_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + payload.data_type().clone(), + true, + )])); + let filter = JoinFilter::new( + Arc::new(IsNotNullExpr::new(col("payload", &filter_schema)?)), + vec![ColumnIndex { + index: 2, + side: JoinSide::Left, + }], + filter_schema, + ); + let scalar_at = |array: &ArrayRef, row: Option| { + row.map_or_else( + || ScalarValue::try_from(array.data_type()), + |row| ScalarValue::try_from_array(array.as_ref(), row), + ) + }; + let cases = [ + ( + JoinType::Inner, + false, + vec![ + (Some(1), Some(0)), + (Some(2), Some(0)), + (Some(1), Some(1)), + (Some(2), Some(1)), + ], + ), + ( + JoinType::Full, + true, + vec![ + (Some(1), Some(0)), + (Some(1), Some(1)), + (Some(0), None), + (Some(2), None), + (Some(3), None), + (None, Some(2)), + ], + ), + ]; + for (join_type, filtered, pairs) in cases { + let expected = pairs + .into_iter() + .map(|(left, right)| { + Ok(vec![ + scalar_at(&build_keys, left)?, + scalar_at(&payload, left)?, + scalar_at(&probe_keys, right)?, + scalar_at(&response, right)?, + ]) + }) + .collect::>>()?; + for use_perfect_hash in [false, true] { + let join = join_builder( + &build, + (Arc::clone(&probe_schema), vec![probe_batch.clone()]), + join_type, + vec![0, 2, 3, 4], + )? + .with_filter(filtered.then(|| filter.clone())) + .build()?; + let output = collect_join(join, 8192, use_perfect_hash).await?; + let mut actual = scalar_rows(&output)?; + assert_eq!(actual.len(), expected.len()); + for row in &expected { + let index = actual.iter().position(|actual| actual == row).unwrap(); + actual.swap_remove(index); + } + assert!(actual.is_empty()); + } + } + Ok(()) +} + +#[tokio::test] +async fn nested_null_payloads_do_not_copy_hidden_children_on_fanout() -> Result<()> { + const MATCHES: usize = 66; + let bytes = Buffer::from_vec(vec![b'p'; PADDING_BYTES]); + assert!(bytes.len() * MATCHES > i32::MAX as usize); + // Hide the bytes at each level: struct, list, and string. + for null_level in 0..3 { + let wrap = |strings: ArrayRef| -> ArrayRef { + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 1])), + strings, + (null_level == 1).then(|| NullBuffer::new_null(1)), + )); + Arc::new(StructArray::new( + vec![Field::new("items", list.data_type().clone(), true)].into(), + vec![list], + (null_level == 0).then(|| NullBuffer::new_null(1)), + )) + }; + let payload = wrap(Arc::new(StringArray::new( + OffsetBuffer::new(ScalarBuffer::from(vec![0, PADDING_BYTES as i32])), + bytes.clone(), + (null_level == 2).then(|| NullBuffer::new_null(1)), + ))); + let expected_payload = ScalarValue::try_from_array( + wrap(new_null_array(&DataType::Utf8, 1)).as_ref(), + 0, + )?; + let schema = Arc::new(Schema::new(vec![ + Field::new("build_key", DataType::Int32, false), + Field::new("payload", payload.data_type().clone(), true), + ])); + let build = [7, 8] + .into_iter() + .map(|key| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![key])), Arc::clone(&payload)], + ) + .map_err(Into::into) + }) + .collect::>>()?; + assert_large_build(&build); + let keys = (0..MATCHES) + .map(|index| 7 + (index % 2) as i32) + .collect::>(); + let expected = keys + .iter() + .map(|&key| { + vec![ + ScalarValue::Int32(Some(key)), + expected_payload.clone(), + ScalarValue::Int32(Some(key)), + ] + }) + .collect::>(); + for use_perfect_hash in [false, true] { + let probe = probe(Arc::new(Int32Array::from(keys.clone())))?; + let join = + join_builder(&build, probe, JoinType::Inner, vec![0, 1, 2])?.build()?; + let output = collect_join(join, 8192, use_perfect_hash).await?; + assert_eq!(scalar_rows(&output)?, expected); + } + } + Ok(()) +} + +#[tokio::test] +async fn final_left_anti_is_bounded_and_emitted_once_with_fetch() -> Result<()> { + const BATCH_SIZE: usize = 4; + const FETCH: usize = 5; + let payload = nested_payload(&(0..34).collect::>()); + let build = with_payload( + padded_build(vec![ + Arc::new(Int32Array::from_iter_values(0..17)), + Arc::new(Int32Array::from_iter_values(17..34)), + ])?, + Arc::clone(&payload), + )?; + let (probe_schema, first_probe) = probe(Arc::new(Int32Array::from(vec![0, 17])))?; + let (_, second_probe) = probe(Arc::new(Int32Array::from(vec![1, 18])))?; + let build_schema = build[0].schema(); + let left: Arc = + TestMemoryExec::try_new_exec(&[build], build_schema, None)?; + let right: Arc = + TestMemoryExec::try_new_exec(&[first_probe, second_probe], probe_schema, None)?; + let expected = (0..34) + .filter(|value| ![0, 1, 17, 18].contains(value)) + .map(Some) + .collect::>(); + + for use_perfect_hash in [false, true] { + let mut unlimited = vec![]; + for fetch in [None, Some(FETCH)] { + let join = join_builder_from_plans( + Arc::clone(&left), + Arc::clone(&right), + JoinType::LeftAnti, + vec![0, 2], + )? + .with_fetch(fetch) + .build()?; + let context = task_context(BATCH_SIZE, use_perfect_hash); + let first = join.execute(0, Arc::clone(&context))?; + let second = join.execute(1, context)?; + let (mut output, second_output) = + futures::try_join!(common::collect(first), common::collect(second))?; + output.extend(second_output); + assert_perfect_hash(&join, use_perfect_hash); + assert_batch_size(&output, BATCH_SIZE); + let actual = int_values(&output); + let rows = scalar_rows(&output)?; + for (row, key) in rows.iter().zip(&actual) { + assert_eq!( + row[1], + ScalarValue::try_from_array(payload.as_ref(), key.unwrap() as usize)? + ); + } + if fetch.is_some() { + assert_eq!(rows, unlimited[..FETCH]); + } else { + let mut sorted = actual; + sorted.sort_unstable(); + assert_eq!(sorted, expected); + unlimited = rows; + } + } + } + Ok(()) +} + +#[tokio::test] +async fn final_null_aware_left_anti_semantics() -> Result<()> { + let build = padded_build(vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(2)])), + Arc::new(Int32Array::from(vec![Some(3), Some(4), None])), + ])?; + let cases = [ + (vec![], vec![None, None, Some(1), Some(2), Some(3), Some(4)]), + (vec![Some(1)], vec![Some(2), Some(3), Some(4)]), + (vec![Some(1), None], vec![]), + ]; + + for use_perfect_hash in [false, true] { + for (probe_values, expected) in &cases { + let probe = probe(Arc::new(Int32Array::from(probe_values.clone())))?; + let join = join_builder(&build, probe, JoinType::LeftAnti, vec![0])? + .with_null_aware(true) + .build()?; + let output = collect_join(join, 2, use_perfect_hash).await?; + assert_batch_size(&output, 2); + let mut actual = int_values(&output); + actual.sort_unstable(); + assert_eq!(&actual, expected); + } + } + Ok(()) +} + +#[tokio::test] +async fn null_aware_left_mark_preserves_unknown_across_batches() -> Result<()> { + let keys = [Some(1), None, Some(2), Some(3), Some(4), None]; + let build = padded_build(vec![ + Arc::new(Int32Array::from(keys[..3].to_vec())), + Arc::new(Int32Array::from(keys[3..].to_vec())), + ])?; + let cases = [ + (vec![], vec![Some(false); 6]), + ( + vec![Some(1)], + vec![ + Some(true), + None, + Some(false), + Some(false), + Some(false), + None, + ], + ), + ( + vec![Some(1), None], + vec![Some(true), None, None, None, None, None], + ), + ]; + for use_perfect_hash in [false, true] { + for (probe_values, marks) in &cases { + let join = join_builder( + &build, + probe(Arc::new(Int32Array::from(probe_values.clone())))?, + JoinType::LeftMark, + vec![0, 2], + )? + .with_null_aware(true) + .build()?; + let output = collect_join(join, 2, use_perfect_hash).await?; + assert_batch_size(&output, 2); + let expected = keys + .iter() + .zip(marks) + .map(|(&key, &mark)| { + vec![ScalarValue::Int32(key), ScalarValue::Boolean(mark)] + }) + .collect::>(); + assert_rows_unordered(scalar_rows(&output)?, &expected); + } + } + Ok(()) +} + +#[tokio::test] +async fn computed_composite_keys_preserve_cross_batch_duplicates() -> Result<()> { + let build_schema = Arc::new(Schema::new(vec![ + Field::new("build_key", DataType::Int32, false), + Field::new("label", DataType::Utf8, false), + Field::new("padding", DataType::Utf8, false), + Field::new("build_id", DataType::Int32, false), + ])); + let build = [ + (vec![1, 2, 2], vec!["a", "a", "b"], vec![0, 1, 2]), + (vec![2, 3, 4], vec!["b", "a", "b"], vec![3, 4, 5]), + ] + .into_iter() + .map(|(keys, labels, ids)| { + RecordBatch::try_new( + Arc::clone(&build_schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(StringArray::from(labels)), + padding(3), + Arc::new(Int32Array::from(ids)), + ], + ) + .map_err(Into::into) + }) + .collect::>>()?; + assert_large_build(&build); + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("probe_key", DataType::Int32, false), + Field::new("label", DataType::Utf8, false), + Field::new("probe_id", DataType::Int32, false), + ])); + let probe_batch = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![ + Arc::new(Int32Array::from(vec![12, 12, 14, 13])), + Arc::new(StringArray::from(vec!["b", "a", "b", "missing"])), + Arc::new(Int32Array::from(vec![10, 11, 12, 13])), + ], + )?; + let on = vec![ + ( + Arc::new(BinaryExpr::new( + col("build_key", &build_schema)?, + Operator::Plus, + lit(10i32), + )) as _, + col("probe_key", &probe_schema)?, + ), + (col("label", &build_schema)?, col("label", &probe_schema)?), + ]; + for use_perfect_hash in [false, true] { + let join = join_builder( + &build, + (Arc::clone(&probe_schema), vec![probe_batch.clone()]), + JoinType::Inner, + vec![3, 6], + )? + .with_on(on.clone()) + .build()?; + let output = + common::collect(join.execute(0, task_context(2, use_perfect_hash))?).await?; + assert_perfect_hash(&join, false); + assert_batch_size(&output, 2); + assert_rows_unordered( + scalar_rows(&output)?, + &[(2, 10), (3, 10), (1, 11), (5, 12)] + .map(|(left, right)| vec![left.into(), right.into()]), + ); + } + Ok(()) +} + +fn assert_rows_unordered( + mut actual: Vec>, + expected: &[Vec], +) { + assert_eq!(actual.len(), expected.len(), "{actual:?} != {expected:?}"); + for row in expected { + let index = actual + .iter() + .position(|actual| actual == row) + .unwrap_or_else(|| panic!("missing {row:?} in {actual:?}")); + actual.swap_remove(index); + } +} + +#[tokio::test] +async fn dictionary_keys_with_distinct_dictionaries_and_logical_nulls() -> Result<()> { + let dictionary = + |keys: Vec>, values: Vec>| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int8Array::from(keys), + Arc::new(StringArray::from(values)), + )?)) + }; + let build = padded_build(vec![ + dictionary( + vec![Some(0), Some(1), None, Some(2)], + vec![Some("red"), None, Some("blue")], + )?, + dictionary( + vec![Some(1), Some(0), Some(2), Some(1)], + vec![Some("blue"), Some("red"), None], + )?, + ])?; + let probe_keys = dictionary( + vec![Some(2), Some(1), Some(0), None], + vec![Some("red"), None, Some("blue")], + )?; + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let join = join_builder( + &build, + probe(Arc::clone(&probe_keys))?, + JoinType::Inner, + vec![0, 2], + )? + .with_null_equality(null_equality) + .build()?; + let output = collect_join(join, 2, false).await?; + assert_batch_size(&output, 2); + let mut actual = vec![]; + for batch in output { + let left = cast(batch.column(0), &DataType::Utf8)?; + let right = cast(batch.column(1), &DataType::Utf8)?; + for row in 0..batch.num_rows() { + actual.push(vec![ + ScalarValue::try_from_array(left.as_ref(), row)?, + ScalarValue::try_from_array(right.as_ref(), row)?, + ]); + } + } + let mut expected = vec![vec!["red".into(), "red".into()]; 3]; + expected.extend(vec![vec!["blue".into(), "blue".into()]; 2]); + if null_equality == NullEquality::NullEqualsNull { + expected.extend(vec![vec![ScalarValue::Utf8(None); 2]; 6]); + } + assert_rows_unordered(actual, &expected); + } + Ok(()) +} + +#[tokio::test] +async fn dictionary_inlist_overflow_keeps_batchwise_join() -> Result<()> { + let dictionary = |start: usize, end: usize| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int8Array::from_iter_values((0..end - start).map(|key| key as i8)), + Arc::new(StringArray::from_iter_values( + (start..end).map(|value| format!("value{value:03}")), + )), + )?)) + }; + // Each input is representable, but a single Int8 dictionary cannot hold + // all 130 distinct keys. Optional IN-list construction must not fail the join. + let build = padded_build(vec![dictionary(0, 65)?, dictionary(65, 130)?])?; + let probe_values = ["value000", "value064", "value065", "value129", "missing"]; + let probe_keys = Arc::new(DictionaryArray::::try_new( + Int8Array::from_iter_values(0..probe_values.len() as i8), + Arc::new(StringArray::from(probe_values.to_vec())), + )?) as ArrayRef; + let join = + join_builder(&build, probe(probe_keys)?, JoinType::Inner, vec![0, 2])?.build()?; + let mut config = task_context(2, false).session_config().clone(); + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_distinct_values = 512; + let context = Arc::new(TaskContext::default().with_session_config(config)); + let output = common::collect(join.execute(0, context)?).await?; + assert_batch_size(&output, 2); + assert_eq!( + string_key_rows(&output)?, + probe_values[..4] + .iter() + .map(|value| (value.to_string(), value.to_string())) + .collect::>(), + ); + Ok(()) +} + +#[tokio::test] +async fn large_build_succeeds_without_room_for_contiguous_copy() -> Result<()> { + const LIMIT: usize = 96 * 1024 * 1024; + for shape in ["plain", "computed", "dictionary", "nested"] { + let keys: Vec = if shape == "dictionary" { + [vec!["one", "two"], vec!["two", "three"]] + .into_iter() + .map(|values| { + Ok(Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1]), + Arc::new(StringArray::from(values)), + )?) as ArrayRef) + }) + .collect::>>()? + } else { + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 3])), + ] + }; + let mut build = padded_build(keys)?; + if shape == "nested" { + build = with_payload(build, nested_payload(&[0, 1, 2, 3]))?; + } + let mut counter = + datafusion_common::utils::memory::RecordBatchMemoryCounter::new(); + for batch in &build { + counter.count_batch(batch); + } + let unique_bytes = counter.memory_usage(); + assert!(unique_bytes > COMPACT_BUILD_BYTES && unique_bytes < LIMIT); + let concat_bytes = build + .iter() + .map(|batch| batch.column(1).to_data().get_slice_memory_size()) + .collect::, _>>()? + .into_iter() + .sum::(); + // Both padding strings use the full backing, so concatenating them + // really needs this destination even though their sources share memory. + assert!(concat_bytes > LIMIT - unique_bytes); + + let probe_keys: ArrayRef = if shape == "dictionary" { + Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0]), + Arc::new(StringArray::from(vec!["two"])), + )?) + } else { + Arc::new(Int32Array::from(vec![2])) + }; + let probe = probe(probe_keys)?; + let mut builder = join_builder(&build, probe.clone(), JoinType::Inner, vec![0])?; + if shape == "computed" { + builder = builder.with_on(vec![( + Arc::new(BinaryExpr::new( + col("build_key", &build[0].schema())?, + Operator::Plus, + lit(1i32), + )), + Arc::new(BinaryExpr::new( + col("probe_key", &probe.0)?, + Operator::Plus, + lit(1i32), + )), + )]); + } + let join = builder.build()?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(LIMIT)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(task_context(2, false).session_config().clone()) + .with_runtime(runtime), + ); + let output = common::collect(join.execute(0, Arc::clone(&context))?).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 2); + for batch in &output { + let keys = cast(batch.column(0), &DataType::Utf8)?; + let keys = keys.as_any().downcast_ref::().unwrap(); + let expected = if shape == "dictionary" { "two" } else { "2" }; + assert!(keys.iter().all(|key| key == Some(expected))); + } + assert!(pool.reserved() <= LIMIT); + drop(output); + drop(join); + drop(context); + assert_eq!(pool.reserved(), 0, "{shape}"); + } + Ok(()) +} + +#[tokio::test] +async fn all_join_types_preserve_rows_across_build_batches() -> Result<()> { + let build = with_payload( + padded_build(vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 4])), + ])?, + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + )?; + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("probe_key", DataType::Int32, false), + Field::new("probe_id", DataType::Int32, false), + ])); + let probe_batch = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 2])), + Arc::new(Int32Array::from(vec![10, 11, 12])), + ], + )?; + let pair = |left: Option, right: Option| { + vec![ScalarValue::Int32(left), ScalarValue::Int32(right)] + }; + let matches = vec![ + pair(Some(1), Some(10)), + pair(Some(2), Some(10)), + pair(Some(1), Some(12)), + pair(Some(2), Some(12)), + ]; + let mut left_outer = matches.clone(); + left_outer.extend([pair(Some(0), None), pair(Some(3), None)]); + let mut right_outer = matches.clone(); + right_outer.push(pair(None, Some(11))); + let mut full = left_outer.clone(); + full.push(pair(None, Some(11))); + let cases = [ + (JoinType::Inner, vec![2, 4], matches), + (JoinType::Left, vec![2, 4], left_outer), + (JoinType::Right, vec![2, 4], right_outer), + (JoinType::Full, vec![2, 4], full), + ( + JoinType::LeftSemi, + vec![2], + vec![vec![1.into()], vec![2.into()]], + ), + ( + JoinType::LeftAnti, + vec![2], + vec![vec![0.into()], vec![3.into()]], + ), + ( + JoinType::RightSemi, + vec![1], + vec![vec![10.into()], vec![12.into()]], + ), + (JoinType::RightAnti, vec![1], vec![vec![11.into()]]), + ( + JoinType::LeftMark, + vec![2, 3], + vec![ + vec![0.into(), false.into()], + vec![1.into(), true.into()], + vec![2.into(), true.into()], + vec![3.into(), false.into()], + ], + ), + ( + JoinType::RightMark, + vec![1, 2], + vec![ + vec![10.into(), true.into()], + vec![11.into(), false.into()], + vec![12.into(), true.into()], + ], + ), + ]; + for use_perfect_hash in [false, true] { + for (join_type, projection, expected) in &cases { + let join = join_builder( + &build, + (Arc::clone(&probe_schema), vec![probe_batch.clone()]), + *join_type, + projection.clone(), + )? + .build()?; + let output = collect_join(join, 2, use_perfect_hash).await?; + // Probe-preserving outer alignment adds the unmatched key 3 to + // the lookup's two matches. The coalescer does not split that batch. + let max_batch_rows = if matches!(join_type, JoinType::Right | JoinType::Full) + { + 3 + } else { + 2 + }; + assert!( + output + .iter() + .all(|batch| batch.num_rows() <= max_batch_rows), + "{join_type:?}, perfect hash {use_perfect_hash}: {:?}", + output.iter().map(RecordBatch::num_rows).collect::>(), + ); + assert_rows_unordered(scalar_rows(&output)?, expected); + } + } + Ok(()) +} + +#[tokio::test] +async fn residual_membership_joins_evaluate_both_build_batches() -> Result<()> { + let build = with_payload( + padded_build(vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 4])), + ])?, + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + )?; + let filter_schema = Arc::new(Schema::new(vec![Field::new( + "build_id", + DataType::Int32, + false, + )])); + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + col("build_id", &filter_schema)?, + Operator::GtEq, + lit(2i32), + )), + vec![ColumnIndex { + index: 2, + side: JoinSide::Left, + }], + filter_schema, + ); + let cases = [ + (JoinType::LeftSemi, vec![2], vec![vec![2.into()]]), + ( + JoinType::LeftAnti, + vec![2], + vec![vec![0.into()], vec![1.into()], vec![3.into()]], + ), + ( + JoinType::RightSemi, + vec![0], + vec![vec![2.into()], vec![2.into()]], + ), + (JoinType::RightAnti, vec![0], vec![vec![3.into()]]), + ( + JoinType::LeftMark, + vec![2, 3], + vec![ + vec![0.into(), false.into()], + vec![1.into(), false.into()], + vec![2.into(), true.into()], + vec![3.into(), false.into()], + ], + ), + ( + JoinType::RightMark, + vec![0, 1], + vec![ + vec![2.into(), true.into()], + vec![3.into(), false.into()], + vec![2.into(), true.into()], + ], + ), + ]; + for use_perfect_hash in [false, true] { + for (join_type, projection, expected) in &cases { + let join = join_builder( + &build, + probe(Arc::new(Int32Array::from(vec![2, 3, 2])))?, + *join_type, + projection.clone(), + )? + .with_filter(Some(filter.clone())) + .build()?; + let output = collect_join(join, 2, use_perfect_hash).await?; + assert_batch_size(&output, 2); + assert_rows_unordered(scalar_rows(&output)?, expected); + } + } + Ok(()) +} + +#[tokio::test] +async fn all_null_build_keys_keep_preserved_rows() -> Result<()> { + let build = with_payload( + padded_build(vec![ + Arc::new(Int32Array::from(vec![None, None])), + Arc::new(Int32Array::from(vec![None])), + ])?, + Arc::new(Int32Array::from(vec![10, 11, 12])), + )?; + let cases = [ + (JoinType::Inner, vec![2], vec![]), + ( + JoinType::Full, + vec![2, 3], + vec![ + vec![10.into(), ScalarValue::Int32(None)], + vec![11.into(), ScalarValue::Int32(None)], + vec![12.into(), ScalarValue::Int32(None)], + vec![ScalarValue::Int32(None), 1.into()], + vec![ScalarValue::Int32(None), ScalarValue::Int32(None)], + ], + ), + ( + JoinType::LeftAnti, + vec![2], + vec![vec![10.into()], vec![11.into()], vec![12.into()]], + ), + ( + JoinType::LeftMark, + vec![2, 3], + vec![ + vec![10.into(), false.into()], + vec![11.into(), false.into()], + vec![12.into(), false.into()], + ], + ), + ( + JoinType::RightMark, + vec![0, 1], + vec![ + vec![1.into(), false.into()], + vec![ScalarValue::Int32(None), false.into()], + ], + ), + ]; + for (join_type, projection, expected) in cases { + let join = join_builder( + &build, + probe(Arc::new(Int32Array::from(vec![Some(1), None])))?, + join_type, + projection, + )? + .build()?; + let output = collect_join(join, 2, false).await?; + assert_batch_size(&output, 2); + assert_rows_unordered(scalar_rows(&output)?, &expected); + } + Ok(()) +} + +#[tokio::test] +async fn fixed_size_list_payloads_preserve_child_and_parent_nulls() -> Result<()> { + let mut values = FixedSizeListBuilder::new(Int32Builder::new(), 2); + for (children, valid) in [ + ([Some(10), None], true), + ([Some(20), Some(21)], false), + ([None, Some(31)], true), + ([Some(40), Some(41)], true), + ] { + for child in children { + values.values().append_option(child); + } + values.append(valid); + } + let payload: ArrayRef = Arc::new(values.finish()); + let build = with_payload( + padded_build(vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 4])), + ])?, + Arc::clone(&payload), + )?; + let expected = [Some(1), Some(2), None] + .into_iter() + .map(|row| { + Ok(vec![row.map_or_else( + || ScalarValue::try_from(payload.data_type()), + |row| ScalarValue::try_from_array(payload.as_ref(), row), + )?]) + }) + .collect::>>()?; + for use_perfect_hash in [false, true] { + let join = join_builder( + &build, + probe(Arc::new(Int32Array::from(vec![2, 3])))?, + JoinType::Right, + vec![2], + )? + .build()?; + let output = collect_join(join, 2, use_perfect_hash).await?; + assert_rows_unordered(scalar_rows(&output)?, &expected); + } + Ok(()) +} diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 03387ea574e63..bfa5897f7cc21 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -27,12 +27,14 @@ use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::joins::Map; use crate::joins::MapOffset; use crate::joins::PartitionMode; -use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode}; +use crate::joins::hash_join::exec::{JoinLeftData, MultiBatchBuildData, NullAwareMode}; 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::{ + OnceFut, equal_rows_arr, equal_rows_arr_multi, matchable_join_keys, +}; use crate::stream::EmptyRecordBatchStream; use crate::{ RecordBatchStream, SendableRecordBatchStream, handle_state, @@ -40,7 +42,8 @@ use crate::{ joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, - build_batch_empty_build_side, build_batch_from_indices, + apply_join_filter_to_indices_multi, build_batch_empty_build_side, + build_batch_from_indices, build_batch_from_indices_multi, build_null_aware_left_mark_column, need_produce_result_in_final, }, }; @@ -467,6 +470,7 @@ impl RecordBatchStream for HashJoinStream { pub(super) fn lookup_join_hashmap( build_hashmap: &dyn JoinHashMapType, build_side_values: &[ArrayRef], + multi_batch: Option<&MultiBatchBuildData>, probe_side_values: &[ArrayRef], null_equality: NullEquality, hashes_buffer: &[u64], @@ -492,13 +496,25 @@ pub(super) fn lookup_join_hashmap( // 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, - )?; + let (build_indices, probe_indices) = if let Some(multi_batch) = multi_batch { + let gather_indices = multi_batch.gather_indices(&build_indices_unfiltered); + equal_rows_arr_multi( + &build_indices_unfiltered, + &probe_indices_unfiltered, + multi_batch.values(), + probe_side_values, + &gather_indices, + null_equality, + )? + } else { + equal_rows_arr( + &build_indices_unfiltered, + &probe_indices_unfiltered, + build_side_values, + probe_side_values, + null_equality, + )? + }; // Reclaim buffers *build_indices_buffer = build_indices_unfiltered.into_parts().1.into(); @@ -638,10 +654,18 @@ impl HashJoinStream { // Use the logical null count: a dictionary key whose entry points at a // NULL dictionary value is a NULL key even though the key bitmap has no // physical nulls (`null_count() == 0` but `logical_null_count() > 0`). - let keys_have_null = left_data - .values() - .iter() - .any(|array| array.logical_null_count() > 0); + let keys_have_null = if let Some(multi_batch) = left_data.multi_batch() { + multi_batch + .values() + .iter() + .flatten() + .any(|array| array.logical_null_count() > 0) + } else { + left_data + .values() + .iter() + .any(|array| array.logical_null_count() > 0) + }; let build_data = match self.mode { PartitionMode::Partitioned => PartitionBuildData::Partitioned { @@ -886,6 +910,7 @@ impl HashJoinStream { Map::HashMap(map) => lookup_join_hashmap( map.as_ref(), build_side.left_data.values(), + build_side.left_data.multi_batch(), &state.values, self.null_equality, &self.hashes_buffer, @@ -923,16 +948,29 @@ impl HashJoinStream { // apply join filter if exists let (left_indices, right_indices) = if let Some(filter) = &self.filter { - apply_join_filter_to_indices( - build_side.left_data.batch(), - &state.batch, - left_indices, - right_indices, - filter, - JoinSide::Left, - None, - self.join_type, - )? + if let Some(multi_batch) = build_side.left_data.multi_batch() { + let gather_indices = multi_batch.gather_indices(&left_indices); + apply_join_filter_to_indices_multi( + multi_batch.batches(), + &gather_indices, + &state.batch, + left_indices, + right_indices, + filter, + self.join_type, + )? + } else { + apply_join_filter_to_indices( + build_side.left_data.batch(), + &state.batch, + left_indices, + right_indices, + filter, + JoinSide::Left, + None, + self.join_type, + )? + } } else { (left_indices, right_indices) }; @@ -989,24 +1027,43 @@ impl HashJoinStream { } // Build output batch and push to coalescer - let (build_batch, probe_batch, join_side) = - if self.join_type == JoinType::RightMark { - (&state.batch, build_side.left_data.batch(), JoinSide::Right) - } else { - (build_side.left_data.batch(), &state.batch, JoinSide::Left) - }; + let batch = if self.join_type != JoinType::RightMark + && let Some(multi_batch) = build_side.left_data.multi_batch() + { + let gather_indices = multi_batch.gather_indices(&left_indices); + build_batch_from_indices_multi( + &self.schema, + multi_batch.batches(), + &gather_indices, + &state.batch, + &left_indices, + &right_indices, + &self.column_indices, + self.join_type, + None, + )? + } else { + // RightMark emits only probe columns and a mark. Its aligned left + // indices refer to the probe batch, not to retained build rows. + let (build_batch, probe_batch, join_side) = + if self.join_type == JoinType::RightMark { + (&state.batch, build_side.left_data.batch(), JoinSide::Right) + } else { + (build_side.left_data.batch(), &state.batch, JoinSide::Left) + }; - let batch = build_batch_from_indices( - &self.schema, - build_batch, - probe_batch, - &left_indices, - &right_indices, - &self.column_indices, - join_side, - self.join_type, - None, - )?; + build_batch_from_indices( + &self.schema, + build_batch, + probe_batch, + &left_indices, + &right_indices, + &self.column_indices, + join_side, + self.join_type, + None, + )? + }; let push_status = self.output_buffer.push_batch(batch)?; @@ -1148,17 +1205,32 @@ impl HashJoinStream { // Push this chunk of final indices to output buffer if !left_side.is_empty() { let empty_right_batch = RecordBatch::new_empty(self.right.schema()); - let batch = build_batch_from_indices( - &self.schema, - build_side.left_data.batch(), - &empty_right_batch, - &left_side, - &right_side, - &self.column_indices, - JoinSide::Left, - self.join_type, - mark_column.as_ref(), - )?; + let batch = if let Some(multi_batch) = build_side.left_data.multi_batch() { + let gather_indices = multi_batch.gather_indices(&left_side); + build_batch_from_indices_multi( + &self.schema, + multi_batch.batches(), + &gather_indices, + &empty_right_batch, + &left_side, + &right_side, + &self.column_indices, + self.join_type, + mark_column.as_ref(), + )? + } else { + build_batch_from_indices( + &self.schema, + build_side.left_data.batch(), + &empty_right_batch, + &left_side, + &right_side, + &self.column_indices, + JoinSide::Left, + self.join_type, + mark_column.as_ref(), + )? + }; let push_status = self.output_buffer.push_batch(batch)?; // If limit reached, finish the coalescer and stop emitting @@ -1560,6 +1632,7 @@ fn for_each_scope_match( let (build_indices, probe_indices, next_offset) = lookup_join_hashmap( scope_map, build_scope_values, + None, probe_scope_values, NullEquality::NullEqualsNothing, hashes_buffer, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 16f61d740985f..0b3f547c3c0d5 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -17,6 +17,13 @@ //! Join related functionality used both on logical and physical plans +mod multi_batch; + +pub(crate) use multi_batch::{ + apply_join_filter_to_indices_multi, build_batch_from_indices_multi, + equal_rows_arr_multi, +}; + use std::cmp::{Ordering, min}; use std::collections::HashSet; use std::fmt::{self, Debug}; diff --git a/datafusion/physical-plan/src/joins/utils/multi_batch.rs b/datafusion/physical-plan/src/joins/utils/multi_batch.rs new file mode 100644 index 0000000000000..e9cc5bb8b33b9 --- /dev/null +++ b/datafusion/physical-plan/src/joins/utils/multi_batch.rs @@ -0,0 +1,633 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Gather and compare logical build rows without concatenating the build side. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::iter::once; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, GenericListArray, ListArray, MapArray, MutableArrayData, + OffsetSizeTrait, RecordBatch, StructArray, UInt32Array, UInt64Array, downcast_array, + make_array, new_empty_array, new_null_array, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::compute::{self, take}; +use arrow::datatypes::{FieldRef, Schema}; +use arrow_schema::{ArrowError, DataType, SortOptions}; +use datafusion_common::cast::as_boolean_array; +use datafusion_common::{JoinSide, JoinType, NullEquality, Result, internal_err}; + +use super::{ColumnIndex, JoinFilter, JoinKeyComparator, new_empty_schema_batch}; + +/// Referenced sources in first-use order. Source zero denotes a synthetic null +/// row; other source IDs are one-based indexes into the retained build batches. +struct SelectedBuildSources { + sources: Vec, + indices: Vec<(usize, usize)>, + single_source_indices: Option, +} + +impl SelectedBuildSources { + fn new(indices: &[(usize, usize)]) -> Self { + let mut sources = Vec::new(); + let mut source_indices = HashMap::new(); + let mut last_source = 0; + let mut last_index = 0; + let indices = indices + .iter() + .map(|&(source, row)| { + if source == 0 { + return (0, 0); + } + if source != last_source { + last_index = *source_indices.entry(source).or_insert_with(|| { + sources.push(source - 1); + sources.len() + }); + last_source = source; + } + (last_index, row) + }) + .collect(); + Self { + sources, + indices, + single_source_indices: None, + } + } + + fn gather<'a>( + &mut self, + data_type: &DataType, + mut source: impl FnMut(usize) -> &'a dyn Array, + ) -> Result { + if self.sources.is_empty() { + return Ok(new_null_array(data_type, self.indices.len())); + } + if self.sources.len() == 1 + && !matches!( + data_type, + DataType::Struct(_) + | DataType::List(_) + | DataType::LargeList(_) + | DataType::Map(_, _) + ) + { + let indices = self.single_source_indices.get_or_insert_with(|| { + self.indices + .iter() + .map(|&(source, row)| (source != 0).then_some(row as u64)) + .collect() + }); + return Ok(take(source(self.sources[0]), indices, None)?); + } + let arrays = self + .sources + .iter() + .map(|&index| source(index)) + .collect::>(); + interleave_payload(data_type, &arrays, &self.indices) + } +} + +/// Gather nested nulls without copying their hidden child values. Other Arrow +/// layouts use the same interleave/take kernels as the single-batch join. +fn interleave_payload( + data_type: &DataType, + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result { + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + let nulls: NullBuffer = indices + .iter() + .map(|&(source, row)| source != 0 && values[source - 1].is_valid(row)) + .collect(); + let indices: Cow<'_, [(usize, usize)]> = if nulls.null_count() == 0 { + Cow::Borrowed(indices) + } else { + Cow::Owned( + indices + .iter() + .enumerate() + .map(|(index, &row)| if nulls.is_valid(index) { row } else { (0, 0) }) + .collect(), + ) + }; + let nulls = (nulls.null_count() != 0).then_some(nulls); + match data_type { + DataType::Struct(fields) => { + let arrays = values.iter().map(|a| a.as_struct()).collect::>(); + let children = fields + .iter() + .enumerate() + .map(|(index, field)| { + let children = arrays + .iter() + .map(|array| array.column(index).as_ref()) + .collect::>(); + interleave_payload(field.data_type(), &children, &indices) + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new_with_length( + fields.clone(), + children, + nulls, + indices.len(), + )?)) + } + DataType::List(field) => { + let arrays = values + .iter() + .map(|a| a.as_list::()) + .collect::>(); + Ok(Arc::new(interleave_list(field, &arrays, &indices, nulls)?)) + } + DataType::LargeList(field) => { + let arrays = values + .iter() + .map(|a| a.as_list::()) + .collect::>(); + Ok(Arc::new(interleave_list(field, &arrays, &indices, nulls)?)) + } + DataType::Map(field, ordered) => { + let lists = values + .iter() + .map(|array| ListArray::from(array.as_map().clone())) + .collect::>(); + let arrays = lists.iter().collect::>(); + let (_, offsets, entries, nulls) = + interleave_list(field, &arrays, &indices, nulls)?.into_parts(); + Ok(Arc::new(MapArray::try_new( + Arc::clone(field), + offsets, + entries.as_struct().clone(), + nulls, + *ordered, + )?)) + } + _ => { + let sentinel = new_null_array(data_type, 1); + let arrays = once(sentinel.as_ref()) + .chain(values.iter().copied()) + .collect::>(); + Ok(compute::interleave(&arrays, &indices)?) + } + } +} + +/// Fixed-width list children can be copied by range without allocating an index +/// for every child, which would dwarf Boolean and zero-width child payloads. +fn fixed_width_max_buffer(data_type: &DataType, rows: usize) -> Result> { + let overflow = + || ArrowError::MemoryError(format!("Gather capacity overflow for {data_type}")); + let bytes = match data_type { + DataType::Struct(fields) => { + let mut maximum = 0; + for field in fields { + let Some(bytes) = fixed_width_max_buffer(field.data_type(), rows)? else { + return Ok(None); + }; + maximum = maximum.max(bytes); + } + maximum + } + DataType::Null | DataType::Boolean => 0, + DataType::FixedSizeBinary(width) => { + let width = usize::try_from(*width).map_err(|_| overflow())?; + let bytes = rows.checked_mul(width).ok_or_else(overflow)?; + if bytes > i32::MAX as usize { + return Err(ArrowError::OffsetOverflowError(bytes).into()); + } + bytes + } + _ => { + let Some(width) = data_type.primitive_width() else { + return Ok(None); + }; + rows.checked_mul(width).ok_or_else(overflow)? + } + }; + let bytes = bytes.max(rows.div_ceil(8)); + if bytes > (isize::MAX as usize & !63) { + return Err(overflow().into()); + } + Ok(Some(bytes)) +} + +fn interleave_list( + field: &FieldRef, + arrays: &[&GenericListArray], + indices: &[(usize, usize)], + nulls: Option, +) -> Result> { + let mut offsets = Vec::with_capacity(indices.len() + 1); + let mut child_count = 0usize; + offsets.push(O::usize_as(0)); + for &(source, row) in indices { + if source != 0 { + let source_offsets = arrays[source - 1].value_offsets(); + let len = source_offsets[row + 1].as_usize() - source_offsets[row].as_usize(); + child_count = child_count + .checked_add(len) + .ok_or(ArrowError::OffsetOverflowError(usize::MAX))?; + } + offsets.push( + O::from_usize(child_count) + .ok_or(ArrowError::OffsetOverflowError(child_count))?, + ); + } + let children = if fixed_width_max_buffer(field.data_type(), child_count)?.is_some() { + let data = arrays + .iter() + .map(|array| array.values().to_data()) + .collect::>(); + let mut output = MutableArrayData::new(data.iter().collect(), false, child_count); + for &(source, row) in indices { + if source != 0 { + let source_offsets = arrays[source - 1].value_offsets(); + output.try_extend( + source - 1, + source_offsets[row].as_usize(), + source_offsets[row + 1].as_usize(), + )?; + } + } + make_array(output.freeze()) + } else { + let mut child_indices = Vec::with_capacity(child_count); + for &(source, row) in indices { + if source != 0 { + let source_offsets = arrays[source - 1].value_offsets(); + let start = source_offsets[row].as_usize(); + let end = source_offsets[row + 1].as_usize(); + child_indices.extend((start..end).map(|row| (source, row))); + } + } + let children = arrays + .iter() + .map(|array| array.values().as_ref()) + .collect::>(); + interleave_payload(field.data_type(), &children, &child_indices)? + }; + Ok(GenericListArray::::try_new( + Arc::clone(field), + OffsetBuffer::new(offsets.into()), + children, + nulls, + )?) +} + +#[expect(clippy::too_many_arguments)] +pub(crate) fn build_batch_from_indices_multi( + schema: &Schema, + build_batches: &[RecordBatch], + gather_indices: &[(usize, usize)], + probe_batch: &RecordBatch, + build_indices: &UInt64Array, + probe_indices: &UInt32Array, + column_indices: &[ColumnIndex], + join_type: JoinType, + mark_column: Option<&ArrayRef>, +) -> Result { + if schema.fields().is_empty() { + let row_count = match join_type { + JoinType::RightAnti | JoinType::RightSemi => probe_indices.len(), + _ => build_indices.len(), + }; + return new_empty_schema_batch(schema, row_count); + } + let mut selection = None; + let columns = column_indices + .iter() + .map(|column| { + if column.side == JoinSide::None { + return match mark_column { + Some(mark) => Ok(Arc::clone(mark)), + None => { + Ok(Arc::new(compute::is_not_null(probe_indices)?) as ArrayRef) + } + }; + } + if column.side == JoinSide::Left { + let data_type = build_batches[0].column(column.index).data_type(); + if build_indices.null_count() == build_indices.len() { + Ok(new_null_array(data_type, build_indices.len())) + } else { + selection + .get_or_insert_with(|| SelectedBuildSources::new(gather_indices)) + .gather(data_type, |index| { + build_batches[index].column(column.index).as_ref() + }) + } + } else { + let array = probe_batch.column(column.index); + if probe_indices.null_count() == probe_indices.len() { + Ok(new_null_array(array.data_type(), probe_indices.len())) + } else { + Ok(take(array.as_ref(), probe_indices, None)?) + } + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) +} + +pub(crate) fn apply_join_filter_to_indices_multi( + build_batches: &[RecordBatch], + gather_indices: &[(usize, usize)], + probe_batch: &RecordBatch, + build_indices: UInt64Array, + probe_indices: UInt32Array, + filter: &JoinFilter, + join_type: JoinType, +) -> Result<(UInt64Array, UInt32Array)> { + if build_indices.is_empty() && probe_indices.is_empty() { + return Ok((build_indices, probe_indices)); + } + let intermediate = build_batch_from_indices_multi( + filter.schema(), + build_batches, + gather_indices, + probe_batch, + &build_indices, + &probe_indices, + filter.column_indices(), + join_type, + None, + )?; + let filter_result = filter + .expression() + .evaluate(&intermediate)? + .into_array(intermediate.num_rows())?; + let mask = as_boolean_array(&filter_result)?; + let left = compute::filter(&build_indices, mask)?; + let right = compute::filter(&probe_indices, mask)?; + Ok(( + downcast_array(left.as_ref()), + downcast_array(right.as_ref()), + )) +} + +/// Compare keys directly in the referenced batches using the existing join +/// comparator, including its floating-point and logical-null semantics. +pub(crate) fn equal_rows_arr_multi( + indices_left: &UInt64Array, + indices_right: &UInt32Array, + left_arrays: &[Vec], + right_arrays: &[ArrayRef], + gather_indices: &[(usize, usize)], + null_equality: NullEquality, +) -> Result<(UInt64Array, UInt32Array)> { + if indices_left.len() != indices_right.len() + || indices_left.len() != gather_indices.len() + { + return internal_err!("Cannot compare join indices with different lengths"); + } + if indices_left.is_empty() || right_arrays.is_empty() { + return Ok((Vec::::new().into(), Vec::::new().into())); + } + let selection = SelectedBuildSources::new(gather_indices); + let sort_options = vec![SortOptions::default(); right_arrays.len()]; + let comparators = selection + .sources + .iter() + .map(|&index| { + let arrays = &left_arrays[index]; + if arrays.len() != right_arrays.len() { + return internal_err!( + "Cannot compare join keys with different column counts" + ); + } + JoinKeyComparator::new(arrays, right_arrays, &sort_options, null_equality) + }) + .collect::>>()?; + let mut left_filtered = Vec::with_capacity(indices_left.len()); + let mut right_filtered = Vec::with_capacity(indices_right.len()); + for ((&left, &right), &(source, row)) in indices_left + .values() + .iter() + .zip(indices_right.values()) + .zip(&selection.indices) + { + // Equality candidates come from the hash table, before outer padding. + debug_assert_ne!(source, 0); + if comparators[source - 1].is_equal(row, right as usize) { + left_filtered.push(left); + right_filtered.push(right); + } + } + Ok((left_filtered.into(), right_filtered.into())) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + DictionaryArray, FixedSizeListArray, Float64Array, Int8Array, Int32Array, + ListViewArray, PrimitiveRunBuilder, StringArray, UnionArray, + }; + use arrow::datatypes::{Field, Int8Type, Int32Type, UnionFields}; + + #[test] + fn gather_only_borrows_selected_sources() -> Result<()> { + let first = StringArray::from(vec!["a", "b"]); + let second = StringArray::from(vec!["c", "d"]); + let mut borrowed = Vec::new(); + let result = SelectedBuildSources::new(&[(4097, 1), (2, 0), (4097, 0), (0, 0)]) + .gather(&DataType::Utf8, |source| { + borrowed.push(source); + match source { + 4096 => &first, + 1 => &second, + _ => panic!("unreferenced source {source}"), + } + })?; + assert_eq!(borrowed, vec![4096, 1]); + assert_eq!( + result.as_ref(), + &StringArray::from(vec![Some("b"), Some("c"), Some("a"), None]) as &dyn Array + ); + Ok(()) + } + + #[test] + fn gather_nested_nulls_omits_hidden_children() -> Result<()> { + let child: ArrayRef = Arc::new(StringArray::from(vec!["hidden", "visible"])); + let field = Arc::new(Field::new("item", DataType::Utf8, true)); + let list = ListArray::try_new( + field, + OffsetBuffer::new(vec![0, 1, 2].into()), + child, + Some(NullBuffer::from(vec![false, true])), + )?; + let result = SelectedBuildSources::new(&[(1, 0), (1, 1), (0, 0), (1, 0)]) + .gather(list.data_type(), |_| &list)?; + let result = result.as_list::(); + assert_eq!(result.value_offsets(), &[0, 0, 1, 1, 1]); + assert_eq!(result.null_count(), 3); + assert_eq!( + result.values().as_ref(), + &StringArray::from(vec!["visible"]) as &dyn Array + ); + Ok(()) + } + + #[test] + fn gather_dictionary_payload_matches_concat_take() -> Result<()> { + let first: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), None, Some(1)]), + Arc::new(StringArray::from(vec!["first", "shared"])), + )?); + let second: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), Some(1), Some(0)]), + Arc::new(StringArray::from(vec!["shared", "second"])), + )?); + let arrays = [first, second]; + let combined = + compute::concat(&arrays.iter().map(AsRef::as_ref).collect::>())?; + let indices = UInt64Array::from(vec![Some(4), Some(2), None, Some(1), Some(3)]); + let expected = take(combined.as_ref(), &indices, None)?; + let actual = SelectedBuildSources::new(&[(2, 1), (1, 2), (0, 0), (1, 1), (2, 0)]) + .gather(arrays[0].data_type(), |source| arrays[source].as_ref())?; + let actual = compute::cast(actual.as_ref(), &DataType::Utf8)?; + let expected = compute::cast(expected.as_ref(), &DataType::Utf8)?; + assert_eq!(actual.as_ref(), expected.as_ref()); + Ok(()) + } + + #[test] + fn gather_encoded_payloads_matches_concat_take() -> Result<()> { + let fixed: ArrayRef = + Arc::new(FixedSizeListArray::from_iter_primitive::( + [ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3), None]), + Some(vec![Some(4), Some(5)]), + ], + 2, + )); + let lists = ListArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5)]), + ]); + let views: ArrayRef = Arc::new(ListViewArray::from(lists)); + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([Some(1), Some(1), None, Some(2)]); + let runs: ArrayRef = Arc::new(builder.finish()); + let union: ArrayRef = Arc::new(UnionArray::try_new( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + )?, + vec![0, 1, 0, 1].into(), + Some(vec![0, 0, 1, 1].into()), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(StringArray::from(vec!["a", "b"])), + ], + )?); + for array in [fixed, views, runs, union] { + let sources = [array.slice(0, 2), array.slice(2, 2)]; + let indices = + UInt64Array::from(vec![Some(3), Some(0), None, Some(2), Some(1)]); + let combined = + compute::concat(&sources.iter().map(AsRef::as_ref).collect::>())?; + let expected = take(combined.as_ref(), &indices, None)?; + let actual = + SelectedBuildSources::new(&[(2, 1), (1, 0), (0, 0), (2, 0), (1, 1)]) + .gather(array.data_type(), |source| sources[source].as_ref())?; + assert_eq!(actual.as_ref(), expected.as_ref(), "{}", array.data_type()); + } + Ok(()) + } + + #[test] + fn equality_matches_contiguous_float_dictionary_and_composite_keys() -> Result<()> { + let floats: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(-0.0), + Some(f64::NAN), + None, + Some(0.0), + Some(-f64::NAN), + None, + ])); + let dictionary: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), Some(1), Some(2), Some(0), Some(1), None]), + Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])), + )?); + let second_key: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 1, 0, 3])); + let build_indices = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 0, 4]); + let probe_indices = UInt32Array::from(vec![3, 4, 5, 0, 1, 2, 0, 4]); + let gather = [ + (1, 0), + (1, 1), + (1, 2), + (2, 0), + (2, 1), + (2, 2), + (1, 0), + (2, 1), + ]; + for first_key in [floats, dictionary] { + for composite in [false, true] { + let mut keys = vec![Arc::clone(&first_key)]; + if composite { + keys.push(Arc::clone(&second_key)); + } + let batches = [0, 3] + .iter() + .map(|&start| { + keys.iter().map(|array| array.slice(start, 3)).collect() + }) + .collect::>(); + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let expected = super::super::equal_rows_arr( + &build_indices, + &probe_indices, + &keys, + &keys, + null_equality, + )?; + let actual = equal_rows_arr_multi( + &build_indices, + &probe_indices, + &batches, + &keys, + &gather, + null_equality, + )?; + assert_eq!(actual, expected); + } + } + } + Ok(()) + } +} From 5e3418baf95f1e7f9e951b9c95f50e02711fcd00 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 00:23:57 +0000 Subject: [PATCH 2/6] fix: bound batched join comparison scratch and account metadata --- .../src/joins/hash_join/exec/build_data.rs | 48 ++++- datafusion/physical-plan/src/joins/utils.rs | 115 +++++++++--- .../src/joins/utils/multi_batch.rs | 176 +++++++++++++++--- 3 files changed, 289 insertions(+), 50 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs index 3446b5cd741e2..27c0ddf0a1609 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs @@ -116,10 +116,16 @@ impl MultiBatchBuildData { reservation.try_grow(metadata_size)?; metrics.build_mem_used.add(metadata_size); - let values = batches - .iter() - .map(|batch| evaluate_expressions_to_arrays(on_left, batch)) - .collect::>>()?; + // Fallible iterator collection can grow geometrically. Allocate the + // capacities admitted above, including the common single-key case. + let mut values = Vec::with_capacity(batches.len()); + for batch in &batches { + let mut keys = Vec::with_capacity(on_left.len()); + for expr in on_left { + keys.push(expr.evaluate(batch)?.into_array_of_size(batch.num_rows())?); + } + values.push(keys); + } let mut batch_offsets = Vec::with_capacity(batches.len() + 1); batch_offsets.push(0); for batch in &batches { @@ -365,6 +371,40 @@ mod tests { counter.memory_usage() } + #[test] + fn key_metadata_capacities_match_reservation() -> Result<()> { + // Non-power-of-two counts expose spare capacity from fallible collect. + for (batch_count, key_count) in [(1, 1), (3, 1), (5, 3)] { + let batches = (0..batch_count) + .map(|_| primitive_batch(1025)) + .collect::>(); + let keys = (0..key_count) + .map(|_| Arc::new(Column::new("a", 0)) as PhysicalExprRef) + .collect::>(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(4096)); + let reservation = MemoryConsumer::new("test").register(&pool); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let data = + MultiBatchBuildData::try_new(batches, &keys, &reservation, &metrics)?; + let metadata_bytes = data.batches.capacity() * size_of::() + + data.values.capacity() * size_of::>() + + data + .values + .iter() + .map(|keys| keys.capacity()) + .sum::() + * size_of::() + + (data.batch_offsets.capacity() + data.row_directory.capacity()) + * size_of::(); + assert_eq!(reservation.size(), metadata_bytes); + assert_eq!(metrics.build_mem_used.value(), metadata_bytes); + drop(data); + drop(reservation); + assert_eq!(pool.reserved(), 0); + } + Ok(()) + } + #[test] fn small_slices_of_large_shared_input_stay_compact() -> Result<()> { let parent = primitive_batch((MAX_COMPACT_BUILD_BYTES + 8) / 8); diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 0b3f547c3c0d5..e611b167c9337 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -2469,6 +2469,31 @@ pub struct JoinKeyComparator { rest: Vec, } +/// Probe-side preprocessing shared when one probe batch is compared against +/// several retained build batches. Ordinary pairwise comparators need no cache. +struct PreparedJoinKeyProbe { + columns: Vec<(ArrayRef, Option)>, + null_equality: NullEquality, +} + +impl PreparedJoinKeyProbe { + fn new(arrays: &[ArrayRef], null_equality: NullEquality) -> Self { + let columns = arrays + .iter() + .map(|array| { + let nulls = (null_equality == NullEquality::NullEqualsNothing) + .then(|| array.logical_nulls().filter(|nulls| nulls.null_count() > 0)) + .flatten(); + (normalize_float_zero(array), nulls) + }) + .collect(); + Self { + columns, + null_equality, + } + } +} + impl JoinKeyComparator { /// Build comparators for each join key column pair. pub fn new( @@ -2480,7 +2505,7 @@ impl JoinKeyComparator { debug_assert_eq!(left_arrays.len(), right_arrays.len()); debug_assert_eq!(left_arrays.len(), sort_options.len()); - let mut iter = left_arrays + let iter = left_arrays .iter() .zip(right_arrays.iter()) .zip(sort_options.iter()) @@ -2494,28 +2519,76 @@ impl JoinKeyComparator { // valid. let l_norm = normalize_float_zero(l); let r_norm = normalize_float_zero(r); - let inner = make_comparator(l_norm.as_ref(), r_norm.as_ref(), *opts)?; - if null_equality == NullEquality::NullEqualsNothing { - let ln = l.logical_nulls().filter(|n| n.null_count() > 0); - let rn = r.logical_nulls().filter(|n| n.null_count() > 0); - match (ln, rn) { - // Both sides have nulls — wrap to override both-null. - (Some(ln), Some(rn)) => Ok(Box::new(move |i, j| { - if ln.is_null(i) && rn.is_null(j) { - Ordering::Less - } else { - inner(i, j) - } - }) - as DynComparator), - // One side has no nulls — both-null impossible, no wrap. - _ => Ok(inner), - } - } else { - Ok(inner) - } + Self::make_key_comparator(&l_norm, &r_norm, *opts, null_equality, || { + ( + l.logical_nulls().filter(|n| n.null_count() > 0), + r.logical_nulls().filter(|n| n.null_count() > 0), + ) + }) + }); + + Self::from_comparators(iter) + } + + fn new_with_prepared_probe( + left_arrays: &[ArrayRef], + probe: &PreparedJoinKeyProbe, + sort_options: &[SortOptions], + ) -> Result { + debug_assert_eq!(left_arrays.len(), probe.columns.len()); + debug_assert_eq!(left_arrays.len(), sort_options.len()); + + let iter = left_arrays + .iter() + .zip(&probe.columns) + .zip(sort_options) + .map(|((left, (right, right_nulls)), opts)| { + let left_normalized = normalize_float_zero(left); + Self::make_key_comparator( + &left_normalized, + right, + *opts, + probe.null_equality, + || { + ( + left.logical_nulls().filter(|n| n.null_count() > 0), + right_nulls.clone(), + ) + }, + ) }); + Self::from_comparators(iter) + } + fn make_key_comparator( + left: &ArrayRef, + right: &ArrayRef, + options: SortOptions, + null_equality: NullEquality, + logical_nulls: impl FnOnce() -> (Option, Option), + ) -> Result { + let inner = make_comparator(left.as_ref(), right.as_ref(), options)?; + if null_equality == NullEquality::NullEqualsNothing { + match logical_nulls() { + // Both sides have nulls: override Arrow's both-null equality. + (Some(left), Some(right)) => Ok(Box::new(move |i, j| { + if left.is_null(i) && right.is_null(j) { + Ordering::Less + } else { + inner(i, j) + } + })), + // One side has no nulls, so both-null is impossible. + _ => Ok(inner), + } + } else { + Ok(inner) + } + } + + fn from_comparators( + mut iter: impl Iterator>, + ) -> Result { let first = iter.next().expect("join must have at least one key")?; let rest = iter.collect::>>()?; Ok(Self { first, rest }) diff --git a/datafusion/physical-plan/src/joins/utils/multi_batch.rs b/datafusion/physical-plan/src/joins/utils/multi_batch.rs index e9cc5bb8b33b9..0a37e35bc24aa 100644 --- a/datafusion/physical-plan/src/joins/utils/multi_batch.rs +++ b/datafusion/physical-plan/src/joins/utils/multi_batch.rs @@ -34,7 +34,10 @@ use arrow_schema::{ArrowError, DataType, SortOptions}; use datafusion_common::cast::as_boolean_array; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, internal_err}; -use super::{ColumnIndex, JoinFilter, JoinKeyComparator, new_empty_schema_batch}; +use super::{ + ColumnIndex, JoinFilter, JoinKeyComparator, PreparedJoinKeyProbe, + new_empty_schema_batch, +}; /// Referenced sources in first-use order. Source zero denotes a synthetic null /// row; other source IDs are one-based indexes into the retained build batches. @@ -213,11 +216,7 @@ fn fixed_width_max_buffer(data_type: &DataType, rows: usize) -> Result 0, DataType::FixedSizeBinary(width) => { let width = usize::try_from(*width).map_err(|_| overflow())?; - let bytes = rows.checked_mul(width).ok_or_else(overflow)?; - if bytes > i32::MAX as usize { - return Err(ArrowError::OffsetOverflowError(bytes).into()); - } - bytes + rows.checked_mul(width).ok_or_else(overflow)? } _ => { let Some(width) = data_type.primitive_width() else { @@ -407,30 +406,46 @@ pub(crate) fn equal_rows_arr_multi( } let selection = SelectedBuildSources::new(gather_indices); let sort_options = vec![SortOptions::default(); right_arrays.len()]; - let comparators = selection - .sources - .iter() - .map(|&index| { - let arrays = &left_arrays[index]; - if arrays.len() != right_arrays.len() { - return internal_err!( - "Cannot compare join keys with different column counts" - ); - } - JoinKeyComparator::new(arrays, right_arrays, &sort_options, null_equality) - }) - .collect::>>()?; + let probe = PreparedJoinKeyProbe::new(right_arrays, null_equality); + + // Link candidate positions by source, then compact in original order. + // Arrow comparators may build their own logical-null masks even with our + // prepared probe metadata. Retaining one comparator at a time bounds that + // scratch memory, though Arrow still computes those masks for each source. + let mut source_heads = vec![usize::MAX; selection.sources.len()]; + let mut next_positions = Vec::with_capacity(selection.indices.len()); + for (position, &(source, _)) in selection.indices.iter().enumerate() { + // Equality candidates come from the hash table, before outer padding. + debug_assert_ne!(source, 0); + next_positions.push(source_heads[source - 1]); + source_heads[source - 1] = position; + } + let mut equal = vec![false; indices_left.len()]; + for (&index, mut position) in selection.sources.iter().zip(source_heads) { + let arrays = &left_arrays[index]; + if arrays.len() != right_arrays.len() { + return internal_err!( + "Cannot compare join keys with different column counts" + ); + } + let comparator = + JoinKeyComparator::new_with_prepared_probe(arrays, &probe, &sort_options)?; + while position != usize::MAX { + let (_, row) = selection.indices[position]; + equal[position] = + comparator.is_equal(row, indices_right.value(position) as usize); + position = next_positions[position]; + } + } let mut left_filtered = Vec::with_capacity(indices_left.len()); let mut right_filtered = Vec::with_capacity(indices_right.len()); - for ((&left, &right), &(source, row)) in indices_left + for ((&left, &right), equal) in indices_left .values() .iter() .zip(indices_right.values()) - .zip(&selection.indices) + .zip(equal) { - // Equality candidates come from the hash table, before outer padding. - debug_assert_ne!(source, 0); - if comparators[source - 1].is_equal(row, right as usize) { + if equal { left_filtered.push(left); right_filtered.push(right); } @@ -445,7 +460,16 @@ mod tests { DictionaryArray, FixedSizeListArray, Float64Array, Int8Array, Int32Array, ListViewArray, PrimitiveRunBuilder, StringArray, UnionArray, }; - use arrow::datatypes::{Field, Int8Type, Int32Type, UnionFields}; + use arrow::datatypes::{Field, Float64Type, Int8Type, Int32Type, UnionFields}; + + #[test] + fn fixed_width_children_do_not_have_a_32_bit_byte_offset_limit() -> Result<()> { + assert_eq!( + fixed_width_max_buffer(&DataType::FixedSizeBinary(16), 1 << 27)?, + Some(1 << 31), + ); + Ok(()) + } #[test] fn gather_only_borrows_selected_sources() -> Result<()> { @@ -630,4 +654,106 @@ mod tests { } Ok(()) } + + #[test] + fn probe_preprocessing_is_shared_across_many_build_sources() -> Result<()> { + const SOURCES: usize = 64; + let dictionary: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1, 0, 2]), + Arc::new(StringArray::from(vec![Some("a"), None, Some("b")])), + )?); + let right = vec![ + Arc::new(Float64Array::from(vec![ + Some(-0.0), + Some(2.0), + None, + Some(0.0), + ])) as ArrayRef, + Arc::clone(&dictionary), + ]; + let left = vec![ + Arc::new(Float64Array::from(vec![ + Some(0.0), + Some(2.0), + None, + Some(-0.0), + ])) as ArrayRef, + dictionary, + ]; + let probe = PreparedJoinKeyProbe::new(&right, NullEquality::NullEqualsNothing); + let normalized_values = probe.columns[0] + .0 + .as_primitive::() + .values() + .inner() + .clone(); + assert!( + !normalized_values + .ptr_eq(right[0].as_primitive::().values().inner()) + ); + let dictionary_nulls = probe.columns[1].1.as_ref().unwrap().buffer().clone(); + let value_refs = normalized_values.strong_count(); + let null_refs = dictionary_nulls.strong_count(); + let options = vec![SortOptions::default(); right.len()]; + for _ in 0..SOURCES { + let comparator = + JoinKeyComparator::new_with_prepared_probe(&left, &probe, &options)?; + assert_eq!(normalized_values.strong_count(), value_refs + 1); + assert_eq!(dictionary_nulls.strong_count(), null_refs + 1); + assert!(comparator.is_equal(0, 0)); + assert!(!comparator.is_equal(1, 1)); + assert!(!comparator.is_equal(2, 2)); + assert!(comparator.is_equal(3, 3)); + drop(comparator); + assert_eq!(normalized_values.strong_count(), value_refs); + assert_eq!(dictionary_nulls.strong_count(), null_refs); + } + + let batches = vec![left; SOURCES]; + let contiguous = (0..right.len()) + .map(|column| { + Ok(compute::concat( + &batches + .iter() + .map(|keys| keys[column].as_ref()) + .collect::>(), + )?) + }) + .collect::>>()?; + let mut build_indices = Vec::new(); + let mut probe_indices = Vec::new(); + let mut gather = Vec::new(); + for round in 0..4 { + for source in (0..SOURCES).rev() { + let row = (source + round) % 4; + build_indices.push((source * 4 + row) as u64); + probe_indices.push(round as u32); + gather.push((source + 1, row)); + } + } + let build_indices = UInt64Array::from(build_indices); + let probe_indices = UInt32Array::from(probe_indices); + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let expected = super::super::equal_rows_arr( + &build_indices, + &probe_indices, + &contiguous, + &right, + null_equality, + )?; + let actual = equal_rows_arr_multi( + &build_indices, + &probe_indices, + &batches, + &right, + &gather, + null_equality, + )?; + assert_eq!(actual, expected); + } + Ok(()) + } } From ad9ba6722579c8e17e727bae7d6148390d4eecba Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 00:37:47 +0000 Subject: [PATCH 3/6] test: limit multi-gigabyte allocation boundary check to 64-bit --- datafusion/physical-plan/src/joins/utils/multi_batch.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/physical-plan/src/joins/utils/multi_batch.rs b/datafusion/physical-plan/src/joins/utils/multi_batch.rs index 0a37e35bc24aa..4b3f7bf694ec2 100644 --- a/datafusion/physical-plan/src/joins/utils/multi_batch.rs +++ b/datafusion/physical-plan/src/joins/utils/multi_batch.rs @@ -463,6 +463,7 @@ mod tests { use arrow::datatypes::{Field, Float64Type, Int8Type, Int32Type, UnionFields}; #[test] + #[cfg(target_pointer_width = "64")] fn fixed_width_children_do_not_have_a_32_bit_byte_offset_limit() -> Result<()> { assert_eq!( fixed_width_max_buffer(&DataType::FixedSizeBinary(16), 1 << 27)?, From caf387c5b9601033da32d26a00840b4eb3ba9236 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 01:11:38 +0000 Subject: [PATCH 4/6] Avoid unnecessary batched join copies and gather overhead --- .../benches/hash_join_batches.rs | 101 ++++++-- .../src/joins/hash_join/exec/build_data.rs | 241 +++++++++++++++++- .../src/joins/utils/multi_batch.rs | 185 ++++++++++++-- 3 files changed, 474 insertions(+), 53 deletions(-) diff --git a/datafusion/physical-plan/benches/hash_join_batches.rs b/datafusion/physical-plan/benches/hash_join_batches.rs index bf4046d6aa135..1f2d1d9f42ccb 100644 --- a/datafusion/physical-plan/benches/hash_join_batches.rs +++ b/datafusion/physical-plan/benches/hash_join_batches.rs @@ -22,9 +22,11 @@ //! plan construction, build, probe and output draining are included. Large cases //! retain over 64 MiB of unique Arrow backing. Cases include tiny independent //! batches, shared slices and oversized backing buffers. The small case measures -//! the compact-build path. Every case probes 4096 rows and checks the matched -//! build-row checksum before timing. Inner joins have 50% matches; the outer case -//! has 6.25%. Perfect-hash selection is enabled and disabled over identical inputs. +//! the compact-build path. Probe batches contain 4096 rows; repeated-probe cases +//! reuse the same batch 64 times (16 for one-row wide inputs). Each case checks +//! the matched build-row checksum before timing. Inner joins have 50% matches; +//! the outer case has 6.25%. Perfect-hash selection is enabled and disabled over +//! identical inputs. //! The separate correctness run reports peak reserved bytes, not process RSS; //! timed runs use the normal memory pool without reservation instrumentation. @@ -74,6 +76,7 @@ enum Payload { enum Layout { Independent, Tiny, + OneRow, Sliced, Overallocated, } @@ -81,6 +84,7 @@ enum Layout { struct Workload { build: Vec, probe: RecordBatch, + probe_repetitions: usize, keys: Keys, join_type: JoinType, expected_rows: usize, @@ -111,6 +115,7 @@ impl Workload { payload: Payload, layout: Layout, low_match_outer: bool, + probe_repetitions: usize, ) -> Self { let value = "x".repeat(width); let batch = |start, end| { @@ -154,10 +159,10 @@ impl Workload { .map(|start| backing.slice(start, (rows - start).min(INPUT_BATCH_ROWS))) .collect::>() } else { - let batch_rows = if matches!(layout, Layout::Tiny) { - 64 - } else { - INPUT_BATCH_ROWS + let batch_rows = match layout { + Layout::Tiny => 64, + Layout::OneRow => 1, + _ => INPUT_BATCH_ROWS, }; (0..rows) .step_by(batch_rows) @@ -192,22 +197,25 @@ impl Workload { Self { build, probe, + probe_repetitions, keys, join_type: if low_match_outer { JoinType::Right } else { JoinType::Inner }, - expected_rows: if low_match_outer { - PROBE_ROWS - } else { - PROBE_ROWS / hit_every - }, + expected_rows: probe_repetitions + * if low_match_outer { + PROBE_ROWS + } else { + PROBE_ROWS / hit_every + }, expected_sum: probe_ids .iter() .step_by(hit_every) .map(|&id| i64::from(id)) - .sum(), + .sum::() + * probe_repetitions as i64, } } @@ -227,7 +235,7 @@ impl Workload { None, )?, TestMemoryExec::try_new_exec( - &[vec![self.probe.clone()]], + &[vec![self.probe.clone(); self.probe_repetitions]], right_schema, None, )?, @@ -263,7 +271,7 @@ impl Workload { fn benchmark(c: &mut Criterion) { let runtime = Builder::new_current_thread().enable_all().build().unwrap(); let mut group = c.benchmark_group("hash_join_batches"); - for (name, rows, width, keys, payload, layout, outer) in [ + for (name, rows, width, keys, payload, layout, outer, probe_repetitions) in [ ( "small_plain", 8192, @@ -272,6 +280,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Independent, false, + 1, ), ( "large_plain", @@ -281,6 +290,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Independent, false, + 1, ), ( "large_computed", @@ -290,6 +300,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Independent, false, + 1, ), ( "large_dictionary_keys", @@ -299,6 +310,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Independent, false, + 1, ), ( "large_dictionary_payload", @@ -308,6 +320,7 @@ fn benchmark(c: &mut Criterion) { Payload::Dictionary, Layout::Independent, false, + 1, ), ( "large_list_payload", @@ -317,6 +330,7 @@ fn benchmark(c: &mut Criterion) { Payload::List, Layout::Independent, false, + 1, ), ( "large_tiny_batches", @@ -326,6 +340,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Tiny, false, + 1, ), ( "large_sliced", @@ -335,6 +350,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Sliced, false, + 1, ), ( "large_overallocated", @@ -344,6 +360,7 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Overallocated, false, + 1, ), ( "large_low_match_outer", @@ -353,9 +370,61 @@ fn benchmark(c: &mut Criterion) { Payload::Plain, Layout::Independent, true, + 1, + ), + ( + "large_plain_many_probe", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Independent, + false, + 64, + ), + ( + "large_tiny_batches_many_probe", + 65536, + 1024, + Keys::Integer, + Payload::Plain, + Layout::Tiny, + false, + 64, + ), + ( + "large_dictionary_keys_many_probe", + 65536, + 1024, + Keys::Dictionary, + Payload::Plain, + Layout::Independent, + false, + 64, + ), + ( + "large_dictionary_payload_many_probe", + 65536, + 1024, + Keys::Integer, + Payload::Dictionary, + Layout::Independent, + false, + 64, + ), + ( + "large_one_row_wide_many_probe", + 4096, + 16384, + Keys::Integer, + Payload::Plain, + Layout::OneRow, + false, + 16, ), ] { - let workload = Workload::new(rows, width, keys, payload, layout, outer); + let workload = + Workload::new(rows, width, keys, payload, layout, outer, probe_repetitions); let mut counter = RecordBatchMemoryCounter::new(); for batch in &workload.build { counter.count_batch(batch); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs index 27c0ddf0a1609..b4fc519f53c55 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs @@ -23,6 +23,8 @@ pub(super) const MAX_COMPACT_BUILD_BYTES: usize = 64 * 1024 * 1024; const TARGET_BUILD_BATCH_BYTES: usize = 8 * 1024 * 1024; const TARGET_BUILD_BATCH_ROWS: usize = 8192; const BUILD_ROW_DIRECTORY_STRIDE: usize = 1024; +// Like OrderedArrayAgg, use average bytes per array to amortize fixed overhead. +const MIN_BUILD_BATCH_BYTES_PER_COLUMN: usize = 4 * 1024; /// Keep the compact path for small logical inputs, even when their slices pin /// much larger allocations. Unsupported estimates only disable this optimization. @@ -168,9 +170,17 @@ impl MultiBatchBuildData { .map(|index| { index.map_or((0, 0), |row| { let row = row as usize; - let mut batch = self.row_directory[row / BUILD_ROW_DIRECTORY_STRIDE]; - while self.batch_offsets[batch + 1] <= row { - batch += 1; + let directory_index = row / BUILD_ROW_DIRECTORY_STRIDE; + let mut batch = self.row_directory[directory_index]; + if self.batch_offsets[batch + 1] <= row { + // Wide one-row batches can put many boundaries in one + // directory bucket. Search only that bucket's offsets. + let end = self + .row_directory + .get(directory_index + 1) + .map_or(self.batches.len(), |next| next + 1); + batch += self.batch_offsets[batch + 1..end] + .partition_point(|offset| *offset <= row); } (batch + 1, row - self.batch_offsets[batch]) }) @@ -240,8 +250,9 @@ impl MultiBatchBuildData { } } -/// Coalesce independent flat inputs a bounded group at a time. Shared buffers -/// stay intact: replacing one slice must not release another slice's charge. +/// Coalesce metadata-heavy independent flat inputs a bounded group at a time. +/// Larger batches stay intact unless their backing allocations need repacking. +/// Shared buffers stay intact: replacing one slice must not release another's charge. pub(super) fn coalesce_build_batches( schema: &SchemaRef, mut batches: Vec, @@ -285,12 +296,18 @@ pub(super) fn coalesce_build_batches( let mut pending_reserved_bytes = 0usize; let mut pending_copy_bytes = 0usize; let mut pending_rows = 0usize; + let min_batch_bytes = + MIN_BUILD_BATCH_BYTES_PER_COLUMN.saturating_mul(schema.fields().len().max(1)); for batch in batches { let reserved_bytes = get_record_batch_memory_size(&batch); let copy_bytes = estimate_batch_concat_allocation(&batch).unwrap_or(usize::MAX); let rows = batch.num_rows(); + let preserve = copy_bytes >= min_batch_bytes + && !should_repack_build_batch(schema, reserved_bytes, copy_bytes); if !pending.is_empty() - && (pending_copy_bytes.saturating_add(copy_bytes) > TARGET_BUILD_BATCH_BYTES + && (preserve + || pending_copy_bytes.saturating_add(copy_bytes) + > TARGET_BUILD_BATCH_BYTES || pending_rows.saturating_add(rows) > TARGET_BUILD_BATCH_ROWS) { output.push(coalesce_build_group( @@ -305,6 +322,10 @@ pub(super) fn coalesce_build_batches( pending_copy_bytes = 0; pending_rows = 0; } + if preserve { + output.push(batch); + continue; + } pending_reserved_bytes += reserved_bytes; pending_copy_bytes = pending_copy_bytes.saturating_add(copy_bytes); pending_rows += rows; @@ -331,12 +352,7 @@ fn coalesce_build_group( reservation: &mut MemoryReservation, metrics: &BuildProbeJoinMetrics, ) -> Result { - if batches.len() == 1 - && copy_bytes <= TARGET_BUILD_BATCH_BYTES - && reserved_bytes > copy_bytes.saturating_mul(2) - && !schema.fields().iter().any(|field| { - matches!(field.data_type(), DataType::Utf8View | DataType::BinaryView) - }) + if batches.len() == 1 && should_repack_build_batch(schema, reserved_bytes, copy_bytes) { // Arrow's single-input concat is zero-copy. A second, empty slice // forces a bounded copy without allocating another input buffer. @@ -345,13 +361,27 @@ fn coalesce_build_group( concat_build_batches(schema, batches, false, reserved_bytes, reservation, metrics) } +fn should_repack_build_batch( + schema: &SchemaRef, + reserved_bytes: usize, + copy_bytes: usize, +) -> bool { + copy_bytes <= TARGET_BUILD_BATCH_BYTES + && reserved_bytes > copy_bytes.saturating_mul(2) + && !schema.fields().iter().any(|field| { + matches!(field.data_type(), DataType::Utf8View | DataType::BinaryView) + }) +} + #[cfg(test)] mod tests { use super::*; use arrow::array::{ - FixedSizeListArray, Int64Array, LargeListArray, ListArray, MapArray, StructArray, + FixedSizeListArray, Int64Array, LargeListArray, ListArray, MapArray, StringArray, + StructArray, }; use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int64Type; use arrow_schema::Field; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryPool, @@ -371,6 +401,191 @@ mod tests { counter.memory_usage() } + fn coalesce_for_test(batches: Vec) -> Result> { + coalesce_with_headroom(batches, TARGET_BUILD_BATCH_BYTES) + } + + fn coalesce_with_headroom( + batches: Vec, + headroom: usize, + ) -> Result> { + let bytes = input_bytes(&batches); + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes + headroom)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + let output = coalesce_build_batches( + &batches[0].schema(), + batches, + bytes, + &mut reservation, + &metrics, + )?; + assert_eq!(reservation.size(), input_bytes(&output)); + assert_eq!(metrics.build_mem_used.value(), input_bytes(&output)); + Ok(output) + } + + #[test] + fn coalesce_byte_threshold_scales_with_columns() -> Result<()> { + for columns in [1, 2] { + for (rows, expected_batches) in [(511, 1), (512, 2)] { + let schema = Arc::new(Schema::new( + (0..columns) + .map(|index| { + Field::new(format!("c{index}"), DataType::Int64, false) + }) + .collect::>(), + )); + let batches = (0..2) + .map(|_| { + RecordBatch::try_new( + Arc::clone(&schema), + (0..columns) + .map(|_| { + Arc::new(Int64Array::from(vec![1; rows])) as ArrayRef + }) + .collect(), + ) + }) + .collect::, _>>()?; + assert_eq!( + estimate_batch_concat_allocation(&batches[0])?, + rows * columns * 8 + ); + let originals = batches.clone(); + let output = coalesce_with_headroom( + batches, + if expected_batches == 2 { + 0 + } else { + TARGET_BUILD_BATCH_BYTES + }, + )?; + assert_eq!(output.len(), expected_batches); + assert_eq!( + output.iter().map(RecordBatch::num_rows).sum::(), + 2 * rows + ); + if expected_batches == 2 { + for (original, retained) in originals.iter().zip(&output) { + assert!(Arc::ptr_eq(original.column(0), retained.column(0))); + } + } + } + } + Ok(()) + } + + #[test] + fn coalesce_small_arrays_but_preserve_wide_batches() -> Result<()> { + let narrow = coalesce_for_test(vec![primitive_batch(64), primitive_batch(64)])?; + assert_eq!(narrow.len(), 1); + assert_eq!(narrow[0].num_rows(), 128); + + for (rows, width, expected_batches) in + [(1, 1024, 1), (64, 1024, 2), (1, 16384, 2)] + { + let value = "x".repeat(width); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])); + let batches = (0..2) + .map(|_| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1; rows])), + Arc::new(StringArray::from_iter_values( + (0..rows).map(|_| value.as_str()), + )), + ], + ) + }) + .collect::, _>>()?; + let output = coalesce_for_test(batches)?; + assert_eq!(output.len(), expected_batches); + assert_eq!( + output.iter().map(RecordBatch::num_rows).sum::(), + 2 * rows + ); + } + Ok(()) + } + + #[test] + fn coalesce_preserves_order_around_retained_batches() -> Result<()> { + let schema = primitive_batch(1).schema(); + let mut offset = 0i64; + let batches = [64, 64, 512, 64, 64, 512, 64] + .into_iter() + .map(|rows| { + let values = Int64Array::from_iter_values(offset..offset + rows); + offset += rows; + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values)]) + }) + .collect::, _>>()?; + let output = coalesce_for_test(batches)?; + assert_eq!( + output.iter().map(RecordBatch::num_rows).collect::>(), + vec![128, 512, 128, 512, 64] + ); + let values = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(values, (0..offset).collect::>()); + Ok(()) + } + + #[test] + fn coalesce_preserves_shared_small_buffers() -> Result<()> { + let parent = primitive_batch(128); + let batches = vec![parent.slice(0, 64), parent.slice(64, 64)]; + let originals = batches.clone(); + let output = coalesce_for_test(batches)?; + assert_eq!(output.len(), 2); + for (original, retained) in originals.iter().zip(&output) { + assert!(Arc::ptr_eq(original.column(0), retained.column(0))); + } + Ok(()) + } + + #[test] + fn gather_indices_searches_small_batch_boundaries() -> Result<()> { + let mut lengths = vec![1; 2 * BUILD_ROW_DIRECTORY_STRIDE + 1]; + lengths.extend([700, 700, 700, 3000, 1, 1, 2048, 2]); + let mut expected = vec![None]; + for (batch, &rows) in lengths.iter().enumerate() { + expected.extend((0..rows).map(|row| Some((batch + 1, row)))); + } + let batches = lengths.into_iter().map(primitive_batch).collect(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let reservation = MemoryConsumer::new("test").register(&pool); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let data = MultiBatchBuildData::try_new(batches, &[], &reservation, &metrics)?; + let indices = UInt64Array::from_iter( + std::iter::once(None).chain((0..data.num_rows()).map(|row| Some(row as u64))), + ); + assert_eq!( + data.gather_indices(&indices), + expected + .into_iter() + .map(|pair| pair.unwrap_or((0, 0))) + .collect::>() + ); + Ok(()) + } + #[test] fn key_metadata_capacities_match_reservation() -> Result<()> { // Non-power-of-two counts expose spare capacity from fallible collect. diff --git a/datafusion/physical-plan/src/joins/utils/multi_batch.rs b/datafusion/physical-plan/src/joins/utils/multi_batch.rs index 4b3f7bf694ec2..43b88d43187e8 100644 --- a/datafusion/physical-plan/src/joins/utils/multi_batch.rs +++ b/datafusion/physical-plan/src/joins/utils/multi_batch.rs @@ -18,7 +18,6 @@ //! Gather and compare logical build rows without concatenating the build side. use std::borrow::Cow; -use std::collections::HashMap; use std::iter::once; use std::sync::Arc; @@ -33,6 +32,7 @@ use arrow::datatypes::{FieldRef, Schema}; use arrow_schema::{ArrowError, DataType, SortOptions}; use datafusion_common::cast::as_boolean_array; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, internal_err}; +use hashbrown::HashMap; use super::{ ColumnIndex, JoinFilter, JoinKeyComparator, PreparedJoinKeyProbe, @@ -120,22 +120,37 @@ fn interleave_payload( if indices.is_empty() { return Ok(new_empty_array(data_type)); } - let nulls: NullBuffer = indices - .iter() - .map(|&(source, row)| source != 0 && values[source - 1].is_valid(row)) - .collect(); - let indices: Cow<'_, [(usize, usize)]> = if nulls.null_count() == 0 { - Cow::Borrowed(indices) + // Arrow handles primitive nulls directly. Variable-width nulls still need + // normalization to avoid copying hidden payload, as do nested null parents. + let normalize_nulls = match data_type { + DataType::Struct(_) + | DataType::List(_) + | DataType::LargeList(_) + | DataType::Map(_, _) => true, + DataType::Boolean => false, + data_type if data_type.primitive_width().is_some() => false, + _ => values.iter().any(|array| array.null_count() > 0), + }; + let (indices, nulls): (Cow<'_, [(usize, usize)]>, _) = if normalize_nulls { + let nulls: NullBuffer = indices + .iter() + .map(|&(source, row)| source != 0 && values[source - 1].is_valid(row)) + .collect(); + let indices = if nulls.null_count() == 0 { + Cow::Borrowed(indices) + } else { + Cow::Owned( + indices + .iter() + .enumerate() + .map(|(index, &row)| if nulls.is_valid(index) { row } else { (0, 0) }) + .collect(), + ) + }; + (indices, (nulls.null_count() != 0).then_some(nulls)) } else { - Cow::Owned( - indices - .iter() - .enumerate() - .map(|(index, &row)| if nulls.is_valid(index) { row } else { (0, 0) }) - .collect(), - ) + (Cow::Borrowed(indices), None) }; - let nulls = (nulls.null_count() != 0).then_some(nulls); match data_type { DataType::Struct(fields) => { let arrays = values.iter().map(|a| a.as_struct()).collect::>(); @@ -188,7 +203,13 @@ fn interleave_payload( )?)) } _ => { - let sentinel = new_null_array(data_type, 1); + // An unused null sentinel would force Arrow to build a null bitmap + // even when all selected arrays and output rows are non-null. + let sentinel = if indices.iter().any(|(source, _)| *source == 0) { + new_null_array(data_type, 1) + } else { + new_empty_array(data_type) + }; let arrays = once(sentinel.as_ref()) .chain(values.iter().copied()) .collect::>(); @@ -457,8 +478,8 @@ pub(crate) fn equal_rows_arr_multi( mod tests { use super::*; use arrow::array::{ - DictionaryArray, FixedSizeListArray, Float64Array, Int8Array, Int32Array, - ListViewArray, PrimitiveRunBuilder, StringArray, UnionArray, + BooleanArray, DictionaryArray, FixedSizeListArray, Float64Array, Int8Array, + Int32Array, ListViewArray, PrimitiveRunBuilder, StringArray, UnionArray, }; use arrow::datatypes::{Field, Float64Type, Int8Type, Int32Type, UnionFields}; @@ -472,6 +493,81 @@ mod tests { Ok(()) } + #[test] + fn gather_flat_nonnull_omits_null_bitmap() -> Result<()> { + let first = Int32Array::from(vec![1, 2]); + let second = Int32Array::from(vec![3, 4]); + let indices = [(2, 1), (1, 0), (2, 0)]; + let result = interleave_payload(&DataType::Int32, &[&first, &second], &indices)?; + assert_eq!( + result.as_ref(), + &Int32Array::from(vec![4, 1, 3]) as &dyn Array + ); + assert!(result.nulls().is_none()); + + let first = StringArray::from(vec!["a", "b"]); + let second = StringArray::from(vec!["c", "d"]); + let result = interleave_payload(&DataType::Utf8, &[&first, &second], &indices)?; + assert_eq!( + result.as_ref(), + &StringArray::from(vec!["d", "a", "c"]) as &dyn Array + ); + assert!(result.nulls().is_none()); + Ok(()) + } + + #[test] + fn gather_flat_physical_nulls_and_padding() -> Result<()> { + let first = Int32Array::from(vec![Some(1), None]); + let second = Int32Array::from(vec![None, Some(4)]); + for indices in [ + vec![(1, 1), (2, 1), (2, 0), (1, 0)], + vec![(1, 1), (2, 1), (0, 0), (1, 0)], + ] { + let result = + interleave_payload(&DataType::Int32, &[&first, &second], &indices)?; + assert_eq!( + result.as_ref(), + &Int32Array::from(vec![None, Some(4), None, Some(1)]) as &dyn Array + ); + } + + let first = BooleanArray::from(vec![Some(true), None]); + let second = BooleanArray::from(vec![Some(false)]); + let result = interleave_payload( + &DataType::Boolean, + &[&first, &second], + &[(1, 1), (2, 0), (0, 0), (1, 0)], + )?; + assert_eq!( + result.as_ref(), + &BooleanArray::from(vec![None, Some(false), None, Some(true)]) as &dyn Array + ); + Ok(()) + } + + #[test] + fn gather_flat_variable_width_nulls_omit_hidden_payload() -> Result<()> { + let first = StringArray::new( + OffsetBuffer::new(vec![0, 6, 13].into()), + arrow::buffer::Buffer::from(b"hiddenvisible".as_slice()), + Some(NullBuffer::from(vec![false, true])), + ); + let second = StringArray::from(vec!["other"]); + let result = interleave_payload( + &DataType::Utf8, + &[&first, &second], + &[(1, 0), (2, 0), (1, 1), (0, 0)], + )?; + assert_eq!( + result.as_ref(), + &StringArray::from(vec![None, Some("other"), Some("visible"), None]) + as &dyn Array + ); + assert_eq!(result.as_string::().value_data(), b"othervisible"); + Ok(()) + } + #[test] fn gather_only_borrows_selected_sources() -> Result<()> { let first = StringArray::from(vec!["a", "b"]); @@ -539,6 +635,33 @@ mod tests { Ok(()) } + #[test] + fn gather_dictionary_value_nulls_without_padding() -> Result<()> { + let first: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1]), + Arc::new(StringArray::from(vec![Some("a"), None])), + )?); + let second: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1]), + Arc::new(StringArray::from(vec![None, Some("b")])), + )?); + for array in [&first, &second] { + assert_eq!(array.null_count(), 0); + assert_eq!(array.logical_nulls().unwrap().null_count(), 1); + } + let actual = interleave_payload( + first.data_type(), + &[first.as_ref(), second.as_ref()], + &[(2, 0), (1, 0), (2, 1), (1, 1)], + )?; + let actual = compute::cast(actual.as_ref(), &DataType::Utf8)?; + assert_eq!( + actual.as_ref(), + &StringArray::from(vec![None, Some("a"), Some("b"), None]) as &dyn Array + ); + Ok(()) + } + #[test] fn gather_encoded_payloads_matches_concat_take() -> Result<()> { let fixed: ArrayRef = @@ -578,15 +701,29 @@ mod tests { )?); for array in [fixed, views, runs, union] { let sources = [array.slice(0, 2), array.slice(2, 2)]; - let indices = - UInt64Array::from(vec![Some(3), Some(0), None, Some(2), Some(1)]); let combined = compute::concat(&sources.iter().map(AsRef::as_ref).collect::>())?; - let expected = take(combined.as_ref(), &indices, None)?; - let actual = - SelectedBuildSources::new(&[(2, 1), (1, 0), (0, 0), (2, 0), (1, 1)]) + // Without padding, logical nulls in encoded arrays and non-null + // selections use an empty sentinel instead of a synthetic null row. + for indices in [ + vec![Some(3), Some(0), None, Some(2), Some(1)], + vec![Some(3), Some(0), Some(2), Some(1)], + vec![Some(3), Some(0)], + ] { + let indices = UInt64Array::from(indices); + let expected = take(combined.as_ref(), &indices, None)?; + let gather = indices + .iter() + .map(|index| { + index.map_or((0, 0), |row| { + (row as usize / 2 + 1, row as usize % 2) + }) + }) + .collect::>(); + let actual = SelectedBuildSources::new(&gather) .gather(array.data_type(), |source| sources[source].as_ref())?; - assert_eq!(actual.as_ref(), expected.as_ref(), "{}", array.data_type()); + assert_eq!(actual.as_ref(), expected.as_ref(), "{}", array.data_type()); + } } Ok(()) } From da836c4de80b2b66592de39e53c225a427976b1f Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 02:08:11 +0000 Subject: [PATCH 5/6] perf: compact eligible dictionary builds for repeated probes --- .../benches/hash_join_batches.rs | 25 +- .../physical-plan/src/joins/hash_join/exec.rs | 50 +- .../src/joins/hash_join/exec/build_data.rs | 598 +++++++++++++++++- 3 files changed, 652 insertions(+), 21 deletions(-) diff --git a/datafusion/physical-plan/benches/hash_join_batches.rs b/datafusion/physical-plan/benches/hash_join_batches.rs index 1f2d1d9f42ccb..6f4926047acd2 100644 --- a/datafusion/physical-plan/benches/hash_join_batches.rs +++ b/datafusion/physical-plan/benches/hash_join_batches.rs @@ -69,6 +69,7 @@ enum Keys { enum Payload { Plain, Dictionary, + UniqueDictionary, List, } @@ -121,12 +122,18 @@ impl Workload { let batch = |start, end| { let ids = (start..end).map(|id| id as i32).collect::>(); let key = key_array(&ids, keys); - let values: ArrayRef = Arc::new(StringArray::from_iter_values( - ids.iter().map(|_| value.as_str()), - )); + let values: ArrayRef = if matches!(payload, Payload::UniqueDictionary) { + Arc::new(StringArray::from_iter_values( + ids.iter().map(|id| format!("{id:08}{}", &value[8..])), + )) + } else { + Arc::new(StringArray::from_iter_values( + ids.iter().map(|_| value.as_str()), + )) + }; let values: ArrayRef = match payload { Payload::Plain => values, - Payload::Dictionary => Arc::new( + Payload::Dictionary | Payload::UniqueDictionary => Arc::new( DictionaryArray::::try_new( Int32Array::from_iter_values(0..ids.len() as i32), values, @@ -412,6 +419,16 @@ fn benchmark(c: &mut Criterion) { false, 64, ), + ( + "large_unique_dictionary_payload_many_probe", + 65536, + 1024, + Keys::Integer, + Payload::UniqueDictionary, + Layout::Independent, + false, + 64, + ), ( "large_one_row_wide_many_probe", 4096, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1b7e383b8582e..38ac456cfdfb2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -108,7 +108,9 @@ mod prepared; pub use prepared::PreparedHashJoinBuild; mod build_data; pub(super) use build_data::MultiBatchBuildData; -use build_data::{coalesce_build_batches, should_preserve_batches}; +use build_data::{ + coalesce_build_batches, should_preserve_batches, try_compact_dictionary_build, +}; /// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions. pub(crate) const HASH_JOIN_SEED: SeededRandomState = @@ -3163,13 +3165,23 @@ async fn collect_left_input( null_equality, )? { let batch = if preserve_batches { - multi_batch = Some(MultiBatchBuildData::try_new( - std::mem::take(&mut batches), - &on_left, - &reservation, + if let Some(batch) = try_compact_dictionary_build( + &schema, + &mut batches, + false, + &mut reservation, &metrics, - )?); - multi_batch.as_ref().unwrap().batches()[0].clone() + ) { + batch + } else { + multi_batch = Some(MultiBatchBuildData::try_new( + std::mem::take(&mut batches), + &on_left, + &reservation, + &metrics, + )?); + multi_batch.as_ref().unwrap().batches()[0].clone() + } } else if prepared { concat_batches(&schema, batches.iter())? } else { @@ -3242,14 +3254,24 @@ async fn collect_left_input( // Match the logical row order used while populating the hash table. let batch = if preserve_batches { - batches.reverse(); - multi_batch = Some(MultiBatchBuildData::try_new( - std::mem::take(&mut batches), - &on_left, - &reservation, + if let Some(batch) = try_compact_dictionary_build( + &schema, + &mut batches, + true, + &mut reservation, &metrics, - )?); - multi_batch.as_ref().unwrap().batches()[0].clone() + ) { + batch + } else { + batches.reverse(); + multi_batch = Some(MultiBatchBuildData::try_new( + std::mem::take(&mut batches), + &on_left, + &reservation, + &metrics, + )?); + multi_batch.as_ref().unwrap().batches()[0].clone() + } } else if prepared { concat_batches(&schema, batches.iter().rev())? } else { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs index b4fc519f53c55..cf19708dea833 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/build_data.rs @@ -26,6 +26,139 @@ const BUILD_ROW_DIRECTORY_STRIDE: usize = 1024; // Like OrderedArrayAgg, use average bytes per array to amortize fixed overhead. const MIN_BUILD_BATCH_BYTES_PER_COLUMN: usize = 4 * 1024; +/// Dictionary output gathers otherwise unify their value domains on every +/// probe batch. Prefer paying that cost once when the existing compact-copy +/// reservation fits, without copying unrelated variable-width payload columns. +/// Ineligible layouts retain the generic batched implementation. +fn prefers_compact_dictionary(schema: &Schema) -> bool { + let mut has_dictionary = false; + for field in schema.fields() { + let supported = match field.data_type() { + DataType::Dictionary(_, values) => { + has_dictionary = true; + values.primitive_width().is_some() + || matches!( + values.as_ref(), + DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8View + | DataType::BinaryView + ) + } + data_type => { + data_type.primitive_width().is_some() + || matches!(data_type, DataType::Boolean | DataType::Null) + } + }; + if !supported { + return false; + } + } + has_dictionary +} + +fn dictionary_concat_preflight(schema: &Schema, batches: &[RecordBatch]) -> bool { + for (column, field) in schema.fields().iter().enumerate() { + let DataType::Dictionary(_, values_type) = field.data_type() else { + continue; + }; + // Arrow's fallback shifts hidden dictionary keys too. Arbitrary + // values under physical nulls can overflow that arithmetic. + if batches + .iter() + .any(|batch| batch.column(column).null_count() > 0) + { + return false; + } + let limit = match values_type.as_ref() { + DataType::Utf8 | DataType::Binary => i32::MAX as usize, + DataType::Utf8View | DataType::BinaryView => u32::MAX as usize, + _ => continue, + }; + let mut total = 0usize; + for batch in batches { + let values = batch.column(column).as_any_dictionary().values(); + let size = match values_type.as_ref() { + DataType::Utf8 => offset_span(values.as_string::().offsets()).1, + DataType::Binary => offset_span(values.as_binary::().offsets()).1, + DataType::Utf8View => values.as_string_view().data_buffers().len(), + DataType::BinaryView => values.as_binary_view().data_buffers().len(), + _ => unreachable!(), + }; + let Some(size) = total.checked_add(size).filter(|size| *size <= limit) else { + return false; + }; + total = size; + } + } + true +} + +/// Attempt the ordinary compact representation without consuming the fallback +/// inputs. As in `concat_build_batches`, admission estimates output buffers, +/// not Arrow's internal kernel scratch. Failure keeps batches and charges intact. +pub(super) fn try_compact_dictionary_build( + schema: &SchemaRef, + batches: &mut Vec, + reverse: bool, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Option { + // Arrow's dictionary concat fallback can panic on child offset overflow. + // Conservatively bound the domains even when deduplication could fit them. + if !prefers_compact_dictionary(schema) + || !dictionary_concat_preflight(schema, batches) + { + return None; + } + let mut input_counter = RecordBatchMemoryCounter::new(); + let mut copy_bytes = 0usize; + for batch in batches.iter() { + input_counter.count_batch(batch); + if batches.len() > 1 { + for array in batch.columns() { + copy_bytes = copy_bytes + .checked_add(estimate_concat_allocation(array.as_ref()).ok()?)?; + } + } + } + // Coalescing may have removed empty inputs and released their buffers. + // Recount only this optional path, without including hash-table charges. + let input_bytes = input_counter.memory_usage(); + reservation.try_grow(copy_bytes).ok()?; + metrics.build_mem_used.add(copy_bytes); + + let compact = if reverse { + concat_batches(schema, batches.iter().rev()) + } else { + concat_batches(schema, batches.iter()) + }; + let Ok(compact) = compact else { + reservation.shrink(copy_bytes); + metrics.build_mem_used.sub(copy_bytes); + return None; + }; + + let new_bytes = input_counter.count_batch(&compact); + let extra_bytes = new_bytes.saturating_sub(copy_bytes); + if reservation.try_grow(extra_bytes).is_err() { + drop(compact); + reservation.shrink(copy_bytes); + metrics.build_mem_used.sub(copy_bytes); + return None; + } + metrics.build_mem_used.add(extra_bytes); + let retained_bytes = get_record_batch_memory_size(&compact); + batches.clear(); + let released_bytes = input_bytes + copy_bytes + extra_bytes - retained_bytes; + reservation.shrink(released_bytes); + metrics.build_mem_used.sub(released_bytes); + Some(compact) +} + /// Keep the compact path for small logical inputs, even when their slices pin /// much larger allocations. Unsupported estimates only disable this optimization. pub(super) fn should_preserve_batches( @@ -377,11 +510,11 @@ fn should_repack_build_batch( mod tests { use super::*; use arrow::array::{ - FixedSizeListArray, Int64Array, LargeListArray, ListArray, MapArray, StringArray, - StructArray, + DictionaryArray, FixedSizeListArray, Int8Array, Int64Array, LargeListArray, + ListArray, MapArray, StringArray, StructArray, }; use arrow::buffer::OffsetBuffer; - use arrow::datatypes::Int64Type; + use arrow::datatypes::{Int8Type, Int64Type}; use arrow_schema::Field; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryPool, @@ -401,6 +534,465 @@ mod tests { counter.memory_usage() } + fn dictionary_batch(values: StringArray, keys: Int8Array, start: i64) -> RecordBatch { + let rows = keys.len(); + let dictionary = Arc::new( + DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("dictionary", dictionary.data_type().clone(), true), + Field::new("id", DataType::Int64, false), + ])); + RecordBatch::try_new( + schema, + vec![ + dictionary, + Arc::new(Int64Array::from_iter_values(start..start + rows as i64)), + ], + ) + .unwrap() + } + + #[test] + fn compact_dictionary_policy_keeps_other_payloads_batched() { + let dictionary = + |values| DataType::Dictionary(Box::new(DataType::Int8), Box::new(values)); + let schema = |types: Vec| { + Schema::new( + types + .into_iter() + .enumerate() + .map(|(index, data_type)| { + Field::new(format!("c{index}"), data_type, true) + }) + .collect::>(), + ) + }; + assert!(prefers_compact_dictionary(&schema(vec![ + dictionary(DataType::Utf8), + DataType::Int64, + DataType::Boolean, + DataType::Null, + ]))); + assert!(!prefers_compact_dictionary(&schema(vec![DataType::Int64]))); + for payload in [ + DataType::Utf8, + DataType::Utf8View, + DataType::FixedSizeBinary(1024), + ] { + assert!(!prefers_compact_dictionary(&schema(vec![ + dictionary(DataType::Utf8), + payload + ]))); + } + for values in [ + dictionary(DataType::Utf8), + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))), + ] { + assert!(!prefers_compact_dictionary(&schema(vec![dictionary( + values + )]))); + } + } + + #[test] + fn compact_dictionary_preserves_order_nulls_and_charges() -> Result<()> { + for reverse in [false, true] { + let mut batches = vec![ + dictionary_batch( + StringArray::from(vec![Some("same"), None, Some("other")]), + Int8Array::from(vec![0, 1, 1, 2]), + 0, + ), + dictionary_batch( + StringArray::from(vec![Some("other"), None, Some("same")]), + Int8Array::from(vec![0, 1, 1, 2]), + 4, + ), + ]; + let schema = batches[0].schema(); + let bytes = input_bytes(&batches); + let unrelated_charge = 4096; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes + unrelated_charge)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes + unrelated_charge); + let expected = if reverse { + concat_batches(&schema, batches.iter().rev())? + } else { + concat_batches(&schema, batches.iter())? + }; + let compact = try_compact_dictionary_build( + &schema, + &mut batches, + reverse, + &mut reservation, + &metrics, + ) + .unwrap(); + assert!(batches.is_empty()); + assert_eq!(compact, expected); + assert_eq!( + reservation.size(), + unrelated_charge + get_record_batch_memory_size(&compact) + ); + assert_eq!(metrics.build_mem_used.value(), reservation.size()); + } + Ok(()) + } + + #[test] + fn compact_dictionary_without_headroom_keeps_inputs() -> Result<()> { + let mut batches = (0..2) + .map(|index| { + dictionary_batch( + StringArray::from(vec!["a", "b"]), + Int8Array::from(vec![0, 1]), + index * 2, + ) + }) + .collect::>(); + let originals = batches.clone(); + let bytes = input_bytes(&batches); + let unrelated_charge = 4096; + let pool: Arc = + Arc::new(GreedyMemoryPool::new(bytes + unrelated_charge)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes + unrelated_charge)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes + unrelated_charge); + assert!( + try_compact_dictionary_build( + &batches[0].schema(), + &mut batches, + true, + &mut reservation, + &metrics, + ) + .is_none() + ); + assert_eq!(reservation.size(), bytes + unrelated_charge); + assert_eq!(metrics.build_mem_used.value(), reservation.size()); + for (batch, original) in batches.iter().zip(originals) { + assert!(Arc::ptr_eq(batch.column(0), original.column(0))); + } + assert_eq!(batches.len(), 2); + Ok(()) + } + + #[test] + fn compact_dictionary_null_keys_keep_original_inputs() -> Result<()> { + let mut batches = vec![ + dictionary_batch( + StringArray::from(vec!["a"]), + Int8Array::from(vec![0; 2]), + 0, + ), + dictionary_batch( + StringArray::from(vec!["b"]), + Int8Array::new( + vec![0, i8::MAX].into(), + Some(arrow::buffer::NullBuffer::from(vec![true, false])), + ), + 2, + ), + ]; + let originals = batches.clone(); + let bytes = input_bytes(&batches); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + assert!( + try_compact_dictionary_build( + &batches[0].schema(), + &mut batches, + false, + &mut reservation, + &metrics, + ) + .is_none() + ); + assert_eq!(batches, originals); + for (batch, original) in batches.iter().zip(originals) { + assert!(Arc::ptr_eq(batch.column(0), original.column(0))); + } + assert_eq!(reservation.size(), bytes); + assert_eq!(metrics.build_mem_used.value(), bytes); + Ok(()) + } + + #[tokio::test] + async fn large_dictionary_join_falls_back_without_copy_headroom() -> Result<()> { + use crate::common; + use crate::joins::{HashJoinExecBuilder, PartitionMode}; + use crate::test::TestMemoryExec; + use arrow::buffer::Buffer; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::expressions::col; + + const DOMAIN_BYTES: usize = 65 * 1024 * 1024; + const LIMIT: usize = 96 * 1024 * 1024; + let values = Buffer::from_vec(vec![b'p'; DOMAIN_BYTES]); + let batches = vec![ + dictionary_batch( + StringArray::new( + OffsetBuffer::new( + vec![0, DOMAIN_BYTES as i32, DOMAIN_BYTES as i32].into(), + ), + values.clone(), + None, + ), + Int8Array::from(vec![0, 1]), + 0, + ), + dictionary_batch( + StringArray::new( + OffsetBuffer::new(vec![0, 0, DOMAIN_BYTES as i32].into()), + values, + None, + ), + Int8Array::from(vec![1, 0]), + 2, + ), + ]; + let schema = batches[0].schema(); + let bytes = input_bytes(&batches); + assert!(bytes > MAX_COMPACT_BUILD_BYTES && bytes < LIMIT); + assert!(should_preserve_batches(&batches, bytes)); + assert!(prefers_compact_dictionary(&schema)); + let copy_bytes = batches.iter().flat_map(RecordBatch::columns).try_fold( + 0usize, + |bytes, array| -> Result { + Ok(bytes + estimate_concat_allocation(array.as_ref())?) + }, + )?; + assert!(copy_bytes > LIMIT - bytes); + + for perfect_hash in [false, true] { + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe", + DataType::Int64, + false, + )])); + let probe = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![Arc::new(Int64Array::from(vec![1, 3]))], + )?; + let join = HashJoinExecBuilder::new( + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, + TestMemoryExec::try_new_exec( + &[vec![probe]], + Arc::clone(&probe_schema), + None, + )?, + vec![(col("id", &schema)?, col("probe", &probe_schema)?)], + JoinType::Inner, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .with_projection(Some(vec![1, 0])) + .build()?; + let mut config = SessionConfig::default().with_batch_size(2); + config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = false; + config + .options_mut() + .execution + .perfect_hash_join_small_build_threshold = + if perfect_hash { usize::MAX } else { 0 }; + config + .options_mut() + .execution + .perfect_hash_join_min_key_density = + if perfect_hash { 0.0 } else { f64::INFINITY }; + let pool: Arc = Arc::new(GreedyMemoryPool::new(LIMIT)); + let context = Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?, + ), + ); + let output = common::collect(join.execute(0, Arc::clone(&context))?).await?; + let ids = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(ids, vec![1, 3]); + for batch in &output { + let payload = arrow::compute::cast(batch.column(1), &DataType::Utf8)?; + assert!( + payload + .as_string::() + .iter() + .all(|value| value == Some("")) + ); + } + { + let mut future = join.left_fut.try_once(|| { + Ok(async { internal_err!("build already initialized") }) + })?; + let build = futures::future::poll_fn(|cx| future.get_shared(cx)).await?; + assert!(build.multi_batch().is_some()); + assert_eq!(build.multi_batch().unwrap().batches().len(), 2); + } + let used_array_map = join + .metrics() + .and_then(|metrics| { + metrics.sum_by_name(ARRAY_MAP_CREATED_COUNT_METRIC_NAME) + }) + .map_or(0, |metric| metric.as_usize()); + assert_eq!(used_array_map > 0, perfect_hash); + drop(output); + drop(join); + drop(context); + assert_eq!(pool.reserved(), 0); + } + Ok(()) + } + + #[test] + fn compact_dictionary_capacity_failure_restores_charges() -> Result<()> { + let mut batches = (0..2) + .map(|batch| { + dictionary_batch( + StringArray::from_iter_values( + (0..100).map(|index| format!("{batch}-{index}")), + ), + Int8Array::from_iter_values(0..100), + batch * 100, + ) + }) + .collect::>(); + let schema = batches[0].schema(); + assert!(concat_batches(&schema, batches.iter()).is_err()); + let originals = batches.clone(); + let bytes = input_bytes(&batches); + let unrelated_charge = 4096; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes + unrelated_charge)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes + unrelated_charge); + assert!( + try_compact_dictionary_build( + &schema, + &mut batches, + false, + &mut reservation, + &metrics, + ) + .is_none() + ); + assert_eq!(reservation.size(), bytes + unrelated_charge); + assert_eq!(metrics.build_mem_used.value(), reservation.size()); + assert_eq!(batches, originals); + for (batch, original) in batches.iter().zip(originals) { + assert!(Arc::ptr_eq(batch.column(0), original.column(0))); + } + Ok(()) + } + + #[test] + fn compact_dictionary_recounts_after_empty_shared_inputs() -> Result<()> { + let parent = dictionary_batch( + StringArray::from(vec!["shared"]), + Int8Array::from(vec![0; 4]), + 0, + ); + let empty_parent = dictionary_batch( + StringArray::from_iter_values(["x".repeat(1024 * 1024)]), + Int8Array::from(vec![0]), + 0, + ); + let batches = vec![ + parent.slice(0, 2), + empty_parent.slice(0, 0), + parent.slice(2, 2), + ]; + drop(parent); + drop(empty_parent); + let bytes = input_bytes(&batches); + let unrelated_charge = 4096; + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes + 8192)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes + unrelated_charge)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes + unrelated_charge); + let schema = batches[0].schema(); + let mut batches = + coalesce_build_batches(&schema, batches, bytes, &mut reservation, &metrics)?; + assert_eq!(batches.len(), 2); + assert!(input_bytes(&batches) < bytes); + let compact = try_compact_dictionary_build( + &schema, + &mut batches, + false, + &mut reservation, + &metrics, + ) + .unwrap(); + assert_eq!(compact.num_rows(), 4); + assert_eq!( + reservation.size(), + unrelated_charge + get_record_batch_memory_size(&compact) + ); + assert_eq!(metrics.build_mem_used.value(), reservation.size()); + Ok(()) + } + + #[test] + fn compact_dictionary_skips_combined_32_bit_offset_overflow() -> Result<()> { + // Reusing a one-MiB domain exceeds the conservative logical concat + // bound without allocating GiBs or attempting an overflowing concat. + let batch = dictionary_batch( + StringArray::from_iter_values(["x".repeat(1024 * 1024)]), + Int8Array::from(vec![0]), + 0, + ); + let schema = batch.schema(); + let mut batches = vec![batch; 2048]; + assert!(!dictionary_concat_preflight(&schema, &batches)); + let bytes = input_bytes(&batches); + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes + 4096)); + let mut reservation = MemoryConsumer::new("test").register(&pool); + reservation.try_grow(bytes)?; + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + metrics.build_mem_used.add(bytes); + assert!( + try_compact_dictionary_build( + &schema, + &mut batches, + false, + &mut reservation, + &metrics, + ) + .is_none() + ); + assert_eq!(batches.len(), 2048); + assert_eq!(reservation.size(), bytes); + assert_eq!(metrics.build_mem_used.value(), bytes); + Ok(()) + } + fn coalesce_for_test(batches: Vec) -> Result> { coalesce_with_headroom(batches, TARGET_BUILD_BATCH_BYTES) } From 0e64452c64659be1c71e90bbec67974119a2ab8a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 30 Sep 2026 03:10:30 +0000 Subject: [PATCH 6/6] perf: avoid recoalescing byte-large batched join output --- datafusion/physical-plan/src/coalesce/mod.rs | 145 +++++++++++++++++ .../physical-plan/src/joins/hash_join/exec.rs | 133 +++++++++++++++ .../src/joins/hash_join/stream.rs | 153 +++++++++++++++++- 3 files changed, 427 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index 70c86368c1477..b403ce5db3cca 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -120,6 +120,33 @@ impl LimitedBatchCoalescer { Ok(PushBatchStatus::Continue) } + /// Emit an already materialized batch without copying it into another batch. + /// Buffered rows are emitted first, preserving order. Fetch-clipped batches + /// keep the ordinary buffering policy because their visible output is smaller. + pub(crate) fn push_batch_without_coalescing( + &mut self, + batch: RecordBatch, + ) -> Result { + assert_or_internal_err!( + !self.finished, + "LimitedBatchCoalescer: cannot push batch after finish" + ); + if batch.num_rows() == 0 + || self.fetch.is_some_and(|fetch| { + fetch.saturating_sub(self.total_rows) < batch.num_rows() + }) + { + return self.push_batch(batch); + } + + self.inner.finish_buffered_batch()?; + let threshold = self.inner.biggest_coalesce_batch_size(); + self.inner.set_biggest_coalesce_batch_size(Some(0)); + let result = self.push_batch(batch); + self.inner.set_biggest_coalesce_batch_size(threshold); + result + } + /// Pushes the next [`RecordBatch`] into the coalescer after applying `filter`, /// avoiding a separate materialization pass compared to calling /// [`filter_record_batch`] followed by [`Self::push_batch`]. @@ -272,6 +299,124 @@ mod tests { .run() } + #[test] + fn bypass_preserves_order_and_restores_coalescing() -> Result<()> { + let prefix = uint32_batch(0..2); + let large = uint32_batch(2..6); + let mut coalescer = LimitedBatchCoalescer::new(prefix.schema(), 16, None); + coalescer.push_batch(prefix)?; + coalescer.push_batch_without_coalescing(large.clone())?; + assert_next_batch_values(&mut coalescer, vec![0, 1]); + let output = coalescer.next_completed_batch().unwrap(); + assert!(Arc::ptr_eq(output.column(0), large.column(0))); + assert_eq!(coalescer.inner.biggest_coalesce_batch_size(), Some(8)); + + coalescer.push_batch(uint32_batch(6..8))?; + coalescer.push_batch(uint32_batch(8..10))?; + assert!(coalescer.next_completed_batch().is_none()); + coalescer.finish()?; + assert_next_batch_values(&mut coalescer, vec![6, 7, 8, 9]); + assert!(coalescer.next_completed_batch().is_none()); + assert!(coalescer.push_batch_without_coalescing(large).is_err()); + assert_eq!(coalescer.inner.biggest_coalesce_batch_size(), Some(8)); + Ok(()) + } + + #[test] + fn bypass_preserves_fetch_boundaries() -> Result<()> { + for fetch in [0, 1, 2, 4, 6, 8] { + let prefix = uint32_batch(0..2); + let mut coalescer = + LimitedBatchCoalescer::new(prefix.schema(), 16, Some(fetch)); + coalescer.push_batch(prefix)?; + if coalescer.push_batch_without_coalescing(uint32_batch(2..6))? + == PushBatchStatus::Continue + { + coalescer.push_batch(uint32_batch(6..10))?; + } + assert_eq!(coalescer.inner.biggest_coalesce_batch_size(), Some(8)); + assert_eq!(coalescer.total_rows, fetch); + coalescer.finish()?; + let mut actual = Vec::new(); + let mut sizes = Vec::new(); + while let Some(batch) = coalescer.next_completed_batch() { + sizes.push(batch.num_rows()); + actual.extend_from_slice( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + assert_eq!(actual, (0..fetch as u32).collect::>()); + let expected_sizes = match fetch { + 0 => vec![], + 1 | 2 | 4 => vec![fetch], + 6 => vec![2, 4], + 8 => vec![2, 4, 2], + _ => unreachable!(), + }; + assert_eq!(sizes, expected_sizes); + } + Ok(()) + } + + #[test] + fn bypass_empty_and_zero_column_batches() -> Result<()> { + use arrow::record_batch::RecordBatchOptions; + + let prefix = uint32_batch(0..2); + let mut coalescer = LimitedBatchCoalescer::new(prefix.schema(), 16, None); + coalescer.push_batch(prefix)?; + coalescer.push_batch_without_coalescing(uint32_batch(2..2))?; + assert!(coalescer.next_completed_batch().is_none()); + coalescer.finish()?; + assert_next_batch_values(&mut coalescer, vec![0, 1]); + + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + )?; + let mut coalescer = LimitedBatchCoalescer::new(schema, 16, None); + coalescer.push_batch_without_coalescing(batch)?; + assert_eq!(coalescer.next_completed_batch().unwrap().num_rows(), 3); + assert_eq!(coalescer.total_rows, 3); + assert_eq!(coalescer.inner.biggest_coalesce_batch_size(), Some(8)); + Ok(()) + } + + #[test] + fn bypass_flush_error_preserves_threshold() -> Result<()> { + use arrow::array::{DictionaryArray, Int8Array, StringArray}; + use arrow::datatypes::Int8Type; + + let batch = |start: i32| { + let dictionary = DictionaryArray::::new( + Int8Array::from_iter_values(0..100), + Arc::new(StringArray::from_iter_values( + (start..start + 100).map(|value| value.to_string()), + )), + ); + RecordBatch::try_from_iter(vec![("d", Arc::new(dictionary) as _)]) + }; + let first = batch(0)?; + let mut coalescer = LimitedBatchCoalescer::new(first.schema(), 1000, None); + coalescer.push_batch(first)?; + coalescer.push_batch(batch(100)?)?; + assert!( + coalescer + .push_batch_without_coalescing(batch(200)?) + .is_err() + ); + assert_eq!(coalescer.inner.biggest_coalesce_batch_size(), Some(500)); + assert_eq!(coalescer.total_rows, 200); + Ok(()) + } + #[test] fn test_push_batch_with_filter_nulls_and_fetch() { let batch = uint32_batch(0..8); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 38ac456cfdfb2..c6aa1326598cc 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3631,6 +3631,139 @@ mod tests { Ok(()) } + #[tokio::test] + async fn retained_build_output_avoids_wide_recoalescing() -> Result<()> { + use datafusion_physical_expr::expressions::col; + + const WIDTH: usize = 128 * 1024; + let payload = "x".repeat(WIDTH); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])); + let batches = [0..256, 256..512] + .into_iter() + .map(|ids| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(ids.clone())), + Arc::new(StringArray::from_iter_values( + ids.map(|_| payload.as_str()), + )), + ], + ) + }) + .collect::, _>>()?; + let mut counter = RecordBatchMemoryCounter::new(); + for batch in &batches { + counter.count_batch(batch); + } + assert!(counter.memory_usage() > 64 * 1024 * 1024); + + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe", + DataType::Int32, + false, + )])); + for use_phj in [false, true] { + // The left join leaves one small probe result buffered before its + // wide unmatched tail. Fetch clips that tail back to ordinary output. + for (join_type, probes, fetch, compact, sizes) in [ + (JoinType::Inner, vec![0..16; 4], None, false, vec![16; 4]), + (JoinType::Inner, vec![0..1; 4], None, false, vec![4]), + (JoinType::Inner, vec![0..16; 4], None, true, vec![64]), + ( + JoinType::Left, + vec![0..478, 478..479], + None, + false, + vec![128, 128, 128, 94, 1, 33], + ), + ( + JoinType::Left, + vec![0..478, 478..479], + Some(483), + false, + vec![128, 128, 128, 94, 5], + ), + ] { + let build_batches = if compact { + batches.iter().map(|batch| batch.slice(0, 16)).collect() + } else { + batches.clone() + }; + let probe_batches = probes + .iter() + .cloned() + .map(|ids| { + RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![Arc::new(Int32Array::from_iter_values(ids))], + ) + }) + .collect::, _>>()?; + let join = HashJoinExecBuilder::new( + TestMemoryExec::try_new_exec( + &[build_batches], + Arc::clone(&schema), + None, + )?, + TestMemoryExec::try_new_exec( + &[probe_batches], + Arc::clone(&probe_schema), + None, + )?, + vec![(col("id", &schema)?, col("probe", &probe_schema)?)], + join_type, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .with_projection(Some(vec![0, 1])) + .with_fetch(fetch) + .build()?; + let context = prepare_task_ctx(128, use_phj); + let output = common::collect(join.execute(0, context)?).await?; + assert_eq!( + output.iter().map(RecordBatch::num_rows).collect::>(), + sizes + ); + let actual = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + let expected = if join_type == JoinType::Left { + (0..fetch.unwrap_or(512) as i32).collect::>() + } else { + probes.into_iter().flatten().collect::>() + }; + assert_eq!(actual, expected); + for batch in &output { + assert!( + batch + .column(1) + .as_string::() + .iter() + .all(|value| value == Some(payload.as_str())) + ); + } + let mut future = join.left_fut.try_once(|| { + Ok(async { internal_err!("build already initialized") }) + })?; + let build = futures::future::poll_fn(|cx| future.get_shared(cx)).await?; + assert_eq!(build.multi_batch().is_some(), !compact); + assert_phj_used(&join.metrics().unwrap(), use_phj); + } + } + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index bfa5897f7cc21..beb5182dc0b9a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -48,9 +48,9 @@ use crate::{ }, }; -use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; +use arrow::array::{Array, ArrayRef, AsArray, UInt32Array, UInt64Array}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, internal_datafusion_err, internal_err, @@ -58,9 +58,58 @@ use datafusion_common::{ use datafusion_physical_expr::PhysicalExprRef; use datafusion_common::hash_utils::RandomState; +use datafusion_common::utils::offset_span_len; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::{Stream, StreamExt, ready}; +/// Avoid a second materialization for byte-large output from retained builds. +const MIN_OUTPUT_COALESCE_COPY_BYTES: usize = 2 * 1024 * 1024; + +fn push_join_output( + output: &mut LimitedBatchCoalescer, + batch: RecordBatch, + multi_batch: bool, +) -> Result { + let mut copy_bytes = 0usize; + let bypass = multi_batch + && batch.columns().iter().any(|array| { + // This allocation-free lower bound omits offsets, null metadata and + // encoded/nested children. View payload buffers remain shared. + let data_type = array.data_type(); + let bytes = if let Some(width) = data_type.primitive_width() { + array.len().saturating_mul(width) + } else { + match data_type { + DataType::Boolean => array.len().div_ceil(8), + DataType::Utf8 => offset_span_len(array.as_string::().offsets()), + DataType::LargeUtf8 => { + offset_span_len(array.as_string::().offsets()) + } + DataType::Binary => { + offset_span_len(array.as_binary::().offsets()) + } + DataType::LargeBinary => { + offset_span_len(array.as_binary::().offsets()) + } + DataType::Utf8View | DataType::BinaryView => { + array.len().saturating_mul(size_of::()) + } + DataType::FixedSizeBinary(width) => array + .len() + .saturating_mul(usize::try_from(*width).unwrap_or_default()), + _ => 0, + } + }; + copy_bytes = copy_bytes.saturating_add(bytes); + copy_bytes >= MIN_OUTPUT_COALESCE_COPY_BYTES + }); + if bypass { + output.push_batch_without_coalescing(batch) + } else { + output.push_batch(batch) + } +} + /// Represents build-side of hash join. pub(super) enum BuildSide { /// Indicates that build-side not collected yet @@ -1065,7 +1114,12 @@ impl HashJoinStream { )? }; - let push_status = self.output_buffer.push_batch(batch)?; + let push_status = push_join_output( + &mut self.output_buffer, + batch, + self.join_type != JoinType::RightMark + && build_side.left_data.multi_batch().is_some(), + )?; timer.done(); @@ -1231,7 +1285,11 @@ impl HashJoinStream { mark_column.as_ref(), )? }; - let push_status = self.output_buffer.push_batch(batch)?; + let push_status = push_join_output( + &mut self.output_buffer, + batch, + build_side.left_data.multi_batch().is_some(), + )?; // If limit reached, finish the coalescer and stop emitting if push_status == PushBatchStatus::LimitReached { @@ -1747,6 +1805,93 @@ mod tests { PushdownStrategy, completed_partitions_for_test, make_partitioned_accumulator_for_test, }; + use arrow::array::{ + DictionaryArray, Int32Array, ListArray, StringArray, StringViewArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Int32Type}; + + #[test] + fn large_multi_batch_output_bypasses_only_at_copy_threshold() -> Result<()> { + for bytes in [ + MIN_OUTPUT_COALESCE_COPY_BYTES - 1, + MIN_OUTPUT_COALESCE_COPY_BYTES, + ] { + for multi_batch in [false, true] { + let array: ArrayRef = + Arc::new(StringArray::from(vec!["x".repeat(bytes)])); + let batch = + RecordBatch::try_from_iter(vec![("payload", Arc::clone(&array))])?; + let mut output = LimitedBatchCoalescer::new(batch.schema(), 128, None); + push_join_output(&mut output, batch, multi_batch)?; + let bypass = multi_batch && bytes >= MIN_OUTPUT_COALESCE_COPY_BYTES; + if bypass { + let batch = output.next_completed_batch().unwrap(); + assert!(Arc::ptr_eq(batch.column(0), &array)); + } else { + assert!(output.next_completed_batch().is_none()); + output.finish()?; + assert_eq!(output.next_completed_batch().unwrap().num_rows(), 1); + } + } + } + Ok(()) + } + + #[test] + fn output_copy_policy_ignores_retained_child_and_view_buffers() -> Result<()> { + let large = "x".repeat(MIN_OUTPUT_COALESCE_COPY_BYTES); + let strings: ArrayRef = + Arc::new(StringArray::from(vec![large.as_str(), "small"])); + let views: ArrayRef = + Arc::new(StringViewArray::from(vec![large.as_str(), "small"])); + let dictionary: ArrayRef = Arc::new(DictionaryArray::::new( + Int32Array::from(vec![1]), + Arc::clone(&strings), + )); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, false)), + OffsetBuffer::new(vec![1i32, 2].into()), + Arc::clone(&strings), + None, + )); + for array in [strings.slice(1, 1), views.slice(1, 1), dictionary, list] { + assert!(array.get_buffer_memory_size() >= MIN_OUTPUT_COALESCE_COPY_BYTES); + let batch = RecordBatch::try_from_iter(vec![("payload", array)])?; + let mut output = LimitedBatchCoalescer::new(batch.schema(), 128, None); + push_join_output(&mut output, batch, true)?; + assert!(output.next_completed_batch().is_none()); + output.finish()?; + assert_eq!(output.next_completed_batch().unwrap().num_rows(), 1); + } + Ok(()) + } + + #[test] + fn output_copy_policy_counts_columns_but_not_zero_column_rows() -> Result<()> { + let payload = + StringArray::from(vec!["x".repeat(MIN_OUTPUT_COALESCE_COPY_BYTES / 2)]); + let batch = RecordBatch::try_from_iter(vec![ + ("a", Arc::new(payload.clone()) as ArrayRef), + ("b", Arc::new(payload) as ArrayRef), + ])?; + let mut output = LimitedBatchCoalescer::new(batch.schema(), 128, None); + push_join_output(&mut output, batch, true)?; + assert_eq!(output.next_completed_batch().unwrap().num_rows(), 1); + + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![], + &arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(1)), + )?; + let mut output = LimitedBatchCoalescer::new(schema, 128, None); + push_join_output(&mut output, batch, true)?; + assert!(output.next_completed_batch().is_none()); + output.finish()?; + assert_eq!(output.next_completed_batch().unwrap().num_rows(), 1); + Ok(()) + } fn empty_build_data(partition_id: usize) -> PartitionBuildData { PartitionBuildData::Partitioned {