Source:
public Result apply(AggregationNode parent, Captures captures, Context context) { if (!SystemSessionProperties.useMarkDistinct(context.getSession())) { return Result.empty(); } Map<Set<Symbol>, Symbol> markers = new HashMap<>(); Map<Symbol, Aggregation> newAggregations = new HashMap<>(); PlanNode subPlan = parent.getSource(); for (Map.Entry<Symbol, Aggregation> entry : parent.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); if (aggregation.isDistinct() && !aggregation.getFilter().isPresent() && !aggregation.getMask().isPresent()) { Set<Symbol> inputs = aggregation.getArguments().stream() .map(MultipleDistinctAggregationToMarkDistinct::toSymbol) .collect(toSet()); Symbol marker = markers.get(inputs); if (marker == null) { marker = context.getSymbolAllocator().newSymbol(Iterables.getLast(inputs).getName(), BOOLEAN, "distinct"); markers.put(inputs, marker); ImmutableSet.Builder<Symbol> distinctSymbols = ImmutableSet.<Symbol>builder() .addAll(parent.getGroupingKeys()) .addAll(inputs); parent.getGroupIdSymbol().ifPresent(distinctSymbols::add); subPlan = new MarkDistinctNode( context.getIdAllocator().getNextId(), subPlan, marker, ImmutableList.copyOf(distinctSymbols.build()), Optional.empty()); } newAggregations.put(entry.getKey(), new Aggregation( new CallExpression( aggregation.getCall().getDisplayName(), aggregation.getCall().getFunctionHandle(), aggregation.getCall().getType(), aggregation.getArguments()), aggregation.getFilter(), aggregation.getOrderBy(), false, Optional.of(marker))); } else { newAggregations.put(entry.getKey(), aggregation); } } return Result.ofPlanNode( new AggregationNode( parent.getId(), subPlan, newAggregations, parent.getGroupingSets(), ImmutableList.of(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol())); }

Target:
public Result apply(AggregationNode parent, Captures captures, Context context) { if (!SystemSessionProperties.useMarkDistinct(context.getSession())) { return Result.empty(); } Map<Set<VariableReferenceExpression>, VariableReferenceExpression> markers = new HashMap<>(); Map<VariableReferenceExpression, Aggregation> newAggregations = new HashMap<>(); PlanNode subPlan = parent.getSource(); for (Map.Entry<VariableReferenceExpression, Aggregation> entry : parent.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); if (aggregation.isDistinct() && !aggregation.getFilter().isPresent() && !aggregation.getMask().isPresent()) { Set<VariableReferenceExpression> inputs = aggregation.getArguments().stream() .map(OriginalExpressionUtils::castToExpression) .map(Symbol::from) .map(context.getSymbolAllocator()::toVariableReference) .collect(toSet()); VariableReferenceExpression marker = markers.get(inputs); if (marker == null) { marker = context.getSymbolAllocator().newVariable(Iterables.getLast(inputs).getName(), BOOLEAN, "distinct"); markers.put(inputs, marker); ImmutableSet.Builder<VariableReferenceExpression> distinctVariables = ImmutableSet.<VariableReferenceExpression>builder() .addAll(parent.getGroupingKeys()) .addAll(inputs); parent.getGroupIdVariable().ifPresent(distinctVariables::add); subPlan = new MarkDistinctNode( context.getIdAllocator().getNextId(), subPlan, marker, ImmutableList.copyOf(distinctVariables.build()), Optional.empty()); } newAggregations.put(entry.getKey(), new Aggregation( aggregation.getCall(), aggregation.getFilter(), aggregation.getOrderBy(), false, Optional.of(marker))); } else { newAggregations.put(entry.getKey(), aggregation); } } return Result.ofPlanNode( new AggregationNode( parent.getId(), subPlan, newAggregations, parent.getGroupingSets(), ImmutableList.of(), parent.getStep(), parent.getHashVariable(), parent.getGroupIdVariable())); }