(datafusion) branch final-aggr-with-lazily-materialize-poc-candidate1 updated: Clarify final output transition states

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

Rachelint pushed a commit to branch final-aggr-with-lazily-materialize-poc-candidate1
in repository https://gitbox.apache.org/repos/asf/datafusion.git


The following commit(s) were added to refs/heads/final-aggr-with-lazily-materialize-poc-candidate1 by this push:
     new c9e11d8fec Clarify final output transition states
c9e11d8fec is described below

commit c9e11d8fecda2688b2fd5e57473492111aa8794a
Author: kamille <ka...@apache.org>
AuthorDate: Sun Jul 5 21:51:19 2026 +0800

    Clarify final output transition states
---
 .../physical-plan/src/aggregates/hash_aggregate.rs | 40 ++++++++++++++++------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
index 77d40ec3f0..e331eb82f7 100644
--- a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
+++ b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
@@ -1263,10 +1263,12 @@ impl FinalHashAggregateStream {
         match result {
             Ok(Some(batch)) => {
                 debug_assert!(batch.num_rows() > 0);
-                let next_state = if original_state.hash_table().is_done()
-                    || original_state.hash_table().is_building()
-                {
-                    if self
+
+                if original_state.hash_table().is_building() {
+                    // The current materialized output has been fully emitted, and
+                    // the reusable buffer has moved back to `Building`. Load the
+                    // next staged partition if one exists.
+                    let next_state = if self
                         .partition_run_state
                         .as_ref()
                         .is_some_and(FinalPartitionRunState::has_buffered_runs)
@@ -1284,16 +1286,32 @@ impl FinalHashAggregateStream {
                         }
                     } else {
                         FinalHashAggregateState::Done
-                    }
-                } else {
-                    let _ =
-                        self.update_memory_reservation(Some(original_state.hash_table()));
-                    original_state
-                };
+                    };
 
+                    return ControlFlow::Break((
+                        Poll::Ready(Some(
+                            Ok(batch.record_output(&self.baseline_metrics)),
+                        )),
+                        next_state,
+                    ));
+                }
+
+                if original_state.hash_table().is_done() {
+                    // The current materialized output has been fully emitted, but
+                    // there is no reusable buffer. This is the terminal output
+                    // path for non-replay final aggregation.
+                    return ControlFlow::Break((
+                        Poll::Ready(Some(
+                            Ok(batch.record_output(&self.baseline_metrics)),
+                        )),
+                        FinalHashAggregateState::Done,
+                    ));
+                }
+
+                let _ = self.update_memory_reservation(Some(original_state.hash_table()));
                 ControlFlow::Break((
                     Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
-                    next_state,
+                    original_state,
                 ))
             }
             Ok(None) => {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org