Code review comment: should include WindowExpression inner trival-agg

Code under review:
public Rule build() {
        return logicalAggregate().whenNot(LogicalAggregate::isNormalized).then(aggregate -> {
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            

            
            
            Set<Expression> groupingByExprs =
                    ImmutableSet.copyOf(aggregate.getGroupByExpressions());

            
            List<NamedExpression> aggregateOutput = aggregate.getOutputExpressions();
            List<AggregateFunction> aggFuncs = Lists.newArrayList();
            aggregateOutput.forEach(o -> o.accept(CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));

            
            Set<SubqueryExpr> subqueryExprs = ExpressionUtils.mutableCollect(aggFuncs.stream()
                    .filter(aggFunc -> !aggFunc.isDistinct()).collect(Collectors.toList()),
                    SubqueryExpr.class::isInstance);

            
            Set<Slot> allAggFuncInputSlots = aggFuncs.stream()
                    .flatMap(agg -> agg.getInputSlots().stream()).collect(Collectors.toSet());

            
            Set<Expression> distinctAggChildExprs = aggFuncs.stream()
                    .filter(agg -> agg.isDistinct()).flatMap(agg -> agg.children().stream())
                    .filter(child -> !(child instanceof SlotReference && child instanceof Literal))
                    .collect(Collectors.toSet());

            Set<Alias> existsAlias =
                    ExpressionUtils.mutableCollect(aggregateOutput, Alias.class::isInstance);

            
            Set<Expression> allPushDownExprs =
                    Sets.union(Sets.union(groupingByExprs, subqueryExprs), distinctAggChildExprs);
            NormalizeToSlotContext bottomSlotContext =
                    NormalizeToSlotContext.buildContext(existsAlias, allPushDownExprs);
            Set<NamedExpression> bottomProjects =
                    bottomSlotContext.pushDownToNamedExpression(allPushDownExprs);

            
            Plan bottomPlan;
            if (!bottomProjects.isEmpty()) {
                bottomPlan = new LogicalProject<>(
                        ImmutableList.copyOf(Sets.union(bottomProjects, allAggFuncInputSlots)),
                        aggregate.child());
            } else {
                bottomPlan = aggregate.child();
            }

            
            
            
            
            
            
            
            
            
            

            
            List<Expression> normalizedGroupExprs =
                    bottomSlotContext.normalizeToUseSlotRef(groupingByExprs);

            
            List<AggregateFunction> normalizedAggFuncs =
                    bottomSlotContext.normalizeToUseSlotRef(aggFuncs);

            
            NormalizeToSlotContext normalizedAggFuncsToSlotContext =
                    NormalizeToSlotContext.buildContext(existsAlias, normalizedAggFuncs);

            
            
            List<NamedExpression> normalizedAggOutput = ImmutableList.<NamedExpression>builder()
                    .addAll(bottomProjects.stream().map(NamedExpression::toSlot).iterator())
                    .addAll(normalizedAggFuncsToSlotContext
                            .pushDownToNamedExpression(normalizedAggFuncs))
                    .build();

            
            LogicalAggregate newAggregate =
                    aggregate.withNormalized(normalizedGroupExprs, normalizedAggOutput, bottomPlan);

            
            List<NamedExpression> upperProjects = normalizeOutput(aggregateOutput,
                    bottomSlotContext, normalizedAggFuncsToSlotContext);

            
            return new LogicalProject<>(upperProjects, newAggregate);
        }).toRule(RuleType.NORMALIZE_AGGREGATE);
    }


Revised code:
public Rule build() {
        return logicalAggregate().whenNot(LogicalAggregate::isNormalized).then(aggregate -> {
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            

            
            
            Set<Expression> groupingByExprs =
                    ImmutableSet.copyOf(aggregate.getGroupByExpressions());

            
            List<NamedExpression> aggregateOutput = aggregate.getOutputExpressions();
            List<AggregateFunction> aggFuncs = Lists.newArrayList();
            aggregateOutput.forEach(o -> o.accept(CollectNonWindowedAggFuncs.INSTANCE, aggFuncs));

            
            
            
            Map<Boolean, Set<Expression>> categorizedNoDistinctAggsChildren = aggFuncs.stream()
                    .filter(aggFunc -> !aggFunc.isDistinct())
                    .flatMap(agg -> agg.children().stream())
                    .collect(Collectors.groupingBy(
                            child -> child.containsType(SubqueryExpr.class, WindowExpression.class),
                            Collectors.toSet()));

            
            
            
            Map<Boolean, Set<Expression>> categorizedDistinctAggsChildren = aggFuncs.stream()
                    .filter(aggFunc -> aggFunc.isDistinct()).flatMap(agg -> agg.children().stream())
                    .collect(Collectors.groupingBy(
                            child -> !(child instanceof SlotReference || child instanceof Literal),
                            Collectors.toSet()));

            Set<Expression> needPushSelf = Sets.union(
                    categorizedNoDistinctAggsChildren.getOrDefault(true, new HashSet<>()),
                    categorizedDistinctAggsChildren.getOrDefault(true, new HashSet<>()));
            Set<Slot> needPushInputSlots = ExpressionUtils.getInputSlotSet(Sets.union(
                    categorizedNoDistinctAggsChildren.getOrDefault(false, new HashSet<>()),
                    categorizedDistinctAggsChildren.getOrDefault(false, new HashSet<>())));

            Set<Alias> existsAlias =
                    ExpressionUtils.mutableCollect(aggregateOutput, Alias.class::isInstance);

            
            
            
            
            Set<Expression> allPushDownExprs =
                    Sets.union(groupingByExprs, Sets.union(needPushSelf, needPushInputSlots));
            NormalizeToSlotContext bottomSlotContext =
                    NormalizeToSlotContext.buildContext(existsAlias, allPushDownExprs);
            Set<NamedExpression> pushedGroupByExprs =
                    bottomSlotContext.pushDownToNamedExpression(groupingByExprs);
            Set<NamedExpression> pushedTrivalAggChildren =
                    bottomSlotContext.pushDownToNamedExpression(needPushSelf);
            Set<NamedExpression> pushedTrivalAggInputSlots =
                    bottomSlotContext.pushDownToNamedExpression(needPushInputSlots);
            Set<NamedExpression> bottomProjects = Sets.union(pushedGroupByExprs,
                    Sets.union(pushedTrivalAggChildren, pushedTrivalAggInputSlots));

            
            Plan bottomPlan;
            if (!bottomProjects.isEmpty()) {
                bottomPlan = new LogicalProject<>(ImmutableList.copyOf(bottomProjects),
                        aggregate.child());
            } else {
                bottomPlan = aggregate.child();
            }

            
            
            
            
            
            
            
            
            
            

            
            List<Expression> normalizedGroupExprs =
                    bottomSlotContext.normalizeToUseSlotRef(groupingByExprs);

            
            List<AggregateFunction> normalizedAggFuncs =
                    bottomSlotContext.normalizeToUseSlotRef(aggFuncs);

            
            NormalizeToSlotContext normalizedAggFuncsToSlotContext =
                    NormalizeToSlotContext.buildContext(existsAlias, normalizedAggFuncs);

            
            
            List<NamedExpression> normalizedAggOutput = ImmutableList.<NamedExpression>builder()
                    .addAll(pushedGroupByExprs.stream().map(NamedExpression::toSlot).iterator())
                    .addAll(normalizedAggFuncsToSlotContext
                            .pushDownToNamedExpression(normalizedAggFuncs))
                    .build();

            
            LogicalAggregate newAggregate =
                    aggregate.withNormalized(normalizedGroupExprs, normalizedAggOutput, bottomPlan);

            
            List<NamedExpression> upperProjects = normalizeOutput(aggregateOutput,
                    bottomSlotContext, normalizedAggFuncsToSlotContext);

            
            return new LogicalProject<>(upperProjects, newAggregate);
        }).toRule(RuleType.NORMALIZE_AGGREGATE);
    }