zhuqi-lucas opened a new issue, #25269:
URL: https://github.com/apache/datafusion/issues/25269
### 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.
<details>
<summary>Full repro test</summary>
```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<i32> = (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::<JoinHashMapU32>())?;
// 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(())
}
}
```
</details>
### 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
