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

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

import java.util.ArrayList;
import java.util.List;

import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.hadoop.hive.ql.exec.ExtractOperator;
import org.apache.hadoop.hive.ql.exec.Operator;
import org.apache.hadoop.hive.ql.exec.ReduceSinkOperator;
import org.apache.hadoop.hive.ql.plan.ExprNodeDesc;
import org.apache.hadoop.hive.ql.plan.ReduceSinkDesc;

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.common.exceptions.AlgebricksException;
import edu.uci.ics.hyracks.algebricks.common.utils.Pair;
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.OrderOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.OrderOperator.IOrder;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.ProjectOperator;
import edu.uci.ics.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
import edu.uci.ics.hyracks.algebricks.core.algebra.properties.OrderColumn;

public class SortVisitor extends DefaultVisitor {

    @SuppressWarnings("rawtypes")
    @Override
    public Mutable<ILogicalOperator> visit(ReduceSinkOperator operator,
            Mutable<ILogicalOperator> AlgebricksParentOperatorRef, Translator t) throws AlgebricksException {
        ReduceSinkDesc desc = (ReduceSinkDesc) operator.getConf();
        Operator downStream = (Operator) operator.getChildOperators().get(0);
        List<ExprNodeDesc> keys = desc.getKeyCols();
        if (!(downStream instanceof ExtractOperator && desc.getNumReducers() == 1 && keys.size() > 0)) {
            return null;
        }

        List<ExprNodeDesc> schema = new ArrayList<ExprNodeDesc>();
        List<ExprNodeDesc> values = desc.getValueCols();
        List<ExprNodeDesc> partitionCols = desc.getPartitionCols();
        for (ExprNodeDesc key : keys) {
            t.rewriteExpression(key);
        }
        for (ExprNodeDesc value : values) {
            t.rewriteExpression(value);
        }
        for (ExprNodeDesc col : partitionCols) {
            t.rewriteExpression(col);
        }

        // add a order-by operator and limit if any
        List<Pair<IOrder, Mutable<ILogicalExpression>>> pairs = new ArrayList<Pair<IOrder, Mutable<ILogicalExpression>>>();
        char[] orders = desc.getOrder().toCharArray();
        int i = 0;
        for (ExprNodeDesc key : keys) {
            Mutable<ILogicalExpression> expr = t.translateScalarFucntion(key);
            IOrder order = orders[i] == '+' ? OrderOperator.ASC_ORDER : OrderOperator.DESC_ORDER;

            Pair<IOrder, Mutable<ILogicalExpression>> pair = new Pair<IOrder, Mutable<ILogicalExpression>>(order, expr);
            pairs.add(pair);
            i++;
        }

        // get input variables
        ArrayList<LogicalVariable> inputVariables = new ArrayList<LogicalVariable>();
        VariableUtilities.getProducedVariables(AlgebricksParentOperatorRef.getValue(), inputVariables);

        ArrayList<LogicalVariable> keyVariables = new ArrayList<LogicalVariable>();
        ILogicalOperator currentOperator;
        ILogicalOperator assignOp = t.getAssignOperator(AlgebricksParentOperatorRef, keys, keyVariables);
        if (assignOp != null) {
            currentOperator = assignOp;
            AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);
        }

        OrderColumn[] keyColumns = new OrderColumn[keyVariables.size()];

        for (int j = 0; j < keyColumns.length; j++)
            keyColumns[j] = new OrderColumn(keyVariables.get(j), pairs.get(j).first.getKind());

        // handle order operator
        currentOperator = new OrderOperator(pairs);
        currentOperator.getInputs().add(AlgebricksParentOperatorRef);
        AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);

        // project back, remove generated sort-key columns if any
        if (assignOp != null) {
            currentOperator = new ProjectOperator(inputVariables);
            currentOperator.getInputs().add(AlgebricksParentOperatorRef);
            AlgebricksParentOperatorRef = new MutableObject<ILogicalOperator>(currentOperator);
        }

        /**
         * a special rule for hive's order by output schema of reduce sink
         * operator only contains the columns
         */
        for (ExprNodeDesc value : values) {
            schema.add(value);
        }

        ArrayList<LogicalVariable> variables = new ArrayList<LogicalVariable>();
        ILogicalOperator assignOperator = t.getAssignOperator(AlgebricksParentOperatorRef, schema, variables);
        t.rewriteOperatorOutputSchema(variables, operator);

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