You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "2010YOUY01 (via GitHub)" <gi...@apache.org> on 2026/07/29 08:56:37 UTC

[PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

2010YOUY01 opened a new pull request, #23965:
URL: https://github.com/apache/datafusion/pull/23965

   ## Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Part of https://github.com/apache/datafusion/issues/22710
   
   ## Rationale for this change
   
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   This PR implements larger-than-memory execution for single mode aggregation.
   
   See top comments at `single_stream.rs` for the high-level idea. The key implementations are in the same file.
   
   See 
   
   ## What changes are included in this PR?
   
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   ## Are these changes tested?
   
   <!--
   We typically require tests for all PRs in order to:
   1. Prevent the code from being accidentally broken by subsequent changes
   2. Serve as another way to document the expected behavior of the code
   
   If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
   -->
   Existing tests.
   - [ ] todo for myself: double check the codecov
   
   ## Are there any user-facing changes?
   No
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->
   


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3672611786


##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -33,6 +33,9 @@
 statement ok
 SET datafusion.execution.target_partitions = 1
 
+statement ok
+SET datafusion.execution.batch_size = 128

Review Comment:
   similar to the above notes on the test change https://github.com/apache/datafusion/pull/23965#discussion_r3672576168



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#issuecomment-5147412675

   I am sorry -- I have been away -- I will try and review this one soon


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#issuecomment-5155274965

   @alamb Thanks for the review.
   
   Yes, we're getting close! I expect the remaining PRs to be: spilling for partial/final aggregate spilling, and some smaller fixes.
   
   I plan to merge this PR tomorrow.


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#issuecomment-5117205744

   ## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/23965?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
   :x: Patch coverage is `83.10627%` with `62 lines` in your changes missing coverage. Please review.
   :white_check_mark: Project coverage is 80.75%. Comparing base ([`88365dd`](https://app.codecov.io/gh/apache/datafusion/commit/88365ddd62b17c1eabd20ed0b064f626f9e77686?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`02636dd`](https://app.codecov.io/gh/apache/datafusion/commit/02636dd3cbb27382f26af7820115c0aaa91c151f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
   :warning: Report is 32 commits behind head on main.
   
   | [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/23965?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
   |---|---|---|
   | [...sion/physical-plan/src/aggregates/single\_stream.rs](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fsingle_stream.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL3NpbmdsZV9zdHJlYW0ucnM=) | 83.84% | [45 Missing and 8 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [datafusion/physical-plan/src/aggregates/mod.rs](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Fmod.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL21vZC5ycw==) | 33.33% | [0 Missing and 6 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...plan/src/aggregates/aggregate\_hash\_table/common.rs](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Faggregates%2Faggregate_hash_table%2Fcommon.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9hZ2dyZWdhdGVzL2FnZ3JlZ2F0ZV9oYXNoX3RhYmxlL2NvbW1vbi5ycw==) | 86.95% | [0 Missing and 3 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/23965?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   
   <details><summary>Additional details and impacted files</summary>
   
   
   
   ```diff
   @@            Coverage Diff             @@
   ##             main   #23965      +/-   ##
   ==========================================
   + Coverage   80.65%   80.75%   +0.09%     
   ==========================================
     Files        1093     1096       +3     
     Lines      371619   373556    +1937     
     Branches   371619   373556    +1937     
   ==========================================
   + Hits       299730   301659    +1929     
   + Misses      53993    53907      -86     
   - Partials    17896    17990      +94     
   ```
   </details>
   
   [:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/23965?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
   :loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   <details><summary> :rocket: New features to boost your workflow: </summary>
   
   - :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
   - :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
   </details>


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3697776463


##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -282,11 +288,48 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         }
     }
 
+    pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {

Review Comment:
   [664b84a](https://github.com/apache/datafusion/pull/23965/commits/664b84a4ce46ee91fe67ddca0baf232e43c49121)



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3672580948


##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5766,9 +5762,9 @@ mod tests {
     #[tokio::test]
     async fn test_aggregate_with_spill_if_necessary() -> Result<()> {
         // test with spill
-        run_test_with_spill_pool_if_necessary(2_000, true).await?;
+        run_test_with_spill_pool_if_necessary(20_000, true).await?;

Review Comment:
   Here is the same as the previous note on test change.



##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5664,9 +5658,11 @@ mod tests {
             Field::new("b", DataType::Float64, false),
         ]));
 
+        let group_keys = [2, 3, 4, 4].repeat(1_000);

Review Comment:
   The old implementation set a very tight memory budget, and the refactored implementation might have minor implementation difference, making this test fail.
   
   Here it just try to scale up the test input size and memory limit, the test target is the same.



##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -6744,11 +6740,6 @@ mod tests {
                     matches!(root, DataFusionError::ResourcesExhausted(_)),
                     "Expected ResourcesExhausted, got: {root}",
                 );
-                let msg = root.to_string();

Review Comment:
   Above check on error type is enough I think, here it also assert the exact error message



##########
datafusion/physical-plan/src/aggregates/single_stream.rs:
##########
@@ -139,162 +312,386 @@ impl SingleHashAggregateStream {
     ) -> Result<Self> {
         debug_assert!(matches!(
             agg.mode,
-            super::AggregateMode::Single | super::AggregateMode::SinglePartitioned
+            AggregateMode::Single | AggregateMode::SinglePartitioned
         ));
         debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear);
 
         let schema = Arc::clone(&agg.schema);
         let input = agg.input.execute(partition, Arc::clone(context))?;
+        let input_schema = input.schema();
         let batch_size = context.session_config().batch_size();
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+        let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+        let state_schema = Arc::new(create_schema(
+            input_schema.as_ref(),
+            &agg.group_by,
+            &agg.aggr_expr,
+            AggregateMode::Partial,
+        )?);
 
         let hash_table = AggregateHashTable::<SingleMarker>::new(
             agg,
             partition,
             Arc::clone(&schema),
+            Arc::clone(&state_schema),
             batch_size,
         )?;
 
+        let can_spill = context.runtime_env().disk_manager.tmp_files_enabled();
+        let spill_context = if can_spill {
+            Some(Box::new(SingleSpillContext::new(
+                agg,
+                context,
+                partition,
+                batch_size,
+                &state_schema,
+                spill_metrics,
+            )?))
+        } else {
+            None
+        };
+
         let reservation =
             MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]"))
+                .with_can_spill(can_spill)
                 .register(context.memory_pool());
 
         Ok(Self {
             schema,
             input,
             baseline_metrics,
             reservation,
-            state: Some(SingleHashAggregateState::ReadingInput { hash_table }),
+            state: Some(SingleHashAggregateState::ReadingInput {
+                hash_table,
+                spill_context,
+            }),
         })
     }
 
-    /// Moves the aggregate hash table's inner state to `Outputting`.
-    ///
-    /// The caller guarantees that input is fully consumed, so this function can
-    /// eagerly release the input stream.
-    fn start_output(
-        &mut self,
-        hash_table: &mut AggregateHashTable<SingleMarker>,
-    ) -> Result<()> {
+    fn close_input(&mut self) {
         let input_schema = self.input.schema();
         self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
-        hash_table.start_output()
     }
 
-    /// Handle ReadingInput state - aggregate input batches into the hash table.
+    fn break_with_err(error: DataFusionError) -> SingleHashAggregateStateTransition {
+        ControlFlow::Break((
+            Poll::Ready(Some(Err(error))),
+            SingleHashAggregateState::Error,
+        ))
+    }
+
+    fn break_with_internal_err(message: &str) -> SingleHashAggregateStateTransition {
+        Self::break_with_err(internal_datafusion_err!("{message}"))
+    }
+
+    /// Reserve memory for the current aggregate table.
+    fn reservation_size_for_table(
+        hash_table: &AggregateHashTable<SingleMarker>,
+        spill_context: Option<&SingleSpillContext>,
+    ) -> usize {
+        let table_size = hash_table.memory_size();
+        if spill_context.is_some() {
+            // See `SingleHashAggregateStream` comments for how this is estimated.
+            table_size.saturating_add(
+                hash_table
+                    .building_group_count()
+                    .saturating_mul(size_of::<u32>()),
+            )
+        } else {
+            table_size
+        }
+    }
+
+    /// Consumes one raw input batch and updates the single-stage hash table.
     ///
     /// See comments at `poll_next()` for details.
     ///
     /// Returns the next operator state with control flow decision.
     fn handle_reading_input(
         &mut self,
         cx: &mut Context<'_>,
-        mut original_state: SingleHashAggregateState,
+        original_state: SingleHashAggregateState,
     ) -> SingleHashAggregateStateTransition {
-        debug_assert!(matches!(
-            &original_state,
-            SingleHashAggregateState::ReadingInput { .. }
-        ));
-        debug_assert!(original_state.hash_table().is_building());
+        let SingleHashAggregateState::ReadingInput {
+            mut hash_table,
+            spill_context,
+        } = original_state
+        else {
+            return Self::break_with_internal_err(
+                "Single hash aggregate stream expected ReadingInput state",
+            );
+        };
 
         match self.input.poll_next_unpin(cx) {
-            Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)),
-            // Get a new input batch, aggregate it in the hash table
+            Poll::Pending => ControlFlow::Break((
+                Poll::Pending,
+                SingleHashAggregateState::ReadingInput {
+                    hash_table,
+                    spill_context,
+                },
+            )),
             Poll::Ready(Some(Ok(batch))) => {
                 let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
                 let timer = elapsed_compute.timer();
-                let result = original_state.hash_table_mut().aggregate_batch(&batch);
+                let result = hash_table.aggregate_batch(&batch);
                 timer.done();
 
                 if let Err(e) = result {
-                    return ControlFlow::Break((
-                        Poll::Ready(Some(Err(e))),
-                        original_state,
-                    ));
+                    return Self::break_with_err(e);
                 }
 
-                if let Err(e) = self
-                    .reservation
-                    .try_resize(original_state.hash_table().memory_size())
-                {
-                    return ControlFlow::Break((
-                        Poll::Ready(Some(Err(e))),
-                        original_state,
-                    ));
+                // Check memory reservation, and potentially spill.
+                let timer = elapsed_compute.timer();
+                let resize_result =
+                    self.reservation
+                        .try_resize(Self::reservation_size_for_table(
+                            &hash_table,
+                            spill_context.as_deref(),
+                        ));
+                timer.done();
+                match resize_result {
+                    Ok(()) => {}
+                    Err(e @ DataFusionError::ResourcesExhausted(_)) => {
+                        let Some(spill_context) = spill_context else {
+                            return Self::break_with_err(e.context(
+                                "Single hash aggregate cannot spill because temporary files are not enabled in the DiskManager",
+                            ));
+                        };
+                        if hash_table.building_group_count() == 0 {
+                            return Self::break_with_internal_err(
+                                "Single hash aggregate ran out of memory with no aggregated groups",
+                            );
+                        }
+                        return ControlFlow::Continue(
+                            SingleHashAggregateState::Spilling {
+                                hash_table,
+                                spill_context,
+                            },
+                        );
+                    }
+                    Err(e) => {
+                        return Self::break_with_err(e);
+                    }
                 }
 
-                ControlFlow::Continue(original_state)
+                ControlFlow::Continue(SingleHashAggregateState::ReadingInput {
+                    hash_table,
+                    spill_context,
+                })
             }
-            Poll::Ready(Some(Err(e))) => {
-                ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
-            }
-            // Input ends, move to output state
+            Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
             Poll::Ready(None) => {
-                let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
-                let timer = elapsed_compute.timer();
-                let result = self.start_output(original_state.hash_table_mut());
-                timer.done();
-
-                match result {
-                    Ok(()) => {
-                        ControlFlow::Continue(original_state.into_producing_output())
+                self.close_input();
+                match spill_context {
+                    Some(spill_context) if spill_context.has_spills() => {
+                        ControlFlow::Continue(
+                            SingleHashAggregateState::PreparingMergeInput {
+                                hash_table,
+                                spill_context,
+                            },
+                        )
                     }
-                    Err(e) => {
-                        ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
+                    _ => {
+                        let elapsed_compute =
+                            self.baseline_metrics.elapsed_compute().clone();
+                        let timer = elapsed_compute.timer();
+                        let result = hash_table.start_output();
+                        timer.done();
+
+                        match result {
+                            Ok(()) => ControlFlow::Continue(
+                                SingleHashAggregateState::ProducingOutput { hash_table },
+                            ),
+                            Err(e) => Self::break_with_err(e),
+                        }
                     }
                 }
             }
         }
     }
 
-    /// Handle ProducingOutput state - emit final aggregate value batches.
+    /// Sorts and spills one complete in-memory state run, then resumes input.
+    ///
+    /// See comments at `poll_next()` for details.
+    ///
+    /// Returns the next operator state with control flow decision.
+    fn handle_spilling(
+        &mut self,
+        original_state: SingleHashAggregateState,
+    ) -> SingleHashAggregateStateTransition {
+        let SingleHashAggregateState::Spilling {
+            mut hash_table,
+            mut spill_context,
+        } = original_state
+        else {
+            return Self::break_with_internal_err(
+                "Single hash aggregate stream expected Spilling state",
+            );
+        };
+
+        // Sanity check: it is impossible to OOM when the table is empty.
+        if hash_table.building_group_count() == 0 {
+            return Self::break_with_internal_err(
+                "Single hash aggregation entered Spilling with an empty table",
+            );
+        }
+
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        let timer = elapsed_compute.timer();
+        let mut result = spill_context.spill_table(&mut hash_table);
+
+        // Spilling shrinks the aggregate table and releases its accumulated
+        // memory. Update the reservation accordingly.
+        if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) {
+            result =
+                Err(e.context("Decreasing allocation after spilling should succeed"));
+        }
+
+        timer.done();
+
+        match result {
+            // Finished spilling the aggregate table, continue aggregating from input.
+            Ok(()) => ControlFlow::Continue(SingleHashAggregateState::ReadingInput {
+                hash_table,
+                spill_context: Some(spill_context),
+            }),
+            Err(e) => Self::break_with_err(e),
+        }
+    }
+
+    /// 1. Spills the last in-memory run.
+    /// 2. Constructs a globally ordered input stream by applying a sort-preserving
+    ///    merge to all spills.
+    /// 3. Constructs a replay stream: an ordered final aggregate stream over the
+    ///    fully ordered input constructed from the spills.
+    ///
+    /// See comments at `poll_next()` for details.
+    ///
+    /// Returns the next operator state with control flow decision.
+    fn handle_preparing_merge_input(
+        &mut self,
+        original_state: SingleHashAggregateState,
+    ) -> SingleHashAggregateStateTransition {
+        let SingleHashAggregateState::PreparingMergeInput {
+            mut hash_table,
+            mut spill_context,
+        } = original_state
+        else {
+            return Self::break_with_internal_err(
+                "Single hash aggregate stream expected PreparingMergeInput state",
+            );
+        };
+
+        let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+        let timer = elapsed_compute.timer();
+        let replay = match spill_context.spill_table(&mut hash_table) {
+            Ok(()) => {
+                let group_by_metrics = hash_table.group_by_metrics();
+                drop(hash_table);
+                match self.reservation.try_resize(0) {
+                    Ok(()) => (*spill_context).into_replay_stream(
+                        &self.baseline_metrics,
+                        group_by_metrics,
+                        self.reservation.new_empty(),
+                    ),
+                    Err(e) => Err(e),
+                }
+            }
+            Err(e) => Err(e),
+        };
+        timer.done();
+
+        match replay {
+            Ok(stream) => {
+                ControlFlow::Continue(SingleHashAggregateState::MergingSpills { stream })
+            }
+            Err(e) => Self::break_with_err(e),
+        }
+    }
+
+    /// Forwards output from the fully ordered stream that consumes the merged
+    /// spill runs.
+    ///
+    /// See comments at `poll_next()` for details.
+    ///
+    /// Returns the next operator state with control flow decision.
+    fn handle_merging_spills(
+        &mut self,
+        cx: &mut Context<'_>,
+        original_state: SingleHashAggregateState,
+    ) -> SingleHashAggregateStateTransition {
+        let SingleHashAggregateState::MergingSpills { mut stream } = original_state
+        else {
+            return Self::break_with_internal_err(
+                "Single hash aggregate stream expected MergingSpills state",
+            );
+        };
+
+        match stream.poll_next_unpin(cx) {
+            Poll::Pending => ControlFlow::Break((
+                Poll::Pending,
+                SingleHashAggregateState::MergingSpills { stream },
+            )),
+            Poll::Ready(Some(Ok(batch))) => ControlFlow::Break((
+                Poll::Ready(Some(Ok(batch))),
+                SingleHashAggregateState::MergingSpills { stream },
+            )),
+            Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
+            Poll::Ready(None) => ControlFlow::Continue(SingleHashAggregateState::Done),
+        }
+    }
+
+    /// Emits one batch after input is exhausted.
     ///
     /// See comments at `poll_next()` for details.
     ///
     /// Returns the next operator state with control flow decision.
     fn handle_producing_output(

Review Comment:
   This function has no functional changes, the diffs are
   1. use internal_err instead of assert
   2. simplify some control flow code with `break_with_err`



##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -33,6 +33,9 @@
 statement ok
 SET datafusion.execution.target_partitions = 1
 
+statement ok
+SET datafusion.execution.batch_size = 128

Review Comment:
   similar to the above notes on the test change



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 merged PR #23965:
URL: https://github.com/apache/datafusion/pull/23965


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3697782873


##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -58,13 +60,9 @@ FROM (
   GROUP BY (v * 7) % 100000
 )
 ----
-Plan with Metrics
-01)ProjectionExec: expr=[count(Int64(1))@0 as count(*), sum(total)@1 as sum(total)], metrics=[<slt:ignore>]
-02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1)), sum(total)], metrics=[<slt:ignore>]
-03)----ProjectionExec: expr=[sum(t.v)@1 as total], metrics=[<slt:ignore>]
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=9,<slt:ignore>]
-05)--------ProjectionExec: expr=[value@0 as v], metrics=[<slt:ignore>]
-06)----------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192], metrics=[<slt:ignore>]
+<slt:ignore>
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spilled_rows=<slt:ignore>K,<slt:ignore>]

Review Comment:
   Updated in [e242c5e](https://github.com/apache/datafusion/pull/23965/commits/e242c5e161124982b95eaf4df299bfb29db33419)
   
   Yes it's better to assert the exact value.
   
   If they're running with multiple partitions, then `spill_count` will be non-deterministic, however we're using `target_partitions=1` here, so it's deterministic, it's better to assert the exact metrics number.



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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


Re: [PR] refactor(hash-aggr): Support spilling for single mode aggregation [datafusion]

Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3696680359


##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -282,11 +288,48 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         }
     }
 
+    pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {

Review Comment:
   minor nit does it makes sense to return this as &GroupByMetrics and let the caller clone it? It probably doesn't matter



##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -58,13 +60,9 @@ FROM (
   GROUP BY (v * 7) % 100000
 )
 ----
-Plan with Metrics
-01)ProjectionExec: expr=[count(Int64(1))@0 as count(*), sum(total)@1 as sum(total)], metrics=[<slt:ignore>]
-02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1)), sum(total)], metrics=[<slt:ignore>]
-03)----ProjectionExec: expr=[sum(t.v)@1 as total], metrics=[<slt:ignore>]
-04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spill_count=9,<slt:ignore>]
-05)--------ProjectionExec: expr=[value@0 as v], metrics=[<slt:ignore>]
-06)----------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192], metrics=[<slt:ignore>]
+<slt:ignore>
+04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spilled_rows=<slt:ignore>K,<slt:ignore>]

Review Comment:
   The changes in aggregate_memory_split.slt seem to have lost the the actual value (now just K), but given the comments this seems intentional -- is there a reason not to encode the actual row count in the test?
   



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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