src/org/apache/pig/newplan/logical/rules/ConstantCalculator.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.pig.newplan.logical.rules;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.datastorage.ConfigurationUtil;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.PigHadoopLogger;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.impl.util.UDFContext;
import org.apache.pig.newplan.BaseOperatorPlan;
import org.apache.pig.newplan.DependencyOrderWalker;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.PlanWalker;
import org.apache.pig.newplan.ReverseDependencyOrderWalkerWOSeenChk;
import org.apache.pig.newplan.logical.expression.AllSameExpressionVisitor;
import org.apache.pig.newplan.logical.expression.ConstantExpression;
import org.apache.pig.newplan.logical.expression.ExpToPhyTranslationVisitor;
import org.apache.pig.newplan.logical.expression.LogicalExpression;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.LogicalExpressionVisitor;
import org.apache.pig.newplan.logical.expression.UserFuncExpression;
import org.apache.pig.newplan.logical.optimizer.AllExpressionVisitor;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.optimizer.Rule;
import org.apache.pig.newplan.optimizer.Transformer;
import org.joda.time.DateTimeZone;

public abstract class ConstantCalculator extends Rule {
    private List<LogicalRelationalOperator> processedOperators = new ArrayList<LogicalRelationalOperator>();
    private PigContext pc;

    public ConstantCalculator(String n, PigContext pc) {
        super(n, false);
        this.pc = pc;
    }

    @Override
    public Transformer getNewTransformer() {
        return new ConstantCalculatorTransformer(processedOperators, pc);
    }

    public static class ConstantCalculatorTransformer extends Transformer {
        private List<LogicalRelationalOperator> processedOperators = new ArrayList<LogicalRelationalOperator>();
        private OperatorPlan plan;
        private PigContext pc;

        public ConstantCalculatorTransformer(List<LogicalRelationalOperator> processedOperators, PigContext pc) {
            this.processedOperators = processedOperators;
            this.pc = pc;
        }

        @Override
        public boolean check(OperatorPlan matched) throws FrontendException {
            Iterator<Operator> operators = matched.getOperators();
            while (operators.hasNext()) {
                LogicalRelationalOperator operator = (LogicalRelationalOperator)operators.next();
    
                // If the operator is already processed, we quit.
                if (processedOperators.contains(operator)) {
                    continue;
                }

                processedOperators.add(operator);
                return true;
            }
            return false;
        }

        public static class ConstantCalculatorLogicalPlanVisitor extends AllExpressionVisitor {
            private PigContext pc;
            public ConstantCalculatorLogicalPlanVisitor(OperatorPlan plan, PigContext pc) throws FrontendException {
                super(plan, new DependencyOrderWalker(plan));
                this.pc = pc;
            }

            @Override
            protected LogicalExpressionVisitor getVisitor(
                    LogicalExpressionPlan expr) throws FrontendException {
                return new ConstantCalculatorExpressionVisitor(expr, currentOp, pc);
            }
        }

        public static class ConstantCalculatorExpressionVisitor extends AllSameExpressionVisitor {
            private LogicalRelationalOperator currentOp;
            private PigContext pc;
            private DateTimeZone currentDTZ = null;
            public ConstantCalculatorExpressionVisitor(OperatorPlan expPlan,
                    LogicalRelationalOperator currentOp, PigContext pc) throws FrontendException {
                super(expPlan, new ReverseDependencyOrderWalkerWOSeenChk(expPlan));
                this.currentOp = currentOp;
                this.pc = pc;
            }

            @Override
            protected void execute(LogicalExpression op) throws FrontendException {
                if (op instanceof UserFuncExpression) {
                    UserFuncExpression udf = (UserFuncExpression)op;
                    if (!udf.getEvalFunc().allowCompileTimeCalculation()) {
                        return;
                    }
                }
                boolean valSet = false;
                Object val = null;
                if (currentWalker.getPlan().getSuccessors(op) != null) {
                    // If has successors and all successors are constant, calculate the constant
                    for (Operator succ : currentWalker.getPlan().getSuccessors(op)) {
                        if (!(succ instanceof ConstantExpression)) {
                            return;
                        }
                    }
                    // All successors are constant, calculate the value
                    OperatorPlan expLogicalPlan = new LogicalExpressionPlan();
                    ((BaseOperatorPlan)currentWalker.getPlan()).moveTree(op, (BaseOperatorPlan)expLogicalPlan);
                    PhysicalPlan expPhysicalPlan = new PhysicalPlan();
                    Map<Operator, PhysicalOperator> logToPhyMap = new HashMap<Operator, PhysicalOperator>();
                    PlanWalker childWalker = new ReverseDependencyOrderWalkerWOSeenChk(expLogicalPlan);

                    // Save the old walker and use childWalker as current Walker
                    pushWalker(childWalker);
                    ExpToPhyTranslationVisitor expTranslationVisitor = new 
                            ExpToPhyTranslationVisitor(expLogicalPlan,
                                    childWalker, currentOp, expPhysicalPlan, logToPhyMap);
                    expTranslationVisitor.visit();
                    popWalker();
                    PhysicalOperator root = expPhysicalPlan.getLeaves().get(0);
                    try {
                        UDFContext.getUDFContext().addJobConf(ConfigurationUtil.toConfiguration(pc.getProperties(), true));
                        PigHadoopLogger pigHadoopLogger = PigHadoopLogger.getInstance();
                        PhysicalOperator.setPigLogger(pigHadoopLogger);
                        setDefaultTimeZone();
                        val = root.getNext(root.getResultType()).result;
                        restoreDefaultTimeZone();
                        UDFContext.getUDFContext().addJobConf(null);
                    } catch (ExecException e) {
                        throw new FrontendException(e);
                    }
                    valSet = true;
                } else if (op instanceof UserFuncExpression) {
                    // If solo UDF, calculate UDF
                    UserFuncExpression udf = (UserFuncExpression)op;
                    try {
                        UDFContext.getUDFContext().addJobConf(ConfigurationUtil.toConfiguration(pc.getProperties(), true));
                        setDefaultTimeZone();
                        val = udf.getEvalFunc().exec(null);
                        restoreDefaultTimeZone();
                        UDFContext.getUDFContext().addJobConf(null);
                    } catch (IOException e) {
                        throw new FrontendException(e);
                    }
                    valSet = true;
                }
                if (valSet) {
                    ConstantExpression constantExpr;
                    constantExpr = new ConstantExpression(currentWalker.getPlan(), val);
                    constantExpr.inheritSchema(op);
                    currentWalker.getPlan().replace(op, constantExpr);
                }
            }

            private void setDefaultTimeZone() {
                String dtzStr = pc.getProperties().getProperty("pig.datetime.default.tz");
                if (dtzStr != null && dtzStr.length() > 0) {
                    currentDTZ = DateTimeZone.getDefault();
                    DateTimeZone.setDefault(DateTimeZone.forID(dtzStr));
                }
            }

            private void restoreDefaultTimeZone() {
                if (currentDTZ != null) {
                    DateTimeZone.setDefault(currentDTZ);
                    currentDTZ = null;
                }
            }
        }

        @Override
        public void transform(OperatorPlan matched) throws FrontendException {
            AllExpressionVisitor expressionVisitor = new ConstantCalculatorLogicalPlanVisitor(matched, pc);
            expressionVisitor.visit();
        }

        @Override
        public OperatorPlan reportChanges() {
            return plan;
        }
    }
}