fix: spill a native PartialMerge whose input is sorted on some of its keys

docs/source/contributor-guide/memory_management.md

@@ -396,7 +396,10 @@ over a native sort, a refusal fails the query, so the pool checks its requests o
 pool, as it does those of a consumer that cannot spill (`limited_to_its_share` in `fair_pool.rs`).
 It still counts as a consumer that can spill when the pool divides the task's share, so a sort
 below it stays at half of the share and leaves the aggregate the other half. The partially sorted
-stream still spills when Spark refuses it.
+stream still spills when Spark refuses it. The exemption goes by the stream's name, so it also
+covers a `PartialMerge` that the planner runs as a final aggregate because its input is sorted on
+some of its grouping keys. That one is not a `CometSortAggregateExec` and its input is only
+partially sorted, but it is likewise checked only against the pool and Spark's grant.
 
 Three details matter for tuning:
 

native/core/src/execution/memory_pools/fair_pool.rs

@@ -163,6 +163,11 @@ impl Display for CometFairMemoryPool {
 /// refused request fails the query. So its requests are checked only against the pool, like those
 /// of a consumer that cannot spill. The partially sorted stream still spills when Spark refuses it.
 ///
+/// The exemption goes by name, so it covers every `OrderedFinalAggregateStream`, including a
+/// pure `PartialMerge` that the planner runs as `Final` over input sorted on some of its keys
+/// (#6880). That one is not a `CometSortAggregateExec` and its input is only partially sorted, but
+/// it is checked only against the pool as well.
+///
 /// Only the check differs. The aggregate still counts as a consumer that can spill everywhere
 /// else, so it is one of `num_spill` and its bytes are not `unspillable`. A sort below it then
 /// stays limited to its fair share and leaves the aggregate the rest.

native/core/src/execution/memory_pools/spill_replay.rs

@@ -21,10 +21,10 @@
 //! replays them through an `OrderedFinalAggregateStream` that has no way to spill, so a refused
 //! memory request there fails the task. `FinalHashAggregateStream` does this, and so does
 //! `OrderedFinalAggregateStream` itself, which DataFusion uses when the input is sorted on some of
-//! the grouping keys. So does `SingleHashAggregateStream`, which Comet runs for a Spark
-//! PartialMerge (#6857). The merge reserves read buffers for as many spill files as fit, and those
-//! buffers belong to the same consumer, so the replay often finds the consumer's share already
-//! taken.
+//! the grouping keys, and which Comet also runs for a Spark PartialMerge over such input (#6880).
+//! So does `SingleHashAggregateStream`, which Comet runs for other Spark PartialMerges (#6857).
+//! The merge reserves read buffers for as many spill files as fit, and those buffers belong to the
+//! same consumer, so the replay often finds the consumer's share already taken.
 //!
 //! [`SpillReplayPool`] wraps a Comet pool. When the pool refuses a request from the replay, it
 //! records the request with the pool's `grow` instead, which ignores `CometFairMemoryPool`'s limits
@@ -182,6 +182,7 @@ pub(super) fn unwrap_spill_replay(pool: &Arc<dyn MemoryPool>) -> Option<&Arc<dyn
 }
 
 /// The prefix of the name DataFusion gives the consumer of an `OrderedFinalAggregateStream`.
+/// DataFusion 56 renames the stream and its consumer `ClusteredFinalAggregateStream` (#6410).
 pub(super) const ORDERED_FINAL_AGGREGATE: &str = "OrderedFinalAggregateStream[";
 
 /// Whether `consumer` belongs to a final or single aggregate, whose spill replay can't spill.

native/core/src/execution/merge_as_partial.rs

@@ -27,6 +27,9 @@
 //! evaluates to its state, packed into a struct. Under DataFusion's `Partial` mode, which
 //! outputs state, that gives merge semantics with state output. Under `Single` mode, which
 //! spills rather than emit a group twice, the planner unpacks the struct into the state columns.
+//! It does the same under `Final` mode, which DataFusion 55 runs with an ordered aggregate that
+//! spills when the input is sorted on some of the grouping keys. There, `Single` falls back to a
+//! stream whose spill can ask for more than its share of memory (#6880).
 
 use std::fmt::Debug;
 use std::hash::{Hash, Hasher};
@@ -51,12 +54,12 @@ use datafusion_comet_spark_expr::GetStructField;
 
 use crate::execution::spark_aggregate_state::PartialMergeStateDecoder;
 
-/// An AggregateUDF wrapper that merges state and outputs state in Partial or Single mode.
+/// An AggregateUDF wrapper that merges state and outputs state in Partial, Single or Final mode.
 ///
 /// When DataFusion runs an AggregateExec in Partial mode, it calls `update_batch`
 /// on each accumulator and outputs `state()`. This wrapper intercepts `update_batch`
 /// and redirects it to `merge_batch` on the inner accumulator, effectively
-/// implementing PartialMerge: merge inputs, output state. In Single mode, DataFusion
+/// implementing PartialMerge: merge inputs, output state. In Single and Final modes, DataFusion
 /// outputs `evaluate()` instead, which returns the inner accumulator's state as a struct
 /// of the state fields.
 ///
@@ -165,7 +168,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF {
 
     fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
         // In Partial mode, return_type isn't used for output schema (state_fields is).
-        // In Single mode, the wrapper evaluates to its state.
+        // In Single and Final modes, the wrapper evaluates to its state.
         Ok(DataType::Struct(self.cached_state_fields.clone()))
     }
 
@@ -278,9 +281,9 @@ impl InnerAccumulatorFactory {
     }
 }
 
-/// Projects the output of a Single-mode `aggregate` whose aggregates are all [`MergeAsPartialUDF`]s
-/// into the output of a Partial-mode one: its `num_grouping` grouping columns, then one column per
-/// state field, unpacked from the struct each aggregate evaluates to.
+/// Projects the output of a Single- or Final-mode `aggregate` whose aggregates are all
+/// [`MergeAsPartialUDF`]s into the output of a Partial-mode one: its `num_grouping` grouping
+/// columns, then one column per state field, unpacked from the struct each aggregate evaluates to.
 pub(crate) fn unpack_state_structs(
     aggregate: Arc<dyn ExecutionPlan>,
     num_grouping: usize,