fullstack/hivesterix/hivesterix-runtime/src/main/java/edu/uci/ics/hivesterix/runtime/evaluator/AggregationFunctionEvaluator.java

package edu.uci.ics.hivesterix.runtime.evaluator;

import java.io.DataOutput;
import java.io.IOException;
import java.util.List;

import org.apache.hadoop.hive.ql.exec.ExprNodeEvaluator;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.plan.ExprNodeDesc;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFCount;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator.AggregationBuffer;
import org.apache.hadoop.hive.serde2.SerDe;
import org.apache.hadoop.hive.serde2.SerDeException;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.io.BytesWritable;

import edu.uci.ics.hivesterix.serde.lazy.LazyObject;
import edu.uci.ics.hyracks.algebricks.common.exceptions.AlgebricksException;
import edu.uci.ics.hyracks.algebricks.runtime.base.ICopyAggregateFunction;
import edu.uci.ics.hyracks.dataflow.common.data.accessors.IFrameTupleReference;

public class AggregationFunctionEvaluator implements ICopyAggregateFunction {

    /**
     * the mode of aggregation function
     */
    private GenericUDAFEvaluator.Mode mode;

    /**
     * an array of evaluators
     */
    private ExprNodeEvaluator[] evaluators;

    /**
     * udaf evaluator partial
     */
    private GenericUDAFEvaluator udafPartial;

    /**
     * udaf evaluator complete
     */
    private GenericUDAFEvaluator udafComplete;

    /**
     * cached parameter objects
     */
    private Object[] cachedParameters;

    /**
     * cached row objects
     */
    private LazyObject<? extends ObjectInspector> cachedRowObject;

    /**
     * the output channel
     */
    private DataOutput out;

    /**
     * aggregation buffer
     */
    private AggregationBuffer aggBuffer;

    /**
     * we only use lazy serde to do serialization
     */
    private SerDe lazySer;

    /**
     * the output object inspector for this aggregation function
     */
    private ObjectInspector outputInspector;

    /**
     * the output object inspector for this aggregation function
     */
    private ObjectInspector outputInspectorPartial;

    /**
     * parameter inspectors
     */
    private ObjectInspector[] parameterInspectors;

    /**
     * output make sure the aggregation functio has least object creation
     * 
     * @param desc
     * @param oi
     * @param output
     */
    public AggregationFunctionEvaluator(List<ExprNodeDesc> inputs, List<TypeInfo> inputTypes, String genericUDAFName,
            GenericUDAFEvaluator.Mode aggMode, boolean distinct, ObjectInspector oi, DataOutput output,
            ExprNodeEvaluator[] evals, ObjectInspector[] pInspectors, Object[] parameterCache, SerDe serde,
            LazyObject<? extends ObjectInspector> row, GenericUDAFEvaluator udafunctionPartial,
            GenericUDAFEvaluator udafunctionComplete, ObjectInspector outputOi, ObjectInspector outputOiPartial) {
        // shared object across threads
        this.out = output;
        this.mode = aggMode;
        this.parameterInspectors = pInspectors;

        // thread local objects
        this.evaluators = evals;
        this.cachedParameters = parameterCache;
        this.cachedRowObject = row;
        this.lazySer = serde;
        this.udafPartial = udafunctionPartial;
        this.udafComplete = udafunctionComplete;
        this.outputInspector = outputOi;
        this.outputInspectorPartial = outputOiPartial;
    }

    @Override
    public void init() throws AlgebricksException {
        try {
            aggBuffer = udafPartial.getNewAggregationBuffer();
        } catch (HiveException e) {
            throw new AlgebricksException(e);
        }
    }

    @Override
    public void step(IFrameTupleReference tuple) throws AlgebricksException {
        readIntoCache(tuple);
        processRow();
    }

    private void processRow() throws AlgebricksException {
        try {
            // get values by evaluating them
            for (int i = 0; i < cachedParameters.length; i++) {
                cachedParameters[i] = evaluators[i].evaluate(cachedRowObject);
            }
            processAggregate();
        } catch (HiveException e) {
            throw new AlgebricksException(e);
        }
    }

    private void processAggregate() throws HiveException {
        /**
         * accumulate the aggregation function
         */
        switch (mode) {
            case PARTIAL1:
            case COMPLETE:
                udafPartial.iterate(aggBuffer, cachedParameters);
                break;
            case PARTIAL2:
            case FINAL:
                if (udafPartial instanceof GenericUDAFCount.GenericUDAFCountEvaluator) {
                    Object parameter = ((PrimitiveObjectInspector) parameterInspectors[0])
                            .getPrimitiveWritableObject(cachedParameters[0]);
                    udafPartial.merge(aggBuffer, parameter);
                } else
                    udafPartial.merge(aggBuffer, cachedParameters[0]);
                break;
            default:
                break;
        }
    }

    /**
     * serialize the result
     * 
     * @param result
     *            the evaluation result
     * @throws IOException
     * @throws AlgebricksException
     */
    private void serializeResult(Object result, ObjectInspector oi) throws IOException, AlgebricksException {
        try {
            BytesWritable outputWritable = (BytesWritable) lazySer.serialize(result, oi);
            out.write(outputWritable.getBytes(), 0, outputWritable.getLength());
        } catch (SerDeException e) {
            throw new AlgebricksException(e);
        }
    }

    /**
     * bind the tuple reference to the cached row object
     * 
     * @param r
     */
    private void readIntoCache(IFrameTupleReference r) {
        cachedRowObject.init(r);
    }

    @Override
    public void finish() throws AlgebricksException {
        // aggregator
        try {
            Object result = null;
            result = udafPartial.terminatePartial(aggBuffer);
            if (mode == GenericUDAFEvaluator.Mode.COMPLETE || mode == GenericUDAFEvaluator.Mode.FINAL) {
                result = udafComplete.terminate(aggBuffer);
                serializeResult(result, outputInspector);
            } else {
                serializeResult(result, outputInspectorPartial);
            }
        } catch (HiveException e) {
            throw new AlgebricksException(e);
        } catch (IOException e) {
            throw new AlgebricksException(e);
        }
    }

    @Override
    public void finishPartial() throws AlgebricksException {
        // aggregator.
        try {
            Object result = null;
            // get aggregations
            result = udafPartial.terminatePartial(aggBuffer);
            serializeResult(result, outputInspectorPartial);
        } catch (HiveException e) {
            throw new AlgebricksException(e);
        } catch (IOException e) {
            throw new AlgebricksException(e);
        }
    }
}