Refine GpuHashAggregateExec.setupReference

[sperlingxx]

Signed-off-by: sperlingxx [email protected]

Current PR refined GpuHashAggregateExec.setupReference to support TypedImperativeAggregate, which is required by #2916 . The primary work of current PR is reorganizing the code of boundInputReferences.

Originally, we workaround PartialMerge modes when we construct boundInputReferences. To be specific, the code piece dealing with aggregations which contain both PartialMerge and Partial mode is not generic enough. It regards GpuPivotFirst and GpuAverage as exceptional cases:

      case Partial =>
        // Partial with distinct case
        val updateExpressionsCudfAggsDistinct =
          updateExpressionsDistinct.filter(_.isInstanceOf[CudfAggregate])
              .map(_.asInstanceOf[CudfAggregate].ref)
        if (inputProjectionsDistinct.exists(p => !p.isInstanceOf[NamedExpression])) {
          // Case of distinct average we need to evaluate the "GpuCast and GpuIsNotNull" columns.
          // Refer to how input projections are setup for GpuAverage.
          // In the case where we have expressions to evaluate, pick the unique attributes
          // references from them as you only have one column for it before you start evaluating.
          distinctExpressions = inputProjectionsDistinct
          distinctAttributes = inputProjectionsDistinct.flatMap(ref =>
            ref.references.toSeq).distinct
        } else {
          distinctAttributes = updateAttributesDistinct
          distinctExpressions = updateExpressionsCudfAggsDistinct
        }

With the new implementation, we avoid all kinds of workaround. Instead, we fit all kinds of aggregation modes into three categories by the stage of AggregateExec in aggregation stack:

1. The first stage of aggregation stack, in terms of aggregation modes, it may contain aggregate expressions with Partial or Complete mode. It may also contain no aggregate expressions. For the first aggregation stage, the input projections are necessary, because it consume the outputs of non-Aggregate plans.

2. The final-like stages, including three conditions:

   • the final (last) stage of aggregation stack
   • the (partial) merge stage merely for nonDistinctAggAttributes (the second stage of aggregation stack with one distinct)

   For final-like stages, we just pass through all output attributes of child plan.

3. The third stage of aggregation stack with one distinct, in terms of aggregation modes, it contains both Partial mode (for distinctAgg) and PartialMerge mode (for nonDistinctAgg). For this stage, we need to switch the position of distinctAttributes and nonDistinctAttributes to match the output schema of the previous stage:

   • the schema of the 2nd stage's outputs: groupingAttributes ++ distinctAttributes ++ nonDistinctAggBufferAttributes
   • the schema of the 3rd stage's aggregate expressions: nonDistinctMergeAggExpressions ++ distinctPartialAggExpressions

[sperlingxx]

build

[abellina]

Could you add to the description for this PR around what exactly was changed and why? Please also note this PR is in flight: #2859, and also makes changes to setupReferences.

[sperlingxx]

build

[sperlingxx]

Could you add to the description for this PR around what exactly was changed and why? Please also note this PR is in flight: #2859, and also makes changes to setupReferences.

I added the description. And I rebased this branch onto latest main branch.

[sperlingxx]

build

[revans2]

@abellina and @jlowe could both of you take a look at this? You have a lot more experience with the aggregate code than I do. It looks okay to me, but I don't feel 100% qualified to approve it.

[abellina]

@revans2 yeap, I am looking.

[abellina]

Thanks for the added details @sperlingxx. I am mostly curious on the update expressions, as we used to use them before in the PartialMerge + Partial case https://github.com/NVIDIA/spark-rapids/blob/branch-21.08/sql-plugin/src/main/scala/com/nvidia/spark/rapids/aggregate.scala#L679. I think the comments need to be adjusted but it would be nice to make sure that you found that input projections is all that is needed in that case (perhaps the shapes match?)

Tagging @kuhushukla and @nartal1 as they may be more familiar with both distinct and pivot.

Had other minor nits.

[abellina]

Comments mention distinct update expressions, but I am not finding the use of update expressions in your change, only input projections.

[sperlingxx]

I updated this part along with comments, in order to verify and clarify "input projections is all that is needed in that case".

[abellina]

Suggested change

	if (modeInfo.hasPartialMerge && modeInfo.uniqueModes.contains(Partial)) {
	if (modeInfo.hasPartialMerge && modeInfo.hasPartialMode) {

[abellina]

bonus points => hasPartialMerge => hasPartialMergeMode

[sperlingxx]

We used modeInfo.uniqueModes.contains(Partial) instead of modeInfo.hasPartialMode here, because hasPartialMode is set as hasPartialMerge || uniqueModes.contains(Partial). I adjusted AggregateModeInfo to make its members more consistent with their names.

[sperlingxx]

build

[abellina]

I had a couple of nits, but they are minor. Approving ahead of time. Thanks for making the tweaks.

[sperlingxx]

build

[abellina]

LGTM