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(); TypeProvider types = context.getSymbolAllocator().getTypes(); for (Map.Entry<Symbol, Aggregation> entry : node.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); String name = metadata.getFunctionManager().getFunctionMetadata(aggregation.getFunctionHandle()).getName(); Type geometryType = metadata.getType(GEOMETRY_TYPE_SIGNATURE); if (name.equals(NAME) && aggregation.getArguments().size() == 1) { RowExpression geometry = getOnlyElement(aggregation.getArguments()); Symbol envelopeSymbol = context.getSymbolAllocator().newSymbol("envelope", geometryType); if (isFunctionNameMatch(metadata.getFunctionManager(), geometry, "ST_Envelope")) { envelopeAssignments.put(envelopeSymbol, castToExpression(geometry)); } else { envelopeAssignments.put(envelopeSymbol, new FunctionCall(QualifiedName.of("ST_Envelope"), ImmutableList.of(castToExpression(geometry)))); } aggregations.put(entry.getKey(), new Aggregation( new CallExpression( name, metadata.getFunctionManager().lookupFunction(NAME, fromTypes(geometryType, INTEGER)), context.getSymbolAllocator().getTypes().get(entry.getKey()), ImmutableList.of( castToRowExpression(envelopeSymbol.toSymbolReference()), castToRowExpression(partitionCountSymbol.toSymbolReference()))), Optional.empty(), Optional.empty(), false, 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())); }

Target:
public Result apply(AggregationNode node, Captures captures, Context context) { ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder(); VariableReferenceExpression partitionCountVariable = context.getSymbolAllocator().newVariable("partition_count", INTEGER); ImmutableMap.Builder<VariableReferenceExpression, Expression> envelopeAssignments = ImmutableMap.builder(); for (Map.Entry<VariableReferenceExpression, Aggregation> entry : node.getAggregations().entrySet()) { Aggregation aggregation = entry.getValue(); String name = metadata.getFunctionManager().getFunctionMetadata(aggregation.getFunctionHandle()).getName(); Type geometryType = metadata.getType(GEOMETRY_TYPE_SIGNATURE); if (name.equals(NAME) && aggregation.getArguments().size() == 1) { RowExpression geometry = getOnlyElement(aggregation.getArguments()); VariableReferenceExpression envelopeVariable = context.getSymbolAllocator().newVariable("envelope", geometryType); if (isFunctionNameMatch(geometry, "ST_Envelope")) { envelopeAssignments.put(envelopeVariable, castToExpression(geometry)); } else { envelopeAssignments.put(envelopeVariable, new FunctionCall(QualifiedName.of("ST_Envelope"), ImmutableList.of(castToExpression(geometry)))); } aggregations.put(entry.getKey(), new Aggregation( new CallExpression( name, metadata.getFunctionManager().lookupFunction(NAME, fromTypes(geometryType, INTEGER)), entry.getKey().getType(), ImmutableList.of( castToRowExpression(asSymbolReference(envelopeVariable)), castToRowExpression(asSymbolReference(partitionCountVariable)))), Optional.empty(), Optional.empty(), false, 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().getOutputVariables()) .put(partitionCountVariable, new LongLiteral(Integer.toString(getHashPartitionCount(context.getSession())))) .putAll(envelopeAssignments.build()) .build()), aggregations.build(), node.getGroupingSets(), node.getPreGroupedVariables(), node.getStep(), node.getHashVariable(), node.getGroupIdVariable())); }