From 9153d794f302c285cf8b8de17b95c02268509c46 Mon Sep 17 00:00:00 2001 From: msaf Date: Sun, 11 Oct 2026 11:53:40 +0100 Subject: [PATCH] perf: stream a converted SortAggregate whose ordering a consumer needs A SortAggregate converted to a native hash aggregate dropped the sort below it, and a consumer that relied on its ordering got a full sort back above it. With high-cardinality keys, such as the first aggregate of a COUNT(DISTINCT) expansion feeding Spark's next SortAggregate, the hash aggregate barely reduces the rows, spills them sorted, and the restored sort sorts them again (gold_device_profiles: 523 GiB aggregate spill plus 452 GiB sort spill, 131 vs 59 task-hours on 1.1-4). When a consumer requires the ordering, the sort now stays and the native aggregate runs over it; DataFusion detects the sorted input and streams the groups, so the output keeps the order and no sort is restored. The converted operator keeps the SortAggregateExec as its original plan, so a revert to Spark stays ordered. The native planner sorts the output if DataFusion does not detect the ordering, so the promised order always holds. Without such a consumer the sort is still dropped, as before. MaxMinByGroupsAccumulator::size walked every group after each batch, which made a hash aggregate with many groups quadratic; it now keeps a running byte count. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/jni_api.rs | 1 + native/core/src/execution/planner.rs | 97 ++++++++++++++-- native/proto/src/proto/operator.proto | 1 + native/spark-expr/src/agg_funcs/max_min_by.rs | 84 +++++++++++--- .../apache/comet/rules/CometExecRule.scala | 60 +++++++++- .../apache/spark/sql/comet/operators.scala | 10 ++ .../comet/exec/CometAggregateSuite.scala | 108 ++++++++++++++++-- 7 files changed, 320 insertions(+), 41 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 2109292025c..0e3d5bb1b59 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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() } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 63a983baef9..6416e315152 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1674,21 +1674,57 @@ impl PhysicalPlanner { }) .collect(); - let aggregate: Arc = 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 = 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 = 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, bool)>) -> Operator { let int_literal = |value: Option| Expr { expr_struct: Some(Literal(spark_expression::Literal { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 4d210b1e380..0d831858605 100644 --- a/native/proto/src/proto/operator.proto +++ b/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 { diff --git a/native/spark-expr/src/agg_funcs/max_min_by.rs b/native/spark-expr/src/agg_funcs/max_min_by.rs index 147f1eb92bf..6bedb196521 100644 --- a/native/spark-expr/src/agg_funcs/max_min_by.rs +++ b/native/spark-expr/src/agg_funcs/max_min_by.rs @@ -339,6 +339,7 @@ struct MaxMinByGroupsAccumulator { best_ordering: Vec, /// Per-group flag: has a non-null ordering been seen yet? has_ordering: Vec, + 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, Vec) { + 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::(); + self.row_bytes -= taken; + (value_rows, ordering_rows) + } +} + +fn row_len(row: &OwnedRow) -> usize { + row.row().as_ref().len() } impl GroupsAccumulator for MaxMinByGroupsAccumulator { @@ -438,9 +464,7 @@ impl GroupsAccumulator for MaxMinByGroupsAccumulator { } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - 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()))?; @@ -448,9 +472,7 @@ impl GroupsAccumulator for MaxMinByGroupsAccumulator { } fn state(&mut self, emit_to: EmitTo) -> Result> { - 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::()` 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| -> usize { - rows.iter().map(|r| r.row().as_ref().len()).sum::() - }; + // 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::() - + 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); + } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 15e0737df44..4ba100ab5e9 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/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 diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 1333c8c567f..8a93cd31730 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -2236,6 +2236,7 @@ trait CometBaseAggregate { resultExpressions, aggregate) } else { + setOrderedInput(hashAggBuilder, aggregate) Some(builder.setHashAgg(hashAggBuilder).build()) } } else { @@ -2254,12 +2255,21 @@ trait CometBaseAggregate { * ops share the caller's `plan_id` so the aggregate's native metrics roll up under the same * Spark operator. */ + private def setOrderedInput( + hashAggBuilder: OperatorOuterClass.HashAggregate.Builder, + aggregate: BaseAggregateExec): Unit = { + if (aggregate.getTagValue(CometExecRule.ORDERED_INPUT_AGGREGATE_TAG).isDefined) { + hashAggBuilder.setOrderedInput(true) + } + } + private def buildAggOp( builder: Operator.Builder, hashAggBuilder: OperatorOuterClass.HashAggregate.Builder, naturalOutput: Seq[Attribute], resultExpressions: Seq[NamedExpression], aggregate: BaseAggregateExec): Option[Operator] = { + setOrderedInput(hashAggBuilder, aggregate) if (resultExpressions.map(_.toAttribute) == naturalOutput) { return Some(builder.setHashAgg(hashAggBuilder).build()) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 04b3746a5d1..eb2819dcd99 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.{ExpandExec, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} @@ -3501,8 +3501,101 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("a converted sort aggregate keeps the ordering a sort-merge join relies on") { + private def orderedAggregates(plan: SparkPlan): Seq[CometHashAggregateExec] = + nativeAggregates(plan).filter { a => + a.originalPlan.isInstanceOf[SortAggregateExec] && a.child.isInstanceOf[CometSortExec] + } + + private def orderedInput(agg: CometHashAggregateExec): Boolean = { + val op = agg.nativeOp + val hashAgg = if (op.hasHashAgg) op.getHashAgg else op.getChildren(0).getHashAgg + hashAgg.getOrderedInput + } + + private def orderSorts(plan: SparkPlan): Seq[CometSortExec] = + collect(plan) { + case s: CometSortExec if s.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => + s + } + + for (aqe <- Seq("false", "true")) { + test(s"a sort aggregate whose ordering a sort-merge join relies on runs ordered (AQE=$aqe)") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withStrings { + withParquetTable((0 until 50).map(i => (i % 6, s"v$i")), "str_right") { + val query = + """SELECT a._1, a.m, r._2 FROM (SELECT _1, max(_2) AS m FROM str_tbl GROUP BY _1) a + |JOIN str_right r ON a._1 = r._1""".stripMargin + val (_, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + val join = collect(cometPlan) { case j: CometSortMergeJoinExec => j } + assert(join.size == 1, s"plan:\n$cometPlan") + assert(orderSorts(cometPlan).isEmpty, s"plan:\n$cometPlan") + val ordered = orderedAggregates(cometPlan) + assert(ordered.size == 1, s"plan:\n$cometPlan") + assert(orderedInput(ordered.head), s"plan:\n$cometPlan") + assert(ordered.head.modes == Seq(Final), s"plan:\n$cometPlan") + } + } + } + } + } + + test("sort aggregates over a pre-partitioned input stream into a sort-merge join") { withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + withParquetTable((0 until 50).map(i => (i % 6, s"v$i")), "str_right") { + val query = + """SELECT a.k, a.m, r._2 FROM + |(SELECT _1 + 1 AS k, max(_2) AS m + |FROM (SELECT * FROM str_tbl DISTRIBUTE BY _1 + 1) GROUP BY _1 + 1) a + |JOIN str_right r ON a.k = r._1""".stripMargin + val (_, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + assert(orderSorts(cometPlan).isEmpty, s"plan:\n$cometPlan") + val aggs = nativeAggregates(cometPlan) + assert(aggs.size == 2, s"plan:\n$cometPlan") + assert(aggs.forall(orderedInput), s"plan:\n$cometPlan") + } + } + } + } + + for (aqe <- Seq("false", "true")) { + test(s"a distinct-expansion aggregate feeding a sort aggregate runs ordered (AQE=$aqe)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val rows = (0 until 600).map { i => + (i % 37, if (i % 11 == 0) null else s"s${i % 13}", i % 17, if (i % 5 == 0) "x" else "y") + } + withParquetTable(rows, "expand_tbl") { + val query = + """SELECT _1, count(DISTINCT _3), count(DISTINCT _2), max(_2), max_by(_4, _3), + |sum(_3), min(_4) FROM expand_tbl GROUP BY _1""".stripMargin + val (sparkPlan, cometPlan) = checkSparkAnswer(sql(query)) + assert(collect(sparkPlan) { case e: ExpandExec => e }.nonEmpty, s"plan:\n$sparkPlan") + assert(orderSorts(cometPlan).isEmpty, s"plan:\n$cometPlan") + val ordered = orderedAggregates(cometPlan) + assert(ordered.size == 1, s"plan:\n$cometPlan") + assert(ordered.head.modes == Seq(Final), s"plan:\n$cometPlan") + assert(orderedInput(ordered.head), s"plan:\n$cometPlan") + assert( + nativeAggregates(cometPlan).exists(a => + a.modes == Seq(Partial) && !a.child.isInstanceOf[CometSortExec]), + s"plan:\n$cometPlan") + } + } + } + } + + test("an ordered sort aggregate reverted by the cost-based choice runs as a sort aggregate") { + withSQLConf( + CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key -> "true", + CometConf.COMET_EXEC_COST_BASED_ENGINES_COST_TABLE.key -> "agg.flat.comet=1000000,0,0", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { withStrings { @@ -3511,15 +3604,8 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { """SELECT a._1, a.m, r._2 FROM (SELECT _1, max(_2) AS m FROM str_tbl GROUP BY _1) a |JOIN str_right r ON a._1 = r._1""".stripMargin val (_, cometPlan) = checkSparkAnswer(sql(query)) - assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") - val join = collect(cometPlan) { case j: CometSortMergeJoinExec => j } - assert(join.size == 1, s"plan:\n$cometPlan") - val orderSorts = collect(cometPlan) { - case s: CometSortExec - if s.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => - s - } - assert(orderSorts.nonEmpty, s"plan:\n$cometPlan") + assert(orderedAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") } } }