src/org/apache/pig/impl/util/UDFContextSeparator.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.impl.util;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;

import org.apache.hadoop.conf.Configuration;
import org.apache.pig.Algebraic;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POCast;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserFunc;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhyPlanVisitor;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLoad;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POStore;
import org.apache.pig.impl.plan.DepthFirstWalker;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.UDFContext.UDFContextKey;

public class UDFContextSeparator extends PhyPlanVisitor {

    public static enum UDFType {
        LOADFUNC,
        STOREFUNC,
        USERFUNC,
    };

    private String planOpKey;
    private DepthFirstWalker<PhysicalOperator, PhysicalPlan> dfw;
    private Map<String, Map<Enum<UDFType>, List<UDFContext.UDFContextKey>>> udfContextsPerPlan;
    private UDFContext udfContext;
    private Set<UDFContext.UDFContextKey> allKeys;
    private Set<UDFContext.UDFContextKey> knownKeys;
    private Set<UDFContext.UDFContextKey> unKnownKeys;
    private Set<UDFContext.UDFContextKey> algebraicUDFKeys;

    public UDFContextSeparator(){
        super(null, null);
        dfw = new DepthFirstWalker<PhysicalOperator, PhysicalPlan>(null);
        udfContext = UDFContext.getUDFContext();
        allKeys = udfContext.getUdfConfs().keySet();
        knownKeys = new HashSet<UDFContext.UDFContextKey>();
        algebraicUDFKeys = new HashSet<UDFContext.UDFContextKey>();
        udfContextsPerPlan = new HashMap<String, Map<Enum<UDFType>, List<UDFContext.UDFContextKey>>>();
    }

    public Set<UDFContext.UDFContextKey> getUnKnownKeys() {
        if (unKnownKeys == null) {
            unKnownKeys = new HashSet<UDFContext.UDFContextKey>(allKeys);
            unKnownKeys.removeAll(knownKeys);
            for (Entry<UDFContextKey, Properties> entry : udfContext.getUdfConfs().entrySet()) {
                if (entry.getValue().isEmpty()) {
                    // Remove empty values
                    unKnownKeys.remove(entry.getKey());
                }
            }
        }
        return unKnownKeys;
    }

    public void setPlan(PhysicalPlan plan, String planOpKey){
        mPlan = plan;
        dfw.setPlan(plan);
        mCurrentWalker = dfw;
        this.planOpKey = planOpKey;
        this.udfContextsPerPlan.put(planOpKey, new HashMap<Enum<UDFType>, List<UDFContextKey>>());
    }

    @Override
    public void visitUserFunc(POUserFunc userFunc) throws VisitorException {
        if (userFunc.getFunc() instanceof Algebraic) {
            for (UDFContext.UDFContextKey key : allKeys) {
                if (key.getClassName().equals(userFunc.getFunc().getClass().getName())) {
                    // If Algebraic handle differently. To be on the safer side
                    // as user might be just accessing properties by base class name
                    // instead of by Initial, Intermediate and Final classes
                    algebraicUDFKeys.add(key);
                    knownKeys.add(key);
                }
            }
        } else {
            findAndAddKeys(userFunc.getFunc().getClass().getName(),
                    userFunc.getSignature(), UDFType.USERFUNC);
        }
    }

    @Override
    public void visitLoad(POLoad ld) throws VisitorException {
        findAndAddKeys(ld.getLoadFunc().getClass().getName(),
                ld.getSignature(), UDFType.LOADFUNC);
    }


    @Override
    public void visitStore(POStore st) throws VisitorException {
        findAndAddKeys(st.getStoreFunc().getClass().getName(),
                st.getSignature(), UDFType.STOREFUNC);
    }

    @Override
    public void visitCast(POCast op) {
        if (op.getFuncSpec() != null) {
            findAndAddKeys(op.getFuncSpec().getClass().getName(),
                    null, UDFType.USERFUNC);
        }
    }

    private void findAndAddKeys(String keyClassName, String signature, UDFType udfType) {
        for (UDFContext.UDFContextKey key : allKeys) {
            if (key.getClassName().equals(keyClassName)
                    && (key.getArgs() == null
                    || signature == null
                    || Arrays.asList(key.getArgs()).contains(signature))) {
                Map<Enum<UDFType>, List<UDFContextKey>> udfKeysByType = udfContextsPerPlan
                        .get(planOpKey);
                List<UDFContextKey> keyList = udfContextsPerPlan.get(planOpKey)
                        .get(udfType);
                if (keyList == null) {
                    keyList = new ArrayList<UDFContext.UDFContextKey>();
                    udfKeysByType.put(udfType, keyList);
                }
                keyList.add(key);
                knownKeys.add(key);
            }
        }
    }

    public void serializeUDFContext(Configuration conf, String planOpKey,
            UDFType... udfTypes) throws IOException {
        Map<UDFContextKey, Properties> udfConfs = udfContext.getUdfConfs();
        HashMap<UDFContextKey, Properties> udfConfsToSerialize = new HashMap<UDFContextKey, Properties>();
        Map<Enum<UDFType>, List<UDFContextKey>> udfKeysByType = udfContextsPerPlan.get(planOpKey);
        if (udfKeysByType != null) {
            for (UDFType udfType : udfTypes) {
                List<UDFContextKey> keyList = udfContextsPerPlan.get(planOpKey).get(udfType);
                if (keyList != null) {
                    for (UDFContextKey key : keyList) {
                        udfConfsToSerialize.put(key, udfConfs.get(key));
                    }
                }
                if (udfType.equals(UDFType.USERFUNC)) {
                    for (UDFContextKey key : algebraicUDFKeys) {
                        udfConfsToSerialize.put(key, udfConfs.get(key));
                    }
                }
            }
        }
        serialize(conf, udfConfsToSerialize);
    }

    public void serializeUDFContext(Configuration conf, String planOpKey,
            POStore store) throws IOException {
        Map<UDFContextKey, Properties> udfConfs = udfContext.getUdfConfs();
        HashMap<UDFContextKey, Properties> udfConfsToSerialize = new HashMap<UDFContextKey, Properties>();
        // Find keys specific to just this StoreFunc
        Map<Enum<UDFType>, List<UDFContextKey>> udfKeysByType = udfContextsPerPlan.get(planOpKey);
        if (udfKeysByType != null) {
            List<UDFContextKey> keyList = udfContextsPerPlan.get(planOpKey).get(
                    UDFType.STOREFUNC);
            if (keyList != null) {
                String keyClassName = store.getStoreFunc().getClass().getName();
                String signature = store.getSignature();
                for (UDFContextKey key : keyList) {
                    if (key.getClassName().equals(keyClassName)
                            && (key.getArgs() == null
                            || Arrays.asList(key.getArgs()).contains(signature))) {
                        udfConfsToSerialize.put(key, udfConfs.get(key));
                    }
                }
            }
        }
        serialize(conf, udfConfsToSerialize);
    }

    private void serialize(Configuration conf,
            HashMap<UDFContextKey, Properties> udfConfsToSerialize)
            throws IOException {
        HashMap<UDFContextKey, Properties> udfConfs = udfContext.getUdfConfs();

        // Save empty values from being serialized unnecessarily
        Iterator<Entry<UDFContextKey, Properties>> iter = udfConfsToSerialize.entrySet().iterator();
        while (iter.hasNext()) {
            Entry<UDFContextKey, Properties> entry = iter.next();
            if (entry.getValue().isEmpty()) {
                iter.remove();
            }
        }

        // Add unknown ones for serialization
        for (UDFContextKey key : getUnKnownKeys()) {
            udfConfsToSerialize.put(key, udfConfs.get(key));
        }
        conf.set(UDFContext.UDF_CONTEXT, ObjectSerializer.serialize(udfConfsToSerialize));
        conf.set(UDFContext.CLIENT_SYS_PROPS, ObjectSerializer.serialize(udfContext.getClientSystemProps()));
    }

}