Source:
public Result apply(AggregationNode node, Captures captures, Context context) { ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder(); Symbol partitionCountSymbol = context.getSymbolAllocator().newSymbol("partition_count", INTEGER); ImmutableMap.Builder<Symbol, Expression> envelopeAssignments = ImmutableMap.builder(); for (Map.Entry<Symbol, Aggregation> entry : node.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); FunctionCall call = aggregation.getCall(); QualifiedName name = call.getName(); if (name.toString().equals(NAME) && call.getArguments().size() == 1) { Expression geometry = getOnlyElement(call.getArguments()); Symbol envelopeSymbol = context.getSymbolAllocator().newSymbol("envelope", metadata.getType(GEOMETRY_TYPE_SIGNATURE)); if (geometry instanceof FunctionCall && ((FunctionCall) geometry).getName().toString().equalsIgnoreCase("ST_Envelope")) { envelopeAssignments.put(envelopeSymbol, geometry); } else { envelopeAssignments.put(envelopeSymbol, new FunctionCall(QualifiedName.of("ST_Envelope"), ImmutableList.of(geometry))); } aggregations.put(entry.getKey(), new Aggregation( new FunctionCall(name, ImmutableList.of(envelopeSymbol.toSymbolReference(), partitionCountSymbol.toSymbolReference())), INTERNAL_SIGNATURE, aggregation.getMask())); } else { aggregations.put(entry); } } return Result.ofPlanNode( new AggregationNode( node.getId(), new ProjectNode( context.getIdAllocator().getNextId(), node.getSource(), Assignments.builder() .putIdentities(node.getSource().getOutputSymbols()) .put(partitionCountSymbol, new LongLiteral(format("%s", getHashPartitionCount(context.getSession())))) .putAll(envelopeAssignments.build()) .build()), aggregations.build(), node.getGroupingSets(), node.getPreGroupedSymbols(), node.getStep(), node.getHashSymbol(), node.getGroupIdSymbol())); }

Target:
public Result apply(AggregationNode node, Captures captures, Context context) { ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder(); Symbol partitionCountSymbol = context.getSymbolAllocator().newSymbol("partition_count", INTEGER); ImmutableMap.Builder<Symbol, Expression> envelopeAssignments = ImmutableMap.builder(); for (Map.Entry<Symbol, Aggregation> entry : node.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); FunctionCall call = aggregation.getCall(); QualifiedName name = call.getName(); if (name.toString().equals(NAME) && call.getArguments().size() == 1) { Expression geometry = getOnlyElement(call.getArguments()); Symbol envelopeSymbol = context.getSymbolAllocator().newSymbol("envelope", metadata.getType(GEOMETRY_TYPE_SIGNATURE)); if (geometry instanceof FunctionCall && ((FunctionCall) geometry).getName().toString().equalsIgnoreCase("ST_Envelope")) { envelopeAssignments.put(envelopeSymbol, geometry); } else { envelopeAssignments.put(envelopeSymbol, new FunctionCall(QualifiedName.of("ST_Envelope"), ImmutableList.of(geometry))); } aggregations.put(entry.getKey(), new Aggregation( new FunctionCall(name, ImmutableList.of(envelopeSymbol.toSymbolReference(), partitionCountSymbol.toSymbolReference())), INTERNAL_SIGNATURE, aggregation.getMask())); } else { aggregations.put(entry); } } return Result.ofPlanNode( new AggregationNode( node.getId(), new ProjectNode( context.getIdAllocator().getNextId(), node.getSource(), Assignments.builder() .putIdentities(node.getSource().getOutputSymbols()) .put(partitionCountSymbol, new LongLiteral(Integer.toString(getHashPartitionCount(context.getSession())))) .putAll(envelopeAssignments.build()) .build()), aggregations.build(), node.getGroupingSets(), node.getPreGroupedSymbols(), node.getStep(), node.getHashSymbol(), node.getGroupIdSymbol())); }