src/org/apache/pig/newplan/logical/rules/ColumnPruneHelper.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.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.pig.data.DataType;
import org.apache.pig.impl.logicalLayer.FrontendException;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.OperatorPlan;
import org.apache.pig.newplan.OperatorSubPlan;
import org.apache.pig.newplan.ReverseDependencyOrderWalker;
import org.apache.pig.newplan.logical.expression.LogicalExpression;
import org.apache.pig.newplan.logical.expression.LogicalExpressionPlan;
import org.apache.pig.newplan.logical.expression.ProjectExpression;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOCross;
import org.apache.pig.newplan.logical.relational.LODistinct;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOInnerLoad;
import org.apache.pig.newplan.logical.relational.LOJoin;
import org.apache.pig.newplan.logical.relational.LOLimit;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.relational.LORank;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LOSplit;
import org.apache.pig.newplan.logical.relational.LOSplitOutput;
import org.apache.pig.newplan.logical.relational.LOStore;
import org.apache.pig.newplan.logical.relational.LOStream;
import org.apache.pig.newplan.logical.relational.LOUnion;
import org.apache.pig.newplan.logical.relational.LogicalRelationalNodesVisitor;
import org.apache.pig.newplan.logical.relational.LogicalRelationalOperator;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;
import org.apache.pig.newplan.logical.relational.SchemaNotDefinedException;

/**
 * Helper class used by ColumnMapKeyPrune to figure out what columns can be pruned.
 * It doesn't make any changes to the operator plan
 *
 */
public class ColumnPruneHelper {
    protected static final String INPUTUIDS = "ColumnPrune:InputUids";
    public static final String OUTPUTUIDS = "ColumnPrune:OutputUids";
    protected static final String REQUIREDCOLS = "ColumnPrune:RequiredColumns";

    private OperatorPlan currentPlan;
    private OperatorSubPlan subPlan;

    public ColumnPruneHelper(OperatorPlan currentPlan) {
        this.currentPlan = currentPlan;
    }

    private OperatorSubPlan getSubPlan() throws FrontendException {
        OperatorSubPlan p = null;
        if (currentPlan instanceof OperatorSubPlan) {
            p = new OperatorSubPlan(((OperatorSubPlan)currentPlan).getBasePlan());
        } else {
            p = new OperatorSubPlan(currentPlan);
        }
        Iterator<Operator> iter = currentPlan.getOperators();

        while(iter.hasNext()) {
            Operator op = iter.next();
            if (op instanceof LOForEach) {
                addOperator(op, p);
            }
        }

        return p;
    }

    private void addOperator(Operator op, OperatorSubPlan subplan) throws FrontendException {
        if (op == null) {
            return;
        }

        subplan.add(op);

        List<Operator> ll = currentPlan.getPredecessors(op);
        if (ll == null) {
            return;
        }

        for(Operator pred: ll) {
            addOperator(pred, subplan);
        }
    }


    @SuppressWarnings("unchecked")
    public boolean check() throws FrontendException {
        List<Operator> sources = currentPlan.getSources();
        // if this rule has run before, just return false
        if (sources.size() > 1 && sources.get(0).getAnnotation(INPUTUIDS) != null) {
            clearAnnotation();
            return false;
        }

        // create sub-plan that ends with foreach
        subPlan = getSubPlan();
        if (subPlan.size() == 0) {
            clearAnnotation();
            return false;
        }

        ColumnDependencyVisitor v = new ColumnDependencyVisitor(currentPlan);
        try {
            v.visit();
        }catch(SchemaNotDefinedException e) {
            // if any operator has an unknown schema, just return false
            clearAnnotation();
            return false;
        }

        List<Operator> ll = subPlan.getSources();
        boolean found = false;
        for(Operator op: ll) {
            if (op instanceof LOLoad) {
                Set<Long> uids = (Set<Long>)op.getAnnotation(INPUTUIDS);
                LogicalSchema s = ((LOLoad) op).getSchema();
                Set<Integer> required = getColumns(s, uids);

                if (required.size() < s.size()) {
                    op.annotate(REQUIREDCOLS, required);
                    found = true;
                }
            }
        }

        if (!found)
            clearAnnotation();

        return found;
    }

    private void clearAnnotation() {
        Iterator<Operator> iter = currentPlan.getOperators();
        while (iter.hasNext()) {
            Operator op = iter.next();
            op.removeAnnotation(INPUTUIDS);
            op.removeAnnotation(OUTPUTUIDS);
            op.removeAnnotation(REQUIREDCOLS);
        }
    }

    // get a set of column indexes from a set of uids
    protected Set<Integer> getColumns(LogicalSchema schema, Set<Long> uids) throws FrontendException {
        if (schema == null) {
            throw new SchemaNotDefinedException("Schema is not defined.");
        }

        Set<Integer> cols = new HashSet<Integer>();
        Iterator<Long> iter = uids.iterator();
        while(iter.hasNext()) {
            long uid = iter.next();
            int index = schema.findField(uid);
            if (index == -1) {
                throw new FrontendException("UID " + uid + " is not found in the schema " + schema, 2241);
            }

            cols.add(index);
        }

        return cols;
    }

    public OperatorPlan reportChanges() {
        return subPlan;
    }

    // Visitor to calculate the input and output uids for each operator
    // It doesn't change the plan, only put calculated info as annotations
    // The input and output uids are not necessarily the top level uids of
    // a schema. They may be the uids of lower level fields of complex fields
    // that have their own schema.
    static private class ColumnDependencyVisitor extends LogicalRelationalNodesVisitor {

        public ColumnDependencyVisitor(OperatorPlan plan) throws FrontendException {
            super(plan, new ReverseDependencyOrderWalker(plan));
        }

        @Override
        public void visit(LOLoad load) throws FrontendException {
            Set<Long> output = setOutputUids(load);

            // for load, input uids are same as output uids
            load.annotate(INPUTUIDS, output);
        }

        @Override
        public void visit(LOFilter filter) throws FrontendException {
            Set<Long> output = setOutputUids(filter);

            // the input uids contains all the output uids and
            // projections in filter conditions
            Set<Long> input = new HashSet<Long>(output);

            LogicalExpressionPlan exp = filter.getFilterPlan();
            collectUids(filter, exp, input);

            filter.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOStore store) throws FrontendException {
            Set<Long> output = setOutputUids(store);

            if (output.isEmpty()) {
                // to deal with load-store-load-store case
                LogicalSchema s = store.getSchema();
                if (s == null) {
                    throw new SchemaNotDefinedException("Schema for " + store.getName() + " is not defined.");
                }

                for(int i=0; i<s.size(); i++) {
                    output.add(s.getField(i).uid);
                }
            }

            // for store, input uids are same as output uids
            store.annotate(INPUTUIDS, output);
        }

        @Override
        public void visit(LOJoin join) throws FrontendException {
            Set<Long> output = setOutputUids(join);

            // the input uids contains all the output uids and
            // projections in join expressions
            Set<Long> input = new HashSet<Long>(output);

            Collection<LogicalExpressionPlan> exps = join.getExpressionPlanValues();
            Iterator<LogicalExpressionPlan> iter = exps.iterator();
            while(iter.hasNext()) {
                LogicalExpressionPlan exp = iter.next();
                collectUids(join, exp, input);
            }

            join.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOCogroup cg) throws FrontendException {
            Set<Long> output = setOutputUids(cg);

            // the input uids contains all the output uids and
            // projections in join expressions
            Set<Long> input = new HashSet<Long>();

            // Add all the uids required for doing cogroup. As in all the
            // keys on which the cogroup is done.
            for( LogicalExpressionPlan plan : cg.getExpressionPlans().values() ) {
                collectUids(cg, plan, input);
            }

            // Now check for the case where the output uid is a generated one
            // If that is the case we need to add the uids which generated it in
            // the input
            long firstUid=-1;
            Map<Integer,Long> generatedInputUids = cg.getGeneratedInputUids();
            for( Map.Entry<Integer, Long> entry : generatedInputUids.entrySet() ) {
                Long uid = entry.getValue();
                LogicalRelationalOperator pred =
                    (LogicalRelationalOperator) cg.getPlan().getPredecessors(cg).get(entry.getKey());
                if( output.contains(uid) ) {
                    // Hence we need to all the full schema of the bag
                    input.addAll( getAllUids( pred.getSchema() ) );
                }
                if (pred.getSchema()!=null)
                    firstUid = pred.getSchema().getField(0).uid;
            }

            if (input.isEmpty() && firstUid!=-1) {
                input.add(firstUid);
            }

            cg.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOLimit limit) throws FrontendException {
            Set<Long> output = setOutputUids(limit);

            // the input uids contains all the output uids and
            // projections in limit expression
            Set<Long> input = new HashSet<Long>(output);

            LogicalExpressionPlan exp = limit.getLimitPlan();
            if (exp != null)
                collectUids(limit, exp, input);

            limit.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOStream stream) throws FrontendException {
            // output is not used, setOutputUids is used to check if it has output schema
            Set<Long> output = setOutputUids(stream);

            // Every field is required
            LogicalRelationalOperator pred = (LogicalRelationalOperator)plan.getPredecessors(stream).get(0);

            Set<Long> input = getAllUids(pred.getSchema());

            stream.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LODistinct distinct) throws FrontendException {
            setOutputUids(distinct);
            
            Set<Long> input = new HashSet<Long>();

            // Every field is required
            LogicalSchema s = distinct.getSchema();
            if (s == null) {
                throw new SchemaNotDefinedException("Schema for " + distinct.getName() + " is not defined.");
            }

            for(int i=0; i<s.size(); i++) {
                input.add(s.getField(i).uid);
            }
            distinct.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOCross cross) throws FrontendException {
            Set<Long> output = setOutputUids(cross);
            // Since we do not change the topology of the plan, we keep
            // at least one input for each predecessor.
            List<Operator> preds = plan.getPredecessors(cross);
            for (Operator pred : preds) {
                LogicalSchema schema = ((LogicalRelationalOperator)pred).getSchema();
                Set<Long> uids = getAllUids(schema);
                boolean allPruned = true;
                for (Long uid : uids) {
                    if (output.contains(uid))
                        allPruned = false;
                }
                if (allPruned)
                    output.add(schema.getField(0).uid);
            }
            cross.annotate(INPUTUIDS, output);
        }

        @Override
        public void visit(LOUnion union) throws FrontendException {
            Set<Long> output = setOutputUids(union);
            Set<Long> input = new HashSet<Long>();
            for (long uid : output) {
                input.addAll(union.getInputUids(uid));
            }
            union.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOSplit split) throws FrontendException {
            Set<Long> output = setOutputUids(split);
            split.annotate(INPUTUIDS, output);
        }

        @Override
        public void visit(LOSplitOutput splitOutput) throws FrontendException {
            Set<Long> output = setOutputUids(splitOutput);

            // the input uids contains all the output uids and
            // projections in splitOutput conditions
            Set<Long> input = new HashSet<Long>();

            for (long uid : output) {
                input.add(splitOutput.getInputUids(uid));
            }

            LogicalExpressionPlan exp = splitOutput.getFilterPlan();
            collectUids(splitOutput, exp, input);

            splitOutput.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOSort sort) throws FrontendException {
            Set<Long> output = setOutputUids(sort);

            Set<Long> input = new HashSet<Long>(output);

            for (LogicalExpressionPlan exp : sort.getSortColPlans()) {
                collectUids(sort, exp, input);
            }

            sort.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LORank rank) throws FrontendException {
            Set<Long> output = setOutputUids(rank);

            Set<Long> input = new HashSet<Long>(output);

            for (LogicalExpressionPlan exp : rank.getRankColPlans()) {
                collectUids(rank, exp, input);
            }

            rank.annotate(INPUTUIDS, input);
        }

        /*
         * This function returns all uids present in the given schema
         */
        private Set<Long> getAllUids( LogicalSchema schema ) {
            Set<Long> uids = new HashSet<Long>();

            if( schema == null ) {
                return uids;
            }

            for( LogicalFieldSchema field : schema.getFields() ) {
                if( ( field.type == DataType.TUPLE || field.type == DataType.BAG )
                        && field.schema != null ) {
                   uids.addAll( getAllUids( field.schema ) );
                }
                uids.add( field.uid );
            }
            return uids;
        }

        @SuppressWarnings("unchecked")
        @Override
        public void visit(LOForEach foreach) throws FrontendException {
            Set<Long> output = setOutputUids(foreach);

            LOGenerate gen = OptimizerUtils.findGenerate(foreach);
            gen.annotate(OUTPUTUIDS, output);

            visit(gen);

            Set<Long> input = (Set<Long>)gen.getAnnotation(INPUTUIDS);

            // Make sure at least one column will retain
            if (input.isEmpty()) {
                LogicalRelationalOperator pred = (LogicalRelationalOperator)plan.getPredecessors(foreach).get(0);
                if (pred.getSchema()!=null)
                    input.add(pred.getSchema().getField(0).uid);
            }
            foreach.annotate(INPUTUIDS, input);
        }

        @Override
        @SuppressWarnings("unchecked")
        public void visit(LOGenerate gen) throws FrontendException {
             Set<Long> output = (Set<Long>)gen.getAnnotation(OUTPUTUIDS);

             Set<Long> input = new HashSet<Long>();

             List<LogicalExpressionPlan> ll = gen.getOutputPlans();

             Iterator<Long> iter = output.iterator();
             while(iter.hasNext()) {
                 long uid = iter.next();
                 for(int i=0; i<ll.size(); i++) {
                     LogicalExpressionPlan exp = ll.get(i);
                     boolean found = false;
                     LogicalSchema planSchema = gen.getOutputPlanSchemas().get(i);
                     for (LogicalFieldSchema fs : planSchema.getFields()) {
                         if (fs.uid == uid) {
                             found = true;
                             break;
                         }
                     }

                     if (found) {
                         List<Operator> srcs = exp.getSinks();
                         for (Operator src : srcs) {
                             if (src instanceof ProjectExpression) {
                                 List<LOInnerLoad> innerLoads = LOForEach.findReacheableInnerLoadFromBoundaryProject((ProjectExpression)src).first;
                                 for (LOInnerLoad innerLoad : innerLoads) {
                                     ProjectExpression prj = innerLoad.getProjection();
                                     if (prj.isProjectStar()) {
                                         if (prj.findReferent().getSchema()!=null) {
                                             for (LogicalSchema.LogicalFieldSchema fs : prj.findReferent().getSchema().getFields()) {
                                                 input.add(fs.uid);
                                             }
                                         }
                                     }
                                     else {
                                         if (prj.findReferent().getSchema()!=null) {
                                             LogicalSchema.LogicalFieldSchema fs = prj.findReferent().getSchema().getField(prj.getColNum()); 
                                             input.add(fs.uid);
                                         }
                                     }
                                 }
                             }
                         }
                     }
                 }
             }

             // for the flatten bag, we need to make sure at least one field is in the input
             for(int i=0; i<ll.size(); i++) {
                 if (!gen.getFlattenFlags()[i]) {
                     continue;
                 }
                 LogicalExpressionPlan exp = ll.get(i);
                 LogicalExpression sink = (LogicalExpression)exp.getSources().get(0);
                 if (sink.getFieldSchema().type!=DataType.TUPLE && sink.getFieldSchema().type!=DataType.BAG)
                     continue;
                 List<Operator> srcs = exp.getSinks();
                 for (Operator src : srcs) {
                     if (!(src instanceof ProjectExpression))
                         continue;
                     List<LOInnerLoad> innerLoads = LOForEach.findReacheableInnerLoadFromBoundaryProject((ProjectExpression)src).first;
                     for (LOInnerLoad innerLoad : innerLoads) {
                         ProjectExpression prj = innerLoad.getProjection();
                         if (prj.isProjectStar()) {
                             if (prj.findReferent().getSchema()!=null) {
                                 for (LogicalSchema.LogicalFieldSchema fs : prj.findReferent().getSchema().getFields()) {
                                     input.add(fs.uid);
                                 }
                             }
                         }
                         else {
                             if (prj.findReferent().getSchema()!=null) {
                                 LogicalSchema.LogicalFieldSchema fs = prj.findReferent().getSchema().getField(prj.getColNum());
                                 input.add(fs.uid);
                             }
                         }
                     }
                 }
             }
             gen.annotate(INPUTUIDS, input);
        }

        @Override
        public void visit(LOInnerLoad load) throws FrontendException {
            Set<Long> output = setOutputUids(load);
            load.annotate(INPUTUIDS, output);
        }

        private void collectUids(LogicalRelationalOperator currentOp, LogicalExpressionPlan exp, Set<Long> uids) throws FrontendException {
            List<Operator> ll = exp.getSinks();
            for(Operator op: ll) {
                if (op instanceof ProjectExpression) {
                    if (!((ProjectExpression)op).isRangeOrStarProject()) {
                        long uid = ((ProjectExpression)op).getFieldSchema().uid;
                        uids.add(uid);
                    } else {
                        LogicalRelationalOperator ref = ((ProjectExpression)op).findReferent();
                        LogicalSchema s = ref.getSchema();
                        if (s == null) {
                            throw new SchemaNotDefinedException("Schema not defined for " + ref.getAlias());
                        }
                        for(LogicalFieldSchema f: s.getFields()) {
                            uids.add(f.uid);
                        }
                    }
                }
            }
        }

        @SuppressWarnings("unchecked")
        // Get output uid from output schema. If output schema does not exist,
        // throw exception
        private Set<Long> setOutputUids(LogicalRelationalOperator op) throws FrontendException {

            List<Operator> ll = plan.getSuccessors(op);
            Set<Long> uids = new HashSet<Long>();

            LogicalSchema s = op.getSchema();
            if (s == null) {
                throw new SchemaNotDefinedException("Schema for " + op.getName() + " is not defined.");
            }

            if (ll != null) {
                // if this is not sink, the output uids are union of input uids of its successors
                for(Operator succ: ll) {
                    Set<Long> inputUids = (Set<Long>)succ.getAnnotation(INPUTUIDS);
                    if (inputUids != null) {
                        Iterator<Long> iter = inputUids.iterator();
                        while(iter.hasNext()) {
                            long uid = iter.next();

                            if (s.findField(uid) != -1) {
                                uids.add(uid);
                            }
                        }
                    }
                }
            } else {
                // if  it's leaf, set to its schema
                for(int i=0; i<s.size(); i++) {
                    uids.add(s.getField(i).uid);
                }
            }

            op.annotate(OUTPUTUIDS, uids);
            return uids;
        }
    }
}