Code review comment: use FunctionSet.SUM instead of literal

Code under review:
public List<OptExpression> transform(OptExpression input, OptimizerContext context) {
        
        /*
         * Before
         *
         * LogicalProjectOperator {projection=[28: sum]}
         * ->  LogicalAggregation {type=GLOBAL ,aggregations={28: sum=sum(27: expr)} ,groupKeys=[]}
         *   ->  LogicalProjectOperator {projection=[add(cast(1: id_int as bigint(20)), 1)]}
         *
         * After
         *
         * LogicalProjectOperator {projection=[add(30: sum, multiply(31: count, 1))]
         * ->  LogicalAggregation {type=GLOBAL ,aggregations={30: sum=sum(29: id_int), 31: count=count()} ,groupKeys=[]}
         *   ->  LogicalProjectOperator {projection=[add(cast(1: id_int as bigint(20)), 1), 1: id_int]}
         */
        LogicalProjectOperator postAggProjectOperator = (LogicalProjectOperator) input.getOp();
        LogicalAggregationOperator aggregationOperator = (LogicalAggregationOperator) input.getInputs().get(0).getOp();
        LogicalProjectOperator preAggProjectOperator =
                (LogicalProjectOperator) input.getInputs().get(0).getInputs().get(0).getOp();

        ColumnRefSet requiredColumns = context.getTaskContext().getRequiredColumns();

        Map<ColumnRefOperator, ScalarOperator> newPreAggProjections = Maps.newHashMap();
        Map<ColumnRefOperator, ScalarOperator> newPostAggProjections = Maps.newHashMap();

        RewriteSumVisitor rewriteSumVisitor = new RewriteSumVisitor(
                aggregationOperator.getAggregations(),
                preAggProjectOperator.getColumnRefMap(),
                context.getColumnRefFactory(),
                newPreAggProjections,
                newPostAggProjections
                );

        
        for (Map.Entry<ColumnRefOperator, ScalarOperator> postAggProjection :
                postAggProjectOperator.getColumnRefMap().entrySet()) {
            ColumnRefOperator columnRefOperator = postAggProjection.getKey();
            ScalarOperator projectExpr = postAggProjection.getValue();
            ScalarOperator newProjectExpr = projectExpr.accept(rewriteSumVisitor, null);
            newPostAggProjections.put(columnRefOperator, newProjectExpr);
        }

        if (rewriteSumVisitor.newAggregations.isEmpty()) {
            return Lists.newArrayList();
        }

        requiredColumns.except(postAggProjectOperator.getColumnRefMap().keySet());
        requiredColumns.union(newPostAggProjections.keySet());

        rewriteSumVisitor.newAggregations.putAll(rewriteSumVisitor.reservedAggregations);
        newPreAggProjections.putAll(rewriteSumVisitor.oldPreAggProjections);

        LogicalProjectOperator newPostAggProjectOperator =
                new LogicalProjectOperator(newPostAggProjections);

        LogicalAggregationOperator newAggOperator = new LogicalAggregationOperator(
                aggregationOperator.getType(),
                aggregationOperator.getGroupingKeys(), rewriteSumVisitor.newAggregations);
        newAggOperator.setProjection(aggregationOperator.getProjection());

        OptExpression newPreAggProjectOpt = OptExpression.create(new LogicalProjectOperator(newPreAggProjections),
                input.getInputs().get(0).getInputs().get(0).getInputs());
        newPreAggProjectOpt.setLogicalProperty(new LogicalProperty(new ColumnRefSet(newPreAggProjections.keySet())));

        OptExpression newAggOpt = OptExpression.create(newAggOperator, newPreAggProjectOpt);
        newAggOpt.setLogicalProperty(new LogicalProperty(newAggOperator.getOutputColumns(null)));

        OptExpression newPostAggProjectOpt = OptExpression.create(newPostAggProjectOperator, newAggOpt);
        newPostAggProjectOpt.setLogicalProperty(new LogicalProperty(new ColumnRefSet(newPostAggProjections.keySet())));

        return Lists.newArrayList(newPostAggProjectOpt);
    }

    private static class RewriteSumVisitor extends ScalarOperatorVisitor<ScalarOperator, Void> {
        public Map<ColumnRefOperator, CallOperator> oldAggregations;
        public Map<ColumnRefOperator, ScalarOperator> oldPreAggProjections;
        public ColumnRefFactory columnRefFactory;
        public Map<ColumnRefOperator, ScalarOperator> newPostAggProjections;
        public Map<ColumnRefOperator, CallOperator> newAggregations;
        public Map<ColumnRefOperator, ScalarOperator> newPreAggProjections;

        public Map<ColumnRefOperator, CallOperator> reservedAggregations;

        private Map<ScalarOperator, ColumnRefOperator> commonArguments;

        public RewriteSumVisitor(Map<ColumnRefOperator, CallOperator> oldAggregations,
                                 Map<ColumnRefOperator, ScalarOperator> oldPreAggProjections,
                                 ColumnRefFactory columnRefFactory,
                                 Map<ColumnRefOperator, ScalarOperator> newPreAggProjections,
                                 Map<ColumnRefOperator, ScalarOperator> newPostAggProjections) {
            this.oldAggregations = oldAggregations;
            this.oldPreAggProjections = oldPreAggProjections;
            this.columnRefFactory = columnRefFactory;
            this.newPostAggProjections = newPostAggProjections;
            this.newAggregations = Maps.newHashMap();
            this.newPreAggProjections = newPreAggProjections;
            this.reservedAggregations = Maps.newHashMap();
            this.commonArguments = Maps.newHashMap();
        }

        @Override
        public ScalarOperator visit(ScalarOperator scalarOperator, Void context) {
            int childrenSize = scalarOperator.getChildren().size();
            for (int i = 0; i < childrenSize; i++) {
                ScalarOperator child = scalarOperator.getChild(i);
                scalarOperator.setChild(i, child.accept(this, null));
            }
            return scalarOperator;
        }

        private ScalarOperator rewriteCastForAggExpr(ScalarOperator expr) {
            if (expr instanceof CastOperator) {
                CastOperator castOperator = (CastOperator) expr;
                Type toType = castOperator.getType();
                ScalarOperator fromOperator = castOperator.getChild(0);
                Type fromType = fromOperator.getType();
                
                
                
                if (ScalarType.isFullyCompatible(fromType, toType)) {
                    AggregateFunction possibleAggregateFunction = AggregateFunction.createBuiltin(
                            FunctionSet.SUM, Lists.newArrayList(fromType), toType, toType,
                            false, true, false);
                    if (GlobalStateMgr.getCurrentState().getFunction(
                            possibleAggregateFunction, Function.CompareMode.IS_IDENTICAL) != null) {
                        return rewriteCastForAggExpr(fromOperator);
                    }
                }
            }
            return expr;
        }

        private ScalarOperator createNewAggFunction(ScalarOperator arg0, ScalarOperator arg1, Type returnType) {
            if (arg0.isConstant()) {
                

                AggregateFunction countFunction = AggregateFunction.createBuiltin(
                        FunctionSet.COUNT, Lists.newArrayList(arg1.getType()), Type.BIGINT, Type.BIGINT,
                        false, true, true);

                
                List<ScalarOperator> countArguments = Lists.newArrayList();
                if (arg1.isNullable()) {
                    countArguments.add(arg1);
                }
                CallOperator newAggFunction = new CallOperator("count",
                        Type.BIGINT, countArguments, countFunction);

                ColumnRefOperator newAggRef = columnRefFactory.create(
                        newAggFunction, newAggFunction.getType(), true);
                newAggregations.put(newAggRef, newAggFunction);

                
                ScalarOperator countOperator = null;
                ScalarOperator constOperator = null;
                PrimitiveType primitiveType = returnType.getPrimitiveType();
                switch (primitiveType) {
                    case DECIMAL128:
                        countOperator = new CastOperator(
                                ScalarType.createDecimalV3Type(PrimitiveType.DECIMAL128, 18, 0),
                                newAggRef, false);
                        int precision = ((ScalarType) arg0.getType()).getScalarPrecision();
                        int scale = ((ScalarType) arg0.getType()).getScalarScale();
                        Type constType = ScalarType.createDecimalV3Type(PrimitiveType.DECIMAL128, precision, scale);
                        constOperator = arg0.isConstantNull() ? ConstantOperator.createNull(constType) :
                                ConstantOperator.createDecimal(new BigDecimal(arg0.toString()), constType);
                        break;
                    case DOUBLE:
                        countOperator = new CastOperator(returnType, newAggRef, false);
                        constOperator = arg0.isConstantNull() ? ConstantOperator.createNull(Type.DOUBLE) :
                                ConstantOperator.createDouble(((ConstantOperator) arg0).getDouble());
                        break;
                    case LARGEINT:
                        countOperator = new CastOperator(returnType, newAggRef, false);
                        constOperator = arg0.isConstantNull() ? ConstantOperator.createNull(Type.LARGEINT) :
                                ConstantOperator.createLargeInt(new BigInteger(arg0.toString()));
                        break;
                    case BIGINT:
                        
                        countOperator = newAggRef;
                        constOperator = arg0.isConstantNull() ? ConstantOperator.createNull(Type.BIGINT) :
                                ConstantOperator.createBigint(Long.parseLong(arg0.toString()));
                        break;
                    default:
                        Preconditions.checkState(false, "unexpected sum function result type");
                }
                CallOperator newMultiply = new CallOperator("multiply", returnType,
                        Lists.newArrayList(
                                countOperator, constOperator));
                return newMultiply;
            } else {
                

                ColumnRefOperator newColumnRef;

                if (arg0.isColumnRef() && oldPreAggProjections.containsKey(arg0)) {
                    
                    
                    newColumnRef = (ColumnRefOperator) arg0;
                } else if (commonArguments.containsKey(arg0)) {
                    
                    
                    newColumnRef = commonArguments.get(arg0);
                } else {
                    newColumnRef = columnRefFactory.create(arg0, arg0.getType(), arg0.isNullable());
                    newPreAggProjections.put(newColumnRef, arg0);
                    commonArguments.put(arg0, newColumnRef);
                }

                Type sumFunctionType = returnType;
                if (returnType.isDecimalV3()) {
                    
                    int argScale = ((ScalarType) arg0.getType()).getScalarScale();
                    sumFunctionType = ScalarType.createDecimalV3Type(PrimitiveType.DECIMAL128, 38, argScale);
                }

                AggregateFunction sumFunction = AggregateFunction.createBuiltin(
                        FunctionSet.SUM, Lists.newArrayList(arg0.getType()), sumFunctionType, sumFunctionType,
                        false, true, false);

                CallOperator newAggFunction = new CallOperator("sum",
                        sumFunctionType, Lists.newArrayList(newColumnRef), sumFunction);

                ColumnRefOperator newAggRef = columnRefFactory.create(
                        newAggFunction, newAggFunction.getType(), true);
                newAggregations.put(newAggRef, newAggFunction);
                return newAggRef;
            }
        }

        @Override
        public ScalarOperator visitVariableReference(ColumnRefOperator op, Void context) {
            if (oldAggregations.containsKey(op)) {
                CallOperator aggFunction = oldAggregations.get(op);
                if (aggFunction.isAggregate() && !aggFunction.isDistinct()
                        && aggFunction.getFnName().equals("sum") &&
                        !aggFunction.getType().isDecimalV2()) {
                    ScalarOperator aggExpr = aggFunction.getArguments().get(0);
                    if (aggExpr.isColumnRef()) {
                        ColumnRefOperator aggColumnRef = (ColumnRefOperator) aggExpr;
                        ScalarOperator aggColumnExpr = oldPreAggProjections.get(aggColumnRef);
                        if (aggColumnExpr.getOpType() == OperatorType.CALL) {
                            CallOperator callOperator = (CallOperator) aggColumnExpr;
                            String functionName = callOperator.getFnName();
                            if (functionName.equals("add") || functionName.equals("subtract")) {
                                List<ScalarOperator> arguments = callOperator.getArguments();
                                Preconditions.checkState(arguments.size() == 2);
                                
                                if (arguments.get(0).isConstant() || arguments.get(1).isConstant()) {
                                    
                                    
                                    
                                    
                                    
                                    ScalarOperator newArg0 = rewriteCastForAggExpr(arguments.get(0));
                                    ScalarOperator newArg1 = rewriteCastForAggExpr(arguments.get(1));

                                    ScalarOperator agg0 = createNewAggFunction(
                                            newArg0, newArg1, aggFunction.getType());
                                    ScalarOperator agg1 = createNewAggFunction(
                                            newArg1, newArg0, aggFunction.getType());

                                    CallOperator newOperator = new CallOperator(functionName,
                                            aggFunction.getType(), Lists.newArrayList(agg0, agg1));
                                    return newOperator;
                                }
                            }
                        }
                    }
                }
                reservedAggregations.put(op, oldAggregations.get(op));
            }
            return op;
        }
    }
}

Line commented on: && aggFunction.getFnName().equals("sum") &&

Revised code:
public List<OptExpression> transform(OptExpression input, OptimizerContext context) {
        
        /*
         * Before
         *
         * LogicalAggregation {type=GLOBAL ,aggregations={28: sum=sum(27: expr)} ,groupKeys=[]}
         * ->  LogicalProjectOperator {projection=[add(cast(1: id_int as bigint(20)), 1)]}
         *
         * After
         *
         * LogicalProjectOperator {projection=[add(30: sum, multiply(31: count, 1))]}
         * -> LogicalAggregation {type=GLOBAL ,aggregations={30: sum=sum(29: id_int), 31: count=count()} ,groupKeys=[]}
         *   ->  LogicalProjectOperator {projection=[add(cast(1: id_int as bigint(20)), 1), 1: id_int]}
         */
        LogicalAggregationOperator aggregationOperator = (LogicalAggregationOperator) input.getOp();
        LogicalProjectOperator preAggProjectOperator = (LogicalProjectOperator) input.getInputs().get(0).getOp();

        ColumnRefSet requiredColumns = context.getTaskContext().getRequiredColumns();

        Map<ColumnRefOperator, ScalarOperator> newPreAggProjections = Maps.newHashMap();
        Map<ColumnRefOperator, ScalarOperator> newPostAggProjections = Maps.newHashMap();

        AggFunctionRewriter aggFunctionRewriter = new AggFunctionRewriter(
                aggregationOperator.getAggregations(),
                preAggProjectOperator.getColumnRefMap(),
                context.getColumnRefFactory(),
                newPreAggProjections,
                newPostAggProjections);

        for (Map.Entry<ColumnRefOperator, CallOperator> aggregation :
                aggregationOperator.getAggregations().entrySet()) {
            ColumnRefOperator columnRefOperator = aggregation.getKey();
            CallOperator aggFunction = aggregation.getValue();
            newPostAggProjections.put(columnRefOperator, aggFunctionRewriter.rewrite(columnRefOperator, aggFunction));
        }

        if (aggFunctionRewriter.newAggregations.isEmpty()) {
            return Lists.newArrayList();
        }

        requiredColumns.union(newPostAggProjections.keySet());
        Map<ColumnRefOperator, CallOperator> newAggregations = Maps.newHashMap();
        newAggregations.putAll(aggFunctionRewriter.newAggregations);
        newAggregations.putAll(aggFunctionRewriter.reservedAggregations);

        
        
        
        if (aggregationOperator.getGroupingKeys() != null && !aggregationOperator.getGroupingKeys().isEmpty()) {
            HashSet<CallOperator> uniqueNewAggregations = new HashSet<>(newAggregations.values());
            HashSet<CallOperator> uniqueOldAggregations = new HashSet<>(aggregationOperator.getAggregations().values());
            if (uniqueNewAggregations.size() >= uniqueOldAggregations.size()) {
                return Lists.newArrayList();
            }
        }

        newPreAggProjections.putAll(aggFunctionRewriter.oldPreAggProjections);

        LogicalProjectOperator newPostAggProjectOperator =
                new LogicalProjectOperator(newPostAggProjections);

        LogicalAggregationOperator newAggOperator = new LogicalAggregationOperator(
                aggregationOperator.getType(),
                aggregationOperator.getGroupingKeys(), newAggregations);
        Preconditions.checkState(aggregationOperator.getProjection() == null,
                "projection in LogicalAggOperator shouldn't be set in logical rewrite phase");

        OptExpression newPreAggProjectOpt = OptExpression.create(new LogicalProjectOperator(newPreAggProjections),
                input.getInputs().get(0).getInputs().get(0).getInputs());

        OptExpression newAggOpt = OptExpression.create(newAggOperator, newPreAggProjectOpt);

        OptExpression newPostAggProjectOpt = OptExpression.create(newPostAggProjectOperator, newAggOpt);
        return Lists.newArrayList(newPostAggProjectOpt);
    }