Source:
private AggregationNode createNonDistinctAggregation( AggregateInfo aggregateInfo, VariableReferenceExpression distinctVariable, VariableReferenceExpression duplicatedDistinctVariable, Set<VariableReferenceExpression> groupByKeys, GroupIdNode groupIdNode, MarkDistinctNode originalNode, ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> aggregationOutputSymbolsMapBuilder) { ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder(); for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregateInfo.getAggregations().entrySet()) { if (!entry.getValue().getMask().isPresent()) { VariableReferenceExpression newVariable = symbolAllocator.newVariable(entry.getKey()); Aggregation aggregation = entry.getValue(); aggregationOutputSymbolsMapBuilder.put(newVariable, entry.getKey()); List<RowExpression> arguments; if (!duplicatedDistinctVariable.equals(distinctVariable) && extractVariables(entry.getValue().getArguments(), symbolAllocator.getTypes()).contains(distinctVariable)) { ImmutableList.Builder<RowExpression> argumentsBuilder = ImmutableList.builder(); for (RowExpression argument : aggregation.getArguments()) { if (getVariable(argument, symbolAllocator.getTypes()).map(a -> a.equals(distinctVariable)).orElse(false)) { argumentsBuilder.add(castToRowExpression(asSymbolReference(duplicatedDistinctVariable))); } else { argumentsBuilder.add(argument); } } arguments = argumentsBuilder.build(); } else { arguments = aggregation.getArguments(); } aggregations.put(newVariable, new Aggregation( new CallExpression( aggregation.getCall().getDisplayName(), aggregation.getCall().getFunctionHandle(), aggregation.getCall().getType(), arguments), Optional.empty(), Optional.empty(), false, Optional.empty())); } } return new AggregationNode( idAllocator.getNextId(), groupIdNode, aggregations.build(), singleGroupingSet(ImmutableList.copyOf(groupByKeys)), ImmutableList.of(), SINGLE, originalNode.getHashVariable(), Optional.empty()); }

Target:
private AggregationNode createNonDistinctAggregation( AggregateInfo aggregateInfo, VariableReferenceExpression distinctVariable, VariableReferenceExpression duplicatedDistinctVariable, Set<VariableReferenceExpression> groupByKeys, GroupIdNode groupIdNode, MarkDistinctNode originalNode, ImmutableMap.Builder<VariableReferenceExpression, VariableReferenceExpression> aggregationOutputSymbolsMapBuilder) { ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder(); for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregateInfo.getAggregations().entrySet()) { if (!entry.getValue().getMask().isPresent()) { VariableReferenceExpression newVariable = symbolAllocator.newVariable(entry.getKey()); Aggregation aggregation = entry.getValue(); aggregationOutputSymbolsMapBuilder.put(newVariable, entry.getKey()); List<RowExpression> arguments; if (!duplicatedDistinctVariable.equals(distinctVariable) && extractVariables(entry.getValue().getArguments(), symbolAllocator.getTypes()).contains(distinctVariable)) { ImmutableList.Builder<RowExpression> argumentsBuilder = ImmutableList.builder(); for (RowExpression argument : aggregation.getArguments()) { if (castToExpression(argument) instanceof SymbolReference && toVariableReference(Symbol.from(castToExpression(argument)), symbolAllocator.getTypes()).equals(distinctVariable)) { argumentsBuilder.add(castToRowExpression(asSymbolReference(duplicatedDistinctVariable))); } else { argumentsBuilder.add(argument); } } arguments = argumentsBuilder.build(); } else { arguments = aggregation.getArguments(); } aggregations.put(newVariable, new Aggregation( new CallExpression( aggregation.getCall().getDisplayName(), aggregation.getCall().getFunctionHandle(), aggregation.getCall().getType(), arguments), Optional.empty(), Optional.empty(), false, Optional.empty())); } } return new AggregationNode( idAllocator.getNextId(), groupIdNode, aggregations.build(), singleGroupingSet(ImmutableList.copyOf(groupByKeys)), ImmutableList.of(), SINGLE, originalNode.getHashVariable(), Optional.empty()); }