andygrove commented on code in PR #4565:
URL: https://github.com/apache/datafusion-comet/pull/4565#discussion_r4169204091
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2051,29 +2043,72 @@ object CometObjectHashAggregateExec
}
}
-case class CometHashAggregateExec(
- override val nativeOp: Operator,
- override val originalPlan: SparkPlan,
- override val output: Seq[Attribute],
- groupingExpressions: Seq[NamedExpression],
- aggregateExpressions: Seq[AggregateExpression],
- resultExpressions: Seq[NamedExpression],
- input: Seq[Attribute],
- child: SparkPlan,
- override val serializedPlanOpt: SerializedPlan)
+object CometSortAggregateExec
+ extends CometOperatorSerde[SortAggregateExec]
+ with CometBaseAggregate {
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+ CometConf.COMET_EXEC_AGGREGATE_ENABLED)
+
+ override def getSupportLevel(op: SortAggregateExec): SupportLevel =
+ baseAggregateSupportLevel(op)
+
+ override def convert(
+ aggregate: SortAggregateExec,
+ builder: Operator.Builder,
+ childOp: OperatorOuterClass.Operator*):
Option[OperatorOuterClass.Operator] = {
+
+ // SortAggregate is planned for TypedImperativeAggregate functions whose
intermediate
+ // buffer formats differ between Spark and Comet (same risk as
ObjectHashAggregate).
+ // Require Comet shuffle so a Partial->Final pair never spans the
JVM/native boundary.
+ if (!isCometShuffleEnabled(aggregate.conf)) {
+ return None
+ }
+
+ doConvert(aggregate, builder, childOp: _*)
+ }
+
+ override def createExec(nativeOp: Operator, op: SortAggregateExec):
CometNativeExec = {
+ // The native AggregateExec auto-detects Sorted input mode from the
child's output ordering
Review Comment:
Agreed, the comment named the wrong mechanism, and sunchao's array-key repro
below shows the first-seen invariant doesn't hold either: DataFusion's
vectorized group interning reorders groups on a hash collision, and a NULL and
an empty array always collide. So instead of documenting that invariant, the
native planner now sorts the aggregate output whenever the AggregateExec's
output ordering doesn't already satisfy the grouping keys (1c16c7d15e). The
explanation lives once, on `CometSortAggregateExec.orderedByGroupingKeys`, and
the `CometBaseAggregateExec` scaladoc no longer claims auto-detection.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2103,17 +2165,40 @@ case class CometHashAggregateExec(
override def hashCode(): Int =
Objects.hashCode(output, groupingExpressions, aggregateExpressions, input,
modes, child)
+}
- override lazy val metrics: Map[String, SQLMetric] = {
- val baseline = CometMetricNode.baselineMetrics(sparkContext)
- if (groupingExpressions.nonEmpty) {
- baseline ++ CometMetricNode.aggregateMetrics(sparkContext)
- } else {
- baseline
+case class CometSortAggregateExec(
+ override val nativeOp: Operator,
+ override val originalPlan: SparkPlan,
+ override val output: Seq[Attribute],
+ groupingExpressions: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ resultExpressions: Seq[NamedExpression],
+ input: Seq[Attribute],
+ child: SparkPlan,
+ override val serializedPlanOpt: SerializedPlan)
+ extends CometBaseAggregateExec {
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ this.copy(child = newChild)
+
+ override def equals(obj: Any): Boolean = {
+ obj match {
+ case other: CometSortAggregateExec =>
Review Comment:
Done in 77259ab685. `equals` and `hashCode` are on `CometBaseAggregateExec`
now, with a same-class check, so both wrappers keep only `withNewChildInternal`.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2051,29 +2043,72 @@ object CometObjectHashAggregateExec
}
}
-case class CometHashAggregateExec(
- override val nativeOp: Operator,
- override val originalPlan: SparkPlan,
- override val output: Seq[Attribute],
- groupingExpressions: Seq[NamedExpression],
- aggregateExpressions: Seq[AggregateExpression],
- resultExpressions: Seq[NamedExpression],
- input: Seq[Attribute],
- child: SparkPlan,
- override val serializedPlanOpt: SerializedPlan)
+object CometSortAggregateExec
Review Comment:
Done in 77259ab685. `CometBaseAggregate[T]` now supplies `enabledConfig`,
`getSupportLevel` and `convert`, and each serde object keeps
`requiresCometShuffle`, its operator-specific check (the decimal SUM guards
that landed on main in the meantime) and `createExec`. The shared
`getSupportLevel` comment says the partial/final test knobs gate sort
aggregates too.
##########
spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql:
##########
@@ -0,0 +1,214 @@
+-- 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.
+
+-- Disabling ObjectHashAggregate forces Spark to plan SortAggregateExec for the
Review Comment:
Done in 7fd7568d96. `collect_set.sql` and
`collect_set_floating_fallback.sql` now carry a `ConfigMatrix` over
`useObjectHashAggregateExec`, and `sort_aggregate.sql` keeps only the
multi-key, expression-key and mixed-aggregate shapes plus a check for the
decimal SUM fallback. The shuffle gate has its own fixture,
`sort_aggregate_without_comet_shuffle.sql`, which asserts the SortAggregate
shuffle reason and fails if the gate is removed. It uses
`spark.comet.shuffle.enabled`, because the `spark.comet.exec.shuffle.enabled`
alias didn't take effect in the SQL file suite.
##########
spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql:
##########
@@ -70,11 +70,16 @@ CREATE TABLE test_types(
grp string
) USING parquet
+-- first/last IGNORE NULLS are non-deterministic when a group has more than
one non-null value,
Review Comment:
Added a Sort Aggregation section to the operator compatibility guide in
abcd4f76d9, and a compatibility note to the PR description. I didn't add a
fixture for a string group with two non-null values, since there's no single
answer to assert against Spark.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -2408,6 +2409,46 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ // useObjectHashAggregateExec=false forces Spark to plan SortAggregateExec
for
+ // TypedImperativeAggregate functions like collect_set. Comet converts those
just like
+ // ObjectHashAggregateExec via the shared CometBaseAggregate path. Broader
data-type and
+ // edge-case coverage lives in the SQL file test
+ //
spark/src/test/resources/sql-tests/expressions/aggregate/sort_aggregate.sql;
these Scala
+ // tests additionally assert that Spark actually planned a
SortAggregateExec, which the SQL
+ // framework cannot check.
+ private def assertSortAggregateRunsNatively(query: String): Unit = {
+ withSQLConf(
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
+ withTempView("tbl") {
+ Seq((1, "a"), (2, "a"), (1, "a"), (3, "b"), (4, "b"), (4, "b"))
+ .toDF("v", "g")
+ .createOrReplaceTempView("tbl")
+ // Spark must actually plan a SortAggregateExec for this query;
otherwise the test
+ // would pass without exercising the new code path.
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val plan = stripAQEPlan(sql(query).queryExecution.executedPlan)
+ assert(
+ plan.find(_.isInstanceOf[SortAggregateExec]).isDefined,
+ s"Expected SortAggregateExec in Spark-only plan but got:\n$plan")
+ }
+ checkSparkAnswerAndOperator(sql(query))
+ }
+ }
+ }
+
+ test("SortAggregate with collect_set is converted to native") {
Review Comment:
Done in 7fd7568d96.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -82,6 +82,7 @@ object CometExecRule {
classOf[GenerateExec] -> CometExplodeExec,
classOf[HashAggregateExec] -> CometHashAggregateExec,
classOf[ObjectHashAggregateExec] -> CometObjectHashAggregateExec,
+ classOf[SortAggregateExec] -> CometSortAggregateExec,
Review Comment:
Thanks, confirmed. Since #5041 lets FIRST run in PartialMerge, your exact
query now reaches the issue #4724 guard instead, which also goes through
`findPartialAggInPlan`, and it failed natively for the same reason. c506925525
makes `findPartialAggInPlan`, `revertUnsafePartialAggregates` and
`hasUnrepairedNativeBuffer` walk through `SortExec` and `CometSortExec`.
`CometAggregateSuite` now covers a Final that can't convert and a distinct
`collect_list` chain, with AQE on and off. Both failed before the change, one
with an NPE in Spark's Final and one with a native schema error. A planner test
in `CometExecRuleSuite` covers the repair pass for both sort node types.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2051,29 +2043,72 @@ object CometObjectHashAggregateExec
}
}
-case class CometHashAggregateExec(
- override val nativeOp: Operator,
- override val originalPlan: SparkPlan,
- override val output: Seq[Attribute],
- groupingExpressions: Seq[NamedExpression],
- aggregateExpressions: Seq[AggregateExpression],
- resultExpressions: Seq[NamedExpression],
- input: Seq[Attribute],
- child: SparkPlan,
- override val serializedPlanOpt: SerializedPlan)
+object CometSortAggregateExec
+ extends CometOperatorSerde[SortAggregateExec]
+ with CometBaseAggregate {
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+ CometConf.COMET_EXEC_AGGREGATE_ENABLED)
+
+ override def getSupportLevel(op: SortAggregateExec): SupportLevel =
+ baseAggregateSupportLevel(op)
+
+ override def convert(
+ aggregate: SortAggregateExec,
+ builder: Operator.Builder,
+ childOp: OperatorOuterClass.Operator*):
Option[OperatorOuterClass.Operator] = {
+
+ // SortAggregate is planned for TypedImperativeAggregate functions whose
intermediate
+ // buffer formats differ between Spark and Comet (same risk as
ObjectHashAggregate).
+ // Require Comet shuffle so a Partial->Final pair never spans the
JVM/native boundary.
+ if (!isCometShuffleEnabled(aggregate.conf)) {
+ return None
+ }
+
+ doConvert(aggregate, builder, childOp: _*)
+ }
+
+ override def createExec(nativeOp: Operator, op: SortAggregateExec):
CometNativeExec = {
+ // The native AggregateExec auto-detects Sorted input mode from the
child's output ordering
+ // and produces output sorted by the grouping keys;
CometExec.outputOrdering defaults to
+ // originalPlan.outputOrdering, which is SortAggregateExec's grouping-key
ordering, so
+ // downstream operators that elided a sort against it still see a
satisfying ordering.
+ CometSortAggregateExec(
Review Comment:
Confirmed with your cached repro: Comet returned `[null,n]`, `[[1],o]`,
`[[],e]` for the ORDER BY query. 1c16c7d15e adds `ordered_by_grouping_keys` to
the HashAggregate proto for sort aggregates, and the native planner sorts the
aggregate output on the grouping columns when the AggregateExec's output
ordering doesn't satisfy them. It also normalizes float grouping keys the way
native sorts do, so a float-keyed final streams over the sort below it rather
than taking the hash path. In the usual shape, with the sort in the same native
plan, nothing extra is added. The cached array case is now a
`CometInMemoryCacheSuite` test, and there are native planner tests for both
shapes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
