Skip to content

Commit 6fceca2

Browse files
andygrove2010YOUY01
andygrove
and
2010YOUY01
authored
[branch-55] perf: Avoid copying when materializing output in OrderedPartialAggregateStream (#25312) (#25843)
## Which issue does this PR close? - Part of #25758 - Related to #25157 ## Rationale for this change On `branch-55`, aggregating input that is sorted on a prefix of the group keys can be several times slower than on 54 (#25157, TPC-DS Q75). The new aggregation path, which is on by default in 55, emits completed groups `batch_size` at a time, and each emit shifts the state of all the remaining groups. The current workaround is `SET datafusion.execution.enable_migration_aggregate = false`. ## What changes are included in this PR? This PR backports #25312 from @2010YOUY01 to the `branch-55` line. `OrderedPartialAggregateStream` now materializes each run of completed groups once and emits slices of it. The cherry-pick conflicts in `common_ordered.rs` and `ordered_partial_table.rs`. On `main`, the new `materialize_groups` records per-accumulator timings through APIs that are not on `branch-55` (`AccumulatorPhase`, `aggregate_accumulator_metrics`, `time_emitting`). Here it takes `is_final: bool`, like the existing `next_output_batch_for_mode`, and records `emitting_time` as before. Otherwise the logic matches `main`. `ordered_partial_stream.rs` applied cleanly and matches `main` at #25312, apart from a doc comment that #25007 removed on `main`. This PR does not include #25639, which does the same for `OrderedFinalAggregateStream`. It builds on the #25538 spill refactor, so it would need a larger port, and most of the regression is recovered without it (see below). ## Are these changes tested? #25312 relies on existing tests. On this branch I ran: - `cargo test -p datafusion-physical-plan --lib` - `cargo test -p datafusion --test core_integration memory_limit` - `cargo test -p datafusion --features extended_tests --test fuzz -- aggregate`, which includes the sorted-input `streaming_aggregate_test` and the limited-memory aggregate fuzz tests - the full sqllogictest suite - `./dev/rust_lint.sh` I also timed the query from #25157 (10M rows, `SELECT count(*) FROM (SELECT DISTINCT ...)` over a Parquet file declared `WITH ORDER (d_year ASC)`), using `datafusion-cli` built with `--profile release-nonlto` on an M3 Max. Median of 5 runs: | Build | Time | |---|---| | `branch-55` | 2.91 s | | `branch-55`, `enable_migration_aggregate = false` | 0.31 s | | this PR | 0.43 s | ## Are there any user-facing changes? Ordered partial aggregation is faster. Output batches are still at most `batch_size` rows, except when the stream cannot reserve memory for a materialized run: then it passes the run downstream as one batch, as on `main`. No public API changes. Co-authored-by: Yongting You <[email protected]>
1 parent c0fa2f2 commit 6fceca2

3 files changed

Lines changed: 272 additions & 206 deletions

File tree

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -287,23 +287,6 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
287287
Ok(Some(batch))
288288
}
289289

290-
/// Returns the [`EmitTo`], clamped to the specified batch size
291-
///
292-
/// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number
293-
/// of groups to emit from `GroupValues` / accumulators, and
294-
/// `should_remove_groups` indicates whether `GroupOrdering` must also shift
295-
/// its tracked indexes.
296-
pub(super) fn clamp_emit_to(
297-
&self,
298-
group_count: usize,
299-
emit_to: EmitTo,
300-
) -> (EmitTo, bool) {
301-
match emit_to {
302-
EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true),
303-
EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false),
304-
EmitTo::All => (EmitTo::First(self.batch_size), false),
305-
}
306-
}
307290
/// Aggregates one evaluated input batch.
308291
///
309292
/// This common utility is used by ordered partial and ordered final aggregation.
@@ -378,19 +361,36 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
378361
let Some(emit_to) = self.buffer.group_ordering.emit_to() else {
379362
return Ok(None);
380363
};
381-
let (emit_to, should_remove_groups) =
382-
self.clamp_emit_to(self.buffer.group_values.len(), emit_to);
364+
let emit_to = match emit_to {
365+
EmitTo::First(n) => EmitTo::First(n.min(self.batch_size)),
366+
EmitTo::All if self.num_groups() > self.batch_size => {
367+
EmitTo::First(self.batch_size)
368+
}
369+
EmitTo::All => EmitTo::All,
370+
};
371+
self.materialize_groups(emit_to, is_final).map(Some)
372+
}
383373

374+
/// Removes the selected groups once and materializes their output columns.
375+
/// The caller chooses the completed prefix and any output-size limit.
376+
///
377+
/// # Argument: `is_final`
378+
///
379+
/// - `true`: output final aggregate values.
380+
/// - `false`: output partial accumulator states.
381+
pub(super) fn materialize_groups(
382+
&mut self,
383+
emit_to: EmitTo,
384+
is_final: bool,
385+
) -> Result<RecordBatch> {
384386
let timer = self.group_by_metrics.emitting_time.timer();
385387
let mut output = self.buffer.group_values.emit(emit_to)?;
386-
if should_remove_groups {
387-
match emit_to {
388-
EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n),
389-
// `EmitTo::All` is only used after `input_done`, when all
390-
// buffered groups are known complete and the ordering state is
391-
// no longer needed.
392-
EmitTo::All => {}
393-
}
388+
// EOF can also emit a prefix when a caller limits its batch size,
389+
// but the completed ordering state no longer tracks group indexes.
390+
if let EmitTo::First(n) = emit_to
391+
&& matches!(self.buffer.group_ordering.emit_to(), Some(EmitTo::First(_)))
392+
{
393+
self.buffer.group_ordering.remove_groups(n);
394394
}
395395

396396
for acc in &mut self.buffer.accumulators {
@@ -405,6 +405,6 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
405405
let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
406406
debug_assert!(batch.num_rows() > 0);
407407

408-
Ok(Some(batch))
408+
Ok(batch)
409409
}
410410
}

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs‎

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -84,23 +84,17 @@ impl OrderedAggregateTable<PartialMarker> {
8484
self.aggregate_evaluated_batch(&evaluated_batch, false)
8585
}
8686

87-
/// Emits the next batch of partial state rows for groups proven complete by
88-
/// the input ordering.
89-
///
90-
/// For example, when the query is `GROUP BY a` and the input is ordered by
91-
/// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3`
92-
/// are complete and safe to emit.
93-
///
94-
/// Key steps:
95-
/// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly.
96-
/// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and
97-
/// all `GroupsAccumulator`s.
98-
///
99-
/// This may output small batches. Avoiding tiny batches is left to future
100-
/// ordered-aggregation optimizations.
101-
pub(in crate::aggregates) fn next_output_batch(
87+
/// Materializes all groups proven complete by the input ordering, leaving
88+
/// the active ordered-key range in the table.
89+
pub(in crate::aggregates) fn take_completed_state_batch(
10290
&mut self,
10391
) -> Result<Option<RecordBatch>> {
104-
self.next_output_batch_for_mode(false)
92+
if self.is_empty() {
93+
return Ok(None);
94+
}
95+
let Some(emit_to) = self.group_ordering().emit_to() else {
96+
return Ok(None);
97+
};
98+
self.materialize_groups(emit_to, false).map(Some)
10599
}
106100
}

0 commit comments

Comments
 (0)