spork-master/src/org/apache/pig/newplan/logical/rules/MapKeysPruneHelper.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.HashMap;
import java.util.HashSet;
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.DependencyOrderWalker;
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.LogicalExpressionVisitor;
import org.apache.pig.newplan.logical.expression.MapLookupExpression;
import org.apache.pig.newplan.logical.expression.UserFuncExpression;
import org.apache.pig.newplan.logical.optimizer.AllExpressionVisitor;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LOGenerate;
import org.apache.pig.newplan.logical.relational.LOJoin;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LOSplitOutput;
import org.apache.pig.newplan.logical.relational.LOStore;
import org.apache.pig.newplan.logical.relational.LOUnion;
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;

/**
 * This filter Marks every Load Operator which has a Map 
 * with MAP_MARKER_ANNOTATION. The annotation value is 
 * <code>Map<Integer,Set<String>><code> where Integer is the column number 
 * of the field and Set is the set of Keys in this field ( field is a map field only ).
 * 
 * It does this for only the top level schema in load. 
 * 
 * Algorithm:
 *  Traverse the Plan in ReverseDependency order ( ie. Sink to Source )
 *      For LogicalRelationalOperators having MapLookupExpression in their 
 *          expressionPlan collect uid and keys related to it. This is
 *          retained in the visitor
 *      For ForEach having nested LogicalPlan use the same visitor hence
 *          there is no distinction required
 *      At Sources find all the uids provided by this source and annotate this 
 *      LogicalRelationalOperator ( load ) with <code>Map<Integer,Set<String>></code>
 *      containing only the column numbers that this LogicalRelationalOperator generates
 *      
 * NOTE: This is a simple Map Pruner. If a map key is mentioned in the script
 *      then this pruner assumes you need the key. This pruner is not as optimized
 *      as column pruner ( which removes a column if it is mentioned but never used )
 *
 */
public class MapKeysPruneHelper {

    public static final String REQUIRED_MAPKEYS = "MapPruner:RequiredKeys";
    
    private OperatorPlan currentPlan;
    private OperatorSubPlan subplan;
    
    public MapKeysPruneHelper(OperatorPlan currentPlan) {
        this.currentPlan = currentPlan;
        
        if (currentPlan instanceof OperatorSubPlan) {
            subplan = new OperatorSubPlan(((OperatorSubPlan)currentPlan).getBasePlan());
        } else {
            subplan = new OperatorSubPlan(currentPlan);
        }
    }
  

    @SuppressWarnings("unchecked")
    public boolean check() throws FrontendException {       
        
        // First check if we have a load with a map in it or not
        List<Operator> sources = currentPlan.getSources();
        
        for( Operator source : sources ) {
            LogicalSchema schema = ((LogicalRelationalOperator)source).getSchema();
            // If any of the loads has a null schema we dont know the ramifications here
            // so we skip this optimization
            if( schema == null ) {
                return false;
            }
        }
                    
        // Now we check what keys are needed
        MapMarker marker = new MapMarker(currentPlan);
        marker.visit();
        
        // If the uid is the input uid of LOStore, LOCogroup, LOUnion, UserFunc, that means
        // the entire map may be used. For simplicity, we do not prune any map key in this case
        Set<Long> fullMapUids = new HashSet<Long>();
        FullMapCollector collector = new FullMapCollector(currentPlan, fullMapUids);
        collector.visit();
        
        // If we have found specific keys which are needed then we return true;
        // Else if we dont have any specific keys we return false
        boolean hasAnnotation = false;
        for( Operator source : sources ) {
            Map<Integer,Set<String>> annotationValue = 
                (Map<Integer, Set<String>>) ((LogicalRelationalOperator)source).getAnnotation(REQUIRED_MAPKEYS);
            
            // Now for all full maps found in sinks we cannot prune them at source
            if( ! fullMapUids.isEmpty() && annotationValue != null && 
                    !annotationValue.isEmpty() ) {
                Integer[] annotationKeyArray = annotationValue.keySet().toArray( new Integer[0] );
                LogicalSchema sourceSchema = ((LogicalRelationalOperator)source).getSchema();
                for( Integer col : annotationKeyArray ) {                	
                    if( fullMapUids.contains(sourceSchema.getField(col).uid)) {
                        annotationValue.remove( col );
                    }
                }
            }
            
            if ( annotationValue != null && annotationValue.isEmpty()) {
                ((LogicalRelationalOperator)source).removeAnnotation(REQUIRED_MAPKEYS);
                annotationValue = null;
            }
            
            // Can we still prune any keys
            if( annotationValue != null ) {
                hasAnnotation = true;
                subplan.add(source);
            }
        }
        
        // If all the sinks dont have any schema, we cant to any optimization
        return hasAnnotation;
    }
    
    /**
     * This function checks if the schema has a map.
     * We dont check for a nested structure.
     * @param schema Schema to be checked
     * @return true if it has a map, else false
     * @throws NullPointerException incase Schema is null
     */
    private boolean hasMap(LogicalSchema schema ) {
        for( LogicalFieldSchema field : schema.getFields() ) {
            if( field.type == DataType.MAP ) {
                return true;
            }
        }
        return false;
    }
    
    /**
     * This function returns a set of Uids corresponding to
     * map datatype in the first level of this schema
     * @param schema Schema having fields
     * @return
     */
    private static Set<Long> getMapUids(LogicalSchema schema ) {
        Set<Long> uids = new HashSet<Long>();
        if( schema != null ) {
            for( LogicalFieldSchema field : schema.getFields() ) {
                uids.add( field.uid );
            }
        }
        return uids;
    }

    public OperatorPlan reportChanges() {
        return subplan;
    }

    /**
     * This class collects all the information required to create
     * the list of keys required for a map
     */
    static public class MapMarker extends AllExpressionVisitor {
        
        Map<Long,Set<String>> inputUids = null;

        protected MapMarker(OperatorPlan plan) throws FrontendException {
            super(plan, new ReverseDependencyOrderWalker(plan));
            inputUids = new HashMap<Long,Set<String>>();
        }
        
        @Override
        public void visit(LOLoad load) throws FrontendException {
            if( load.getSchema() != null ) {
                Map<Integer,Set<String>> annotation = new HashMap<Integer,Set<String>>();
                for( int i=0; i<load.getSchema().size(); i++) {
                    LogicalFieldSchema field = load.getSchema().getField(i);
                    if( inputUids.containsKey( field.uid ) ) {
                        annotation.put(i, inputUids.get( field.uid ) );
                    }
                }
                load.annotate(REQUIRED_MAPKEYS, annotation);
            }
        }

        @Override
        public void visit(LOFilter filter) throws FrontendException {
            currentOp = filter;
            MapExprMarker v = (MapExprMarker) getVisitor(filter.getFilterPlan());
            v.visit();
            mergeUidKeys( v.inputUids );
        }
        
        @Override
        public void visit(LOJoin join) throws FrontendException {
            currentOp = join;
            Collection<LogicalExpressionPlan> c = join.getExpressionPlanValues();
            for (LogicalExpressionPlan plan : c) {
                MapExprMarker v = (MapExprMarker) getVisitor(plan);
                v.visit();
                mergeUidKeys( v.inputUids );
            }
        }
        
        @Override
        public void visit(LOGenerate gen) throws FrontendException {
            currentOp = gen;
            Collection<LogicalExpressionPlan> plans = gen.getOutputPlans();
            for( LogicalExpressionPlan plan : plans ) {
                MapExprMarker v = (MapExprMarker) getVisitor(plan);
                v.visit();
                mergeUidKeys( v.inputUids );
            }
        }
        
        @Override
        public void visit(LOSort sort) throws FrontendException {
            currentOp = sort;
            Collection<LogicalExpressionPlan> c = sort.getSortColPlans();
            for (LogicalExpressionPlan plan : c) {
                MapExprMarker v = (MapExprMarker) getVisitor(plan);
                v.visit();
                mergeUidKeys( v.inputUids );
            }
        }
        
        
        @Override
        public void visit(LOSplitOutput splitOutput) throws FrontendException {
            currentOp = splitOutput;
            MapExprMarker v = (MapExprMarker) getVisitor(splitOutput.getFilterPlan());
            v.visit();
            mergeUidKeys( v.inputUids );
            if (splitOutput.getSchema()!=null) {
                for (LogicalFieldSchema fs : splitOutput.getSchema().getFields()) {
                    long inputUid = splitOutput.getInputUids(fs.uid);
                    if( inputUid!=-1) {
                        Set<String> mapKeySet = inputUids.get(fs.uid);
                        if (mapKeySet!=null) {
                            if (inputUids.containsKey(inputUid))
                                inputUids.get(inputUid).addAll(mapKeySet);
                            else
                                inputUids.put(inputUid, mapKeySet);
                        }
                    }
                }
            }
        }
        
        private void mergeUidKeys( Map<Long, Set<String> > inputMap ) {
            for( Map.Entry<Long, Set<String>> entry : inputMap.entrySet() ) {
                if( inputUids.containsKey(entry.getKey()) ) {
                    Set<String> mapKeySet = inputUids.get(entry.getKey());
                    mapKeySet.addAll(entry.getValue());
                } else {
                    inputUids.put(entry.getKey(), inputMap.get(entry.getKey()));
                }
            }
        }

        @Override
        protected LogicalExpressionVisitor getVisitor(LogicalExpressionPlan expr) throws FrontendException {
            return new MapExprMarker(expr );
        }
        
        static class MapExprMarker extends LogicalExpressionVisitor {

            Map<Long,Set<String>> inputUids = null;
            
            protected MapExprMarker(OperatorPlan p) throws FrontendException {
                super(p, new DependencyOrderWalker(p));
                inputUids = new HashMap<Long,Set<String>>();
            }

            @Override
            public void visit(MapLookupExpression op) throws FrontendException {
                Long uid = op.getMap().getFieldSchema().uid;
                String key = op.getLookupKey();
                
                HashSet<String> mapKeySet = null;
                if( inputUids.containsKey(uid) ) {
                    mapKeySet = (HashSet<String>) inputUids.get(uid);                                        
                } else {
                    mapKeySet = new HashSet<String>();
                    inputUids.put(uid, mapKeySet);
                }
                mapKeySet.add(key);
            }
        }
    }
    
    static public class FullMapCollector extends AllExpressionVisitor {
        Set<Long> fullMapUids = new HashSet<Long>();

        protected FullMapCollector(OperatorPlan plan, Set<Long> fullMapUids) throws FrontendException {
            super(plan, new ReverseDependencyOrderWalker(plan));
            this.fullMapUids = fullMapUids;
        }
        
        @Override
        public void visit(LOStore store) throws FrontendException {
            super.visit(store);
            Set<Long> uids = getMapUids(store.getSchema());
            fullMapUids.addAll(uids);
        }
        
        @SuppressWarnings("unchecked")
        @Override
        public void visit(LOUnion union) throws FrontendException {
            super.visit(union);
            List<Operator> preds = plan.getPredecessors(union);
            if (preds!=null) {
                for (Operator pred : preds) {
                    LogicalSchema schema = ((LogicalRelationalOperator)pred).getSchema();
                    Set<Long> uids = getMapUids(schema);
                    fullMapUids.addAll(uids);
                }
            }
        }
        
        @SuppressWarnings("unchecked")
        @Override
        public void visit(LOCogroup cogroup) throws FrontendException {
            super.visit(cogroup);
            List<Operator> preds = plan.getPredecessors(cogroup);
            if (preds!=null) {
                for (Operator pred : preds) {
                    LogicalSchema schema = ((LogicalRelationalOperator)pred).getSchema();
                    Set<Long> uids = getMapUids(schema);
                    fullMapUids.addAll(uids);
                }
            }
        }
        
        @Override
        public void visit(LOSplitOutput splitOutput) throws FrontendException {
            super.visit(splitOutput);
            if (splitOutput.getSchema()!=null) {
                for (LogicalFieldSchema fs : splitOutput.getSchema().getFields()) {
                    if (fullMapUids.contains(fs.uid) && splitOutput.getInputUids(fs.uid)!=-1)
                        fullMapUids.add(splitOutput.getInputUids(fs.uid));
                }
            }
        }

        @Override
        protected LogicalExpressionVisitor getVisitor(LogicalExpressionPlan expr)
                throws FrontendException {
            return new FullMapExpCollector(expr, fullMapUids);
        }
        
        static class FullMapExpCollector extends LogicalExpressionVisitor {
            Set<Long> fullMapUids = new HashSet<Long>();
            protected FullMapExpCollector(OperatorPlan plan, Set<Long> fullMapUids)
                    throws FrontendException {
                super(plan, new DependencyOrderWalker(plan));
                this.fullMapUids = fullMapUids;
            }
            
            @Override
            public void visit(UserFuncExpression userFunc) throws FrontendException {
                List<Operator> succs = userFunc.getPlan().getSuccessors(userFunc);
                if (succs==null) return;
                LogicalExpression succ = (LogicalExpression)succs.get(0);
                if (succ.getFieldSchema()!=null && succ.getFieldSchema().type==DataType.MAP)
                    fullMapUids.add(succ.getFieldSchema().uid);
            }
        }
    }
}
