fullstack/hivesterix/hivesterix-translator/src/main/java/edu/uci/ics/hivesterix/logical/plan/visitor/MapJoinVisitor.java

package edu.uci.ics.hivesterix.logical.plan.visitor;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.hadoop.hive.ql.exec.ColumnInfo;
import org.apache.hadoop.hive.ql.exec.MapJoinOperator;
import org.apache.hadoop.hive.ql.exec.Operator;
import org.apache.hadoop.hive.ql.plan.ExprNodeConstantDesc;
import org.apache.hadoop.hive.ql.plan.ExprNodeDesc;
import org.apache.hadoop.hive.ql.plan.ExprNodeGenericFuncDesc;
import org.apache.hadoop.hive.ql.plan.MapJoinDesc;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPAnd;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPEqual;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;

import edu.uci.ics.hivesterix.logical.plan.visitor.base.DefaultVisitor;
import edu.uci.ics.hivesterix.logical.plan.visitor.base.Translator;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalExpression;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.ILogicalOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.base.LogicalVariable;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.InnerJoinOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.ProjectOperator;

@SuppressWarnings("rawtypes")
public class MapJoinVisitor extends DefaultVisitor {

    /**
     * map a join operator (in hive) to its parent operators (in asterix)
     */
    private HashMap<Operator, List<Mutable<ILogicalOperator>>> opMap = new HashMap<Operator, List<Mutable<ILogicalOperator>>>();

    @Override
    public Mutable<ILogicalOperator> visit(MapJoinOperator operator,
            Mutable<ILogicalOperator> AlgebricksParentOperatorRef, Translator t) {
        List<Operator<? extends Serializable>> joinSrc = operator.getParentOperators();
        List<Mutable<ILogicalOperator>> parents = opMap.get(operator);
        if (parents == null) {
            parents = new ArrayList<Mutable<ILogicalOperator>>();
            opMap.put(operator, parents);
        }
        parents.add(AlgebricksParentOperatorRef);
        if (joinSrc.size() != parents.size())
            return null;

        ILogicalOperator currentOperator;
        // make an map join operator
        // TODO: will have trouble for n-way joins
        MapJoinDesc joinDesc = (MapJoinDesc) operator.getConf();

        Map<Byte, List<ExprNodeDesc>> keyMap = joinDesc.getKeys();
        // get the projection expression (already re-written) from each source
        // table
        Map<Byte, List<ExprNodeDesc>> exprMap = joinDesc.getExprs();

        int inputSize = operator.getParentOperators().size();
        // get a list of reduce sink descs (input descs)

        // get the parent operator
        List<Mutable<ILogicalOperator>> parentOps = parents;

        List<String> fieldNames = new ArrayList<String>();
        List<TypeInfo> types = new ArrayList<TypeInfo>();
        for (Operator ts : joinSrc) {
            List<ColumnInfo> columns = ts.getSchema().getSignature();
            for (ColumnInfo col : columns) {
                fieldNames.add(col.getInternalName());
                types.add(col.getType());
            }
        }

        // get number of equality conjunctions in the final join condition
        Set<Entry<Byte, List<ExprNodeDesc>>> keyEntries = keyMap.entrySet();
        Iterator<Entry<Byte, List<ExprNodeDesc>>> entry = keyEntries.iterator();

        int size = 0;
        if (entry.hasNext())
            size = entry.next().getValue().size();

        // make up the join conditon expression
        List<ExprNodeDesc> joinConditionChildren = new ArrayList<ExprNodeDesc>();
        for (int i = 0; i < size; i++) {
            // create a join key pair
            List<ExprNodeDesc> keyPair = new ArrayList<ExprNodeDesc>();
            for (int j = 0; j < inputSize; j++) {
                keyPair.add(keyMap.get(Byte.valueOf((byte) j)).get(i));
            }
            // create a hive equal condition
            ExprNodeDesc equality = new ExprNodeGenericFuncDesc(TypeInfoFactory.booleanTypeInfo,
                    new GenericUDFOPEqual(), keyPair);
            // add the equal condition to the conjunction list
            joinConditionChildren.add(equality);
        }
        // get final conjunction expression
        ExprNodeDesc conjunct = null;

        if (joinConditionChildren.size() > 1)
            conjunct = new ExprNodeGenericFuncDesc(TypeInfoFactory.booleanTypeInfo, new GenericUDFOPAnd(),
                    joinConditionChildren);
        else if (joinConditionChildren.size() == 1)
            conjunct = joinConditionChildren.get(0);
        else {
            // there is no join equality condition, full outer join
            conjunct = new ExprNodeConstantDesc(TypeInfoFactory.booleanTypeInfo, new Boolean(true));
        }
        // get an ILogicalExpression from hive's expression
        Mutable<ILogicalExpression> expression = t.translateScalarFucntion(conjunct);

        ArrayList<LogicalVariable> left = new ArrayList<LogicalVariable>();
        ArrayList<LogicalVariable> right = new ArrayList<LogicalVariable>();

        Set<Entry<Byte, List<ExprNodeDesc>>> kentries = keyMap.entrySet();
        Iterator<Entry<Byte, List<ExprNodeDesc>>> kiterator = kentries.iterator();
        int iteration = 0;
        ILogicalOperator assignOperator = null;
        while (kiterator.hasNext()) {
            List<ExprNodeDesc> outputExprs = kiterator.next().getValue();

            if (iteration == 0)
                assignOperator = t.getAssignOperator(AlgebricksParentOperatorRef, outputExprs, left);
            else
                assignOperator = t.getAssignOperator(AlgebricksParentOperatorRef, outputExprs, right);

            if (assignOperator != null) {
                currentOperator = assignOperator;
                AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);
            }
            iteration++;
        }

        List<Mutable<ILogicalOperator>> inputs = parentOps;

        // get the join operator
        currentOperator = new InnerJoinOperator(expression);

        // set the inputs from asterix join operator
        for (Mutable<ILogicalOperator> input : inputs)
            currentOperator.getInputs().add(input);
        AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);

        // add assign and project operator
        // output variables
        ArrayList<LogicalVariable> variables = new ArrayList<LogicalVariable>();
        Set<Entry<Byte, List<ExprNodeDesc>>> entries = exprMap.entrySet();
        Iterator<Entry<Byte, List<ExprNodeDesc>>> iterator = entries.iterator();
        while (iterator.hasNext()) {
            List<ExprNodeDesc> outputExprs = iterator.next().getValue();
            assignOperator = t.getAssignOperator(AlgebricksParentOperatorRef, outputExprs, variables);

            if (assignOperator != null) {
                currentOperator = assignOperator;
                AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);
            }
        }

        currentOperator = new ProjectOperator(variables);
        currentOperator.getInputs().add(AlgebricksParentOperatorRef);
        t.rewriteOperatorOutputSchema(variables, operator);
        // opMap.clear();
        return new MutableObject<ILogicalOperator>(currentOperator);
    }
}