From 3adfb1eccbbd1aebd4185f634296b74777d09541 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 08:42:10 +0100 Subject: [PATCH 1/8] fix: keep a columnar-to-row over a Comet broadcast out of whole-stage codegen A Spark BroadcastHashJoinExec can read a Comet broadcast through CometColumnarToRowExec, whose doExecuteBroadcast turns the Arrow batches into the join's relation. This happens under AQE when a re-optimization leaves the join in Spark after its build side was already planned as a Comet broadcast stage, for example when OptimizeSkewedJoin puts an AQEShuffleRead under the sort-merge join below it. When the join itself is outside whole-stage codegen (more output fields than spark.sql.codegen.maxFields), CollapseCodegenStages wrapped the CometColumnarToRowExec into a WholeStageCodegenExec of its own, and the join failed with "WholeStageCodegen (n) does not implement doExecuteBroadcast". CometColumnarToRowExec over a Comet broadcast exchange, a broadcast query stage or a reused exchange of one now reports supportCodegen = false, so it stays the join's build plan, or sits behind an InputAdapter that forwards executeBroadcast when the join is codegen'd. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../sql/comet/CometColumnarToRowExec.scala | 9 +++ .../apache/comet/exec/CometJoinSuite.scala | 57 ++++++++++++++++++- 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala index 98c1105edc1..de3ecb31206 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala @@ -65,6 +65,15 @@ case class CometColumnarToRowExec(child: SparkPlan) override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override def supportCodegen: Boolean = !readsBroadcast(child) + + private def readsBroadcast(plan: SparkPlan): Boolean = plan match { + case _: CometBroadcastExchangeExec => true + case stage: BroadcastQueryStageExec => readsBroadcast(stage.plan) + case reused: ReusedExchangeExec => readsBroadcast(reused.child) + case _ => false + } + // `ColumnarToRowExec` processes the input RDD directly, which is kind of a leaf node in the // codegen stage and needs to do the limit check. protected override def canCheckLimitNotReached: Boolean = true diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 8bb8d5ddf70..392ad516a00 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,11 +31,12 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometColumnarToRowExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} +import org.apache.spark.sql.execution.{InputAdapter, LocalTableScanExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} @@ -1479,6 +1480,56 @@ class CometJoinSuite extends CometTestBase { } } + for (costBased <- Seq("true", "false"); maxFields <- Seq("3", "100")) { + test( + "a Spark broadcast hash join over a skewed join reads a Comet broadcast stage " + + s"(costBasedEngines=$costBased, codegen.maxFields=$maxFields)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.SKEW_JOIN_ENABLED.key -> "true", + SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR.key -> "2", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100", + SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> maxFields, + CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key -> costBased) { + withParquetTable( + (0 until 4000).map(i => (if (i % 10 < 8) 0 else i % 50, i % 10, i)), + "skew_f") { + withParquetTable((0 until 500).map(i => (i % 50, i)), "skew_g") { + withParquetTable((0 until 10).map(i => (i, s"name_$i")), "skew_d") { + val query = + """SELECT /*+ BROADCAST(d), MERGE(f, g) */ f._1, f._2, f._3, g._2, d._2 + |FROM skew_f f JOIN skew_g g ON f._1 = g._1 + |LEFT JOIN skew_d d ON f._2 = d._1 + |DISTRIBUTE BY f._3""".stripMargin + val (_, plan) = checkSparkAnswer(sql(query)) + val skewed = collect(plan) { case j: SortMergeJoinExec if j.isSkewJoin => j } + assert(skewed.nonEmpty, s"expected a skewed Spark sort-merge join:\n$plan") + val sparkBroadcastJoins = collect(plan) { case j: BroadcastHashJoinExec => j } + assert(sparkBroadcastJoins.size == 1, s"plan:\n$plan") + val build = sparkBroadcastJoins.head.right + val readsCometBroadcast = build.collectFirst { case c: CometColumnarToRowExec => + c.child match { + case InputAdapter(stage) => stage + case stage => stage + } + } + assert( + readsCometBroadcast.exists(_.isInstanceOf[BroadcastQueryStageExec]), + s"expected the join to read a Comet broadcast stage:\n$plan") + assert(!build.isInstanceOf[WholeStageCodegenExec], s"plan:\n$plan") + } + } + } + } + } + } + test("BroadcastNestedLoopJoin with unequal filter") { withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { // Include NULL keys: predicate `_1 > _1` returns NULL with a NULL operand, so From 764adf614b0f3748bc9b35bc44585147a7549804 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 10:49:34 +0100 Subject: [PATCH 2/8] perf: evaluate JVM-dispatched filter conjuncts only on rows passing preceding conjuncts DataFusion's AND evaluates its right side on every row once more than 20% of rows pass the left side. A conjunct that runs through the JVM codegen dispatcher (ScalaUDF, regex family outside the native subset, other dispatched expressions) then costs a JNI round trip and Arrow export for rows Spark would never evaluate it on. CometFilterExec now splits the serialized deterministic predicate into its top-level conjuncts and rewrites `a AND b`, where b contains a JvmScalarUdf, into `CASE WHEN a THEN b END`. DataFusion's CaseExpr projects the columns b needs, filters them by a and evaluates b on the selected rows only. Null results reject rows just like false, so the filter output is unchanged. Scan pushdown is untouched because only the native filter predicate changes. Controlled by spark.comet.exec.filter.shortCircuitJvmDispatch.enabled (default true). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scala/org/apache/comet/CometConf.scala | 9 + .../apache/spark/sql/comet/operators.scala | 82 ++++++++- .../exec/CometFilterShortCircuitSuite.scala | 168 ++++++++++++++++++ 3 files changed, 256 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 03bd2701754..38f9129142a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -486,6 +486,15 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.filter.shortCircuitJvmDispatch.enabled") + .category(CATEGORY_EXEC) + .doc("Whether a native filter evaluates a top-level conjunct that runs through the JVM " + + "codegen dispatcher only on rows where all preceding conjuncts are true, by rewriting " + + "`a AND b` into `CASE WHEN a THEN b END`. Applies to deterministic conditions only.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") 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 0718697a794..1333c8c567f 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 @@ -53,13 +53,14 @@ import org.apache.spark.util.SerializableConfiguration import org.apache.spark.util.io.ChunkedByteBuffer import com.google.common.base.Objects -import com.google.protobuf.CodedOutputStream +import com.google.protobuf.{CodedOutputStream, Message} import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, ConfigEntry, ContribServices} import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.rules.CometExecRule -import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.{CometOperatorSerde, Compatible, ExprOuterClass, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} @@ -1428,15 +1429,90 @@ object CometFilterExec extends CometOperatorSerde[FilterExec] { val cond = exprToProto(op.condition, op.child.output) if (cond.isDefined && childOp.nonEmpty) { + val predicate = + if (CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.get() && + op.condition.deterministic) { + shortCircuitJvmDispatch(cond.get) + } else { + cond.get + } val filterBuilder = OperatorOuterClass.Filter .newBuilder() - .setPredicate(cond.get) + .setPredicate(predicate) Some(builder.setFilter(filterBuilder).build()) } else { None } } + def shortCircuitJvmDispatch(predicate: Expr): Expr = { + val conjuncts = flattenAnd(predicate) + val expensive = conjuncts.map(containsJvmScalarUdf) + if (!expensive.drop(1).contains(true)) { + predicate + } else { + val segments = ArrayBuffer(ArrayBuffer(conjuncts.head)) + conjuncts.zip(expensive).tail.foreach { case (conjunct, isExpensive) => + if (isExpensive) segments += ArrayBuffer(conjunct) else segments.last += conjunct + } + segments.tail.foldLeft(balancedAnd(segments.head.toIndexedSeq)) { (guard, segment) => + Expr + .newBuilder() + .setCaseWhen( + ExprOuterClass.CaseWhen + .newBuilder() + .addWhen(guard) + .addThen(balancedAnd(segment.toIndexedSeq))) + .build() + } + } + } + + private def flattenAnd(predicate: Expr): Seq[Expr] = { + val leaves = ArrayBuffer.empty[Expr] + var stack: List[Expr] = predicate :: Nil + while (stack.nonEmpty) { + val current = stack.head + stack = stack.tail + if (current.hasAnd) { + stack = current.getAnd.getLeft :: current.getAnd.getRight :: stack + } else { + leaves += current + } + } + leaves.toSeq + } + + private def balancedAnd(conjuncts: IndexedSeq[Expr]): Expr = { + if (conjuncts.length == 1) { + conjuncts.head + } else { + val mid = conjuncts.length / 2 + Expr + .newBuilder() + .setAnd( + ExprOuterClass.BinaryExpr + .newBuilder() + .setLeft(balancedAnd(conjuncts.slice(0, mid))) + .setRight(balancedAnd(conjuncts.slice(mid, conjuncts.length)))) + .build() + } + } + + private def containsJvmScalarUdf(message: Message): Boolean = message match { + case e: Expr if e.hasJvmScalarUdf => true + case _ => + message.getAllFields.values.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case children: java.util.List[_] => + children.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case _ => false + } + case _ => false + } + } + override def createExec(nativeOp: Operator, op: FilterExec): CometNativeExec = { CometFilterExec(nativeOp, op, op.output, op.condition, op.child, SerializedPlan(None)) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala new file mode 100644 index 00000000000..e95e8455f72 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala @@ -0,0 +1,168 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.comet.{CometFilterExec, CometNativeScanExec, CometScanExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + +import org.apache.comet.CometConf +import org.apache.comet.serde.ExprOuterClass.Expr + +class CometFilterShortCircuitSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val backreference = "^(\\\\d)\\\\1" + + override def beforeAll(): Unit = { + super.beforeAll() + spark.udf.register( + "strict_even", + (x: Int) => { + if (x % 3 != 0) throw new IllegalStateException(s"strict_even evaluated on $x") + x % 2 == 0 + }) + spark.udf.register("is_even", (x: Int) => x % 2 == 0) + } + + private def withData(f: => Unit): Unit = { + withTable("t") { + sql("""CREATE TABLE t USING parquet AS + |SELECT + | IF(id % 7 = 0, NULL, CAST(id AS INT)) AS x, + | IF(id % 5 = 0, NULL, id % 4 = 0) AS b, + | IF(id % 11 = 0, NULL, CAST(id * 37 % 1000 AS STRING)) AS s + |FROM range(4000)""".stripMargin) + f + } + } + + private def filterPredicates(df: DataFrame): Seq[Expr] = + collect(df.queryExecution.executedPlan) { case f: CometFilterExec => + f.nativeOp.getFilter.getPredicate + } + + private def caseWhenDepth(e: Expr): Int = + if (e.hasCaseWhen) 1 + caseWhenDepth(e.getCaseWhen.getWhen(0)) else 0 + + private def scanDataFilters(df: DataFrame): Seq[Expression] = + collect(df.queryExecution.executedPlan) { + case s: CometNativeScanExec => s.dataFilters + case s: CometScanExec => s.dataFilters + }.flatten + + private def checkRewritten(query: String, depth: Int = 1): Unit = { + val df = sql(query) + checkSparkAnswerAndOperator(df) + val predicates = filterPredicates(df) + assert(predicates.nonEmpty, df.queryExecution.executedPlan) + assert(predicates.map(caseWhenDepth).max == depth, predicates) + } + + private def checkNotRewritten(query: String): Unit = { + val df = sql(query) + checkSparkAnswer(df) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen), filterPredicates(df)) + } + + test("expensive conjunct is evaluated only on rows passing preceding conjuncts") { + withData { + checkRewritten("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + checkRewritten("SELECT x FROM t WHERE x > 10 AND x % 3 = 0 AND strict_even(x) AND x < 3000") + } + } + + test("disabled rewrite evaluates the expensive conjunct on every row") { + withData { + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + val e = intercept[Exception](df.collect()) + assert( + e.toString.contains("strict_even evaluated on") || + Option(e.getCause).exists(_.toString.contains("strict_even evaluated on"))) + } + } + } + + test("three-valued logic matches Spark") { + withData { + checkRewritten("SELECT x, b, s FROM t WHERE b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE NOT b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE (b OR x IS NULL) AND is_even(x)") + checkRewritten(s"SELECT x, b, s FROM t WHERE b AND s RLIKE '$backreference'") + checkRewritten(s"SELECT x, b, s FROM t WHERE x % 3 = 0 AND NOT (s RLIKE '$backreference')") + checkRewritten("SELECT x, b, s FROM t WHERE x % 2 = 0 AND (NOT is_even(x) OR b)") + } + } + + test("nested AND/OR and multiple expensive conjuncts") { + withData { + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND (b OR s RLIKE '$backreference') AND x > 100") + checkRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) AND (is_even(x) AND s RLIKE '$backreference')", + depth = 2) + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND strict_even(x) AND b AND s RLIKE '$backreference'", + depth = 2) + checkNotRewritten(s"SELECT x, s FROM t WHERE x % 3 = 0 OR s RLIKE '$backreference'") + checkNotRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) OR (is_even(x) AND s RLIKE '$backreference')") + } + } + + test("no rewrite without a preceding conjunct or JVM dispatch") { + withData { + checkNotRewritten("SELECT x FROM t WHERE is_even(x)") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND b AND s IS NOT NULL") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND s LIKE '1%'") + } + } + + test("nondeterministic condition is not rewritten") { + withData { + val df = sql( + "SELECT x FROM t WHERE monotonically_increasing_id() % 5 < 4 AND x % 3 = 0 AND is_even(x)") + assert(filterPredicates(df).nonEmpty, df.queryExecution.executedPlan) + assert( + filterPredicates(df).forall(p => !p.hasCaseWhen), + df.queryExecution.optimizedPlan.toString + df.queryExecution.executedPlan) + checkSparkAnswer(df) + } + } + + test("scan pushdown is unchanged") { + withData { + val query = "SELECT x FROM t WHERE x > 10 AND b AND is_even(x)" + val rewritten = sql(query) + var baselineFilters: Seq[Expression] = Nil + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql(query) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + baselineFilters = scanDataFilters(df) + } + assert(filterPredicates(rewritten).exists(_.hasCaseWhen)) + assert(scanDataFilters(rewritten).nonEmpty) + assert(scanDataFilters(rewritten) == baselineFilters) + checkSparkAnswerAndOperator(rewritten) + } + } +} From 6f7a970eb61fbbd872404a6fb32d3061e6498593 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 11:58:01 +0100 Subject: [PATCH 3/8] feat: native MIN/MAX over strings and SortAggregate as a native hash aggregate Spark plans MIN or MAX over a string as a SortAggregate, since the aggregation buffer is not mutable, and Comet did not convert SortAggregate, so such an aggregate and often the joins around it ran in Spark. MIN and MAX now accept StringType with the default UTF8_BINARY collation. The native min/max compare UTF-8 bytes, which is Spark's binary order; collated strings stay in Spark. A SortAggregateExec is converted to a CometHashAggregateExec built from an ObjectHashAggregateExec with the same fields. Spark runs that operator over any input order, so an aggregate reverted to Spark later (cost-based engine choice, unsafe partial aggregates) stays correct, and the cost-based choice prices it as an object hash aggregate. A sort by exactly the grouping keys directly below the aggregate is dropped; without one only order-insensitive aggregates (MIN, MAX, COUNT, SUM, AVG) are converted. A native sort above the aggregate restores the output ordering a consumer such as a sort-merge join may rely on, and is dropped again when a shuffle reads it directly. Aggregates Comet does not support keep the SortAggregate. spark.comet.exec.sortAggregate.enabled (default true) turns the conversion off. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scala/org/apache/comet/CometConf.scala | 7 + .../apache/comet/rules/CometExecRule.scala | 104 +++++++++++++- .../org/apache/comet/serde/aggregates.scala | 10 +- .../expressions/aggregate/first_last.sql | 4 +- .../expressions/aggregate/min_max.sql | 2 +- .../comet/exec/CometAggregateSuite.scala | 127 +++++++++++++++++- .../rules/CostBasedEngineChoiceSuite.scala | 24 +++- .../rules/WideRowSortFallbackSuite.scala | 1 + 8 files changed, 267 insertions(+), 12 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 03bd2701754..4c8ca3b402c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -237,6 +237,13 @@ object CometConf extends ShimCometConf { createExecEnabledConfig("sortMergeJoin", defaultValue = true) val COMET_EXEC_AGGREGATE_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("aggregate", defaultValue = true) + val COMET_EXEC_SORT_AGGREGATE_ENABLED: ConfigEntry[Boolean] = + createExecEnabledConfig( + "sortAggregate", + defaultValue = true, + notes = Some( + "When enabled, a SortAggregate that Comet can run is converted to a native hash " + + "aggregate, with a native sort restoring its output ordering")) val COMET_EXEC_COLLECT_LIMIT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("collectLimit", defaultValue = true) val COMET_EXEC_COALESCE_ENABLED: ConfigEntry[Boolean] = 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 65b85540e41..2cc826e0eef 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -23,7 +23,7 @@ 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} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction, AggregateMode, Average, Count, Final, Max, Min, Partial, PartialMerge, Sum} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag @@ -35,7 +35,7 @@ import org.apache.spark.sql.comet.shims.ShimCometEmptyRelation import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} @@ -150,6 +150,14 @@ object CometExecRule { */ val ENGINE_CHOICE_SPARK_TAG: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.engineChoiceSpark") + /** + * Tag set on the native sort that restores the output ordering of a `SortAggregateExec` + * converted to a native hash aggregate. A shuffle directly above it discards the ordering, so + * the sort is dropped there. + */ + val SORT_AGGREGATE_ORDER_TAG: TreeNodeTag[Unit] = + TreeNodeTag[Unit]("comet.sortAggregateOrder") + /** * 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 @@ -587,6 +595,9 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) convertToComet(s, CometShuffleExchangeExec) .getOrElse(preserveSparkAggregateBuffers(s)) + case agg: SortAggregateExec if CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.get(conf) => + convertSortAggregate(agg).getOrElse(agg) + case op => // if all children are native (or if this is a leaf node) then see if there is a // registered handler for creating a fully native plan @@ -644,7 +655,7 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } plan.transformUp { case op => - val converted = convertNode(op) + val converted = convertNode(dropSortAggregateOrderUnderShuffle(op)) // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions // when the broadcast child has a Comet plan underneath. This enables exchange reuse // between the DPP subquery and the join's CometBroadcastExchangeExec because both @@ -928,6 +939,93 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } } + private def dropSortAggregateOrderUnderShuffle(op: SparkPlan): SparkPlan = op match { + case shuffle: ShuffleExchangeLike => + shuffle.child match { + case sort: CometSortExec + if sort.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => + shuffle.withNewChildren(Seq(sort.child)) + case _ => shuffle + } + case other => other + } + + private def orderInsensitive(fn: AggregateFunction): Boolean = fn match { + case _: Min | _: Max | _: Count | _: Sum | _: Average => true + case _ => false + } + + /** + * A `SortAggregateExec` as a native hash aggregate. Spark plans one when an aggregation buffer + * is not mutable, such as MIN or MAX over strings, and runs it over input sorted by the + * grouping keys. The hash aggregate starts from an `ObjectHashAggregateExec` with the same + * fields, which Spark runs correctly over any input order, so any operator reverted to Spark + * later stays correct. The sort by the grouping keys directly below the aggregate is dropped, + * as rows within a group were in no particular order anyway. Without such a sort the input + * order could carry meaning, so only aggregates that do not depend on it are converted. A + * native sort above restores the ordering consumers may rely on. + */ + private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = { + val required = agg.requiredChildOrdering.headOption.getOrElse(Nil) + val sortedByKeys = agg.child match { + case sort: CometSortExec => + sort.sortOrder.length == required.length && + sort.sortOrder.zip(required).forall { case (a, b) => + a.semanticEquals(b) + } && + (sort.originalPlan match { + case s: SortExec => !s.global + case _ => false + }) + case _ => false + } + val input = if (sortedByKeys) agg.child.children.head else agg.child + if (!sortedByKeys && + !agg.aggregateExpressions.forall(e => orderInsensitive(e.aggregateFunction))) { + withFallbackReason( + agg, + "SortAggregate over input not sorted by its grouping keys alone may depend on the " + + "input order") + return None + } + if (!input.isInstanceOf[CometNativeExec]) { + return None + } + val hashAgg = ObjectHashAggregateExec( + agg.requiredChildDistributionExpressions, + agg.isStreaming, + agg.numShufflePartitions, + agg.groupingExpressions, + agg.aggregateExpressions, + agg.aggregateAttributes, + agg.initialInputBufferOffset, + agg.resultExpressions, + input) + hashAgg.copyTagsFrom(agg) + val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec).flatMap { native => + val ordering = agg.outputOrdering + if (ordering.isEmpty) { + Some(native) + } else { + val sort = SortExec(ordering, global = false, native) + agg.logicalLink.foreach(sort.setLogicalLink) + tryConvertToComet(sort, CometSortExec).map { nativeSort => + nativeSort.setTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG, ()) + nativeSort + } + } + } + if (converted.isEmpty) { + hashAgg + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .foreach(reasons => withFallbackReasons(agg, reasons)) + if (!hasFallbackReason(agg)) { + withFallbackReason(agg, "SortAggregate could not be converted to a native hash aggregate") + } + } + converted + } + /** Convert a Spark plan to a Comet plan using the specified serde handler */ private def convertToComet(op: SparkPlan, handler: CometOperatorSerde[_]): Option[SparkPlan] = { val converted = tryConvertToComet(op, handler) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index ad53cceb31b..bc45c605025 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -1395,7 +1395,7 @@ object CometMode extends CometAggregateExpressionSerde[Mode] with CometTypeShim } } -object AggSerde { +object AggSerde extends CometTypeShim { import org.apache.spark.sql.types._ def minMaxDataTypeSupported(dt: DataType): Boolean = { @@ -1436,7 +1436,13 @@ object AggSerde { /** Shared support level for `Min` / `Max` based on the result data type. */ def minMaxSupportLevel(dt: DataType): SupportLevel = { - if (!minMaxDataTypeSupported(dt)) { + if (dt.isInstanceOf[StringType]) { + if (isStringCollationType(dt)) { + Unsupported(Some(s"Unsupported collated string type: $dt")) + } else { + Compatible() + } + } else if (!minMaxDataTypeSupported(dt)) { Unsupported(Some(s"Unsupported data type: $dt")) } else if ((dt == FloatType || dt == DoubleType) && COMET_EXEC_STRICT_FLOATING_POINT.get()) { diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql index 8aedc4f5f17..7fd5dcde731 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql @@ -183,7 +183,7 @@ SELECT first(val) IGNORE NULLS FROM test_single_null -- first IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +query SELECT grp, first(i_val) IGNORE NULLS, first(l_val) IGNORE NULLS, @@ -327,7 +327,7 @@ SELECT last(val) IGNORE NULLS FROM test_single_null -- last IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +query SELECT grp, last(i_val) IGNORE NULLS, last(l_val) IGNORE NULLS, diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql index dbec88e2f66..9fcde06186e 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql @@ -21,7 +21,7 @@ CREATE TABLE test_min_max(i int, d double, s string, grp string) USING parquet statement INSERT INTO test_min_max VALUES (1, 1.5, 'b', 'x'), (3, 3.5, 'a', 'x'), (2, 2.5, 'c', 'y'), (NULL, NULL, NULL, 'y'), (-1, -1.5, 'z', 'x') -query expect_fallback(SortAggregate is not supported) +query SELECT min(i), max(i), min(d), max(d), min(s), max(s) FROM test_min_max query 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 e82831ec426..20cbd42042b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,11 +31,11 @@ import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} 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} +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.SQLExecution +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.functions.{avg, col, count_distinct, expr, sum} import org.apache.spark.sql.internal.SQLConf @@ -3448,4 +3448,125 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + private val stringRows: Seq[(Integer, String)] = Seq( + (1, "b"), + (1, ""), + (1, null), + (1, "a"), + (2, null), + (2, null), + (3, "\u00e9t\u00e9"), + (3, "ete"), + (3, "\u65e5\u672c"), + (4, "\uff61"), + (4, "\ud83d\ude00"), + (4, "z"), + (5, ""), + (null, "x"), + (null, "\u00ff")) + + private def withStrings(f: => Unit): Unit = + withParquetTable(stringRows, "str_tbl")(f) + + private def sortAggregates(plan: SparkPlan): Seq[SortAggregateExec] = + collect(plan) { case a: SortAggregateExec => a } + + private def nativeAggregates(plan: SparkPlan): Seq[CometHashAggregateExec] = + collect(plan) { case a: CometHashAggregateExec => a } + + for (aqe <- Seq("false", "true")) { + test(s"min and max over strings run natively, grouped and ungrouped (AQE=$aqe)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withStrings { + for (query <- Seq( + "SELECT min(_2), max(_2), count(_2) FROM str_tbl", + "SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1", + "SELECT _1, max(_2) FROM str_tbl WHERE _1 = 2 GROUP BY _1", + "SELECT max(_2) FROM str_tbl WHERE _1 > 100", + "SELECT _1, min(_2), sum(_1) FROM str_tbl GROUP BY _1")) { + val (sparkPlan, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(sparkPlan).nonEmpty, s"$query:\n$sparkPlan") + assert(sortAggregates(cometPlan).isEmpty, s"$query:\n$cometPlan") + assert(nativeAggregates(cometPlan).nonEmpty, s"$query:\n$cometPlan") + } + } + } + } + } + + test("min and max over strings compare UTF-8 bytes as Spark does") { + withStrings { + checkSparkAnswerAndOperator(sql("SELECT min(_2), max(_2) FROM str_tbl WHERE _1 = 4")) + checkSparkAnswerAndOperator(sql("SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1")) + } + } + + test("a converted sort aggregate keeps the ordering a sort-merge join relies on") { + 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._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") + } + } + } + } + + test("the ordering sort of a converted partial sort aggregate is dropped under its shuffle") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + val (_, cometPlan) = checkSparkAnswer(sql("SELECT _1, max(_2) FROM str_tbl GROUP BY _1")) + val shuffles = collect(cometPlan) { case s: CometShuffleExchangeExec => s } + assert(shuffles.nonEmpty, s"plan:\n$cometPlan") + assert(shuffles.forall(s => !s.child.isInstanceOf[CometSortExec]), s"plan:\n$cometPlan") + assert(nativeAggregates(cometPlan).size == 2, s"plan:\n$cometPlan") + } + } + } + + test("a sort aggregate Comet cannot run natively stays in Spark") { + withStrings { + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT _1, max_by(_2, _2) FROM str_tbl GROUP BY _1")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } + + test("an order-sensitive sort aggregate over input sorted beyond its keys stays in Spark") { + withStrings { + val (_, cometPlan) = checkSparkAnswer( + sql("SELECT _1, first(_2) FROM (SELECT * FROM str_tbl DISTRIBUTE BY _1 SORT BY _1, _2) " + + "GROUP BY _1")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } + + test("a converted sort aggregate reverted by the cost-based choice runs as an object hash") { + 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.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1")) + assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + assert( + collect(cometPlan) { case a: ObjectHashAggregateExec => a }.nonEmpty, + s"plan:\n$cometPlan") + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala index 25ec4c2e0eb..2e2ff2673e6 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala @@ -19,6 +19,7 @@ package org.apache.comet.rules +import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.expressions.aggregate.Partial @@ -27,7 +28,7 @@ import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.{ExpandExec, ProjectExec, SortExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} -import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{col, sum} @@ -43,6 +44,9 @@ import org.apache.comet.rules.EngineCostTable.CostClass._ class CostBasedEngineChoiceSuite extends CometTestBase { + override protected def sparkConf: SparkConf = + super.sparkConf.set(CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key, "false") + private val flag = CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key private val costTable = CometConf.COMET_EXEC_COST_BASED_ENGINES_COST_TABLE.key private val expensiveNativeSort = costTable -> "sort.flat.comet=1000,0,0" @@ -1669,4 +1673,22 @@ class CostBasedEngineChoiceSuite extends CometTestBase { } } } + + test("a sort aggregate converted to a native hash aggregate is priced as an object hash") { + withTables { + withSQLConf( + CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + flag -> "true") { + val plan = run("SELECT k, max(s) FROM t GROUP BY k") + assert(count(plan) { case a: SortAggregateExec => a } == 0, s"plan:\n$plan") + val aggs = nodes(plan).collect { case a: CometHashAggregateExec => a } + assert(aggs.size == 2, s"plan:\n$plan") + aggs.foreach { agg => + assert(agg.originalPlan.isInstanceOf[ObjectHashAggregateExec], s"plan:\n$plan") + assert(model.costClasses(agg) == Seq(Agg, AggObjectHash), s"plan:\n$plan") + } + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala index 37253b2485c..5de75ae186b 100644 --- a/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala @@ -43,6 +43,7 @@ class WideRowSortFallbackSuite extends CometTestBase { override protected def sparkConf: SparkConf = super.sparkConf + .set(CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key, "false") .set(shuffleMinLeaves, "0") .set(CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key, "false") From a675e548904726c87f12f12df3991f2d48608619 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 12:28:30 +0100 Subject: [PATCH 4/8] fix: spill a native sort before output only when the JVM reads that output spark.comet.exec.sort.spillBeforeOutputThreshold (568955d21) protects a Spark sort, sort-merge join or window that reads a native sort's output through a columnar-to-row transition in the same task: the native sort cannot be made to release memory, so it spills its in-memory input before producing output and holds only the merge buffers. The threshold was a session extension, so every ExternalSorter in the native plan applied it, including sorts whose output never leaves the plan: under a window, a window group limit, a sort-merge join, an aggregate or the native shuffle writer. Those sorts wrote and read back their whole input once per task for nothing, because their consumer reserves memory from the same native pool and the sort releases it as its output is merged. SortExec now carries a spill_before_output flag, kept by with_fetch, projection swaps and filter pushdown, and ExternalSorter applies the threshold only when it is set. The planner sets it for a Sort reached from the root of the native plan through streaming operators only (Projection, Filter, Limit, Sample, Expand, Explode), whose output is what the JVM consumes. Window, partition-aggregate and sorted window operators do not use ExternalSorter and are unaffected. On a 1/40 sample of sat_product_info_loader's main stage (scan, window, sort-merge join, union, sort, native shuffle) with the joom-1.1-6 image and the patched library loaded from java.library.path, the query took 0.736 task-h with no spill, the same as threshold 0 (0.733-0.756), against 0.824-0.839 task-h and 29.2 GB spilled before. On the 2026-10-10 night, CometSort operators under native consumers spilled 21.6 TB, about one spill per task, mostly under window group limit, window, sort-merge join and native shuffle. Tests: the planner marks only sorts reached through streaming operators, and a planned SortExec carries the flag; the native sort with a JVM consumer still spills before output at the threshold, and with a native consumer it keeps its input in memory; in CometExecSuite, with a 1 byte threshold, sorts read through columnar-to-row spill and sorts under a window and a sort-merge join do not, all matching Spark. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/jni_api.rs | 36 +++- native/core/src/execution/planner.rs | 166 +++++++++++++++++- .../src/sorts/sort.rs | 30 +++- .../apache/comet/exec/CometExecSuite.scala | 35 ++++ 4 files changed, 256 insertions(+), 11 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 8fcaf964ca4..2109292025c 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1241,6 +1241,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) .with_sql_text_pool(&exec_context.spark_plan) + .with_jvm_output_sorts(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) .with_class_loader(exec_context.class_loader.clone()) .with_shuffle_partition_pusher( @@ -3657,6 +3658,7 @@ mod native_sort_spill_tests { share: usize, executor_cores: usize, spill_before_output: Option, + jvm_consumer: bool, ) -> OutputTrace { let shape = source.shape; let rows = source.rows; @@ -3708,7 +3710,8 @@ mod native_sort_spill_tests { }, )]) .unwrap(); - let sort: Arc = Arc::new(SortExec::new(ordering, child)); + let sort: Arc = + Arc::new(SortExec::new(ordering, child).with_spill_before_output(jvm_consumer)); let mut stream = sort.execute(0, session.task_ctx()).unwrap(); let mut trace = OutputTrace { reserved: vec![], @@ -3744,8 +3747,14 @@ mod native_sort_spill_tests { let share = SPILL_BEFORE_OUTPUT_SHARE; for rows_per_batch in [8192, 3] { let case = format!("{shape:?} rows={rows} rows_per_batch={rows_per_batch}"); - let held = - sort_and_trace(KibRows::new(shape, rows, rows_per_batch), share, 8, None).await; + let held = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + None, + true, + ) + .await; eprintln!( "{case} off: spill_count={} peak={}MiB reserved at 0%={}MiB 50%={}MiB 90%={}MiB", held.spill_count(), @@ -3762,6 +3771,7 @@ mod native_sort_spill_tests { share, 8, Some(share / 4), + true, ) .await; eprintln!( @@ -3782,11 +3792,31 @@ mod native_sort_spill_tests { share, 8, Some(held.peak_reserved), + true, ) .await; assert_eq!(below.spill_count(), 0, "{case}"); assert_eq!(below.reserved.len(), held.reserved.len(), "{case}"); assert!(below.reserved_at(0.9) >= input_bytes(rows) / 2, "{case}"); + + let native_consumer = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + Some(share / 4), + false, + ) + .await; + assert_eq!(native_consumer.spill_count(), 0, "{case}"); + assert_eq!( + native_consumer.reserved.len(), + held.reserved.len(), + "{case}" + ); + assert!( + native_consumer.reserved_at(0.9) >= input_bytes(rows) / 2, + "{case}" + ); } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1febe577337..63a983baef9 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -314,6 +314,29 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + jvm_output_sorts: std::collections::HashSet, +} + +pub(crate) fn jvm_output_sorts(root: &Operator) -> std::collections::HashSet { + let mut sorts = std::collections::HashSet::new(); + let mut pending = vec![root]; + while let Some(op) = pending.pop() { + match op.op_struct.as_ref() { + Some(OpStruct::Sort(_)) => { + sorts.insert(op.plan_id); + } + Some( + OpStruct::Projection(_) + | OpStruct::Filter(_) + | OpStruct::Limit(_) + | OpStruct::Sample(_) + | OpStruct::Expand(_) + | OpStruct::Explode(_), + ) => pending.extend(op.children.iter()), + _ => {} + } + } + sorts } impl Default for PhysicalPlanner { @@ -333,12 +356,18 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + jvm_output_sorts: std::collections::HashSet::new(), } } /// Load the SQL text pool from the root operator of the plan about to be planned. Must be /// called with the *root* operator: the JVM only populates the pool there, and /// `QueryContext.sql_text_idx` values are indices into it. + pub fn with_jvm_output_sorts(mut self, root: &Operator) -> Self { + self.jvm_output_sorts = jvm_output_sorts(root); + self + } + pub fn with_sql_text_pool(mut self, root: &Operator) -> Self { self.sql_text_pool = root .sql_text_pool @@ -1783,7 +1812,8 @@ impl PhysicalPlanner { LexOrdering::new(exprs?).unwrap(), Arc::clone(&child.native_plan), ) - .with_fetch(fetch), + .with_fetch(fetch) + .with_spill_before_output(self.jvm_output_sorts.contains(&spark_plan.plan_id)), ); if let Some(skip) = sort.skip.filter(|&n| n > 0).map(|n| n as usize) { @@ -5391,7 +5421,7 @@ mod tests { use crate::execution::operators::{ExecutionError, PartitionedRankLimitExec, WindowFnKind}; use crate::execution::planner::{ - convert_spark_types_to_arrow_schema, literal_to_array_ref, + convert_spark_types_to_arrow_schema, jvm_output_sorts, literal_to_array_ref, parse_file_scan_tasks_from_common, }; use crate::execution::shuffle::CometPartitioning; @@ -6755,6 +6785,138 @@ mod tests { } } + fn op(plan_id: u32, op_struct: OpStruct, children: Vec) -> Operator { + Operator { + plan_id, + sql_text_pool: vec![], + children, + op_struct: Some(op_struct), + } + } + + fn sort_op(plan_id: u32, child: Operator) -> Operator { + let order = spark_expression::Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(spark_expression::Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(create_proto_datatype()), + })), + query_context: None, + expr_id: None, + })), + direction: spark_expression::SortDirection::Ascending as i32, + null_ordering: spark_expression::NullOrdering::NullsFirst as i32, + }, + ))), + query_context: None, + expr_id: None, + }; + op( + plan_id, + OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![order], + fetch: None, + skip: None, + }), + vec![child], + ) + } + + fn leaf(plan_id: u32) -> Operator { + let mut scan = create_scan(); + scan.plan_id = plan_id; + scan + } + + #[test] + fn jvm_output_sorts_are_sorts_reached_from_the_root_through_streaming_operators() { + use std::collections::HashSet; + let ids = |v: &[u32]| v.iter().copied().collect::>(); + + assert_eq!(jvm_output_sorts(&sort_op(1, leaf(0))), ids(&[1])); + + let projected = op( + 3, + OpStruct::Projection(Default::default()), + vec![op( + 2, + OpStruct::Filter(Default::default()), + vec![sort_op(1, leaf(0))], + )], + ); + assert_eq!(jvm_output_sorts(&projected), ids(&[1])); + + let nested = sort_op( + 2, + op( + 3, + OpStruct::Window(Default::default()), + vec![sort_op(1, leaf(0))], + ), + ); + assert_eq!(jvm_output_sorts(&nested), ids(&[2])); + + for consumer in [ + OpStruct::ShuffleWriter(Default::default()), + OpStruct::Window(Default::default()), + OpStruct::HashAgg(Default::default()), + OpStruct::ParquetWriter(Default::default()), + OpStruct::WindowGroupLimit(Default::default()), + ] { + let plan = op(2, consumer, vec![sort_op(1, leaf(0))]); + assert!(jvm_output_sorts(&plan).is_empty(), "{plan:?}"); + } + + let join = op( + 5, + OpStruct::SortMergeJoin(Default::default()), + vec![sort_op(1, leaf(0)), sort_op(2, leaf(3))], + ); + assert!(jvm_output_sorts(&join).is_empty()); + assert!( + jvm_output_sorts(&op(6, OpStruct::Projection(Default::default()), vec![join])) + .is_empty() + ); + } + + fn sort_execs(plan: &Arc, out: &mut Vec) { + if let Some(sort) = plan.downcast_ref::() { + out.push(sort.spill_before_output()); + } + for child in plan.children() { + sort_execs(child, out); + } + } + + fn planned_sort_flags(root: &Operator) -> Vec { + let planner = PhysicalPlanner::default().with_jvm_output_sorts(root); + let (_, _, plan) = planner.create_plan(root, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + flags + } + + #[test] + fn planner_lets_only_a_sort_read_by_the_jvm_spill_before_output() { + assert_eq!(planned_sort_flags(&sort_op(1, leaf(0))), vec![true]); + assert_eq!( + planned_sort_flags(&create_filter(sort_op(1, leaf(0)), 3)), + vec![true] + ); + assert_eq!( + planned_sort_flags(&sort_op(2, sort_op(1, leaf(0)))), + vec![true, false] + ); + let sort = sort_op(1, leaf(0)); + let planner = PhysicalPlanner::default(); + let (_, _, plan) = planner.create_plan(&sort, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + assert_eq!(flags, vec![false]); + } + fn create_scan() -> Operator { Operator { plan_id: 0, diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index 51514d63b94..ce3173f5f73 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -1241,6 +1241,7 @@ pub struct SortExec { /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. /// If `fetch` is `None`, this will be `None`. filter: Option>>, + spill_before_output: bool, } impl SortExec { @@ -1260,9 +1261,19 @@ impl SortExec { common_sort_prefix: sort_prefix, cache: Arc::new(cache), filter: None, + spill_before_output: false, } } + pub fn with_spill_before_output(mut self, spill_before_output: bool) -> Self { + self.spill_before_output = spill_before_output; + self + } + + pub fn spill_before_output(&self) -> bool { + self.spill_before_output + } + /// Whether this `SortExec` preserves partitioning of the children pub fn preserve_partitioning(&self) -> bool { self.preserve_partitioning @@ -1336,6 +1347,7 @@ impl SortExec { fetch: self.fetch, cache: Arc::clone(&self.cache), filter: self.filter.clone(), + spill_before_output: self.spill_before_output, } } @@ -1724,10 +1736,14 @@ impl ExecutionPlan for SortExec { execution_options.sort_spill_reservation_bytes; let in_place_bytes = execution_options.sort_in_place_threshold_bytes; - let spill_before_output = context - .session_config() - .get_extension::() - .map_or(0, |threshold| threshold.0); + let spill_before_output = if self.spill_before_output { + context + .session_config() + .get_extension::() + .map_or(0, |threshold| threshold.0) + } else { + 0 + }; let compression = context.session_config().spill_compression(); let runtime = context.runtime_env(); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -1867,7 +1883,8 @@ impl ExecutionPlan for SortExec { Ok(Some(Arc::new( SortExec::new(updated_exprs, make_with_child(projection, self.input())?) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ))) } @@ -1947,7 +1964,8 @@ impl ExecutionPlan for SortExec { let new_sort = Arc::new( SortExec::new(self.expr.clone(), new_child) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ) as Arc; Ok(FilterPushdownPropagation { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index b79b412cc03..36e2a8c25e8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -136,6 +136,41 @@ class CometExecSuite extends CometTestBase { } } + test("a native sort spills before output only when the JVM reads its output") { + def sortSpills(plan: SparkPlan): Seq[Long] = + collect(plan) { case s: CometSortExec => s.metrics("spill_count").value } + withSQLConf( + CometConf.COMET_EXEC_SORT_SPILL_BEFORE_OUTPUT_THRESHOLD.key -> "1", + CometConf.COMET_EXEC_WINDOW_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withParquetTable((0 until 20000).map(i => (i % 37, s"v${i * 7919 % 20000}", i)), "a") { + withParquetTable((0 until 5000).map(i => (i % 37, s"w$i")), "b") { + val (_, jvmRead) = checkSparkAnswer(sql("SELECT _1, _2, _3 FROM a ORDER BY _2, _3")) + assert(sortSpills(jvmRead).nonEmpty) + assert(sortSpills(jvmRead).forall(_ > 0), jvmRead) + + val (_, filtered) = + checkSparkAnswer(sql("SELECT _2, _3 + 1 FROM a WHERE _1 > 3 ORDER BY _2, _3")) + assert(sortSpills(filtered).nonEmpty) + assert(sortSpills(filtered).forall(_ > 0), filtered) + + val (_, windowRead) = checkSparkAnswer( + sql("SELECT _1, _2, row_number() OVER (PARTITION BY _1 ORDER BY _2, _3) FROM a")) + assert(collect(windowRead) { case w: CometWindowExec => w }.nonEmpty, windowRead) + assert(sortSpills(windowRead).nonEmpty) + assert(sortSpills(windowRead).forall(_ == 0), windowRead) + + val (_, joinRead) = checkSparkAnswer( + sql("SELECT a._1, a._2, b._2 FROM a JOIN b ON a._1 = b._1 AND a._3 < 300")) + assert(collect(joinRead) { case j: CometSortMergeJoinExec => j }.nonEmpty, joinRead) + assert(sortSpills(joinRead).nonEmpty) + assert(sortSpills(joinRead).forall(_ == 0), joinRead) + } + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) From 092d70af28b3ea4bbded9121335cfc54dbd07cfc Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 12:52:00 +0100 Subject: [PATCH 5/8] ci: run CometFilterShortCircuitSuite Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 1c11bd8de6d..473006aa378 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -529,6 +529,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 1cf2f06d213..d00c7f6d2e2 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -177,6 +177,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite From 4d7eec2d2965545fac326680da2f41b05ff7a99a Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 13:17:25 +0100 Subject: [PATCH 6/8] fix: sort a converted SortAggregate's output only for a consumer that needs it The native sort above a SortAggregate converted to a hash aggregate was added unconditionally and dropped only under a shuffle. It is now added the way EnsureRequirements would: when a consumer is visited, before it is converted, each child whose required ordering is no longer met and that leads through unary operators to a converted SortAggregate gets a sort on that edge. A partial aggregate under a shuffle, a final aggregate read by a hash aggregate, a project or a write, and any consumer without an ordering requirement get none. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apache/comet/rules/CometExecRule.scala | 71 ++++++++++++------- .../comet/exec/CometAggregateSuite.scala | 10 ++- 2 files changed, 53 insertions(+), 28 deletions(-) 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 2cc826e0eef..c461f43abfb 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} +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.aggregate.{AggregateFunction, AggregateMode, Average, Count, Final, Max, Min, Partial, PartialMerge, Sum} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule @@ -150,10 +150,13 @@ object CometExecRule { */ val ENGINE_CHOICE_SPARK_TAG: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.engineChoiceSpark") + /** Tag set on a native hash aggregate converted from a `SortAggregateExec`. */ + val CONVERTED_SORT_AGGREGATE_TAG: TreeNodeTag[Unit] = + TreeNodeTag[Unit]("comet.convertedSortAggregate") + /** - * Tag set on the native sort that restores the output ordering of a `SortAggregateExec` - * converted to a native hash aggregate. A shuffle directly above it discards the ordering, so - * the sort is dropped there. + * Tag set on the sort inserted above a converted `SortAggregateExec` for a consumer that + * requires the ordering the sort aggregate provided. */ val SORT_AGGREGATE_ORDER_TAG: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.sortAggregateOrder") @@ -655,7 +658,7 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } plan.transformUp { case op => - val converted = convertNode(dropSortAggregateOrderUnderShuffle(op)) + val converted = convertNode(restoreSortAggregateOrdering(op)) // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions // when the broadcast child has a Comet plan underneath. This enables exchange reuse // between the DPP subquery and the join's CometBroadcastExchangeExec because both @@ -939,15 +942,39 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } } - private def dropSortAggregateOrderUnderShuffle(op: SparkPlan): SparkPlan = op match { - case shuffle: ShuffleExchangeLike => - shuffle.child match { - case sort: CometSortExec - if sort.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => - shuffle.withNewChildren(Seq(sort.child)) - case _ => shuffle + private def leadsToConvertedSortAggregate(plan: SparkPlan): Boolean = + plan.getTagValue(CometExecRule.CONVERTED_SORT_AGGREGATE_TAG).isDefined || + (plan.children.size == 1 && leadsToConvertedSortAggregate(plan.children.head)) + + /** + * Spark's EnsureRequirements left out the sort a consumer needs when a `SortAggregateExec` + * below it already provided the ordering. Once that aggregate is a native hash aggregate, the + * sort goes back on exactly the edges whose requirement is no longer met, before the consumer + * itself is converted. + */ + private def restoreSortAggregateOrdering(op: SparkPlan): SparkPlan = { + if (op.isInstanceOf[CometPlan]) return op + val required = op.requiredChildOrdering + if (required.length != op.children.length || required.forall(_.isEmpty)) return op + var changed = false + val children = op.children.zip(required).map { case (child, ordering) => + if (ordering.nonEmpty && leadsToConvertedSortAggregate(child) && + !SortOrder.orderingSatisfies(child.outputOrdering, ordering)) { + changed = true + val sort = SortExec(ordering, global = false, child) + child.logicalLink.foreach(sort.setLogicalLink) + val restored = if (child.isInstanceOf[CometNativeExec]) { + tryConvertToComet(sort, CometSortExec).getOrElse(sort) + } else { + sort + } + restored.setTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG, ()) + restored + } else { + child } - case other => other + } + if (changed) op.withNewChildren(children) else op } private def orderInsensitive(fn: AggregateFunction): Boolean = fn match { @@ -963,7 +990,8 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) * later stays correct. The sort by the grouping keys directly below the aggregate is dropped, * as rows within a group were in no particular order anyway. Without such a sort the input * order could carry meaning, so only aggregates that do not depend on it are converted. A - * native sort above restores the ordering consumers may rely on. + * consumer that relied on the ordering of the sort aggregate gets a sort back, see + * [[restoreSortAggregateOrdering]]. */ private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = { val required = agg.requiredChildOrdering.headOption.getOrElse(Nil) @@ -1002,19 +1030,8 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) agg.resultExpressions, input) hashAgg.copyTagsFrom(agg) - val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec).flatMap { native => - val ordering = agg.outputOrdering - if (ordering.isEmpty) { - Some(native) - } else { - val sort = SortExec(ordering, global = false, native) - agg.logicalLink.foreach(sort.setLogicalLink) - tryConvertToComet(sort, CometSortExec).map { nativeSort => - nativeSort.setTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG, ()) - nativeSort - } - } - } + val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec) + converted.foreach(_.setTagValue(CometExecRule.CONVERTED_SORT_AGGREGATE_TAG, ())) if (converted.isEmpty) { hashAgg .getTagValue(CometExplainInfo.FALLBACK_REASONS) 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 20cbd42042b..298d2a73eb5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -3525,7 +3525,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("the ordering sort of a converted partial sort aggregate is dropped under its shuffle") { + test("a converted sort aggregate gets no sort without a consumer that needs its ordering") { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { withStrings { val (_, cometPlan) = checkSparkAnswer(sql("SELECT _1, max(_2) FROM str_tbl GROUP BY _1")) @@ -3533,6 +3533,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { assert(shuffles.nonEmpty, s"plan:\n$cometPlan") assert(shuffles.forall(s => !s.child.isInstanceOf[CometSortExec]), s"plan:\n$cometPlan") assert(nativeAggregates(cometPlan).size == 2, s"plan:\n$cometPlan") + assert( + collect(cometPlan) { + case s: CometSortExec + if s.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => + s + }.isEmpty, + s"plan:\n$cometPlan") + assert(collect(cometPlan) { case s: CometSortExec => s }.isEmpty, s"plan:\n$cometPlan") } } } From e930327f406b071e102b25272a7ce7b81a5dda8f Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 14:57:06 +0100 Subject: [PATCH 7/8] fix: keep a SortAggregate with an order-sensitive aggregate in Spark Spark's sort is stable, so a SortAggregate over input sorted by its grouping keys sees each group's rows in their input order, and FIRST, LAST, ANY_VALUE and FIRST_VALUE return the first row of a group as read. The native hash aggregate does not guarantee that order, so converting such a SortAggregate broke the exact first/last results the fork promises on an ordered single split (CometFirstLastBoolAggFuzzSuite, 240 mismatches on Spark 4.1 and 559 on 3.5 for first, any_value and first_value over strings). A SortAggregate is now converted only when every aggregate function is order-insensitive (MIN, MAX, COUNT, SUM, AVG), whatever its input. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../org/apache/comet/rules/CometExecRule.scala | 16 +++++++--------- .../expressions/aggregate/first_last.sql | 4 ++-- 2 files changed, 9 insertions(+), 11 deletions(-) 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 c461f43abfb..7b33aad7b0b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -987,11 +987,11 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) * is not mutable, such as MIN or MAX over strings, and runs it over input sorted by the * grouping keys. The hash aggregate starts from an `ObjectHashAggregateExec` with the same * fields, which Spark runs correctly over any input order, so any operator reverted to Spark - * later stays correct. The sort by the grouping keys directly below the aggregate is dropped, - * as rows within a group were in no particular order anyway. Without such a sort the input - * order could carry meaning, so only aggregates that do not depend on it are converted. A - * consumer that relied on the ordering of the sort aggregate gets a sort back, see - * [[restoreSortAggregateOrdering]]. + * later stays correct. Only aggregates whose result does not depend on the input order are + * converted: Spark's sort is stable, so FIRST, LAST and the like see the rows of a group in + * their input order, which the native hash aggregate does not guarantee. 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]]. */ private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = { val required = agg.requiredChildOrdering.headOption.getOrElse(Nil) @@ -1008,12 +1008,10 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) case _ => false } val input = if (sortedByKeys) agg.child.children.head else agg.child - if (!sortedByKeys && - !agg.aggregateExpressions.forall(e => orderInsensitive(e.aggregateFunction))) { + if (!agg.aggregateExpressions.forall(e => orderInsensitive(e.aggregateFunction))) { withFallbackReason( agg, - "SortAggregate over input not sorted by its grouping keys alone may depend on the " + - "input order") + "SortAggregate with an aggregate function that depends on the input order stays in Spark") return None } if (!input.isInstanceOf[CometNativeExec]) { diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql index 7fd5dcde731..b67a4247eb8 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql @@ -183,7 +183,7 @@ SELECT first(val) IGNORE NULLS FROM test_single_null -- first IGNORE NULLS: multiple data types -- ============================================================ -query +query expect_fallback(depends on the input order) SELECT grp, first(i_val) IGNORE NULLS, first(l_val) IGNORE NULLS, @@ -327,7 +327,7 @@ SELECT last(val) IGNORE NULLS FROM test_single_null -- last IGNORE NULLS: multiple data types -- ============================================================ -query +query expect_fallback(depends on the input order) SELECT grp, last(i_val) IGNORE NULLS, last(l_val) IGNORE NULLS, From 92f4b72974664c3eaed7e338ea1b31ea29f89ed2 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 16:04:56 +0100 Subject: [PATCH 8/8] feat: native MAX_BY and MIN_BY with a string value MAX_BY and MIN_BY accept a StringType value with the default UTF8_BINARY collation; the ordering stays limited to fixed-length types. The native MaxMinBy accumulators already store the value as Arrow row bytes and only compare the ordering, so no native change is needed. A string value makes Spark plan a SortAggregate, which is now converted when its other aggregates are order-insensitive. MAX_BY and MIN_BY depend on the input order only among rows tied on the ordering, where Spark is non-deterministic too and the native path already documents that it may pick a different row. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apache/comet/rules/CometExecRule.scala | 12 ++-- .../org/apache/comet/serde/aggregates.scala | 27 ++++---- .../expressions/aggregate/max_by.sql | 41 ++++++++---- .../expressions/aggregate/min_by.sql | 38 ++++++++--- .../comet/exec/CometAggregateSuite.scala | 65 +++++++++++++++++++ 5 files changed, 146 insertions(+), 37 deletions(-) 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 7b33aad7b0b..15e0737df44 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -23,7 +23,7 @@ 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.aggregate.{AggregateFunction, AggregateMode, Average, Count, Final, Max, Min, Partial, PartialMerge, Sum} +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 import org.apache.spark.sql.catalyst.trees.TreeNodeTag @@ -978,7 +978,7 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } private def orderInsensitive(fn: AggregateFunction): Boolean = fn match { - case _: Min | _: Max | _: Count | _: Sum | _: Average => true + case _: Min | _: Max | _: Count | _: Sum | _: Average | _: MaxMinBy => true case _ => false } @@ -989,9 +989,11 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) * fields, which Spark runs correctly over any input order, so any operator reverted to Spark * later stays correct. Only aggregates whose result does not depend on the input order are * converted: Spark's sort is stable, so FIRST, LAST and the like see the rows of a group in - * their input order, which the native hash aggregate does not guarantee. 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]]. + * 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]]. */ private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = { val required = agg.requiredChildOrdering.headOption.getOrElse(Nil) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index bc45c605025..4583c2bce00 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -129,27 +129,28 @@ abstract class CometMaxMinBy[T <: MaxMinBy] extends CometAggregateExpressionSerd " Results may differ from Spark in that case.") override def getUnsupportedReasons(): Seq[String] = Seq( - "The value and ordering must both be fixed-length types (boolean, integral, floating-point," + - " decimal, date, or timestamp). A variable-length or nested type such as string, binary, or" + - " struct falls back to Spark.") + "The ordering must be a fixed-length type (boolean, integral, floating-point, decimal, date," + + " or timestamp). The value may also be a string with the default UTF8_BINARY collation." + + " Other variable-length or nested types such as binary or struct fall back to Spark.") override def getSupportLevel(expr: T): SupportLevel = { - // Both the value and ordering must be fixed-length types. + // The ordering must be a fixed-length type. The value may also be a UTF8_BINARY string: the + // native side only stores it as Arrow row bytes and never compares it. // - // On its own a variable-length type never reaches here: Spark only uses HashAggregate (the - // aggregate operator Comet accelerates) when the aggregation buffer is mutable, and the buffer - // holds both the running value and the running ordering, so a StringType in either position - // forces SortAggregate, which Comet does not convert. + // A string in the buffer makes Spark plan a SortAggregate, which Comet converts to a native + // hash aggregate (see `CometExecRule.convertSortAggregate`), or an ObjectHashAggregate when a + // TypedImperativeAggregate sits in the same aggregate. // - // The check is still load-bearing, because a TypedImperativeAggregate elsewhere in the same - // aggregate switches Spark to ObjectHashAggregate, which Comet does convert. In that shape a - // string ordering would otherwise be compared by Arrow's row format as raw UTF-8 bytes, while + // A string ordering stays in Spark: Arrow's row format would compare raw UTF-8 bytes, while // Spark compares collation sort keys. See the fallback cases in max_by.sql. // // The native side compares the ordering column via Arrow's row format, which supports all of // the fixed-length orderable types allowed below. - if (!AggSerde.minMaxDataTypeSupported(expr.valueExpr.dataType)) { - Unsupported(Some(s"Unsupported value data type: ${expr.valueExpr.dataType}")) + val valueType = expr.valueExpr.dataType + val stringValue = + valueType.isInstanceOf[StringType] && !AggSerde.isStringCollationType(valueType) + if (!stringValue && !AggSerde.minMaxDataTypeSupported(valueType)) { + Unsupported(Some(s"Unsupported value data type: $valueType")) } else if (!AggSerde.minMaxDataTypeSupported(expr.orderingExpr.dataType)) { Unsupported(Some(s"Unsupported ordering data type: ${expr.orderingExpr.dataType}")) } else { diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql index 9c8354f06ab..b8270db9e1a 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql @@ -17,9 +17,8 @@ -- max_by(x, y) returns the value of x associated with the maximum value of y. -- --- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate --- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length --- value types such as string force SortAggregate and fall back to Spark. +-- The value (x) must be a fixed-length type or a string. A string value makes Spark plan a +-- SortAggregate, which Comet converts to a native hash aggregate. -- -- Ordering values are kept unique within each group so results are deterministic (max_by is -- non-deterministic when several rows tie on the maximum ordering). @@ -260,14 +259,11 @@ query SELECT grp, max_by(v, ord) FROM mb_signed_zero GROUP BY grp ORDER BY grp -- ============================================================ --- Variable-length value or ordering falls back to Spark +-- Variable-length ordering falls back to Spark -- --- A plain max_by over a string is planned as SortAggregate, which Comet never converts, so the --- serde's type check is not what stops it. Pairing it with a TypedImperativeAggregate switches --- Spark to ObjectHashAggregate, which Comet does convert, and then getSupportLevel is the only --- thing that keeps the aggregate off the native path. That matters most for a string *ordering*: --- Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys, so this --- check is load-bearing for correctness and not just an optimisation. +-- A string value runs natively: as a plain aggregate through the converted SortAggregate, and +-- next to a TypedImperativeAggregate through ObjectHashAggregate. A string ordering stays in +-- Spark: Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys. -- ============================================================ statement @@ -277,8 +273,31 @@ statement INSERT INTO mb_varlen VALUES (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2') -query expect_fallback(Unsupported value data type) +query SELECT grp, max_by(s, v), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER BY grp query expect_fallback(Unsupported ordering data type) SELECT grp, max_by(v, s), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER BY grp + +-- ============================================================ +-- String value with an integer ordering +-- ============================================================ + +statement +CREATE TABLE mb_str(s string, ord int, grp string) USING parquet + +statement +INSERT INTO mb_str VALUES + ('b', 1, 'g1'), ('', 5, 'g1'), (NULL, 3, 'g1'), + ('é日本', 2, 'g2'), (NULL, 9, 'g2'), ('😀', 4, 'g2'), + ('z', NULL, 'g3'), ('y', NULL, 'g3'), + ('x', 6, 'g4') + +query +SELECT grp, max_by(s, ord) FROM mb_str GROUP BY grp ORDER BY grp + +query +SELECT max_by(s, ord) FROM mb_str + +query +SELECT grp, max_by(s, ord), max(s), count(*) FROM mb_str GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql index d015acc35d3..55828e5c65c 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql @@ -17,9 +17,8 @@ -- min_by(x, y) returns the value of x associated with the minimum value of y. -- --- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate --- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length --- value types such as string force SortAggregate and fall back to Spark. +-- The value (x) must be a fixed-length type or a string. A string value makes Spark plan a +-- SortAggregate, which Comet converts to a native hash aggregate. -- -- Ordering values are kept unique within each group so results are deterministic (min_by is -- non-deterministic when several rows tie on the minimum ordering). @@ -245,11 +244,11 @@ query SELECT grp, min_by(v, ord) FROM mnb_signed_zero GROUP BY grp ORDER BY grp -- ============================================================ --- Variable-length value or ordering falls back to Spark +-- Variable-length ordering falls back to Spark -- --- See the equivalent section in max_by.sql for why this needs a TypedImperativeAggregate --- alongside it: without one, Spark plans SortAggregate and Comet never sees the aggregate at all, --- so the serde's type check would go untested. +-- A string value runs natively: as a plain aggregate through the converted SortAggregate, and +-- next to a TypedImperativeAggregate through ObjectHashAggregate. A string ordering stays in +-- Spark: Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys. -- ============================================================ statement @@ -259,8 +258,31 @@ statement INSERT INTO mnb_varlen VALUES (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2') -query expect_fallback(Unsupported value data type) +query SELECT grp, min_by(s, v), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER BY grp query expect_fallback(Unsupported ordering data type) SELECT grp, min_by(v, s), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER BY grp + +-- ============================================================ +-- String value with an integer ordering +-- ============================================================ + +statement +CREATE TABLE mnb_str(s string, ord int, grp string) USING parquet + +statement +INSERT INTO mnb_str VALUES + ('b', 1, 'g1'), ('', 5, 'g1'), (NULL, 3, 'g1'), + ('é日本', 2, 'g2'), (NULL, 9, 'g2'), ('😀', 4, 'g2'), + ('z', NULL, 'g3'), ('y', NULL, 'g3'), + ('x', 6, 'g4') + +query +SELECT grp, min_by(s, ord) FROM mnb_str GROUP BY grp ORDER BY grp + +query +SELECT min_by(s, ord) FROM mnb_str + +query +SELECT grp, min_by(s, ord), max(s), count(*) FROM mnb_str GROUP BY grp ORDER BY grp 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 298d2a73eb5..04b3746a5d1 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -3577,4 +3577,69 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } } + + private val byRows: Seq[(Integer, String, Integer)] = Seq( + (1, "b", 1), + (1, "", 5), + (1, null, 3), + (2, "\u00e9\u65e5\u672c", 2), + (2, null, 9), + (2, "\ud83d\ude00", 4), + (3, "z", null), + (3, "y", null), + (4, "x", 6), + (null, "w", 7)) + + for (aqe <- Seq("false", "true"); fn <- Seq("max_by", "min_by")) { + test(s"$fn over a string value runs natively, grouped and ungrouped (AQE=$aqe)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withParquetTable(byRows, "by_tbl") { + for (query <- Seq( + s"SELECT _1, $fn(_2, _3) FROM by_tbl GROUP BY _1", + s"SELECT $fn(_2, _3) FROM by_tbl", + s"SELECT _1, $fn(_2, _3), max(_2), count(*) FROM by_tbl GROUP BY _1", + s"SELECT $fn(_2, _3) FROM by_tbl WHERE _1 = 3")) { + val (sparkPlan, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(sparkPlan).nonEmpty, s"$query:\n$sparkPlan") + assert(sortAggregates(cometPlan).isEmpty, s"$query:\n$cometPlan") + assert(nativeAggregates(cometPlan).nonEmpty, s"$query:\n$cometPlan") + } + } + } + } + } + + test("max_by and min_by over a string value pick one of the values tied on the ordering") { + val rows = (0 until 400).map(i => (i % 4, s"v${i % 7}", if (i % 3 == 0) 10 else i % 3)) + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3") { + withParquetTable(rows, "by_ties") { + for ((fn, tiedOrd) <- Seq("max_by" -> 10, "min_by" -> 1)) { + val df = sql(s"SELECT _1, $fn(_2, _3) FROM by_ties GROUP BY _1") + val result = df.collect().map(r => r.getInt(0) -> r.getString(1)).toMap + assert(nativeAggregates(df.queryExecution.executedPlan).nonEmpty) + val allowed = rows.filter(_._3 == tiedOrd).groupBy(_._1).map { case (k, v) => + k -> v.map(_._2).toSet + } + assert(result.keySet == allowed.keySet, s"$fn: $result") + result.foreach { case (k, v) => assert(allowed(k).contains(v), s"$fn group $k: $v") } + } + } + } + } + + test("max_by over a string value keeps a native partial with a Spark final") { + withSQLConf(CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false") { + withParquetTable(byRows, "by_tbl") { + checkSparkAnswer(sql("SELECT _1, max_by(_2, _3), min_by(_2, _3) FROM by_tbl GROUP BY _1")) + } + } + } + + test("an ungrouped order-sensitive sort aggregate over an ordered input stays in Spark") { + withParquetTable(byRows, "by_tbl") { + val (_, cometPlan) = checkSparkAnswer( + sql("SELECT max_by(_2, _1), first(_2) FROM (SELECT * FROM by_tbl ORDER BY _3, _2)")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } }