You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/13 16:06:03 UTC
(datafusion) branch gh-readonly-queue/main/pr-26182-071b4ddc14d904ae81df20735b2bfe821d58c169 created (now 1d165c5cca)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-26182-071b4ddc14d904ae81df20735b2bfe821d58c169
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 1d165c5cca refactor: retain hash join build batches after allocation failure (#26182)
This branch includes the following new commits:
new 1d165c5cca refactor: retain hash join build batches after allocation failure (#26182)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: refactor: retain hash join build batches after allocation failure (#26182)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26182-071b4ddc14d904ae81df20735b2bfe821d58c169
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1d165c5ccadc3a02a5e482b32315ed757b7dca91
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Tue Oct 13 14:47:42 2026 +0000
refactor: retain hash join build batches after allocation failure (#26182)
## Which issue does this PR close?
Part of #24768. Splits the build refactor out of #25217 and #26015 for
easier review.
## Rationale for this change
The sort-merge fallback needs the build rows after a hash table
allocation fails.
## What changes are included in this PR?
Separate hash-table construction from input collection and retain the
batches when building or concatenating them fails. Also adds Clippy
allowances for `Expr` hash-set keys.
## What is the testing strategy for this PR?
Tests cover retained batches, memory accounting, and existing
memory-limit behavior for ordered and unbounded inputs.
## Are there any user-facing changes?
None. The fallback comes in a follow-up PR.
Co-authored-by: Jay Zhan <ja...@gmail.com>
---
datafusion/optimizer/src/push_down_filter.rs | 4 +
.../physical-plan/src/joins/hash_join/exec.rs | 251 +++++++++++++++++----
2 files changed, 216 insertions(+), 39 deletions(-)
diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs
index 999cddfb62..68ef37ece5 100644
--- a/datafusion/optimizer/src/push_down_filter.rs
+++ b/datafusion/optimizer/src/push_down_filter.rs
@@ -1135,10 +1135,13 @@ impl OptimizerRule for PushDownFilter {
// multiple window functions, each with potentially different partition keys.
// Therefore, we need to ensure that any potential partition key returned is used in
// ALL window functions. Otherwise, filters cannot be pushed by through that column.
+ #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // Expr hashing ignores Arrow buffer reservation tracking.
fn extract_partition_keys(func: &WindowFunction) -> HashSet<&Expr> {
func.params.partition_by.iter().collect()
}
+ #[allow(clippy::allow_attributes, clippy::mutable_key_type)]
+ // Expr hashing ignores Arrow buffer reservation tracking.
let potential_partition_keys = window
.window_expr
.iter()
@@ -1601,6 +1604,7 @@ fn with_filters(predicates: Vec<Expr>, plan: LogicalPlan) -> LogicalPlan {
///
/// A predicate containing a subquery is never pushed: what the subquery reads is
/// not visible from the expression tree.
+#[allow(clippy::allow_attributes, clippy::mutable_key_type)] // Expr hashing ignores Arrow buffer reservation tracking.
fn reads_only_partition_keys(
expr: &Expr,
partition_keys: &HashSet<&Expr>,
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 5c68018bda..eb14d76e8b 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -3154,9 +3154,10 @@ fn estimate_concat_allocation(array: &dyn Array) -> Result<usize> {
/// `inputs_reserved` is the memory `reservation` already holds for `batches`.
/// The copy is reserved before it is made, and once `batches` are dropped the
/// reservation is trimmed to the memory retained by the returned batch.
+/// On success, `batches` contains that batch; errors retain all input rows.
fn concat_build_batches(
schema: &SchemaRef,
- batches: Vec<RecordBatch>,
+ batches: &mut Vec<RecordBatch>,
reverse: bool,
inputs_reserved: usize,
reservation: &mut MemoryReservation,
@@ -3165,7 +3166,7 @@ fn concat_build_batches(
// Concatenating a single batch is zero-copy
let copy_size = if batches.len() > 1 {
let mut copy_size = 0;
- for batch in &batches {
+ for batch in batches.iter() {
for array in batch.columns() {
copy_size += estimate_concat_allocation(array.as_ref())?;
}
@@ -3182,7 +3183,8 @@ fn concat_build_batches(
} else {
concat_batches(schema, batches.iter())?
};
- drop(batches);
+ batches.clear();
+ batches.push(batch.clone());
// The inputs are gone: only hold on to what the concatenated batch retains,
// which includes any buffers it still shares with the inputs.
@@ -3256,11 +3258,6 @@ async fn collect_left_input(
let schema = left_stream.schema();
let prepared = mode == BuildMode::Prepared;
- // The extra scope maps + null bitmap are only built for correlated
- // null-aware joins (see `NullAwareMode`).
- let with_null_aware_row_state = null_aware.is_some_and(NullAwareMode::is_correlated);
- let null_aware_value_keys = null_aware.map_or(1, NullAwareMode::value_keys);
-
let is_phj_candidate = is_perfect_hash_join_candidate(&on_left, &schema)?;
let mut state = BuildSideState::try_new(
@@ -3300,7 +3297,7 @@ async fn collect_left_input(
// Bind the reservation first so error paths release allocations before their charge.
let BuildSideState {
- mut reservation,
+ reservation,
mut batches,
num_rows,
metrics,
@@ -3319,7 +3316,7 @@ async fn collect_left_input(
};
// Compute bounds
- let mut bounds = match bounds_accumulators {
+ let bounds = match bounds_accumulators {
Some(accumulators) if num_rows > 0 => {
let bounds = accumulators
.into_iter()
@@ -3330,30 +3327,89 @@ async fn collect_left_input(
_ => None,
};
+ build_in_memory(
+ &random_state,
+ &schema,
+ &mut batches,
+ num_rows,
+ max_batch_rows,
+ input_bytes,
+ copy_bytes,
+ &on_left,
+ &metrics,
+ reservation,
+ bounds,
+ with_visited_indices_bitmap,
+ probe_threads_count,
+ should_compute_dynamic_filters,
+ &config,
+ null_equality,
+ null_aware,
+ &array_map_created_count,
+ prepared,
+ )
+}
+
+/// Builds the hash table (or perfect-hash [`ArrayMap`]) and the bitmaps over
+/// the collected build `batches`, the second half of [`collect_left_input`].
+///
+/// `reservation` already covers `batches`; the structures built here grow it
+/// further. On error the reservation is dropped, releasing everything, while
+/// `batches` keeps every build row (possibly concatenated into one batch) for
+/// callers that need to recover from a failed build.
+#[expect(clippy::too_many_arguments)]
+fn build_in_memory(
+ random_state: &RandomState,
+ schema: &SchemaRef,
+ batches: &mut Vec<RecordBatch>,
+ num_rows: usize,
+ max_batch_rows: usize,
+ input_bytes: usize,
+ copy_bytes: usize,
+ on_left: &[PhysicalExprRef],
+ metrics: &BuildProbeJoinMetrics,
+ mut reservation: MemoryReservation,
+ mut bounds: Option<PartitionBounds>,
+ with_visited_indices_bitmap: bool,
+ probe_threads_count: usize,
+ should_compute_dynamic_filters: bool,
+ config: &ConfigOptions,
+ null_equality: NullEquality,
+ null_aware: Option<NullAwareMode>,
+ array_map_created_count: &Count,
+ prepared: bool,
+) -> Result<JoinLeftData> {
+ // The extra scope maps + null bitmap are only built for correlated
+ // null-aware joins (see `NullAwareMode`).
+ let with_null_aware_row_state = null_aware.is_some_and(NullAwareMode::is_correlated);
+ let null_aware_value_keys = null_aware.map_or(1, NullAwareMode::value_keys);
+
+ let is_phj_candidate = is_perfect_hash_join_candidate(on_left, schema)?;
+
let (join_hash_map, batch, left_values) = if let Some((min_val, max_val)) =
array_map_key_range(
bounds.as_ref(),
- &schema,
- &batches,
- &on_left,
+ schema,
+ batches,
+ on_left,
&mut reservation,
config.execution.perfect_hash_join_small_build_threshold,
config.execution.perfect_hash_join_min_key_density,
null_equality,
)? {
let batch = if prepared {
- concat_batches(&schema, batches.iter())?
+ concat_batches(schema, batches.iter())?
} else {
concat_build_batches(
- &schema,
- std::mem::take(&mut batches),
+ schema,
+ batches,
false,
input_bytes,
&mut reservation,
- &metrics,
+ metrics,
)?
};
- let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
+ let left_values = evaluate_expressions_to_arrays(on_left, &batch)?;
let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?;
array_map_created_count.add(1);
@@ -3366,7 +3422,7 @@ async fn collect_left_input(
// Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the
// `u64` indice variant
// Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown
- let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?;
+ let mut hashmap = new_join_hashmap(num_rows, &mut reservation, metrics)?;
let scratch_reservation = reservation.new_empty();
if prepared {
@@ -3385,11 +3441,11 @@ async fn collect_left_input(
hashes_buffer.clear();
hashes_buffer.resize(batch.num_rows(), 0);
update_hash(
- &on_left,
+ on_left,
batch,
&mut *hashmap,
offset,
- &random_state,
+ random_state,
&mut hashes_buffer,
0,
true,
@@ -3400,19 +3456,19 @@ async fn collect_left_input(
// Merge all batches into a single batch, so we can directly index into the arrays
let batch = if prepared {
- concat_batches(&schema, batches.iter().rev())?
+ concat_batches(schema, batches.iter().rev())?
} else {
concat_build_batches(
- &schema,
- std::mem::take(&mut batches),
+ schema,
+ batches,
true,
input_bytes,
&mut reservation,
- &metrics,
+ metrics,
)?
};
- let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
+ let left_values = evaluate_expressions_to_arrays(on_left, &batch)?;
(Map::HashMap(hashmap), batch, left_values)
};
@@ -3469,7 +3525,7 @@ async fn collect_left_input(
// Scope-only NULL marking uses a HashMap (the primary join map may
// use ArrayMap for full-key matches, but scope keys have arbitrary
// shape).
- let mut scope_map = new_join_hashmap(num_rows, &mut reservation, &metrics)?;
+ let mut scope_map = new_join_hashmap(num_rows, &mut reservation, metrics)?;
let mut hashes_buffer = vec![0; batch.num_rows()];
update_hash(
@@ -3477,7 +3533,7 @@ async fn collect_left_input(
&batch,
&mut *scope_map,
0,
- &random_state,
+ random_state,
&mut hashes_buffer,
0,
true,
@@ -3513,9 +3569,9 @@ async fn collect_left_input(
None
} else {
let null_rows = build_indices.len();
- let mut map = new_join_hashmap(null_rows, &mut reservation, &metrics)?;
+ let mut map = new_join_hashmap(null_rows, &mut reservation, metrics)?;
let mut hashes_buffer = vec![0; null_rows];
- create_hashes(&scope_values, &random_state, &mut hashes_buffer)?;
+ create_hashes(&scope_values, random_state, &mut hashes_buffer)?;
map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0);
Some(map)
};
@@ -3623,7 +3679,7 @@ async fn collect_left_input(
}
if prepared {
- drop(batches);
+ batches.clear();
// Prepared keys are direct columns. IN-list arrays share these batch
// buffers, including through multi-key StructArray children.
let retained = RecordBatchMemoryCounter::new().count_batch(&batch);
@@ -8522,7 +8578,7 @@ mod tests {
};
let batch = concat_build_batches(
&schema,
- batches.clone(),
+ &mut batches.clone(),
reverse,
inputs_reserved,
&mut reservation,
@@ -8535,7 +8591,7 @@ mod tests {
#[test]
fn concat_build_batches_reserves_copy() -> Result<()> {
- let batches = concat_test_batches(4, 1000);
+ let mut batches = concat_test_batches(4, 1000);
let schema = batches[0].schema();
let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new());
let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum();
@@ -8544,9 +8600,10 @@ mod tests {
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs * 5 / 4));
let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?;
assert_eq!(inputs_reserved, inputs);
+ let original = batches.clone();
let err = concat_build_batches(
&schema,
- batches.clone(),
+ &mut batches,
false,
inputs_reserved,
&mut reservation,
@@ -8554,6 +8611,8 @@ mod tests {
)
.unwrap_err();
assert_contains!(err.to_string(), "Resources exhausted");
+ assert_eq!(batches, original);
+ drop(original);
drop(reservation);
// With room for the copy, the reservation ends up at what is retained
@@ -8561,7 +8620,7 @@ mod tests {
let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?;
let batch = concat_build_batches(
&schema,
- batches,
+ &mut batches,
false,
inputs_reserved,
&mut reservation,
@@ -8569,13 +8628,14 @@ mod tests {
)?;
assert_eq!(reservation.size(), get_record_batch_memory_size(&batch));
assert_eq!(pool.reserved(), reservation.size());
+ assert_eq!(batches, vec![batch]);
Ok(())
}
/// Concatenating a single batch is zero-copy, so nothing more is reserved.
#[test]
fn concat_build_batches_single_batch_not_reserved_twice() -> Result<()> {
- let batches = concat_test_batches(1, 1000);
+ let mut batches = concat_test_batches(1, 1000);
let schema = batches[0].schema();
let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new());
let inputs = get_record_batch_memory_size(&batches[0]);
@@ -8584,7 +8644,7 @@ mod tests {
let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?;
let batch = concat_build_batches(
&schema,
- batches,
+ &mut batches,
true,
inputs_reserved,
&mut reservation,
@@ -8602,7 +8662,7 @@ mod tests {
) -> Result<()> {
let batches = concat_test_batches(4, 1000);
let column = batches[0].schema().index_of(column)?;
- let batches = batches
+ let mut batches = batches
.into_iter()
.map(|batch| batch.project(&[column]))
.collect::<Result<Vec<_>, _>>()?;
@@ -8617,7 +8677,7 @@ mod tests {
let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?;
let batch = concat_build_batches(
&schema,
- batches,
+ &mut batches,
false,
inputs_reserved,
&mut reservation,
@@ -11729,4 +11789,117 @@ mod tests {
assert!(join.set_dynamic_filter(df).is_err());
Ok(())
}
+ fn fallback_inputs() -> Result<(Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>)> {
+ let side = |offset: i32| {
+ let batches = (0..32)
+ .map(|i| {
+ let ids: Vec<i32> = (i * 16..(i + 1) * 16).collect();
+ let keys: Vec<i32> = ids.iter().map(|id| id % 47 + offset).collect();
+ build_table_i32(("id", &ids), ("key", &keys), ("value", &ids))
+ })
+ .collect::<Vec<_>>();
+ TestMemoryExec::try_new_exec(
+ std::slice::from_ref(&batches),
+ batches[0].schema(),
+ None,
+ )
+ .map(|exec| exec as Arc<dyn ExecutionPlan>)
+ };
+ Ok((side(0)?, side(5)?))
+ }
+
+ fn fallback_context(limit: Option<usize>) -> Arc<TaskContext> {
+ let mut runtime = RuntimeEnvBuilder::new();
+ if let Some(limit) = limit {
+ runtime = runtime.with_memory_limit(limit, 1.0);
+ }
+ Arc::new(
+ TaskContext::default()
+ .with_session_config(
+ SessionConfig::default()
+ .with_batch_size(16)
+ .with_sort_spill_reservation_bytes(0),
+ )
+ .with_runtime(runtime.build_arc().unwrap()),
+ )
+ }
+
+ fn fallback_join(
+ left: Arc<dyn ExecutionPlan>,
+ right: Arc<dyn ExecutionPlan>,
+ join_type: JoinType,
+ ) -> Result<HashJoinExec> {
+ HashJoinExecBuilder::new(
+ left,
+ right,
+ vec![(
+ Arc::new(Column::new("key", 1)),
+ Arc::new(Column::new("key", 1)),
+ )],
+ join_type,
+ )
+ .with_partition_mode(PartitionMode::Partitioned)
+ .build()
+ }
+
+ #[tokio::test]
+ async fn partitioned_join_memory_limit() -> Result<()> {
+ let (left, right) = fallback_inputs()?;
+ let reference =
+ fallback_join(Arc::clone(&left), Arc::clone(&right), JoinType::Inner)?;
+ let expected =
+ common::collect(reference.execute(0, fallback_context(None))?).await?;
+ assert!(!expected.is_empty());
+ let join = fallback_join(left, right, JoinType::Inner)?;
+ let error = common::collect(join.execute(0, fallback_context(Some(4 * 1024)))?)
+ .await
+ .unwrap_err();
+ assert_contains!(error.to_string(), "Resources exhausted");
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn sort_merge_fallback_preserves_promised_ordering() -> Result<()> {
+ let (left, right) = fallback_inputs()?;
+ let right = Arc::new(crate::sorts::sort::SortExec::new(
+ [PhysicalSortExpr::new_default(Arc::new(Column::new(
+ "id", 0,
+ )))]
+ .into(),
+ right,
+ ));
+ let join = fallback_join(left, right, JoinType::Inner)?;
+ assert!(join.properties().output_ordering().is_some());
+ let error = common::collect(join.execute(0, fallback_context(Some(4096)))?)
+ .await
+ .unwrap_err();
+ assert_contains!(error.to_string(), "Resources exhausted");
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn memory_limited_join_keeps_streaming_inputs() -> Result<()> {
+ let (left, right) = fallback_inputs()?;
+ let batches = common::collect(right.execute(0, fallback_context(None))?).await?;
+ let right: Arc<dyn ExecutionPlan> =
+ Arc::new(crate::streaming::StreamingTableExec::try_new(
+ right.schema(),
+ vec![Arc::new(
+ crate::test::TestPartitionStream::new_with_batches(batches),
+ )],
+ None,
+ [],
+ true,
+ None,
+ )?);
+ // A hash join can stream this probe. A fallback must not turn it into a
+ // sort that waits forever for an unbounded source to finish.
+ assert!(right.boundedness().is_unbounded());
+ let join = fallback_join(left, right, JoinType::Inner)?;
+ let error = common::collect(join.execute(0, fallback_context(Some(4096)))?)
+ .await
+ .unwrap_err();
+ assert_contains!(error.to_string(), "Resources exhausted");
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org