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") } } }