[VL] Improve FlushableHashAggregateRule

[zhztheplayer]

The patch improves FlushableHashAggregateRule to apply flushable optimizations in additional cases.

The rule code also gets simplified.

[Copilot]

🟡 Changes recommended

Fix protected aggregate tracking and preserve ordering for order-sensitive aggregates.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This pull request broadens FlushableHashAggregateRule optimizations and adds single-partition regression coverage.

Changes:

• Expands eligible aggregate rewrites.
• Simplifies transformation logic.
• Adds flushable aggregation tests without a shuffle.

File summaries

File	Summary
backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala	Adds single-partition flushable aggregation coverage.
backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala	Broadens and simplifies aggregate conversion logic.

Review details

Suppressed comments (1)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:42

• This broadens the rewrite to intermediate SortHashAggregateExecTransformers even when no shuffle separates the partial and final stages. In the existing test collect_list with ordering path (the suite uses one shuffle partition), this can turn the sorted partial aggregate into FlushableHashAggregateExecTransformer; allowFlush may emit multiple buffers and the final collect_list merge is concatenation, so the input ordering is no longer preserved (and the test's two-SortHash assertion can fail). Keep sort-preserving aggregates regular for order-sensitive functions, or otherwise restrict this new no-exchange case to aggregates whose flush semantics are order-independent.

      case agg: SortHashAggregateExecTransformer if isEligible(agg, protectedAggIds) =>
        toFlushableAgg(agg)

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

A critical issue affecting join-pushed partial aggregates requires correction and regression coverage.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:83

• Because transformUp now visits aggregates that are not below an exchange, this guard is insufficient for grouping-only complete aggregates. MergeTwoPhasesHashBaseAggregate can collapse a no-exchange SortAggregateExec into one node with empty aggregateExpressions and requiredChildDistributionExpressions = None (gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MergeTwoPhasesHashBaseAggregate.scala:143-152); forall is then vacuously true and isGroupingOnlyFinalAgg returns false. The complete deduplication node is converted to allowFlush=true, which can emit duplicate grouping keys and make SELECT DISTINCT or grouping without aggregate functions incorrect. Keep merged complete nodes regular or determine that a downstream merge exists, and add a no-exchange grouping-only regression test.

    !isGroupingOnlyFinalAgg(agg) &&
    agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge) &&

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

Address the grouping-partitioning guard and preserve protected distinct-aggregation nodes during rewriting.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:85

• protectedAggIds is collected from the original tree, but transformUp rewrites children before evaluating their parents. When the partial aggregate below this protected PartialMerge node is converted, Spark rebuilds the parent aggregate with a new plan-node id, so its original id is no longer in protectedAggIds and it can be converted to flushable. This breaks the one-distinct invariant that the existing test below expects (the protected partial merge must remain regular). Preserve the protection structurally during the rewrite (for example, use the prior top-down/boundary-aware traversal or carry the protection marker through rebuilt nodes) rather than matching stale ids.

  private def isEligible(
      agg: HashAggregateExecTransformer,
      protectedAggIds: Set[Int]): Boolean = {
    !isGroupingOnlyFinalAgg(agg) &&
    agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge) &&
    !protectedAggIds.contains(agg.id) &&
    !aggregatesNotSupportFlush(agg.aggregateExpressions)

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

One or more issues must be addressed before approval.

Review details

Suppressed comments (1)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:86

• The new top-down traversal now converts aggregates even when their child is already partitioned by the aggregate keys. In that case a flushable aggregate can emit duplicate grouping keys from a partition that Spark has deliberately arranged to contain the complete grouping set, changing the semantics of downstream partial aggregation (for example, when grouping keys are a superset of the input partitioning keys). The analogous Bolt rule explicitly skips this case via child.outputPartitioning.satisfies(ClusteredDistribution(groupingExpressions)); retain the same guard here before converting the aggregate.

  private def isEligible(
      agg: HashAggregateExecTransformer,
      protectedAggIds: Set[Int]): Boolean = {
    !isGroupingOnlyFinalAgg(agg) &&
    agg.aggregateExpressions.forall(p => p.mode == Partial || p.mode == PartialMerge) &&
    !protectedAggIds.contains(agg.id) &&
    !aggregatesNotSupportFlush(agg.aggregateExpressions)

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

The rewrite can duplicate join-aggregate partial rows and multiply join matches without a downstream merge.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

The aggregate rewrite may violate distribution assumptions and affect join-side partial aggregates.

Review details

Suppressed comments (2)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:43

• This traversal now converts every eligible partial/partial-merge aggregate, including one whose child already satisfies ClusteredDistribution on its grouping keys. Such an aggregate can be the first stage after an exchange; making it flushable allows duplicate grouping keys to be emitted even though Spark/native planning relies on that input being clustered, which can make downstream partial counts or joins overcount. The existing Bolt implementation explicitly skips this case in isAggInputAlreadyDistributedWithAggKeys (backends-bolt/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:92-95,127-139). Retain the distribution guard here (while preserving the new no-exchange case) before converting the aggregate.

      case agg: RegularHashAggregateExecTransformer if isEligible(agg, protectedAggIds) =>
        toFlushableAgg(agg)
      case agg: SortHashAggregateExecTransformer if isEligible(agg, protectedAggIds) =>
        toFlushableAgg(agg)

backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala:43

• This traversal now rewrites every eligible partial aggregate, including the ordinary HashAggregateExec created by ImplementJoinAggregate.planPartialPhase. That phase is explicitly packed in a ProjectExec and then consumed by the join (ImplementJoinAggregate.scala:107-115,126-137,194), so allowing it to flush can emit duplicate grouping keys that multiply join rows; the previous exchange-bounded traversal did not reach this join-side aggregate. Preserve an exclusion/boundary for join-aggregate partial phases (or otherwise prove their outputs are re-aggregated before the join).

    plan.transformDown {
      case agg: RegularHashAggregateExecTransformer if isEligible(agg, protectedAggIds) =>
        toFlushableAgg(agg)
      case agg: SortHashAggregateExecTransformer if isEligible(agg, protectedAggIds) =>
        toFlushableAgg(agg)

• Files reviewed: 42/42 changed files
• Comments generated: 0 new
• Review effort level: Lite