presto-main/src/main/java/io/prestosql/sql/planner/iterative/rule/RewriteSpatialPartitioningAggregation.java/RewriteSpatialPartitioningAggregation/apply
class RewriteSpatialPartitioningAggregation:    @Override
    public Result apply(AggregationNode node, Captures captures, Context context)
    {
        ResolvedFunction spatialPartitioningFunction = metadata.resolveFunction(QualifiedName.of(NAME), fromTypeSignatures(GEOMETRY_TYPE_SIGNATURE, INTEGER.getTypeSignature()));
        ResolvedFunction stEnvelopeFunction = metadata.resolveFunction(QualifiedName.of("ST_Envelope"), fromTypeSignatures(GEOMETRY_TYPE_SIGNATURE));

        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();
            String name = aggregation.getResolvedFunction().getSignature().getName();
            if (name.equals(NAME) && aggregation.getArguments().size() == 1) {
                Expression geometry = getOnlyElement(aggregation.getArguments());
                Symbol envelopeSymbol = context.getSymbolAllocator().newSymbol("envelope", metadata.getType(GEOMETRY_TYPE_SIGNATURE));
                if (isStEnvelopeFunctionCall(geometry, stEnvelopeFunction)) {
                    envelopeAssignments.put(envelopeSymbol, geometry);
                }
                else {
                    envelopeAssignments.put(envelopeSymbol, new FunctionCallBuilder(metadata)
                            .setName(QualifiedName.of("ST_Envelope"))
                            .addArgument(GEOMETRY_TYPE_SIGNATURE, geometry)
                            .build());
                }
                aggregations.put(entry.getKey(),
                        new Aggregation(
                                spatialPartitioningFunction,
                                ImmutableList.of(envelopeSymbol.toSymbolReference(), partitionCountSymbol.toSymbolReference()),
                                false,
                                Optional.empty(),
                                Optional.empty(),
                                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()));
    }