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

[zhuqi-lucas]

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

#[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(())
    }
}

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.