comphead commented on code in PR #25188:
URL: https://github.com/apache/datafusion/pull/25188#discussion_r4088096774
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs:
##########
@@ -106,6 +107,65 @@ impl AggregateHashTable<PartialMarker> {
})
}
+ /// Starts a bounded-memory drain of partial aggregate states.
+ pub(in crate::aggregates) fn start_early_emit(&mut self) {
+ self.start_outputting();
+ }
+
+ /// Emits at most one output batch while releasing its groups from the
table.
+ ///
+ /// Unlike terminal output, this must not materialize all states: early
+ /// emission can be triggered precisely because the complete state does not
+ /// fit in the memory pool. Once drained, rebuild an empty table so raw
input
+ /// aggregation can resume.
+ pub(in crate::aggregates) fn next_early_emit_batch(
+ &mut self,
+ ) -> Result<Option<RecordBatch>> {
+ let state_schema = Arc::clone(&self.state_schema);
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
+ let group_by_metrics = self.group_by_metrics.clone();
+ let AggregateHashTableState::Outputting(mut state) =
+ std::mem::replace(&mut self.state, AggregateHashTableState::Done)
+ else {
+ return Ok(None);
+ };
+
+ let emit_to =
EmitTo::First(self.batch_size.min(state.group_values.len()));
Review Comment:
Repeated `EmitTo::First(batch_size)` makes this drain O(N²/batch_size). Each
call does O(remaining) work: `GroupValuesPrimitive`, `GroupValuesRows` and
`GroupValuesColumn` run `map.retain` over the whole table and shift or copy the
remaining values, `GroupValuesBytes`/`GroupValuesBytesView` materialize the map
and re-intern the tail, and `ArrayAggGroupsAccumulator` compacts every retained
entry. #23250 removed this pattern from partial terminal output for the same
reason, and `MaterializedAggregateOutput` documents it.
Local numbers, this head vs its merge base: `datafusion-cli -m <limit>
--mem-pool-type fair`, 4 partitions, partial-skip probe disabled, 20M rows, 5M
groups (15M for the two-column key). `emitting_time` is the Partial
`AggregateExec` metric summed over partitions.
| key | limit | `emitting_time` main | `emitting_time` PR | wall main | wall
PR |
|---|---|---|---|---|---|
| `Int64` | 128M | 0.5 ms | 1.42 s | 0.37 s | 0.65 s |
| `Int64` | 512M | 0.5 ms | 4.5 s | 0.33 s | 1.44 s |
| string | 512M | 0.15 ms | 10.3 s | 0.59 s | 3.17 s |
| `(Int64, Int64)` | 512M | 0.1 ms | 2.8 s | 0.75 s | 1.42 s |
The cost grows with the memory limit, so larger pools get slower.
As far as I can tell, the rewrite is driven by the 500 B grouping-set test.
With only the `size()` changes, `aggregate_grouping_sets_*_with_spill` fails
because the emptied three-column table now accounts 648 B (`Failed to allocate
additional 648.0 B ... pool_size: 500.0 B`). Keeping main's `take_state_batch`
path and raising that budget to 700 passes all `aggregates::` tests with the
snapshot unchanged, including your new `early_emit_count > 0` check.
`nested_nullability`, `aggregate_memory_spill.slt` and the original
`memory_limit` expectations also pass that way. Could we do that here and move
bounded early-emit batches to a separate PR, ideally on top of #7065?
<details><summary>Queries</summary>
```sql
SET datafusion.execution.target_partitions = 4;
SET datafusion.execution.skip_partial_aggregation_probe_rows_threshold =
1000000000000;
-- Int64 key
EXPLAIN ANALYZE SELECT count(*), sum(c) FROM (SELECT (value % 500000) +
(value / 2000000) * 500000 AS k, count(*) AS c FROM generate_series(1,
20000000) GROUP BY k);
-- string key
EXPLAIN ANALYZE SELECT count(*), sum(c) FROM (SELECT concat('key-',
CAST((value % 500000) + (value / 2000000) * 500000 AS VARCHAR)) AS k, count(*)
AS c FROM generate_series(1, 20000000) GROUP BY k);
-- (Int64, Int64) key
EXPLAIN ANALYZE SELECT count(*), sum(c) FROM (SELECT (value % 500000) +
(value / 2000000) * 500000 AS k1, value % 3 AS k2, count(*) AS c FROM
generate_series(1, 20000000) GROUP BY k1, k2);
```
</details>
##########
datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs:
##########
@@ -551,12 +556,16 @@ impl GroupedHashAggregateStream {
};
let group_values = new_group_values(group_schema, &group_ordering)?;
- let reservation = MemoryConsumer::new(name)
- // We interpret 'can spill' as 'can handle memory back pressure'.
- // This value needs to be set to true for the default memory pool
implementations
- // to ensure fair application of back pressure amongst the memory
consumers.
- .with_can_spill(oom_mode != OutOfMemoryMode::ReportError)
- .register(context.memory_pool());
+ let merge_pool = Arc::new(MergeMemoryPool::new(
Review Comment:
I confirmed that `legacy_aggregate_spill_merge_leaves_memory_for_replay`
needs this once descriptors are counted. `FairSpillPool::try_grow` checks each
reservation against the fair share on its own, so the state reservation and its
`new_empty()` merge sibling can together exceed that share.
`FinalHashAggregateStream::into_replay_stream` (the default path) still
creates the merge reservation with `reservation.new_empty()`. Does it need the
same shared pool? `migrated_aggregate_spill_merge_leaves_memory_for_replay`
runs with a 2 MiB budget, so it may not surface this.
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -589,74 +581,41 @@ impl PartialHashAggregateStream {
}
}
- /// emit a materialized partial-state on memory pressure
- /// batch in `batch_size`(from configuration) slices
+ /// Drain partial aggregate states in bounded output batches after memory
+ /// pressure. Each batch is removed from the table before the next one is
+ /// materialized, so this path never requires the complete state to fit.
async fn emit_on_memory_pressure(
&mut self,
- // After each incremental emitting step, the `remaining_groups` will
be updated
- // with batch slicing.
- mut remaining_groups: RecordBatch,
+ hash_table: &mut AggregateHashTable<PartialMarker>,
emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
- hash_table_mem_size: usize,
) -> Result<()> {
- let remaining_groups_memory = remaining_groups.get_array_memory_size();
-
- // Emitting clears the aggregate table and releases its
- // accumulated memory. Update the reservation accordingly.
- // We account here for the remaining groups memory to see if we can
return batch size states
- // if there is not enough memory, fallback to emit large batch
- match self
- .reservation
- .try_resize(hash_table_mem_size + remaining_groups_memory)
- {
- Ok(_) => {
- // Continue with slicing
- }
- Err(DataFusionError::ResourcesExhausted(_)) => {
- // Fail to reserve memory for the hash table + state batch
while slicing so emit a huge batch
-
- // Try resize without holding the state batch, if it fails
there is nothing we can do
- self.reservation.try_resize(hash_table_mem_size)?;
-
- self.reduction_factor.add_part(remaining_groups.num_rows());
- emitter
-
.emit(remaining_groups.record_output(&self.baseline_metrics))
- .await;
+ hash_table.start_early_emit();
+ loop {
+ let batch = hash_table.next_early_emit_batch()?.ok_or_else(|| {
+ internal_datafusion_err!(
+ "Partial hash aggregate exhausted early-emission state
unexpectedly"
+ )
+ })?;
- return Ok(());
+ self.reduction_factor.add_part(batch.num_rows());
+ // The reservation may already be above the pool limit that caused
+ // early emission. As with terminal output, make progress by
+ // releasing table state rather than requiring this output batch to
+ // fit alongside all remaining groups. A failed resize is expected
+ // until enough groups have been released; no materialized output
+ // batch is retained across the next iteration.
+ match self.reservation.try_resize(hash_table.memory_size()) {
+ Ok(()) | Err(DataFusionError::ResourcesExhausted(_)) => {}
Review Comment:
If the incremental drain stays, a few things here:
- Once the table is rebuilt, the resize error is still ignored, so this
stream no longer returns `ResourcesExhausted` at all. That is why the two
`memory_limit` expectations moved to `FinalHashAggregateStream[0]`. Main keeps
this resize strict after the groups are released.
- The drain runs after `timer.done()`, so it is no longer counted in
`elapsed_compute`. In the runs above, Partial `elapsed_compute` stays at about
455 ms on both branches while `emitting_time` goes from 0.5 ms to 1.42 s.
`produce_output` restarts the timer around each `emit().await`, and the same
pattern would work here.
- Tests: the `assert_ne!` message still says the batches "should be slices
of the same materialized state batch", and the comments still describe slicing.
The three `test_partial_hash_stream_*` tests share one setup and could be a
single test parameterized over the memory limit.
`partial_stream_under_memory_limit` no longer needs to return the `RuntimeEnv`.
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs:
##########
@@ -1342,6 +1386,170 @@ mod tests {
GroupIndexView, group_column_supported_type, make_group_column,
supported_schema,
};
+ fn expected_size(group_values: &GroupValuesColumn<false>) -> usize {
Review Comment:
`expected_size` restates `size()` field by field, and the same pattern is in
`row.rs`, `single_group_by/primitive.rs` and `row_backed.rs`. Those assertions
pass for any formula that matches itself, so they would not catch a missed
allocation. I would keep the behavioral checks, such as
`clear_shrink_releases_vectorized_and_emit_scratch_capacity` and "size grows
after `intern`/`emit`", and drop the mirrored formulas.
`size_includes_collision_emit_and_vectorized_buffers` and
`size_retains_vectorized_and_emit_scratch_capacity` also overlap.
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -396,9 +396,16 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
let batch = RecordBatch::try_new(state_schema, output)?;
debug_assert!(batch.num_rows() > 0);
- // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
- // key/index buffers too so the memory reservation can be released
- // before the batch is sorted for spilling.
+ // State emission should reset accumulators, but spill recovery must
+ // release every emitted allocation even for an accumulator that
retains
+ // capacity. Rebuild the accumulator set before returning the state
batch.
+ state.accumulators = state
Review Comment:
Which accumulator keeps capacity after `state(EmitTo::All)`? If it is a
specific one, fixing its `state` or `size` would be more targeted. Otherwise a
small shared helper would help, since this rebuild now appears in `common.rs`,
`common_ordered.rs`, `partial_table.rs` and the existing `partial_skip_table`.
--
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]
