2010YOUY01 commented on code in PR #24061:
URL: https://github.com/apache/datafusion/pull/24061#discussion_r3702935348
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -107,13 +116,38 @@ use crate::{InputOrderMode, RecordBatchStream,
SendableRecordBatchStream, metric
/// accumulated groups, then switches to a skip state. In that state, each
/// remaining input batch is converted directly to partial aggregate state rows
/// without inserting the rows into the grouped hash table.
+///
+/// # Feature: Memory-limited Execution
+///
+/// ## Partial Aggregation
+///
+/// Partial aggregation can emit incomplete results because the final stage
merges
+/// all intermediate states for the same group. If the memory reservation
exceeds
+/// its limit after aggregating an input batch, this stream emits all
accumulated
+/// states and continues aggregating the remaining input with an empty table.
+///
+/// ## Final Aggregation
+///
+/// During final aggregation, group keys and states accumulate. If memory usage
+/// exceeds the budget, spilling is triggered as follows:
+/// 1. After aggregating a new input batch, if the memory reservation exceeds
its
+/// limit, spill all accumulated groups and states.
+/// - Sort all groups by the group keys before spilling.
+/// 2. Repeat until the input is exhausted.
+/// 3. Perform a sort-preserving merge of all spill files and feed the merged
output
+/// into an ordered streaming aggregation, which ensures bounded memory
usage and
+/// evaluates the final result.
+/// - [`OrderedFinalAggregateStream`] is reused for the streaming
aggregation.
pub(crate) struct PartialHashAggregateStream {
Review Comment:
We could put partial and final stream to two different files, I plan to this
after this PR.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -3427,8 +3418,8 @@ mod tests {
assert!(final_stats.total_byte_size.get_value().is_some());
let task_ctx = if spill {
- // enlarge memory limit to let the final aggregation finish
- new_spill_ctx(2, 2600)
+ // Enlarge the memory limit enough to replay spilled states.
+ new_spill_ctx(2, 4640)
Review Comment:
The legacy implementation underestimate the memory usage somehow, so here it
relaxes the memory budget
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -3390,7 +3380,8 @@ mod tests {
| 2 | 1 | 1.0 |
| 3 | 1 | 2.0 |
| 3 | 2 | 5.0 |
- | 4 | 3 | 11.0 |
+ | 4 | 1 | 4.0 |
Review Comment:
they're both valid partial aggregation output, due to early emitting under
memory limit.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -3457,17 +3448,12 @@ mod tests {
let spilled_bytes = metrics.spilled_bytes().unwrap();
let spilled_rows = metrics.spilled_rows().unwrap();
+ assert_eq!(3, output_rows);
if spill {
- // When spilling, the output rows metrics become partial output
size + final output size
Review Comment:
This should be a bug in the legacy implementation. Final aggregation should
return the same `output_rows` regardless of the spilling condition.
--
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]
