sunchao commented on code in PR #4565:
URL: https://github.com/apache/datafusion-comet/pull/4565#discussion_r4104683220
##########
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:
[P1] Traverse `SortExec` when tagging incompatible aggregate producers. With
Comet shuffle enabled, AQE disabled and
`spark.sql.execution.useObjectHashAggregateExec=false`, `SELECT g, first(v),
collect_list(v), count(DISTINCT v) FROM t GROUP BY g` creates sort-aggregate
merge stages that must fall back because `FIRST` does not support
`PartialMerge`. However, `findPartialAggInPlan` stops at their intervening
`SortExec`, so the native partial remains untagged. Spark then consumes Comet's
array-valued `collect_list` buffer where it requires serialized binary, causing
query failure instead of safe fallback. Before this registration, these sort
aggregates remained on Spark. Extend the producer search through row-preserving
sorts and test this mixed aggregate query.
Evidence: On Spark 4.1.3, used an eight-row Parquet table with `g=id%2` and
integer `v=id`. A harness containing the exact-head `findPartialAggInPlan`
returned `None` for all three buffer-consuming sort aggregates, while the
object-hash equivalent found the producer. The exact-head buffer adjustment
changed `collect_list` state from `BinaryType` to
`ArrayType(IntegerType,true)`. A bounded Spark merge-contract check rejected
that array state with `ClassCastException: GenericArrayData cannot be cast to
[B`. `unsupportedPartialMergeFallbackReason` explicitly rejects the query's
`FIRST` merge. Reproduction source and output:
`/tmp/pr4565-current-validation/PR4565CurrentValidation.scala` and
`validation-final.log`.
##########
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:
[P1] Preserve the advertised ordering for array grouping keys. With native
caching enabled at startup, cache a single-partition relation sorted by array
key `k`, containing `(NULL,'n'), (array(),'e'), (array(1),'o')`, then run
`SELECT k, first(v) FROM cached_sorted GROUP BY k ORDER BY k`. Spark eliminates
the sorts because the cache and sort aggregates advertise ascending order.
Comet's native scan does not carry that ordering, so DataFusion selects
`Linear`. Its vectorized group interning can reorder these keys, producing
`NULL, [1], []` instead of `NULL, [], [1]`. This violates the explicit `ORDER
BY`. Carry the guaranteed ordering into native planning, enforce it with a
native sort, or fall back. The earlier retraction tested primitive/string keys
and does not cover this reproduced array-key case.
Evidence: Spark 4.1.3 produced `SortAggregate[Final] ->
SortAggregate[Partial] -> InMemoryTableScan` with zero executable `SortExec`
nodes and returned `NULL, [], [1]`. Compiled
`/tmp/pr4565-validation/aggregate_order.rs` against this checkout's native
dependencies: both DataFusion 55 aggregate stages reported `Linear` and
returned `NULL, [1], []`. `new_group_values` selects `GroupValuesColumn<false>`
for this array schema, whose vectorized interning does not preserve first-seen
order. The reachable Comet cache path requires
`spark.comet.exec.inMemoryCache.enabled=true` with the Comet plugin and shuffle
configured. Its serializer supports arrays, while `ScanExec` constructs empty
ordering metadata.
--
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]
