Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,11 @@ case class MergeTwoPhasesHashBaseAggregate(session: SparkSession)
private def isPartialAgg(partialAgg: BaseAggregateExec, finalAgg: BaseAggregateExec): Boolean = {
// Aggregates with a FILTER clause can be merged as long as the FILTER predicate is carried
// over to the Complete mode aggregate. Note the physical final aggregate has its FILTER
// stripped (Spark's AggUtils.mayRemoveAggFilters only keeps FILTER in Partial/Complete modes),
// so the FILTER must be restored from the partial aggregate when merging. Spark's aggregate
// planning produces the partial and final aggregate expression lists together and keeps them
// positionally aligned (the final phase reads the partial buffer by position), so a
// partial/final pair can be matched by position. We cannot match by `resultId`: for a single
// distinct aggregate, `AggUtils.planAggregateWithOneDistinct` builds the partial and final
// distinct expressions with fresh `AggregateExpression` instances, so their `resultId`s differ.
// stripped (Spark's AggUtils.mayRemoveAggFilters only keeps FILTER in Partial/Complete
// modes), so the FILTER must be restored from the partial aggregate when merging. Spark's
// aggregate planning produces the partial and final aggregate expression lists together and
// keeps them positionally aligned (the final phase reads the partial buffer by position), so
// a partial/final pair can be matched by position.
if (
partialAgg.aggregateExpressions.forall(x => x.mode == Partial) &&
finalAgg.aggregateExpressions.forall(x => x.mode == Final) &&
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,43 @@ abstract class BaseMergeTwoPhasesHashBaseAggregateSuite extends WholeStageTransf
compareResult = true,
df => checkHashAggregateCount(df, 1)
)

// pure distinct + FILTER: Spark's planAggregateWithOneDistinct inserts an exchange
// between the partial and final stages, so the merge rule's direct parent-child pattern
// match does NOT fire. This records that boundary: distinct + FILTER is not merged into
// a single aggregate, and the FILTER is preserved through the unmerged stages.
compareResultsAgainstVanillaSpark(
"""
|SELECT count(DISTINCT key) FILTER (WHERE key > 50) AS pc
|FROM v1
|""".stripMargin,
compareResult = true,
df => {
df.collect()
val plans = collect(df.queryExecution.executedPlan) {
case agg: HashAggregateExecBaseTransformer => agg
}
assert(plans.size > 1, "distinct + FILTER should not be merged into a single aggregate")
}
Comment on lines +127 to +133
)

// mixed distinct + non-distinct + FILTER: same exchange boundary; the merge rule does
// not fire. Verifies FILTER is preserved on both sides through the unmerged stages.
compareResultsAgainstVanillaSpark(
"""
|SELECT count(DISTINCT key) FILTER (WHERE key > 50) AS dc,
| count(key) FILTER (WHERE key LIKE '%1%') AS pc
|FROM v1
|""".stripMargin,
compareResult = true,
df => {
df.collect()
val plans = collect(df.queryExecution.executedPlan) {
case agg: HashAggregateExecBaseTransformer => agg
}
assert(plans.size > 1, "distinct + FILTER should not be merged into a single aggregate")
}
)
}

// with exchange hash aggregate
Expand Down
Loading