CollectLeft hash join build holds ~2x build side while charging the memory pool 1x: concat_batches copy is never reserved and superseded batches are never shrunk
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
In `collect_left_input` (CollectLeft hash join build), the input batches are charged to the `MemoryReservation` incrementally as they stream in, but the `concat_batches` result — a **full second copy of the entire build side** — is never charged, and the original batches' reservation is never shrunk after the concatenated batch supersedes them:
- each incoming batch: `state.reservation.try_grow(batch_size)?` ✔️
- hash table estimate: `try_grow(estimated_hashtable_size)?` ✔️
- `let batch = concat_batches(&schema, batches_iter.clone())?;` — **not charged** ✘
- there is no `shrink`/`try_shrink` anywhere in the file; `batches` stays alive until the function returns, so both copies coexist through the bitmap/scope-map/InList-pushdown construction that follows.
The join-key arrays evaluated from the concatenated batch (for non-`Column` on-expressions) are also unaccounted. The ArrayMap (perfect-hash) build path has the identical pattern: it reserves the map's vectors, then concatenates unreserved.
Consequence: near the pool limit a CollectLeft build transiently allocates ~2× the build side while the pool sees 1×. Instead of a clean `ResourcesExhausted` error — the reason the reservation exists — the process can abort with a real OOM. #24558 (single-partition threshold raised to 4MB) makes CollectLeft builds more common, widening the exposure.
### To Reproduce
Test appended to the `hash_join/exec.rs` tests module: 8 batches × 8192 rows × 3 Int32 columns on the build side, two join keys (defeats the ArrayMap path), Inner CollectLeft join, `GreedyMemoryPool` sized to
```
pool limit = charged bytes (batches 786,432 B + hash table 2,228,280 B) + 393,216 B headroom
```
The headroom (393 KB) is **smaller than the uncharged concat copy (786 KB)**, so if accounting were correct the query MUST fail with `ResourcesExhausted` during the concat. Observed:
```
VERDICT CONFIRMED: query succeeded (0 output rows) with pool limit 3,407,928 B
although concat_batches allocated a second ~786,432 B copy of the build side
that was never charged to the reservation
```
Sanity leg: with pool = build/2 the same query fails with `Resources exhausted`, proving the pool is watching this operator.
Full repro test
```rust
#[tokio::test]
async fn collect_left_concat_copy_not_charged_to_pool() -> Result<()> {
const NUM_BATCHES: usize = 8;
const ROWS_PER_BATCH: usize = 8192;
let mut left_batches = Vec::new();
let mut build_side_size = 0usize;
for i in 0..NUM_BATCHES {
let start = (i * ROWS_PER_BATCH) as i32;
let vals: Vec = (start..start + ROWS_PER_BATCH as i32).collect();
let batch = build_table_i32(("a1", &vals), ("b1", &vals), ("c1", &vals));
build_side_size += get_record_batch_memory_size(&batch);
left_batches.push(batch);
}
let schema = left_batches[0].schema();
let num_rows = NUM_BATCHES * ROWS_PER_BATCH;
let right_batch = build_table_i32(
("a2", &vec![10, 11]),
("b2", &vec![12, 13]),
("c2", &vec![14, 15]),
);
// Two join keys => the ArrayMap (perfect hash join) path declines,
// exercising the JoinHashMap path and its `concat_batches` call.
let on = vec![
(
Arc::new(Column::new_with_schema("a1", &schema)?) as _,
Arc::new(Column::new_with_schema("a2", &right_batch.schema())?) as _,
),
(
Arc::new(Column::new_with_schema("b1", &schema)?) as _,
Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _,
),
];
let hashtable_size =
estimate_memory_size::<(u32, u64)>(num_rows, size_of::())?;
// Everything collect_left_input charges for an Inner join:
// input batches (incrementally) + estimated hash table (no bitmap).
let charged = build_side_size + hashtable_size;
let run = |limit: usize| {
let left_batches = left_batches.clone();
let schema = Arc::clone(&schema);
let right_batch = right_batch.clone();
let on: JoinOn = on.clone();
async move {
let runtime = RuntimeEnvBuilder::new()
.with_memory_limit(limit, 1.0)
.build_arc()?;
let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));
let left =
TestMemoryExec::try_new_exec(&[left_batches], schema, None).unwrap();
let right = TestMemoryExec::try_new_exec(
&[vec![right_batch.clone()]],
right_batch.schema(),
None,
)
.unwrap();
let join = join(
left,
right,
on,
&JoinType::Inner,
NullEquality::NullEqualsNothing,
)?;
common::collect(join.execute(0, task_ctx)?).await
}
};
// Sanity leg: half the build side must overflow while collecting batches,
// proving the pool is watching this operator.
let err = run(build_side_size / 2).await.unwrap_err();
assert_contains!(err.to_string(), "Resources exhausted");
// Main leg: charged amounts fit with build_side/2 headroom; the concat
// copy (~build_side_size) does not.
let limit = charged + build_side_size / 2;
match run(limit).await {
Ok(batches) => {
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
eprintln!(
"VERDICT CONFIRMED: query succeeded ({rows} output rows) with pool \
limit {limit} B = charged ({charged} B: {build_side_size} B batches \
+ {hashtable_size} B hash table) + {} B headroom, although \
concat_batches allocated a second ~{build_side_size} B copy of the \
build side that was never charged to the reservation",
build_side_size / 2
);
}
Err(e) => {
eprintln!(
"VERDICT REFUTED: accounting held, query failed under limit {limit} B: {e}"
);
panic!("concat copy appears to be charged: {e}");
}
}
Ok(())
}
}
```
### Expected behavior
`try_grow(concat_size)` before the concat; after the concat succeeds, drop `batches` and `shrink` their reservation (the transient 2× window is then visible to the pool and fails cleanly when it doesn't fit). Same two-line change on the ArrayMap path. Non-`Column` join-key arrays should be counted too.
### Additional context
Verified present on current `main` (the tip adds `RecordBatchMemoryCounter` to dedup shared buffers on the incremental leg, but the concat copy is still uncharged and the file still contains no shrink). Existing memory-limit tests in the same file assert the error-message shape and keep passing. Found during a joins audit; repro test available. I plan to follow up with a fix PR.
Contributor guide
Assessment
This issue has not been assessed yet.