This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24012-e8015e41a7d38098a61c3776debf052ae03ece96 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f6d4f04dc419a5f7b7fcc374e468a43026a66221 Author: Raz Luvaton <[email protected]> AuthorDate: Tue Aug 4 23:43:11 2026 +0300 chore: cleanup `OrderedPartialAggregateStream` more (#24012) ## Which issue does this PR close? N/A ## Rationale for this change My original refactor PR did this but I can just move table and cleanup more - #23951 ## What changes are included in this PR? remove weird non optimizing thing ## Are these changes tested? existing tests ## Are there any user-facing changes? no Co-authored-by: Andrew Lamb <[email protected]> --- .../src/aggregates/ordered_partial_stream.rs | 32 +++++++++------------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 975acc1980..9e93a111a6 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -221,18 +221,7 @@ impl OrderedPartialAggregateStream { self.close_input(); table.input_done(); - let last_batch = self.handle_draining_final(&mut table, &mut emitter).await?; - - // Clear memory before emitting last batch so we don't have to wait for next poll to clear - { - // Clear memory - drop(table); - let _ = self.reservation.try_resize(0); - } - - if let Some(last_batch) = last_batch { - emitter.emit(last_batch).await; - } + self.handle_draining_final(table, &mut emitter).await?; Ok(()) }) @@ -326,23 +315,28 @@ impl OrderedPartialAggregateStream { /// `table.input_done()` has already made every remaining group safe to emit, /// so this state keeps draining until the table is empty. /// - /// Returns the last batch to emit so we can free all the state and memory before emitting, - /// and we won't need to hold while waiting for the next poll. - /// /// See comments at [`Self::create_stream`] for details. /// async fn handle_draining_final( &mut self, - table: &mut OrderedAggregateTable<PartialMarker>, + mut table: OrderedAggregateTable<PartialMarker>, emitter: &mut TryEmitter<RecordBatch, DataFusionError>, - ) -> Result<Option<RecordBatch>> { + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); + while let Some(batch) = table.next_output_batch()? { self.reduction_factor.add_part(batch.num_rows()); if table.is_empty() { - return Ok(Some(batch)); + // Clear memory before emitting last batch so we don't have to wait for next poll to clear + drop(table); + let _ = self.reservation.try_resize(0); + drop(timer); + + emitter.emit(batch).await; + + return Ok(()); } self.reservation.try_resize(table.memory_size())?; @@ -353,6 +347,6 @@ impl OrderedPartialAggregateStream { } // was empty - Ok(None) + Ok(()) } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
