Skip to content

GroupedHashAggregateStream releases its memory reservation at emit while the output batch is still alive #26165

Closed

Description

@pmoust

Describe the bug

GroupedHashAggregateStream shrinks its MemoryReservation to the emptied group state as soon as it emits, while the emitted RecordBatch, which holds those same bytes, is still alive and is handed out in batch_size slices. From that point until the last slice is dropped downstream, the memory pool believes the bytes are free and can grant them to other consumers or queries.

The relevant code in DataFusion 54.1.0 (datafusion-physical-plan/src/aggregates/row_hash.rs):

  1. emit (around line 1103) moves the group state out with group_values.emit(...) (around line 1114). For string keys, GroupValuesBytesView::emit takes the whole ArrowBytesViewMap and converts it zero-copy with into_state.
  2. emit builds the aggregate output arrays with evaluate / state (around lines 1120-1127). These are new allocations.
  3. update_memory_reservation (around line 1133) then resizes the reservation to the now-empty state.
  4. set_input_done_and_produce_output parks the batch in ExecutionState::ProducingOutput, and poll_next hands it out in batch_size slices that share the batch's buffers. So all of it stays live, uncharged, until the last slice is dropped.

To Reproduce

We measured this with a MemoryPool that records per-consumer reservations, sampling jemalloc's stats.allocated every 20 ms. Each run was a single statement over 20M rows of ClickBench hits, with nothing else running.

Statement Reservation at emit Allocated at emit Uncharged Duration
ClickBench Q33, GROUP BY WatchID, ClientIP fell from 2.22 GB to 0.47 GB stayed at 1.68 GB about 1.2 GB 60-80 ms
ClickBench Q19 about 0.77 GB
ClickBench Q17 about 0.39 GB

One statement alone survives this. Under concurrency, the pool hands those bytes to other queries while they are still live. In our case, 10 concurrent connections drove the process 2-3 GB past its pool limit and into the kernel OOM killer.

Expected behavior

While the stream is in ProducingOutput, its reservation should cover the emitted batch, for example get_array_memory_size of what is still to be sliced out. The reservation should shrink as slices are handed out, or when the batch is dropped, rather than immediately at emit.

Additional context

We work around it in our own MemoryPool: shrinks from a GroupedHashAggregateStream[..] consumer that cannot spill (!consumer.can_spill()) stay charged until the consumer unregisters. That over-charges until the stream ends, and it keys on the consumer's name, which is fragile. A fix in the stream itself would be exact.

Related: #23393 (size() under-reporting), which covers a different part of the same gap.

Activity

  1. 2010YOUY01 commented on Oct 10, 2026

    @2010YOUY01
    Contributor

    Thank you for reporting it, GroupedHashAggregateStream has been deprecated in df 55.0, and it will be removed in 56.0

    If this bug can be reproduced in 55.0, we should direct the fix to new implementation

    To Reproduce

    I don't quite get how to reproduce, could you describe the end behavior rather than internal mechanism, for example `in Clickbench Q33, we set DataFusion memory pool to 2GB, but measured peak RSS is 3GB'

  2. pmoust commented on Oct 10, 2026

    @pmoust
    Author

    Checked on main (55.1.0). The new implementation keeps the emitted batch reserved. Only the legacy stream (enable_migration_aggregate = false) shrinks early.

    Unit test: SUM(v) GROUP BY g, 16,384 Utf8View keys, batch_size = 1024, pool.reserved() vs the first slice's get_array_memory_size():

    Path Emitted Reserved while slicing After
    default 753,872 753,872 0
    legacy 753,872 165,128 248

    The legacy stream is being removed, so closing as fixed in 55.

    End behavior on 54.1: Q33 alone over 20M rows, allocated exceeded total pool reservation by ~1.2 GB for 60-80 ms at emit. The concurrent overshoot in the report also included allocator retention, so I don't have a clean "pool X, RSS Y" number for this bug alone.

  3. pmoust commented on Oct 10, 2026

    @pmoust
    Author

    @2010YOUY01 sorry for the cryptic reply. This is a test from ClickHouse/ClickBench#2471

  4. pmoust commented on Oct 10, 2026

    @pmoust
    Author

    One correction to the above: that result is main, not the 55.2 release. In 55.2.0 AggregateExec::execute_typed still falls back to GroupedHashAggregateStream for any shape the should_use_* checks do not claim, with enable_migration_aggregate at its default. A single-stage SELECT DISTINCT ... LIMIT over Utf8View keys takes that path and shows the early release: 65,840 bytes reserved against a 2,490,480-byte first batch. On main the fallback is config-only, so this is fixed there and in whatever release ships that change. Leaving the issue closed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions