Skip to content

[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate non-flushable #12960

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhztheplayer merged 2 commits into from
Sep 4, 2026
Merged

[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate non-flushable #12960

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7d45a4e
[GLUTEN-12959][VL] Keep the final stage of a grouping-only aggregate …
Dhruv-meesho Sep 2, 2026
67ed775
[GLUTEN-12959][VL] Update TPC-DS plan stability files
Dhruv-meesho Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,20 @@ case class FlushableHashAggregateRule(session: SparkSession) extends Rule[SparkP
aggExprs.exists(isUnsupportedAggregation)
}

/**
* Returns true if the aggregate applies no aggregate functions and is the final (or complete)
* stage, e.g. the last step of `SELECT DISTINCT`, or of a `GROUP BY` without aggregate functions.
*/
Comment on lines +73 to +76

@zhztheplayer zhztheplayer Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Can we only keep the 1st paragraph of the comment?

@Dhruv-meesho Dhruv-meesho Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Trimmed, thanks. @zhztheplayer

private def isGroupingOnlyFinalAgg(agg: HashAggregateExecTransformer): Boolean = {
agg.aggregateExpressions.isEmpty && agg.requiredChildDistributionExpressions.isDefined
}

/**
* Walks the plan downward, applying func to each RegularHashAggregateExecTransformer or
* SortHashAggregateExecTransformer that is eligible for flushable conversion. An aggregate is
* eligible when all expressions are Partial/PartialMerge, it is not the protected PartialMerge
* aggregate directly below a distinct-partial aggregate, and no aggregate function disallows
* flushing.
* eligible when all expressions are Partial/PartialMerge, it is not the final stage of a
* grouping-only aggregate, it is not the protected PartialMerge aggregate directly below a
* distinct-partial aggregate, and no aggregate function disallows flushing.
*/
private def replaceEligibleAggregates(
plan: SparkPlan,
Expand All @@ -93,6 +101,9 @@ case class FlushableHashAggregateRule(session: SparkSession) extends Rule[SparkP
}

def transformDown: SparkPlan => SparkPlan = {
case agg: RegularHashAggregateExecTransformer if isGroupingOnlyFinalAgg(agg) =>
// Final stage of a grouping-only aggregate. It must fully aggregate. Skip.
agg
case agg: RegularHashAggregateExecTransformer
if !agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge) =>
// Not an intermediate agg. Skip.
Expand All @@ -110,6 +121,9 @@ case class FlushableHashAggregateRule(session: SparkSession) extends Rule[SparkP
case agg: RegularHashAggregateExecTransformer =>
// All guards passed; replace with the flushable variant.
toFlushableAgg(agg)
case agg: SortHashAggregateExecTransformer if isGroupingOnlyFinalAgg(agg) =>
// See the RegularHashAggregateExecTransformer branch above.
agg
case agg: SortHashAggregateExecTransformer
if !agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge) =>
// Not an intermediate agg. Skip.
Expand Down
33 changes: 33 additions & 0 deletions backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1281,6 +1281,39 @@ class VeloxAggregateFunctionsFlushSuite extends VeloxAggregateFunctionsSuite {
}
}

test("flushable aggregate rule - distinct feeding a join keeps final agg regular") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1k") {
// The join key is a strict subset of the distinct keys, so the distinct's final aggregate
// is repartitioned by a following exchange and thus visited by FlushableHashAggregateRule.
// Flushing it would emit duplicate distinct keys and inflate the join output.
runQueryAndCompare("""
|select count(*) from orders
|left join (select distinct l_orderkey, l_partkey from lineitem) d
| on o_orderkey = d.l_orderkey
|""".stripMargin) {
df =>
val executedPlan = getExecutedPlan(df)
val groupingOnlyFinalAggs = executedPlan.collect {
case agg: HashAggregateExecTransformer
if agg.aggregateExpressions.isEmpty &&
agg.requiredChildDistributionExpressions.isDefined =>
agg
}
assert(
groupingOnlyFinalAggs.nonEmpty,
"expected the distinct's final aggregate in the plan")
assert(
groupingOnlyFinalAggs.forall(
agg => !agg.isInstanceOf[FlushableHashAggregateExecTransformer]),
"the final aggregate of a grouping-only aggregation must not be flushable"
)
}
}
}

test("flushable aggregate decimal sum") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
Expand Down
4 changes: 2 additions & 2 deletions ...ces/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q14a.sf100/explain.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ VeloxColumnarToRow (187)
: : : : +- ColumnarExchange (58)
: : : : +- VeloxResizeBatches (57)
: : : : +- ^ ProjectExecTransformer (55)
: : : : +- ^ FlushableHashAggregateExecTransformer (54)
: : : : +- ^ RegularHashAggregateExecTransformer (54)
: : : : +- ^ InputIteratorTransformer (53)
: : : : +- ColumnarExchange (51)
: : : : +- VeloxResizeBatches (50)
Expand Down Expand Up @@ -370,7 +370,7 @@ Input [3]: [brand_id#27, class_id#28, category_id#29]
(53) InputIteratorTransformer
Input [3]: [brand_id#27, class_id#28, category_id#29]

(54) FlushableHashAggregateExecTransformer
(54) RegularHashAggregateExecTransformer
Input [3]: [brand_id#27, class_id#28, category_id#29]
Keys [3]: [brand_id#27, class_id#28, category_id#29]
Functions: []
Expand Down
2 changes: 1 addition & 1 deletion .../backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q14a.sf100/simplified.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (22)
ProjectExecTransformer [brand_id,class_id,category_id]
FlushableHashAggregateExecTransformer [brand_id,class_id,category_id]
RegularHashAggregateExecTransformer [brand_id,class_id,category_id]
InputIteratorTransformer
InputAdapter
ColumnarExchange [brand_id,class_id,category_id] #7
Expand Down
4 changes: 2 additions & 2 deletions ...ces/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q14b.sf100/explain.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ VeloxColumnarToRow (150)
: : : : +- ColumnarExchange (58)
: : : : +- VeloxResizeBatches (57)
: : : : +- ^ ProjectExecTransformer (55)
: : : : +- ^ FlushableHashAggregateExecTransformer (54)
: : : : +- ^ RegularHashAggregateExecTransformer (54)
: : : : +- ^ InputIteratorTransformer (53)
: : : : +- ColumnarExchange (51)
: : : : +- VeloxResizeBatches (50)
Expand Down Expand Up @@ -342,7 +342,7 @@ Input [3]: [brand_id#27, class_id#28, category_id#29]
(53) InputIteratorTransformer
Input [3]: [brand_id#27, class_id#28, category_id#29]

(54) FlushableHashAggregateExecTransformer
(54) RegularHashAggregateExecTransformer
Input [3]: [brand_id#27, class_id#28, category_id#29]
Keys [3]: [brand_id#27, class_id#28, category_id#29]
Functions: []
Expand Down
2 changes: 1 addition & 1 deletion .../backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q14b.sf100/simplified.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (24)
ProjectExecTransformer [brand_id,class_id,category_id]
FlushableHashAggregateExecTransformer [brand_id,class_id,category_id]
RegularHashAggregateExecTransformer [brand_id,class_id,category_id]
InputIteratorTransformer
InputAdapter
ColumnarExchange [brand_id,class_id,category_id] #6
Expand Down
12 changes: 6 additions & 6 deletions ...rces/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q38.sf100/explain.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ VeloxColumnarToRow (104)
: : +- ColumnarExchange (34)
: : +- VeloxResizeBatches (33)
: : +- ^ ProjectExecTransformer (31)
: : +- ^ FlushableHashAggregateExecTransformer (30)
: : +- ^ RegularHashAggregateExecTransformer (30)
: : +- ^ InputIteratorTransformer (29)
: : +- ColumnarExchange (27)
: : +- VeloxResizeBatches (26)
Expand All @@ -35,7 +35,7 @@ VeloxColumnarToRow (104)
: +- ColumnarExchange (65)
: +- VeloxResizeBatches (64)
: +- ^ ProjectExecTransformer (62)
: +- ^ FlushableHashAggregateExecTransformer (61)
: +- ^ RegularHashAggregateExecTransformer (61)
: +- ^ InputIteratorTransformer (60)
: +- ColumnarExchange (58)
: +- VeloxResizeBatches (57)
Expand All @@ -58,7 +58,7 @@ VeloxColumnarToRow (104)
+- ColumnarExchange (97)
+- VeloxResizeBatches (96)
+- ^ ProjectExecTransformer (94)
+- ^ FlushableHashAggregateExecTransformer (93)
+- ^ RegularHashAggregateExecTransformer (93)
+- ^ InputIteratorTransformer (92)
+- ColumnarExchange (90)
+- VeloxResizeBatches (89)
Expand Down Expand Up @@ -197,7 +197,7 @@ Input [3]: [c_last_name#9, c_first_name#8, d_date#5]
(29) InputIteratorTransformer
Input [3]: [c_last_name#9, c_first_name#8, d_date#5]

(30) FlushableHashAggregateExecTransformer
(30) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#9, c_first_name#8, d_date#5]
Keys [3]: [c_last_name#9, c_first_name#8, d_date#5]
Functions: []
Expand Down Expand Up @@ -320,7 +320,7 @@ Input [3]: [c_last_name#20, c_first_name#19, d_date#16]
(60) InputIteratorTransformer
Input [3]: [c_last_name#20, c_first_name#19, d_date#16]

(61) FlushableHashAggregateExecTransformer
(61) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#20, c_first_name#19, d_date#16]
Keys [3]: [c_last_name#20, c_first_name#19, d_date#16]
Functions: []
Expand Down Expand Up @@ -449,7 +449,7 @@ Input [3]: [c_last_name#30, c_first_name#29, d_date#26]
(92) InputIteratorTransformer
Input [3]: [c_last_name#30, c_first_name#29, d_date#26]

(93) FlushableHashAggregateExecTransformer
(93) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#30, c_first_name#29, d_date#26]
Keys [3]: [c_last_name#30, c_first_name#29, d_date#26]
Functions: []
Expand Down
6 changes: 3 additions & 3 deletions ...s/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q38.sf100/simplified.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (6)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #2
Expand Down Expand Up @@ -52,7 +52,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (12)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #7
Expand Down Expand Up @@ -84,7 +84,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (18)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #10
Expand Down
4 changes: 2 additions & 2 deletions ...urces/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q8.sf100/explain.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ VeloxColumnarToRow (85)
+- ColumnarExchange (69)
+- VeloxResizeBatches (68)
+- ^ ProjectExecTransformer (66)
+- ^ FlushableHashAggregateExecTransformer (65)
+- ^ RegularHashAggregateExecTransformer (65)
+- ^ InputIteratorTransformer (64)
+- ColumnarExchange (62)
+- VeloxResizeBatches (61)
Expand Down Expand Up @@ -332,7 +332,7 @@ Input [1]: [ca_zip#11]
(64) InputIteratorTransformer
Input [1]: [ca_zip#11]

(65) FlushableHashAggregateExecTransformer
(65) RegularHashAggregateExecTransformer
Input [1]: [ca_zip#11]
Keys [1]: [ca_zip#11]
Functions: []
Expand Down
2 changes: 1 addition & 1 deletion ...es/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q8.sf100/simplified.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (10)
ProjectExecTransformer [ca_zip]
FlushableHashAggregateExecTransformer [ca_zip]
RegularHashAggregateExecTransformer [ca_zip]
InputIteratorTransformer
InputAdapter
ColumnarExchange [ca_zip] #6
Expand Down
12 changes: 6 additions & 6 deletions ...rces/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q87.sf100/explain.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ VeloxColumnarToRow (104)
: : +- ColumnarExchange (34)
: : +- VeloxResizeBatches (33)
: : +- ^ ProjectExecTransformer (31)
: : +- ^ FlushableHashAggregateExecTransformer (30)
: : +- ^ RegularHashAggregateExecTransformer (30)
: : +- ^ InputIteratorTransformer (29)
: : +- ColumnarExchange (27)
: : +- VeloxResizeBatches (26)
Expand All @@ -35,7 +35,7 @@ VeloxColumnarToRow (104)
: +- ColumnarExchange (65)
: +- VeloxResizeBatches (64)
: +- ^ ProjectExecTransformer (62)
: +- ^ FlushableHashAggregateExecTransformer (61)
: +- ^ RegularHashAggregateExecTransformer (61)
: +- ^ InputIteratorTransformer (60)
: +- ColumnarExchange (58)
: +- VeloxResizeBatches (57)
Expand All @@ -58,7 +58,7 @@ VeloxColumnarToRow (104)
+- ColumnarExchange (97)
+- VeloxResizeBatches (96)
+- ^ ProjectExecTransformer (94)
+- ^ FlushableHashAggregateExecTransformer (93)
+- ^ RegularHashAggregateExecTransformer (93)
+- ^ InputIteratorTransformer (92)
+- ColumnarExchange (90)
+- VeloxResizeBatches (89)
Expand Down Expand Up @@ -197,7 +197,7 @@ Input [3]: [c_last_name#9, c_first_name#8, d_date#5]
(29) InputIteratorTransformer
Input [3]: [c_last_name#9, c_first_name#8, d_date#5]

(30) FlushableHashAggregateExecTransformer
(30) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#9, c_first_name#8, d_date#5]
Keys [3]: [c_last_name#9, c_first_name#8, d_date#5]
Functions: []
Expand Down Expand Up @@ -320,7 +320,7 @@ Input [3]: [c_last_name#20, c_first_name#19, d_date#16]
(60) InputIteratorTransformer
Input [3]: [c_last_name#20, c_first_name#19, d_date#16]

(61) FlushableHashAggregateExecTransformer
(61) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#20, c_first_name#19, d_date#16]
Keys [3]: [c_last_name#20, c_first_name#19, d_date#16]
Functions: []
Expand Down Expand Up @@ -449,7 +449,7 @@ Input [3]: [c_last_name#30, c_first_name#29, d_date#26]
(92) InputIteratorTransformer
Input [3]: [c_last_name#30, c_first_name#29, d_date#26]

(93) FlushableHashAggregateExecTransformer
(93) RegularHashAggregateExecTransformer
Input [3]: [c_last_name#30, c_first_name#29, d_date#26]
Keys [3]: [c_last_name#30, c_first_name#29, d_date#26]
Functions: []
Expand Down
6 changes: 3 additions & 3 deletions ...s/backends-velox/tpcds-plan-stability/gluten-approved-plans-v1_4/q87.sf100/simplified.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (6)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #2
Expand Down Expand Up @@ -52,7 +52,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (12)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #7
Expand Down Expand Up @@ -84,7 +84,7 @@ VeloxColumnarToRow
VeloxResizeBatches
WholeStageCodegenTransformer (18)
ProjectExecTransformer [c_last_name,c_first_name,d_date]
FlushableHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
RegularHashAggregateExecTransformer [c_last_name,c_first_name,d_date]
InputIteratorTransformer
InputAdapter
ColumnarExchange [c_last_name,c_first_name,d_date] #10
Expand Down
Loading