Skip to content

Commit 26883b5

Browse files
2010YOUY012010YOUY01
committed
review: more comments
1 parent 2302f6e commit 26883b5

1 file changed

Lines changed: 92 additions & 43 deletions

File tree

‎datafusion/physical-plan/src/aggregates/hash_aggregate.rs‎

Lines changed: 92 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -45,12 +45,12 @@ use crate::metrics::{
4545
use crate::stream::EmptyRecordBatchStream;
4646
use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics};
4747

48-
/// Hash aggregation uses a 2-stage (partial and final) hash aggregation, this stream
49-
/// is for the partial stage.
48+
/// Hash aggregation is implemented in two stages: partial and final. This
49+
/// stream implements the partial stage.
5050
///
5151
/// # Example
5252
///
53-
/// select k, avg(v) from t group by k;
53+
/// SELECT k, AVG(v) FROM t GROUP BY k;
5454
///
5555
/// ## Plan
5656
/// AggregateExec(stage=final)
@@ -59,15 +59,18 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric
5959
///
6060
/// ## Partial Stage Behavior
6161
/// Input: raw rows
62-
/// Output: partial states for all groups (e.g. for avg(x), it's sum(x), count(x))
62+
/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)`
63+
/// and `COUNT(x)`)
6364
///
6465
/// ## Final Stage Behavior
6566
/// Input: partial states
66-
/// Output: results for all groups (e.g. for avg(x), it's avg(x) calculated from the state)
67+
/// Output: results for all groups (for example, `AVG(x)` calculated from the
68+
/// state)
6769
///
6870
/// # Optimization: DISTINCT LIMIT Soft Limit
6971
///
70-
/// This optimization applies to both [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`]
72+
/// This optimization applies to both [`PartialHashAggregateStream`] and
73+
/// [`FinalHashAggregateStream`].
7174
///
7275
/// Unordered distinct queries such as:
7376
///
@@ -179,8 +182,8 @@ impl PartialHashAggregateState {
179182
}
180183
}
181184

182-
/// Hash aggregation uses a 2-stage (partial and final) hash aggregation, this stream
183-
/// is for the final stage.
185+
/// Hash aggregation is implemented in two stages: partial and final. This
186+
/// stream implements the final stage.
184187
///
185188
/// See [`PartialHashAggregateStream`] for details.
186189
pub(crate) struct FinalHashAggregateStream {
@@ -196,7 +199,7 @@ pub(crate) struct FinalHashAggregateStream {
196199
/// Memory reservation for group keys and accumulators.
197200
reservation: MemoryReservation,
198201

199-
/// See comments for the same variable in [`PartialHashAggregateStream`]
202+
/// See comments for the same variable in [`PartialHashAggregateStream`].
200203
group_values_soft_limit: Option<usize>,
201204

202205
/// Tracks the high-level stream lifecycle. The hash table owns the lower-level
@@ -205,7 +208,8 @@ pub(crate) struct FinalHashAggregateStream {
205208
}
206209

207210
/// States for final hash aggregation processing.
208-
// Typestate pattern is used, in case the inner logic become more complex in the future.
211+
// The typestate pattern is used in case the inner logic becomes more complex in
212+
// the future.
209213
enum FinalHashAggregateState {
210214
ReadingInput {
211215
hash_table: AggregateHashTable<Final>,
@@ -362,6 +366,8 @@ impl PartialHashAggregateStream {
362366

363367
/// Handle ReadingInput state - aggregate input batches into the hash table.
364368
///
369+
/// See comments at `poll_next()` for details.
370+
///
365371
/// Returns the next operator state with control flow decision.
366372
fn handle_reading_input(
367373
&mut self,
@@ -506,6 +512,8 @@ impl PartialHashAggregateStream {
506512

507513
/// Handle ProducingOutput state - emit partial aggregate state batches.
508514
///
515+
/// See comments at `poll_next()` for details.
516+
///
509517
/// Returns the next operator state with control flow decision.
510518
fn handle_producing_output(
511519
&mut self,
@@ -554,8 +562,8 @@ impl PartialHashAggregateStream {
554562
}
555563
Ok(None) => {
556564
let _ = self.reservation.try_resize(0);
557-
// If in the previous `Aggregating` stage it has decided to skip partial
558-
// aggregation, go the `SkipAggregation` stage; otherwise finish.
565+
// If the previous `Aggregating` stage decided to skip partial
566+
// aggregation, go to the `SkippingAggregation` stage; otherwise finish.
559567
let next_state = match original_state {
560568
PartialHashAggregateState::ProducingOutput {
561569
skip_hash_table: Some(hash_table),
@@ -575,6 +583,8 @@ impl PartialHashAggregateStream {
575583

576584
/// Handle SkippingAggregation state - convert raw input directly to partial states.
577585
///
586+
/// See comments at `poll_next()` for details.
587+
///
578588
/// Returns the next operator state with control flow decision.
579589
fn handle_skipping_aggregation(
580590
&mut self,
@@ -634,31 +644,52 @@ impl Stream for PartialHashAggregateStream {
634644
///
635645
/// See comments in [`PartialHashAggregateStream`] for high-level ideas.
636646
///
637-
/// ============================
638647
/// State transition graph:
639-
/// ============================
640648
///
641-
/// (start) --> ReadingInput
642-
/// ----------------------------
643-
/// ReadingInput -> ReadingInput (after aggregating an input batch)
644-
/// ReadingInput -> ProducingOutput(skip=None) (input exhausted or soft
645-
/// limit reached)
646-
/// ReadingInput -> ProducingOutput(skip=Some) (partial skip triggered)
649+
/// ```text
650+
/// (start)
651+
/// -> ReadingInput
652+
/// The stream starts by polling input and aggregating batches into the
653+
/// in-memory hash table.
647654
///
648-
/// ProducingOutput(skip=None) -> ProducingOutput(skip=None)
649-
/// (after yielding one accumulated output batch)
650-
/// ProducingOutput(skip=None) -> Done (all accumulated output emitted)
655+
/// ReadingInput
656+
/// -> ReadingInput
657+
/// Aggregate one batch, update the inner aggregate hash table, and
658+
/// continue with the next input batch.
659+
/// -> ProducingOutput(skip=None)
660+
/// Input was exhausted, or the soft group limit was reached. Move to
661+
/// the next state to start outputting.
662+
/// -> ProducingOutput(skip=Some)
663+
/// Partial skip aggregation was triggered. First move to the
664+
/// `ProducingOutput` state to drain the accumulated state, then move to
665+
/// the `SkippingAggregation` state to convert input directly to partial
666+
/// state without aggregation.
651667
///
652-
/// ProducingOutput(skip=Some) -> ProducingOutput(skip=Some)
653-
/// (after yielding one accumulated output batch)
654-
/// ProducingOutput(skip=Some) -> SkippingAggregation (all accumulated output
655-
/// emitted)
668+
/// ProducingOutput(skip=None)
669+
/// -> ProducingOutput(skip=None)
670+
/// One accumulated output batch was yielded, repeat to continue producing
671+
/// output incrementally.
672+
/// -> Done
673+
/// All accumulated output was emitted.
656674
///
657-
/// SkippingAggregation -> SkippingAggregation (after yielding one
658-
/// `convert_to_state` batch)
659-
/// SkippingAggregation -> Done (input exhausted)
660-
/// ----------------------------
661-
/// Done -> (end)
675+
/// ProducingOutput(skip=Some)
676+
/// -> ProducingOutput(skip=Some)
677+
/// One accumulated output batch was yielded, repeat to continue producing
678+
/// output incrementally.
679+
/// -> SkippingAggregation
680+
/// All accumulated output was emitted. Continue by converting raw
681+
/// input batches directly to partial aggregate state.
682+
///
683+
/// SkippingAggregation
684+
/// -> SkippingAggregation
685+
/// One `convert_to_state` batch was yielded; repeat to continue
686+
/// processing.
687+
/// -> Done
688+
/// Input was exhausted.
689+
///
690+
/// Done
691+
/// -> (end)
692+
/// ```
662693
fn poll_next(
663694
mut self: std::pin::Pin<&mut Self>,
664695
cx: &mut Context<'_>,
@@ -761,6 +792,8 @@ impl FinalHashAggregateStream {
761792

762793
/// Handle ReadingInput state - aggregate partial state batches into the hash table.
763794
///
795+
/// See comments at `poll_next()` for details.
796+
///
764797
/// Returns the next operator state with control flow decision.
765798
fn handle_reading_input(
766799
&mut self,
@@ -838,6 +871,8 @@ impl FinalHashAggregateStream {
838871

839872
/// Handle ProducingOutput state - emit final aggregate value batches.
840873
///
874+
/// See comments at `poll_next()` for details.
875+
///
841876
/// Returns the next operator state with control flow decision.
842877
fn handle_producing_output(
843878
&mut self,
@@ -887,20 +922,34 @@ impl Stream for FinalHashAggregateStream {
887922
///
888923
/// See comments in [`FinalHashAggregateStream`] for high-level ideas.
889924
///
890-
/// ============================
891925
/// State transition graph:
892-
/// ============================
893926
///
894-
/// (start) --> ReadingInput
895-
/// ----------------------------
896-
/// ReadingInput -> ReadingInput (after aggregating one partial-state input
897-
/// batch)
898-
/// ReadingInput -> ProducingOutput (input exhausted or soft limit reached)
927+
/// ```text
928+
/// (start)
929+
/// -> ReadingInput
930+
/// The stream starts by polling partial-state input and aggregating
931+
/// those states into the final hash table.
932+
///
933+
/// ReadingInput
934+
/// -> ReadingInput
935+
/// Aggregate one partial-state input batch, update the inner aggregate
936+
/// hash table, and continue with the next input batch.
937+
///
938+
/// -> ProducingOutput
939+
/// Input was exhausted, or the soft group limit was reached. Move to
940+
/// the next state to start outputting final aggregate values.
941+
///
942+
/// ProducingOutput
943+
/// -> ProducingOutput
944+
/// One final output batch was yielded; repeat to continue producing
945+
/// output incrementally.
946+
///
947+
/// -> Done
948+
/// All final output was emitted.
899949
///
900-
/// ProducingOutput -> ProducingOutput (after yielding one final output batch)
901-
/// ProducingOutput -> Done (all final output emitted)
902-
/// ----------------------------
903-
/// Done -> (end)
950+
/// Done
951+
/// -> (end)
952+
/// ```
904953
fn poll_next(
905954
mut self: std::pin::Pin<&mut Self>,
906955
cx: &mut Context<'_>,

0 commit comments

Comments
 (0)