perf: stream a converted SortAggregate whose ordering a consumer needs

native/core/src/execution/jni_api.rs

@@ -4302,6 +4302,7 @@ mod partial_merge_spill_tests {
                 mode: AggregateMode::PartialMerge as i32,
                 expr_modes: vec![AggregateMode::PartialMerge as i32],
                 initial_input_buffer_offset: 2,
+                ordered_input: false,
             })),
             ..Default::default()
         }

native/core/src/execution/planner.rs

@@ -1674,21 +1674,57 @@ impl PhysicalPlanner {
                     })
                     .collect();
 
-                let aggregate: Arc<dyn ExecutionPlan> = Arc::new(
-                    datafusion::physical_plan::aggregates::AggregateExec::try_new(
-                        mode,
-                        group_by,
-                        aggr_expr,
-                        filter_exprs?,
-                        Arc::clone(&child.native_plan),
-                        Arc::clone(&schema),
-                    )?,
-                );
+                let num_group_exprs = group_by.expr().len();
+                let aggregate = datafusion::physical_plan::aggregates::AggregateExec::try_new(
+                    mode,
+                    group_by,
+                    aggr_expr,
+                    filter_exprs?,
+                    Arc::clone(&child.native_plan),
+                    Arc::clone(&schema),
+                )?;
+
+                if agg.ordered_input && *aggregate.input_order_mode() != InputOrderMode::Sorted {
+                    warn!(
+                        "Aggregate {} expected input sorted by its grouping expressions, \
+                         sorting its output instead",
+                        spark_plan.plan_id
+                    );
+                    let aggregate: Arc<dyn ExecutionPlan> = Arc::new(aggregate);
+                    let output_schema = aggregate.schema();
+                    let ordering = (0..num_group_exprs).map(|idx| {
+                        PhysicalSortExpr::new(
+                            Arc::new(Column::new(output_schema.field(idx).name(), idx)),
+                            SortOptions {
+                                descending: false,
+                                nulls_first: true,
+                            },
+                        )
+                    });
+                    let sorted: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+                        LexOrdering::new(ordering).unwrap(),
+                        Arc::clone(&aggregate),
+                    ));
+                    return Ok((
+                        scans,
+                        shuffle_scans,
+                        Arc::new(SparkPlan::new_with_additional(
+                            spark_plan.plan_id,
+                            sorted,
+                            vec![child],
+                            vec![aggregate],
+                        )),
+                    ));
+                }
 
                 Ok((
                     scans,
                     shuffle_scans,
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, aggregate, vec![child])),
+                    Arc::new(SparkPlan::new(
+                        spark_plan.plan_id,
+                        Arc::new(aggregate),
+                        vec![child],
+                    )),
                 ))
             }
 
@@ -6638,6 +6674,7 @@ mod tests {
                 mode: spark_operator::AggregateMode::Partial as i32,
                 expr_modes: vec![],
                 initial_input_buffer_offset: 0,
+                ordered_input: false,
             })),
         };
         let projection = Operator {
@@ -6663,6 +6700,44 @@ mod tests {
         assert_eq!("ScanExec", projection_exec.children[0].native_plan.name());
     }
 
+    #[test]
+    fn ordered_input_aggregate_streams_over_a_sort_and_sorts_its_output_otherwise() {
+        let aggregate = |child: Operator| {
+            op(
+                2,
+                OpStruct::HashAgg(spark_operator::HashAggregate {
+                    grouping_exprs: vec![create_bound_reference(0)],
+                    agg_exprs: vec![],
+                    mode: spark_operator::AggregateMode::Final as i32,
+                    expr_modes: vec![],
+                    initial_input_buffer_offset: 0,
+                    ordered_input: true,
+                }),
+                vec![child],
+            )
+        };
+        let planner = PhysicalPlanner::default();
+
+        let (_, _, over_sort) = planner
+            .create_plan(&aggregate(sort_op(1, leaf(0))), &mut vec![], 1)
+            .unwrap();
+        assert_eq!("AggregateExec", over_sort.native_plan.name());
+        assert!(over_sort.additional_native_plans.is_empty());
+        let line = datafusion::physical_plan::displayable(over_sort.native_plan.as_ref())
+            .one_line()
+            .to_string();
+        assert!(line.contains("ordering_mode=Sorted"), "{line}");
+
+        let (_, _, over_scan) = planner
+            .create_plan(&aggregate(leaf(0)), &mut vec![], 1)
+            .unwrap();
+        assert_eq!("SortExec", over_scan.native_plan.name());
+        assert_eq!(1, over_scan.additional_native_plans.len());
+        assert_eq!("AggregateExec", over_scan.additional_native_plans[0].name());
+        assert_eq!(1, over_scan.children.len());
+        assert_eq!("ScanExec", over_scan.children[0].native_plan.name());
+    }
+
     fn window_operator(functions: Vec<(&str, Vec<Expr>, bool)>) -> Operator {
         let int_literal = |value: Option<i32>| Expr {
             expr_struct: Some(Literal(spark_expression::Literal {

native/proto/src/proto/operator.proto

@@ -891,6 +891,7 @@ message HashAggregate {
   // Offset in the child's output where aggregate buffer attributes start.
   // Used by PartialMerge to locate state fields in the input.
   int32 initial_input_buffer_offset = 7;
+  bool ordered_input = 9;
 }
 
 message Limit {

native/spark-expr/src/agg_funcs/max_min_by.rs

@@ -339,6 +339,7 @@ struct MaxMinByGroupsAccumulator {
     best_ordering: Vec<OwnedRow>,
     /// Per-group flag: has a non-null ordering been seen yet?
     has_ordering: Vec<bool>,
+    row_bytes: usize,
 }
 
 impl MaxMinByGroupsAccumulator {
@@ -364,10 +365,14 @@ impl MaxMinByGroupsAccumulator {
             best_value: Vec::new(),
             best_ordering: Vec::new(),
             has_ordering: Vec::new(),
+            row_bytes: 0,
         })
     }
 
     fn resize(&mut self, total_num_groups: usize) {
+        let added = total_num_groups.saturating_sub(self.best_value.len());
+        self.row_bytes +=
+            added * (row_len(&self.null_value_row) + row_len(&self.null_ordering_row));
         self.best_value
             .resize(total_num_groups, self.null_value_row.clone());
         self.best_ordering
@@ -407,13 +412,34 @@ impl MaxMinByGroupsAccumulator {
             let take = !self.has_ordering[group_index]
                 || candidate >= self.best_ordering[group_index].row();
             if take {
+                let value = value_rows.row(idx);
+                self.row_bytes = self.row_bytes + candidate.as_ref().len() + value.as_ref().len()
+                    - row_len(&self.best_ordering[group_index])
+                    - row_len(&self.best_value[group_index]);
                 self.best_ordering[group_index] = candidate.owned();
-                self.best_value[group_index] = value_rows.row(idx).owned();
+                self.best_value[group_index] = value.owned();
                 self.has_ordering[group_index] = true;
             }
         }
         Ok(())
     }
+
+    fn take_rows(&mut self, emit_to: EmitTo) -> (Vec<OwnedRow>, Vec<OwnedRow>) {
+        let value_rows = emit_to.take_needed(&mut self.best_value);
+        let ordering_rows = emit_to.take_needed(&mut self.best_ordering);
+        let _ = emit_to.take_needed(&mut self.has_ordering);
+        let taken = value_rows
+            .iter()
+            .chain(ordering_rows.iter())
+            .map(row_len)
+            .sum::<usize>();
+        self.row_bytes -= taken;
+        (value_rows, ordering_rows)
+    }
+}
+
+fn row_len(row: &OwnedRow) -> usize {
+    row.row().as_ref().len()
 }
 
 impl GroupsAccumulator for MaxMinByGroupsAccumulator {
@@ -438,19 +464,15 @@ impl GroupsAccumulator for MaxMinByGroupsAccumulator {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let value_rows = emit_to.take_needed(&mut self.best_value);
-        let _ = emit_to.take_needed(&mut self.best_ordering);
-        let _ = emit_to.take_needed(&mut self.has_ordering);
+        let (value_rows, _) = self.take_rows(emit_to);
         let arrays = self
             .value_converter
             .convert_rows(value_rows.iter().map(|r| r.row()))?;
         Ok(Arc::clone(&arrays[0]))
     }
 
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        let value_rows = emit_to.take_needed(&mut self.best_value);
-        let ordering_rows = emit_to.take_needed(&mut self.best_ordering);
-        let _ = emit_to.take_needed(&mut self.has_ordering);
+        let (value_rows, ordering_rows) = self.take_rows(emit_to);
         let value_arrays = self
             .value_converter
             .convert_rows(value_rows.iter().map(|r| r.row()))?;
@@ -473,16 +495,12 @@ impl GroupsAccumulator for MaxMinByGroupsAccumulator {
 
     fn size(&self) -> usize {
         // `size_of::<OwnedRow>()` only covers the inline `Box<[u8]>` pointer/len, not the row
-        // bytes behind it, so the per-group payloads have to be summed separately. Leaving them
-        // out under-reports to the memory pool that drives spill decisions.
-        let row_bytes = |rows: &Vec<OwnedRow>| -> usize {
-            rows.iter().map(|r| r.row().as_ref().len()).sum::<usize>()
-        };
+        // bytes behind it, so the per-group payloads are counted separately in `row_bytes`.
+        // Leaving them out under-reports to the memory pool that drives spill decisions.
         size_of_val(self)
             + (self.best_value.capacity() + self.best_ordering.capacity())
                 * std::mem::size_of::<OwnedRow>()
-            + row_bytes(&self.best_value)
-            + row_bytes(&self.best_ordering)
+            + self.row_bytes
             + self.has_ordering.capacity()
     }
 }
@@ -892,4 +910,42 @@ mod tests {
             acc.size()
         );
     }
+
+    #[test]
+    fn groups_accumulator_row_bytes_track_updates_and_emits() {
+        let walked = |acc: &MaxMinByGroupsAccumulator| -> usize {
+            acc.best_value
+                .iter()
+                .chain(acc.best_ordering.iter())
+                .map(row_len)
+                .sum()
+        };
+        let mut acc = max_by_groups(DataType::Utf8, DataType::Int64);
+        let values: ArrayRef = Arc::new(StringArray::from(vec![
+            Some("a"),
+            Some("bbbbbbbbbbbbbbbb"),
+            None,
+            Some("cc"),
+            Some("dddddddddddddddddddddddddddddd"),
+        ]));
+        let ordering: ArrayRef = Arc::new(Int64Array::from(vec![
+            Some(1),
+            Some(5),
+            Some(9),
+            None,
+            Some(2),
+        ]));
+        acc.update_batch(&[values, ordering], &[0, 1, 1, 2, 3], None, 6)
+            .unwrap();
+        assert_eq!(acc.row_bytes, walked(&acc));
+        let values: ArrayRef = Arc::new(StringArray::from(vec![Some("eeeeeeeeeeee"), Some("f")]));
+        let ordering: ArrayRef = Arc::new(Int64Array::from(vec![Some(7), Some(0)]));
+        acc.update_batch(&[values, ordering], &[0, 3], None, 6)
+            .unwrap();
+        assert_eq!(acc.row_bytes, walked(&acc));
+        acc.state(EmitTo::First(2)).unwrap();
+        assert_eq!(acc.row_bytes, walked(&acc));
+        acc.evaluate(EmitTo::All).unwrap();
+        assert_eq!(acc.row_bytes, 0);
+    }
 }

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -22,7 +22,7 @@ package org.apache.comet.rules
 import scala.collection.mutable.ListBuffer
 
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction, AggregateMode, Average, Count, Final, Max, MaxMinBy, Min, Partial, PartialMerge, Sum}
 import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -161,6 +161,12 @@ object CometExecRule {
   val SORT_AGGREGATE_ORDER_TAG: TreeNodeTag[Unit] =
     TreeNodeTag[Unit]("comet.sortAggregateOrder")
 
+  val SORT_AGGREGATE_ORDER_CONSUMED_TAG: TreeNodeTag[Unit] =
+    TreeNodeTag[Unit]("comet.sortAggregateOrderConsumed")
+
+  val ORDERED_INPUT_AGGREGATE_TAG: TreeNodeTag[Unit] =
+    TreeNodeTag[Unit]("comet.orderedInputAggregate")
+
   /**
    * Serializes the native plan of each block of adjacent native operators into its topmost
    * operator. Blocks that already hold a serialized plan are left as they are, so this can run
@@ -657,6 +663,9 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false)
         }
     }
 
+    if (CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.get(conf)) {
+      markOrderConsumedSortAggregates(plan)
+    }
     plan.transformUp { case op =>
       val converted = convertNode(restoreSortAggregateOrdering(op))
       // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions
@@ -977,6 +986,32 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false)
     if (changed) op.withNewChildren(children) else op
   }
 
+  private def markOrderConsumedSortAggregates(plan: SparkPlan): Unit = plan.foreach { op =>
+    if (!op.isInstanceOf[CometPlan] && op.requiredChildOrdering.length == op.children.length) {
+      op.children.zip(op.requiredChildOrdering).foreach { case (child, ordering) =>
+        if (ordering.nonEmpty) {
+          sortAggregateProviding(child, ordering)
+            .foreach(_.setTagValue(CometExecRule.SORT_AGGREGATE_ORDER_CONSUMED_TAG, ()))
+        }
+      }
+    }
+  }
+
+  private def sortAggregateProviding(
+      plan: SparkPlan,
+      ordering: Seq[SortOrder]): Option[SortAggregateExec] = {
+    if (!SortOrder.orderingSatisfies(plan.outputOrdering, ordering)) {
+      None
+    } else {
+      plan match {
+        case agg: SortAggregateExec => Some(agg)
+        case _: SortExec | _: CometSortExec => None
+        case p if p.children.size == 1 => sortAggregateProviding(p.children.head, ordering)
+        case _ => None
+      }
+    }
+  }
+
   private def orderInsensitive(fn: AggregateFunction): Boolean = fn match {
     case _: Min | _: Max | _: Count | _: Sum | _: Average | _: MaxMinBy => true
     case _ => false
@@ -992,8 +1027,10 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false)
    * their input order, which the native hash aggregate does not guarantee. MAX_BY and MIN_BY
    * depend on it only among rows tied on the ordering, where they are non-deterministic in Spark
    * too and already differ from it on the native hash aggregate path. The sort by the grouping
-   * keys directly below the aggregate is then dropped. A consumer that relied on the ordering of
-   * the sort aggregate gets a sort back, see [[restoreSortAggregateOrdering]].
+   * keys directly below the aggregate is then dropped, unless a consumer relies on the ordering
+   * of the sort aggregate: then the sort stays and the native aggregate streams over it, keeping
+   * the ordering. Otherwise such a consumer gets a sort back, see
+   * [[restoreSortAggregateOrdering]].
    */
   private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = {
     val required = agg.requiredChildOrdering.headOption.getOrElse(Nil)
@@ -1009,7 +1046,11 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false)
         })
       case _ => false
     }
-    val input = if (sortedByKeys) agg.child.children.head else agg.child
+    val keepOrder = agg.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_CONSUMED_TAG).isDefined &&
+      agg.groupingExpressions.forall(_.isInstanceOf[Attribute]) &&
+      agg.child.isInstanceOf[CometNativeExec] &&
+      SortOrder.orderingSatisfies(agg.child.outputOrdering, required)
+    val input = if (sortedByKeys && !keepOrder) agg.child.children.head else agg.child
     if (!agg.aggregateExpressions.forall(e => orderInsensitive(e.aggregateFunction))) {
       withFallbackReason(
         agg,
@@ -1030,7 +1071,16 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false)
       agg.resultExpressions,
       input)
     hashAgg.copyTagsFrom(agg)
-    val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec)
+    if (keepOrder) {
+      hashAgg.setTagValue(CometExecRule.ORDERED_INPUT_AGGREGATE_TAG, ())
+    }
+    val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec).map {
+      case native: CometHashAggregateExec if keepOrder =>
+        val ordered = native.copy(originalPlan = agg)
+        ordered.copyTagsFrom(native)
+        ordered
+      case native => native
+    }
     converted.foreach(_.setTagValue(CometExecRule.CONVERTED_SORT_AGGREGATE_TAG, ()))
     if (converted.isEmpty) {
       hashAgg