kosiew commented on code in PR #25188:
URL: https://github.com/apache/datafusion/pull/25188#discussion_r4101107267
##########
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:
[8b69b81](https://github.com/apache/datafusion/pull/25188/commits/8b69b81c03a2c794d6f818b6e8b86093e890d58b)
- remove incremental drain
--
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]
