1. Background

一句话概括 UDAF 的背景就是系统自带的聚合函数无法满足用户需求。

2. Basic

2.1 什么是 UDAF ?

UDAF 即自定义聚合函数。首先看一下什么是聚合函数：聚合函数即是指0行到多行的0个到多个列作为参数输入，返回单一值的函数。通俗的说就是多进一出。经常和 group by 子句一起用。比如 sum,count,avg 等都是很常见的系统聚合函数。

2.2 什么是 ObjectInspector ？

这个概念较为复杂，感兴趣的可以自己深入了解一下。简单的说，ObjectInspector 接口使得Hive可以不拘泥于一种特定数据格式，使得数据流 1）在输入端和输出端切换不同的输入/输出格式 2）在不同的Operator上使用不同的数据格式。 hive 的 UDAF 中只要会使用即可。

3. Deep

3.1 两个核心类

AbstractGenericUDAFResolver
通过 getEvaluator 返回自定义的Evaluator
GenericUDAFEvaluator
通过7个函数，4个步骤完成UDAF全部逻辑

3.2 GenericUDAFEvaluator

7个函数：

init()
初始化输入和输出的数据结构
getNewAggregationBuffer()
返回用于存储中间聚合结果的对象
reset()
重置聚合结果
iterate()
将一行数据放入聚合buffer中
terminatePartial()
返回部分聚合结果
merge()
合并terminatePartial返回的部分聚合结果
terminate()
返回最终结果

4个步骤：

partial1
input -> output: original -> partial aggregation
method invoke: init() -> iterate() -> terminatePartial()
partial2
input -> output: partial aggregation -> partial aggregation
method invoke: init() -> merge() -> terminatePartial()
final
input -> output: partial aggregation -> full aggregation
method invoke: init() -> merge() -> terminate()
complete
input -> output: original -> full aggregation
method invoke: init() -> iterate() -> terminate()

4. Best Practice

4.1 需求：

input:
a1:5 c1:10
a2:3 c2:40
a3:8 c3:100
output:
(a1 * c1 + a2 * c2 + a3 * c3) / (c1 + c2 + c3) 类似于加权平均数

4.2 代码：

public class WeightedAverage extends AbstractGenericUDAFResolver {
    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
            throws SemanticException {
        if (parameters.length != 2) {
            throw new UDFArgumentTypeException(parameters.length - 1,
                    "Exactly two argument is expected.");
        }

        if (parameters[0].getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new UDFArgumentTypeException(0,
                    "Only primitive type arguments are accepted but "
                            + parameters[0].getTypeName() + " is passed.");
        }
        switch (((PrimitiveTypeInfo) parameters[0]).getPrimitiveCategory()) {
            case INT:
            case LONG:
                return new GenericUDAFAverageEvaluator();
            case BYTE:
            case SHORT:
            case FLOAT:
            case DOUBLE:
            case STRING:
            case TIMESTAMP:
            case BOOLEAN:
            default:
                throw new UDFArgumentTypeException(0,
                        "Only int or long type arguments are accepted but "
                                + parameters[0].getTypeName() + " is passed.");
        }
    }


    public static class GenericUDAFAverageEvaluator extends GenericUDAFEvaluator {

        // input For iterate()
        PrimitiveObjectInspector avgOriginalInputOI;
        PrimitiveObjectInspector weightOriginalInputOI;

        // output For terminatePartial()
        Object[] partialAggregationResult;

        // input For merge()
        StructObjectInspector soi;
        StructField countField;
        StructField sumField;
        LongObjectInspector countFieldOI;
        LongObjectInspector sumFieldOI;

        // output For terminate()
        LongWritable fullAggregationResult;

        @Override
        public ObjectInspector init(Mode mode, ObjectInspector[] parameters)
                throws HiveException {
            super.init(mode, parameters);

    // init input
            // Mode.PARTIAL1 || mode == Mode.COMPLETE
            // input:original, method:iterate()
            if (mode == Mode.PARTIAL1 || mode == Mode.COMPLETE) {
                avgOriginalInputOI = (PrimitiveObjectInspector) parameters[0];
                weightOriginalInputOI = (PrimitiveObjectInspector) parameters[1];
            }
            // Mode.PARTIAL2 || Mode.FINAL
            // input:partial aggregation, method:merge()
            else {
                //部分数据作为输入参数时，用到的struct的OI实例，指定输入数据类型，用于解析数据
                soi = (StructObjectInspector) parameters[0];
                countField = soi.getStructFieldRef("count");
                sumField = soi.getStructFieldRef("sum");
                //数组中的每个数据，需要其各自的基本类型OI实例解析
                countFieldOI = (LongObjectInspector) countField.getFieldObjectInspector();
                sumFieldOI = (LongObjectInspector) sumField.getFieldObjectInspector();
            }

    // init output
            // Mode.PARTIAL1 || mode == Mode.PARTIAL2
            // output:partial aggregation, method:terminatePartial()
            if (mode == Mode.PARTIAL1 || mode == Mode.PARTIAL2) {
                partialAggregationResult = new Object[2];
                partialAggregationResult[0] = new LongWritable(0);
                partialAggregationResult[1] = new LongWritable(0);
                /*
                 * 构造Struct的OI实例，用于设定聚合结果数组的类型
                 * 需要字段名List和字段类型List作为参数来构造
                 */
                ArrayList<String> fname = new ArrayList<String>();
                fname.add("count");
                fname.add("sum");
                ArrayList<ObjectInspector> foi = new ArrayList<ObjectInspector>();
                //注：此处的两个OI类型 描述的是 partialResult[] 的两个类型，故需一致
                foi.add(PrimitiveObjectInspectorFactory.writableLongObjectInspector);
                foi.add(PrimitiveObjectInspectorFactory.writableLongObjectInspector);
                return ObjectInspectorFactory.getStandardStructObjectInspector(fname, foi);
            }
           // Mode.COMPLETE || Mode. FINAL
           // output:full aggregation, method:terminate()
            else {
                //FINAL COMPLETE 最终聚合结果为一个数值，并用基本类型OI设定其类型
                fullAggregationResult = new LongWritable(0);
                return PrimitiveObjectInspectorFactory.writableLongObjectInspector;
            }
        }

        /*
         * 聚合数据缓存存储结构
         */
        static class AverageAgg implements AggregationBuffer {
            long count;
            long sum;
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer() throws HiveException {
            AverageAgg result = new AverageAgg();
            reset(result);
            return result;
        }

        @Override
        public void reset(AggregationBuffer agg) throws HiveException {
            AverageAgg myagg = (AverageAgg) agg;
            myagg.count = 0;
            myagg.sum = 0;
        }

        /*
         * 遍历原始数据(将一行数据（Object[] parameters）放入聚合buffer中)
         * input: original
         */
        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters)
                throws HiveException {
            Object p1 = parameters[0];
            Object p2 = parameters[1];
            if (p1 != null && p2 != null) {
                AverageAgg myagg = (AverageAgg) agg;
                try {
                    long avg = PrimitiveObjectInspectorUtils.getLong(p1, avgOriginalInputOI);
                    long count = PrimitiveObjectInspectorUtils.getLong(p2, weightOriginalInputOI);
                    myagg.count += count;
                    myagg.sum += avg*count;
                } catch (NumberFormatException e) {
                    throw new HiveException("NumberFormatException: get value failed");
                }
            }
        }

        /*
         * 得出部分聚合结果
         * output: partial aggregation
         */
        @Override
        public Object terminatePartial(AggregationBuffer agg) throws HiveException {
            AverageAgg myagg = (AverageAgg) agg;
            ((LongWritable) partialAggregationResult[0]).set(myagg.count);
            ((LongWritable) partialAggregationResult[1]).set(myagg.sum);
            return partialAggregationResult;
        }

        /*
         * 合并部分聚合结果(注：Object[] 是 Object 的子类，此处 partial 为 Object[]数组)
         * input: partial aggregation
         */
        @Override
        public void merge(AggregationBuffer agg, Object partial)
                throws HiveException {
            if (partial != null) {
                AverageAgg myagg = (AverageAgg) agg;
                //通过StandardStructObjectInspector实例，分解出 partial 数组元素值
                Object partialCount = soi.getStructFieldData(partial, countField);
                Object partialSum = soi.getStructFieldData(partial, sumField);
                //通过基本数据类型的OI实例解析Object的值
                myagg.count += countFieldOI.get(partialCount);
                myagg.sum += sumFieldOI.get(partialSum);
            }
        }

        /*
         * 得出最终聚合结果
         * output: full aggregation
         */
        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            AverageAgg myagg = (AverageAgg) agg;
            if (myagg.count == 0) {
                return null;
            } else {
                fullAggregationResult.set(myagg.sum / myagg.count);
                return fullAggregationResult;
            }
        }
    }

}

5. Ref

《Hive 编程指南》
Hive之自定义聚合函数UDAF

Hive 之 UDAF