Copilot commented on code in PR #13328:
URL: https://github.com/apache/gluten/pull/13328#discussion_r4295325371
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -150,16 +179,75 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
} else {
childPlan
}
+ val stageChild = if (stage == AggregationStage.TWO_STAGE) {
+ buildMapperAggregate(grouping, preparedChild, rewrittenAggExprs)
+ } else {
+ preparedChild
+ }
+ val stageExpressions = if (stage == AggregationStage.TWO_STAGE) {
+ rewrittenAggExprs.map {
+ ae =>
+ val mode = ae.mode match {
+ case Partial => PartialMerge
+ case Complete => Final
+ case other => other
+ }
+ // Like Spark's AggUtils.mayRemoveAggFilters, remove filters from
buffer-merge stages.
+ ae.copy(mode = mode, filter = None)
+ }
+ } else {
+ rewrittenAggExprs
+ }
+ val distribution = stage match {
+ case AggregationStage.LOCAL => None
+ case AggregationStage.DISTRIBUTED | AggregationStage.TWO_STAGE =>
+ Some(grouping.map(_.toAttribute))
+ }
if (packOutputBuffers) {
- planBufferPhase(
+ buildBufferAggregate(
grouping,
aggregateExpressions,
resultExpressions,
- preparedChild,
- rewrittenAggExprs)
+ stageChild,
+ stageExpressions,
+ distribution)
} else {
- planFinalOutput(grouping, resultExpressions, preparedChild,
rewrittenAggExprs)
+ buildFinalAggregate(grouping, resultExpressions, stageChild,
stageExpressions, distribution)
Review Comment:
With `twoStage.mode=TOP`, the mapper outputs grouping attributes and
aggregate buffers, but the reducer receives the original grouping expressions.
For a computed key such as `substr(w_warehouse_name, 1, 20)`, the reducer
references `w_warehouse_name` after the mapper has dropped it, causing an
input-binding failure. Pass `grouping.map(_.toAttribute)` to the reducer in
TWO_STAGE mode so only the mapper evaluates the original expressions.
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala:
##########
@@ -150,16 +179,75 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
} else {
childPlan
}
+ val stageChild = if (stage == AggregationStage.TWO_STAGE) {
+ buildMapperAggregate(grouping, preparedChild, rewrittenAggExprs)
+ } else {
+ preparedChild
+ }
+ val stageExpressions = if (stage == AggregationStage.TWO_STAGE) {
+ rewrittenAggExprs.map {
+ ae =>
+ val mode = ae.mode match {
+ case Partial => PartialMerge
+ case Complete => Final
Review Comment:
With `twoStage.mode=TOP`, an unwrapped imperative aggregate such as Spark's
`approx_count_distinct` changes from `Complete` to `Final`. The mapper prefixes
its buffers with grouping columns, but `buildFinalAggregate` still uses
`initialInputBufferOffset = 0`. Spark's imperative merge reads buffers
positionally, so grouped queries consume the wrong columns when executed in
Spark, including fallback. Pass `grouping.size` as the reducer's input-buffer
offset for TWO_STAGE mode, preserving the existing offset for other stages. Add
a grouped SUM + approx_count_distinct regression test that exercises Spark
aggregation without native HLL rewriting.
--
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]
