Source:
<code>public Optional<DecorrelationResult> visitAggregation(AggregationNode node, Void context) { if (node.hasEmptyGroupingSet()) { return Optional.empty(); } Optional<DecorrelationResult> childDecorrelationResultOptional = lookup.resolve(node.getSource()).accept(this, null); if (!childDecorrelationResultOptional.isPresent()) { return Optional.empty(); } DecorrelationResult childDecorrelationResult = childDecorrelationResultOptional.get(); Set<Symbol> constantSymbols = childDecorrelationResult.getConstantSymbols(); AggregationNode decorrelatedAggregation = new SymbolMapper(childDecorrelationResult.correlatedSymbolsMapping) .map(node, childDecorrelationResult.node); ImmutableList.Builder<List<Symbol>> newGroupingSets = ImmutableList.builder(); for (List<Symbol> groupingSet : decorrelatedAggregation.getGroupingSets()) { List<Symbol> symbolsToAdd = childDecorrelationResult.symbolsToPropagate.stream() <START> .filter(symbol -> !groupingSet.contains(symbol)) <END> .collect(toImmutableList()); if (!constantSymbols.containsAll(symbolsToAdd)) { return Optional.empty(); } newGroupingSets.add(ImmutableList.<Symbol>builder() .addAll(groupingSet) .addAll(symbolsToAdd).build()); } AggregationNode newAggregation = new AggregationNode( decorrelatedAggregation.getId(), decorrelatedAggregation.getSource(), decorrelatedAggregation.getAggregations(), newGroupingSets.build(), decorrelatedAggregation.getStep(), decorrelatedAggregation.getHashSymbol(), decorrelatedAggregation.getGroupIdSymbol()); boolean atMostSingleRow = newAggregation.getGroupingSets().size() == 1 && constantSymbols.containsAll(newAggregation.getGroupingKeys()); return Optional.of(new DecorrelationResult( newAggregation, childDecorrelationResult.symbolsToPropagate, childDecorrelationResult.correlatedPredicates, childDecorrelationResult.correlatedSymbolsMapping, atMostSingleRow)); }</code><technical_language>nit. groupinSet is a list, calling contains in a loop is O(n^2)</technical_language>

Target:
public Optional<DecorrelationResult> visitAggregation(AggregationNode node, Void context) { if (node.hasEmptyGroupingSet()) { return Optional.empty(); } Optional<DecorrelationResult> childDecorrelationResultOptional = lookup.resolve(node.getSource()).accept(this, null); if (!childDecorrelationResultOptional.isPresent()) { return Optional.empty(); } DecorrelationResult childDecorrelationResult = childDecorrelationResultOptional.get(); Set<Symbol> constantSymbols = childDecorrelationResult.getConstantSymbols(); AggregationNode decorrelatedAggregation = childDecorrelationResult.getCorrelatedSymbolMapper() .map(node, childDecorrelationResult.node); ImmutableList.Builder<List<Symbol>> newGroupingSets = ImmutableList.builder(); for (List<Symbol> groupingSet : decorrelatedAggregation.getGroupingSets()) { Set<Symbol> set = ImmutableSet.copyOf(groupingSet); List<Symbol> symbolsToAdd = childDecorrelationResult.symbolsToPropagate.stream() .filter(symbol -> !set.contains(symbol)) .collect(toImmutableList()); if (!constantSymbols.containsAll(symbolsToAdd)) { return Optional.empty(); } newGroupingSets.add(ImmutableList.<Symbol>builder() .addAll(groupingSet) .addAll(symbolsToAdd) .build()); } AggregationNode newAggregation = new AggregationNode( decorrelatedAggregation.getId(), decorrelatedAggregation.getSource(), decorrelatedAggregation.getAggregations(), newGroupingSets.build(), decorrelatedAggregation.getStep(), decorrelatedAggregation.getHashSymbol(), decorrelatedAggregation.getGroupIdSymbol()); boolean atMostSingleRow = newAggregation.getGroupingSets().size() == 1 && constantSymbols.containsAll(newAggregation.getGroupingKeys()); return Optional.of(new DecorrelationResult( newAggregation, childDecorrelationResult.symbolsToPropagate, childDecorrelationResult.correlatedPredicates, childDecorrelationResult.correlatedSymbolsMapping, atMostSingleRow)); }