You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/04/21 04:40:55 UTC
[PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
andygrove opened a new pull request, #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015
## Which issue does this PR close?
Closes #1389.
Part of #1267.
## Rationale for this change
When a Spark query's Final-mode aggregate cannot be converted to Comet (for example because its result expressions are not supported, as in `concat(flatten(collect_set(col)))`), Comet would still convert the upstream Partial-mode aggregate. The Partial produces intermediate buffers in a format the Spark Final cannot interpret (different encodings for `CollectSet`, `Average`, decimal `Sum`, variance, etc.), which crashes at runtime with errors such as `Not supported on CometListVector`.
Conversely, most aggregates block even a safe Spark-Partial + Comet-Final combination, where the buffer formats are in fact compatible (`MIN`, `MAX`, `COUNT`, bitwise).
This change prevents the crash for unsafe aggregates and unlocks the mixed execution for the safe ones.
## What changes are included in this PR?
- New `supportsMixedPartialFinal` flag on `CometAggregateExpressionSerde`, defaulting to `false`. Set to `true` for `MIN`, `MAX`, `COUNT`, `BitAndAgg`, `BitOrAgg`, `BitXorAgg`, whose intermediate buffer formats match between Spark and Comet.
- `QueryPlanSerde.allAggsSupportMixedExecution` checks the flag across an aggregate's expressions.
- `CometExecRule.tagUnsafePartialAggregates` runs before bottom-up transformation. For each Final-mode aggregate whose expressions are not all mixed-safe, it conservatively checks whether the Final itself is convertible via the new `canFinalAggregateBeConverted` (mirrors the predicates in `CometBaseAggregate.doConvert`). If not, the corresponding Partial (looked up by `findPartialAggInPlan`, traversing through `AQEShuffleReadExec` and `ShuffleQueryStageExec`) is tagged with `COMET_UNSAFE_PARTIAL`.
- `CometBaseAggregate.doConvert` honours the new tag, and now permits the Spark-Partial + Comet-Final case when all aggregates are mixed-safe.
## How are these changes tested?
`CometExecRuleSuite`:
- Existing test for Comet-Partial + Spark-Final with `SUM` (unsafe) is un-ignored; asserts neither side is converted.
- New test for Spark-Partial + Comet-Final with `SUM`; asserts neither side is converted.
- New test for Comet-Partial + Spark-Final with `MIN`/`MAX`/`COUNT`; asserts partial converts to Comet, final stays Spark.
- New test for Spark-Partial + Comet-Final with `MIN`/`MAX`/`COUNT`; asserts partial stays Spark, final converts to Comet.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4481953414
This PR requires conflicts to be resolved, and this will also fix the PR failed CI check
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#discussion_r3261968192
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -1362,10 +1364,24 @@ trait CometBaseAggregate {
// We support {Partial, PartialMerge} mix; other combinations are rejected.
val multiMode = modes.size > 1 && modeSet != Set(Partial, PartialMerge)
// For a final mode HashAggregate, we only need to transform the HashAggregate
- // if there is Comet partial aggregation.
+ // if there is Comet partial aggregation, unless all aggregates have compatible
+ // intermediate buffer formats (safe for mixed Spark/Comet execution).
val sparkFinalMode = modes.contains(Final) && findCometPartialAgg(aggregate.child).isEmpty
- if (multiMode || sparkFinalMode) {
+ if (multiMode) {
+ return None
Review Comment:
`withInfo` ?
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -1362,10 +1364,24 @@ trait CometBaseAggregate {
// We support {Partial, PartialMerge} mix; other combinations are rejected.
val multiMode = modes.size > 1 && modeSet != Set(Partial, PartialMerge)
// For a final mode HashAggregate, we only need to transform the HashAggregate
- // if there is Comet partial aggregation.
+ // if there is Comet partial aggregation, unless all aggregates have compatible
+ // intermediate buffer formats (safe for mixed Spark/Comet execution).
val sparkFinalMode = modes.contains(Final) && findCometPartialAgg(aggregate.child).isEmpty
- if (multiMode || sparkFinalMode) {
+ if (multiMode) {
+ return None
+ }
+
+ if (sparkFinalMode &&
+ !QueryPlanSerde.allAggsSupportMixedExecution(aggregate.aggregateExpressions)) {
+ return None
Review Comment:
`withInfo`
?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#discussion_r3261939377
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -778,4 +794,127 @@ case class CometExecRule(session: SparkSession)
}
}
+ /**
+ * Walk the plan to find Final-mode aggregates that cannot be converted to Comet. For each such
+ * Final, if the aggregate functions have incompatible intermediate buffer formats, tag the
+ * corresponding Partial-mode aggregate so it will also be skipped during conversion.
+ *
+ * This prevents the crash described in issue #1389 where a Comet Partial produces intermediate
+ * data in a format that the Spark Final cannot interpret.
+ */
+ private def tagUnsafePartialAggregates(plan: SparkPlan): Unit = {
+ plan.foreach {
+ case agg: BaseAggregateExec =>
+ // Only consider single-mode Final aggregates. Multi-mode Finals come from Spark's
+ // distinct-aggregate rewrite, where the Comet partial (if any) feeds into a Spark
+ // PartialMerge rather than directly into a Final, which is a different code path
+ // than the Comet-Partial → Spark-Final crash scenario from issue #1389.
+ val modes = agg.aggregateExpressions.map(_.mode).distinct
+ if (modes == Seq(Final) &&
+ !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) &&
+ !canAggregateBeConverted(agg, Final)) {
+ findPartialAggInPlan(agg.child).foreach { partial =>
+ // Only tag if the Partial would otherwise have been converted. If the Partial
+ // itself cannot be converted (e.g. the aggregate function is incompatible for the
+ // input type), there is no buffer-format mismatch to guard against, and tagging
+ // would mask the natural, more specific fallback reason.
+ if (canAggregateBeConverted(partial, Partial)) {
+ partial.setTagValue(
+ CometExecRule.COMET_UNSAFE_PARTIAL,
+ "Partial aggregate disabled: corresponding final aggregate " +
+ "cannot be converted to Comet and intermediate buffer formats are incompatible")
+ }
+ }
+ }
+ case _ =>
+ }
+ }
+
+ /**
+ * Conservative check for whether an aggregate could be converted to Comet. Checks operator
+ * enablement, grouping expressions, aggregate expressions, and result expressions.
+ * Intentionally skips the sparkFinalMode / child-native checks since those depend on
+ * transformation state.
+ *
+ * WARNING: this intentionally mirrors the predicate checks in `CometBaseAggregate.doConvert`
+ * (operators.scala). Any change to the convertibility rules there must be reflected here or
+ * this tagging pass will drift and either crash (missed tag) or over-disable (spurious tag). A
+ * shared predicate helper would be preferable.
+ */
+ private def canAggregateBeConverted(
+ agg: BaseAggregateExec,
+ expectedMode: AggregateMode): Boolean = {
+ val handler = allExecs.get(agg.getClass)
+ if (handler.isEmpty) return false
+ val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]]
+ if (!isOperatorEnabled(serde, agg.asInstanceOf[SparkPlan])) return false
+
+ // ObjectHashAggregate has an extra shuffle-enabled guard in its convert method
+ agg match {
+ case _: ObjectHashAggregateExec if !isCometShuffleEnabled(agg.conf) => return false
+ case _ =>
+ }
+
+ val aggregateExpressions = agg.aggregateExpressions
+ val groupingExpressions = agg.groupingExpressions
+
+ if (groupingExpressions.isEmpty && aggregateExpressions.isEmpty) return false
+
+ if (groupingExpressions.exists(e => QueryPlanSerde.containsMapType(e.dataType))) return false
+
+ if (!groupingExpressions.forall(e =>
+ QueryPlanSerde.exprToProto(e, agg.child.output).isDefined)) {
+ return false
+ }
+
+ if (aggregateExpressions.isEmpty) {
+ // Result expressions always checked when there are no aggregate expressions
+ val attributes =
+ groupingExpressions.map(_.toAttribute) ++ agg.aggregateAttributes
+ return agg.resultExpressions.forall(e =>
+ QueryPlanSerde.exprToProto(e, attributes).isDefined)
+ }
+
+ val modes = aggregateExpressions.map(_.mode).distinct
+ if (modes.size != 1 || modes.head != expectedMode) return false
+
+ // In Final mode, exprToProto resolves against the child's output; in Partial/non-Final mode
+ // it must bind to input attributes. This mirrors the `binding` calculation in
+ // `CometBaseAggregate.doConvert`.
+ val binding = expectedMode != Final
+ if (!aggregateExpressions.forall(e =>
+ QueryPlanSerde.aggExprToProto(e, agg.child.output, binding, agg.conf).isDefined)) {
+ return false
+ }
+
+ // doConvert only checks resultExpressions in Final mode when aggregate expressions exist
+ // (Partial emits the buffer directly). Mirror that here to avoid false negatives.
+ if (expectedMode == Final) {
+ val attributes =
+ groupingExpressions.map(_.toAttribute) ++ agg.aggregateAttributes
+ agg.resultExpressions.forall(e => QueryPlanSerde.exprToProto(e, attributes).isDefined)
+ } else {
+ true
+ }
+ }
+
+ /**
+ * Look for a Partial-mode aggregate that feeds directly into the given plan (the child of a
+ * Final). Walks through exchanges and AQE stages only, stopping at anything else including
+ * other aggregate stages. This avoids tagging unrelated Partials found deeper in the plan (e.g.
+ * the non-distinct Partial in a distinct-aggregate rewrite, which is separated from the Final
+ * by intermediate PartialMerge stages). Requires `aggregateExpressions.nonEmpty` so that
+ * group-by-only dedup stages are not mistaken for the partial we want to tag.
+ */
+ private def findPartialAggInPlan(plan: SparkPlan): Option[BaseAggregateExec] = plan match {
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.nonEmpty &&
+ agg.aggregateExpressions.forall(e => e.mode == Partial) =>
+ Some(agg)
+ case a: AQEShuffleReadExec => findPartialAggInPlan(a.child)
+ case s: ShuffleQueryStageExec => findPartialAggInPlan(s.plan)
+ case e: ShuffleExchangeExec => findPartialAggInPlan(e.child)
+ case _ => None
Review Comment:
does it make sense to log unsupported exec?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4290590901
might be related to https://github.com/apache/datafusion-comet/pull/4003
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4500017923
Thanks for the review @comphead. I addressed your feedback and will go ahead and merge when CI is green.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4285987435
@Shekharrajak This PR draws some inspiration from https://github.com/apache/datafusion-comet/pull/2994. Thanks for the early work towards this.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4322093250
In https://github.com/apache/datafusion-comet/pull/2994 we saw improvements to some Spark plans in the stability suite and I do not see the same with this PR. Moving this to draft until I have time to check this.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4504619503
One workflow keeps failing - sql core-1 4.1.1 ... will keep retrying .. failures are unrelated to PR
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#discussion_r3261964327
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -368,6 +386,19 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
false
}
+ /**
+ * Returns true if the given data type is or contains a `MapType` at any nesting level. Arrow's
+ * row format (used by DataFusion's grouped hash aggregate for composite group keys) does not
+ * support `Map`, so grouping on any type that transitively contains a map would crash in native
+ * execution.
+ */
+ def containsMapType(dt: DataType): Boolean = dt match {
Review Comment:
we already have this method in `operators.scala` we prob need to make this method generic. I'll check it in followup PR
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: allow safe mixed Spark/Comet partial/final aggregate execution [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4015:
URL: https://github.com/apache/datafusion-comet/pull/4015#issuecomment-4390385992
> In #2994 we saw improvements to some Spark plans in the stability suite and I do not see the same with this PR. Moving this to draft until I have time to check this.
This is because #2994 supported `COUNT` and this PR does not due to issues that came up. Follow on issue is https://github.com/apache/datafusion-comet/issues/4242
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org