Skip to content

Commit cab63e4

Browse files
committed
fix: preserve prepared hash join admission checks
1 parent 18b679f commit cab63e4

3 files changed

Lines changed: 198 additions & 20 deletions

File tree

‎datafusion/physical-plan/src/joins/hash_join/exec.rs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3039,16 +3039,11 @@ async fn collect_left_input(
30393039

30403040
let scratch_reservation = reservation.new_empty();
30413041
if prepared {
3042-
// Allow one logical null mask per key plus the combined mask.
3043-
let masks = if null_equality == NullEquality::NullEqualsNothing {
3044-
on_left.len() + 1
3045-
} else {
3046-
0
3047-
};
3048-
scratch_reservation.try_grow(
3049-
max_batch_rows * size_of::<u64>()
3050-
+ (max_batch_rows.div_ceil(8) + 64) * masks,
3051-
)?;
3042+
scratch_reservation.try_grow(prepared::prepared_scratch_bytes(
3043+
max_batch_rows,
3044+
on_left.len(),
3045+
null_equality,
3046+
)?)?;
30523047
}
30533048
// The maximum is known: avoid geometric growth and its excess capacity.
30543049
let mut hashes_buffer = Vec::with_capacity(max_batch_rows);
@@ -3213,10 +3208,15 @@ async fn collect_left_input(
32133208

32143209
if prepared {
32153210
drop(batches);
3211+
// Prepared keys are direct columns. IN-list arrays share these batch
3212+
// buffers, including through multi-key StructArray children.
32163213
let retained = RecordBatchMemoryCounter::new().count_batch(&batch);
3217-
let allowance = input_bytes + copy_bytes;
3218-
debug_assert!(retained <= allowance);
3219-
reservation.shrink(allowance - retained);
3214+
prepared::reconcile_prepared_copy_reservation(
3215+
&reservation,
3216+
input_bytes,
3217+
copy_bytes,
3218+
retained,
3219+
)?;
32203220
}
32213221

32223222
let data = JoinLeftData {

‎datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs‎

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
2020
use super::*;
2121
use crate::spill::spill_manager::GetSlicedSize;
22+
use datafusion_common::exec_datafusion_err;
2223
use datafusion_execution::memory_pool::MemoryPool;
2324

2425
/// An immutable, fully prepared broadcast build, independent of any probe task.
@@ -210,7 +211,8 @@ impl HashJoinExec {
210211
/// Bound copy allocations, including validity, offsets and alignment. Aliased
211212
/// columns count separately because concatenation materializes each column.
212213
pub(super) fn prepared_copy_bytes(batches: &[RecordBatch]) -> Result<usize> {
213-
let mut bytes = 0;
214+
let overflow = || exec_datafusion_err!("Prepared hash-join copy size overflow");
215+
let mut bytes = 0usize;
214216
for batch in batches {
215217
// Concat can materialize validity for an all-valid input when another input
216218
// contains nulls. Arrow's slice measurement only counts existing bitmaps.
@@ -219,17 +221,73 @@ pub(super) fn prepared_copy_bytes(batches: &[RecordBatch]) -> Result<usize> {
219221
.iter()
220222
.filter(|array| array.nulls().is_none())
221223
.count();
222-
bytes +=
223-
batch.get_sliced_size()? + batch.num_rows().div_ceil(8) * missing_validity;
224+
let validity_bytes = batch
225+
.num_rows()
226+
.div_ceil(8)
227+
.checked_mul(missing_validity)
228+
.ok_or_else(overflow)?;
229+
bytes = bytes
230+
.checked_add(batch.get_sliced_size()?)
231+
.and_then(|bytes| bytes.checked_add(validity_bytes))
232+
.ok_or_else(overflow)?;
224233
}
225234
// These are flat arrays: allow one alignment unit per output buffer,
226235
// including potential validity. Concat allocates each buffer only once.
227236
let buffers = batches[0]
228237
.columns()
229238
.iter()
230-
.map(|array| array.to_data().buffers().len() + 1)
231-
.sum::<usize>();
232-
Ok(bytes + buffers * 64)
239+
.try_fold(0usize, |buffers, array| {
240+
buffers
241+
.checked_add(array.to_data().buffers().len())
242+
.and_then(|buffers| buffers.checked_add(1))
243+
.ok_or_else(overflow)
244+
})?;
245+
buffers
246+
.checked_mul(64)
247+
.and_then(|padding| bytes.checked_add(padding))
248+
.ok_or_else(overflow)
249+
}
250+
251+
/// Bound hashes and flat-key null masks for the largest input batch.
252+
pub(super) fn prepared_scratch_bytes(
253+
max_batch_rows: usize,
254+
num_keys: usize,
255+
null_equality: NullEquality,
256+
) -> Result<usize> {
257+
let overflow = || exec_datafusion_err!("Prepared hash-join scratch size overflow");
258+
let hash_bytes = max_batch_rows
259+
.checked_mul(size_of::<u64>())
260+
.ok_or_else(overflow)?;
261+
if null_equality == NullEquality::NullEqualsNull {
262+
return Ok(hash_bytes);
263+
}
264+
// Allow one logical null mask per key plus the combined mask.
265+
let masks = num_keys.checked_add(1).ok_or_else(overflow)?;
266+
let mask_bytes = max_batch_rows
267+
.div_ceil(8)
268+
.checked_add(64)
269+
.and_then(|bytes| bytes.checked_mul(masks))
270+
.ok_or_else(overflow)?;
271+
hash_bytes.checked_add(mask_bytes).ok_or_else(overflow)
272+
}
273+
274+
/// Release the unused copy allowance only after validating the retained size.
275+
pub(super) fn reconcile_prepared_copy_reservation(
276+
reservation: &MemoryReservation,
277+
input_bytes: usize,
278+
copy_bytes: usize,
279+
retained: usize,
280+
) -> Result<()> {
281+
let allowance = input_bytes.checked_add(copy_bytes).ok_or_else(|| {
282+
exec_datafusion_err!("Prepared hash-join copy allowance size overflow")
283+
})?;
284+
if retained > allowance {
285+
return internal_err!(
286+
"Prepared hash-join concat exceeded its admitted copy bound"
287+
);
288+
}
289+
reservation.shrink(allowance - retained);
290+
Ok(())
233291
}
234292

235293
impl HashJoinExecBuilder {

‎datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs‎

Lines changed: 121 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ use crate::{
3030
test::TestMemoryExec,
3131
};
3232
use arrow::array::{
33-
BooleanArray, FixedSizeBinaryArray, Int64Array, StringArray, StringViewArray,
33+
BooleanArray, FixedSizeBinaryArray, Int64Array, NullArray, StringArray,
34+
StringViewArray, StructArray,
3435
};
3536
use arrow::buffer::Buffer;
3637
use arrow::compute::kernels::sort::SortOptions;
@@ -698,6 +699,125 @@ fn prepared_concat_admits_validity_for_non_nullable_input_arrays() -> Result<()>
698699
Ok(())
699700
}
700701

702+
#[tokio::test]
703+
async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> {
704+
let build = batch(vec![Some(1), Some(1_000_000), None, Some(1)]);
705+
for keys in [1, 2] {
706+
let on = (0..keys)
707+
.map(|index| {
708+
let name = build.schema().field(index).name().clone();
709+
let key: PhysicalExprRef = Arc::new(Column::new(&name, index));
710+
(Arc::clone(&key), key)
711+
})
712+
.collect();
713+
let base = join(build.schema(), build.clone())?
714+
.builder()
715+
.with_on(on)
716+
.build()?;
717+
for batches in [
718+
vec![build.clone()],
719+
vec![build.slice(0, 2), build.slice(2, 2)],
720+
] {
721+
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
722+
let prepared = prepare(&base, batches.clone(), Arc::clone(&pool)).await?;
723+
let PushdownStrategy::InList(values) = &prepared.build.membership else {
724+
panic!("expected IN-list membership");
725+
};
726+
let columns = prepared.build.batch.columns();
727+
if keys == 1 {
728+
assert!(Arc::ptr_eq(values, &columns[0]));
729+
} else {
730+
let values = values.as_any().downcast_ref::<StructArray>().unwrap();
731+
for (value, column) in values.columns().iter().zip(columns) {
732+
assert!(Arc::ptr_eq(value, column));
733+
}
734+
}
735+
let bytes = prepared.reserved_bytes();
736+
let mut counter = RecordBatchMemoryCounter::new();
737+
let batch_bytes = counter.count_batch(&prepared.build.batch);
738+
assert!(bytes > batch_bytes, "hash storage must also be reserved");
739+
let membership_batch =
740+
RecordBatch::try_from_iter([("membership", Arc::clone(values))])?;
741+
assert_eq!(counter.count_batch(&membership_batch), 0);
742+
drop(membership_batch);
743+
744+
// Disabling IN-list membership changes no retained key buffers.
745+
let mut config = ConfigOptions::default();
746+
config.optimizer.hash_join_inlist_pushdown_max_size = 0;
747+
let map = base
748+
.prepare_build(
749+
input(batches, build.schema()),
750+
Arc::clone(&pool),
751+
Arc::new(config),
752+
)
753+
.await?;
754+
assert!(matches!(map.build.membership, PushdownStrategy::Map(_)));
755+
assert_eq!(map.reserved_bytes(), bytes);
756+
assert_eq!(pool.reserved(), 2 * bytes);
757+
drop(map);
758+
let lease = Arc::clone(&prepared);
759+
drop(prepared);
760+
assert_eq!(pool.reserved(), bytes);
761+
drop(lease);
762+
assert_eq!(pool.reserved(), 0);
763+
}
764+
}
765+
Ok(())
766+
}
767+
768+
#[test]
769+
fn prepared_concat_rejects_under_admission_without_shrinking() -> Result<()> {
770+
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024));
771+
let reservation = MemoryConsumer::new("prepared concat test").register(&pool);
772+
// 128 bytes of map storage plus 128 bytes of input/copy allowance.
773+
reservation.try_grow(256)?;
774+
let error =
775+
reconcile_prepared_copy_reservation(&reservation, 64, 64, 129).unwrap_err();
776+
assert!(matches!(error, DataFusionError::Internal(_)));
777+
assert!(
778+
error
779+
.to_string()
780+
.contains("exceeded its admitted copy bound")
781+
);
782+
assert_eq!(reservation.size(), 256);
783+
assert!(reconcile_prepared_copy_reservation(&reservation, usize::MAX, 1, 0).is_err());
784+
assert_eq!(reservation.size(), 256);
785+
reconcile_prepared_copy_reservation(&reservation, 64, 64, 128)?;
786+
assert_eq!(reservation.size(), 256);
787+
reconcile_prepared_copy_reservation(&reservation, 64, 64, 80)?;
788+
assert_eq!(reservation.size(), 208);
789+
drop(reservation);
790+
assert_eq!(pool.reserved(), 0);
791+
Ok(())
792+
}
793+
794+
#[test]
795+
fn prepared_sizing_rejects_overflow() -> Result<()> {
796+
use NullEquality::{NullEqualsNothing, NullEqualsNull};
797+
for (rows, keys, equality) in [
798+
(usize::MAX / 8 + 1, 1, NullEqualsNull),
799+
(0, usize::MAX, NullEqualsNothing),
800+
(0, usize::MAX / 64, NullEqualsNothing),
801+
(usize::MAX / 8, 0, NullEqualsNothing),
802+
] {
803+
let error = prepared_scratch_bytes(rows, keys, equality).unwrap_err();
804+
assert!(error.to_string().contains("scratch size overflow"));
805+
}
806+
assert!(prepared_scratch_bytes(usize::MAX / 8, 1, NullEqualsNull).is_ok());
807+
808+
// Null arrays can describe huge row counts without allocating data buffers.
809+
// Exercise both per-batch validity multiplication and cumulative addition.
810+
let array: ArrayRef = Arc::new(NullArray::new(usize::MAX));
811+
for (columns, batches) in [(8, 1), (4, 2)] {
812+
let batch = RecordBatch::try_from_iter(
813+
(0..columns).map(|i| (format!("key{i}"), Arc::clone(&array))),
814+
)?;
815+
let error = prepared_copy_bytes(&vec![batch; batches]).unwrap_err();
816+
assert!(error.to_string().contains("copy size overflow"));
817+
}
818+
Ok(())
819+
}
820+
701821
#[tokio::test]
702822
async fn prepared_byte_keys_use_hash_membership() -> Result<()> {
703823
for array in [

0 commit comments

Comments
 (0)