Source:
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator, Session session) { if (!(node instanceof AggregationNode)) { return Optional.empty(); } AggregationNode parent = (AggregationNode) node; PlanNode source = lookup.resolve(parent.getSource()); if (!(source instanceof MarkDistinctNode)) { return Optional.empty(); } MarkDistinctNode child = (MarkDistinctNode) source; boolean hasFilters = parent.getAssignments().values().stream() .map(Aggregation::getCall) .map(FunctionCall::getFilter) .anyMatch(Optional::isPresent); if (hasFilters) { return Optional.empty(); } Collection<Aggregation> aggregations = parent.getAssignments().values(); List<Symbol> masks = aggregations.stream() .map(Aggregation::getMask).filter(Optional::isPresent).map(Optional::get).collect(toImmutableList()); Set<Symbol> uniqueMasks = ImmutableSet.copyOf(masks); if (uniqueMasks.size() != 1 || masks.size() != aggregations.size()) { return Optional.empty(); } Symbol mask = Iterables.getOnlyElement(uniqueMasks); if (!child.getMarkerSymbol().equals(mask)) { return Optional.empty(); } return Optional.of( new AggregationNode( idAllocator.getNextId(), new AggregationNode( idAllocator.getNextId(), child.getSource(), Collections.emptyMap(), ImmutableList.of(child.getDistinctSymbols()), SINGLE, child.getHashSymbol(), Optional.empty()), parent.getAssignments() .entrySet().stream() .collect(Collectors.toMap( Map.Entry::getKey, e -> removeDistinct(e.getValue()))), parent.getGroupingSets(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol())); }

Target:
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator, Session session) { if (!(node instanceof AggregationNode)) { return Optional.empty(); } AggregationNode parent = (AggregationNode) node; PlanNode source = lookup.resolve(parent.getSource()); if (!(source instanceof MarkDistinctNode)) { return Optional.empty(); } MarkDistinctNode child = (MarkDistinctNode) source; boolean hasFilters = parent.getAggregations().values().stream() .map(Aggregation::getCall) .map(FunctionCall::getFilter) .anyMatch(Optional::isPresent); if (hasFilters) { return Optional.empty(); } Collection<Aggregation> aggregations = parent.getAggregations().values(); List<Symbol> masks = aggregations.stream() .map(Aggregation::getMask) .filter(Optional::isPresent) .map(Optional::get) .collect(toImmutableList()); Set<Symbol> uniqueMasks = ImmutableSet.copyOf(masks); if (uniqueMasks.size() != 1 || masks.size() != aggregations.size()) { return Optional.empty(); } Symbol mask = Iterables.getOnlyElement(uniqueMasks); if (!child.getMarkerSymbol().equals(mask)) { return Optional.empty(); } return Optional.of( new AggregationNode( idAllocator.getNextId(), new AggregationNode( idAllocator.getNextId(), child.getSource(), Collections.emptyMap(), ImmutableList.of(child.getDistinctSymbols()), SINGLE, child.getHashSymbol(), Optional.empty()), parent.getAggregations() .entrySet().stream() .collect(Collectors.toMap( Map.Entry::getKey, e -> removeDistinct(e.getValue()))), parent.getGroupingSets(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol())); }