Spark SQL

import com.google.gson.GsonBuilder;
import lombok.Data;
import lombok.experimental.Accessors;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils;

import java.util.Comparator;
import java.util.Objects;


/**
 * 在 SparkSQL 中，first_value() 和 last_value() 是一种窗口函数，依赖 over 开窗的 order by，
 * 但想要分组取第一个元素函数的时候就必须需要再写子查询进行过滤
 * 在 FlinkSQL 中，first_value() 和 last_value() 是一种 UDAF 函数，函数中可以传递排序字段
 * <p>
 * 该函数就是在 SparkSQL 中实现对齐 FlinkSQL 的语义，降低 SQL 复杂度
 * <p>
 * Usage: sorted_first_value(value_col, order_col, [asc|desc])
 * <p>
 * ----------------------------------------------------
 * id   configuration_code  zh_name            price_version
 * 1	API_Usage           API调用小时         20230703
 * 2	API_Usage           API调用次数_小时结   20230412
 * 3	API_Usage           API调用小时         20230321
 * 4	API_Usage           API调用次数_小时结   20230302
 * 5	API_Usage           API调用次数按小时    20230117
 * 6	API_Usage           API调用次数按小时    20230117
 * 7	API_Usage           API调用次数按小时    20230117
 * 8	API_Usage           API调用小时	      20221209
 * 9	API_Usage           API调用小时	      20220930
 * ----------------------------------------------------
 * <p>
 * -- 按价格版本取每个 configuration_code 的最新的中文名
 * SELECT  product,
 *         configuration_code,
 *         sorted_first_value(zh_name, price_version, 'desc') AS zh_name
 * FROM    eps_tech.dwd_trade_common_configuration_df
 * WHERE   date = '${date}'
 * AND     configuration_code = 'API_Usage'
 * GROUP BY
 *         product,
 *         configuration_code;
 */
public class SortedFirstValue extends AbstractGenericUDAFResolver {

    /**
     * 校验多个入参的参数类型，并生成 GenericUDAFEvaluator
     *
     * @param parameterInfos
     * @return
     * @throws SemanticException
     */
    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameterInfos) throws SemanticException {
        // Type-checking goes here!
        if (parameterInfos == null || parameterInfos.length != 3) {
            throw new UDFArgumentLengthException("Exactly 3 params is expected");
        }

        // 解析第一个参数类型，最终需要取返回值的列
        ObjectInspector valueColOI = TypeInfoUtils.getStandardJavaObjectInspectorFromTypeInfo(parameterInfos[0]);
        if (valueColOI.getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new UDFArgumentTypeException(0, "Must be a primitive type");
        }

        // 解析第二个参数类型，排序字段
        ObjectInspector orderColOI = TypeInfoUtils.getStandardJavaObjectInspectorFromTypeInfo(parameterInfos[1]);
        if (orderColOI.getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new UDFArgumentTypeException(1, "Must be a primitive type");
        }

        // 解析第三个参数类型，ASC｜DESC String 类型
        ObjectInspector orderOI = TypeInfoUtils.getStandardJavaObjectInspectorFromTypeInfo(parameterInfos[2]);
        if (orderColOI.getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new UDFArgumentTypeException(1, "Must be a primitive type");
        }

        PrimitiveObjectInspector orderPOI = (PrimitiveObjectInspector) orderOI;
        if (orderPOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.STRING) {
            throw new UDFArgumentTypeException(2, "Must be a string type, but pass " + orderPOI.getPrimitiveCategory().name());
        }

        return new SortedFirstValueEvaluator();
    }

    /**
     * UDAF 计算逻辑
     */
    static class SortedFirstValueEvaluator extends GenericUDAFEvaluator {

        private transient ObjectInspector internalMergeOI;

        @Data
        @Accessors(chain = true)
        static class Row {
            private Object value;
            private Object orderBy;
            private String order = "ASC";

            public String toJson() {
                return new GsonBuilder().create().toJson(this);
            }

            public static Row parseJson(String json) {
                return new GsonBuilder().create().fromJson(json, Row.class);
            }
        }

        /**
         * 用于缓存中间结果的数据结构
         */
        static class SortedFirstValueBuffer extends AbstractAggregationBuffer {

            Object container;

            public SortedFirstValueBuffer() {
                container = null;
            }

        }

        /**
         * 实例化 Evaluator 类的时候调用，初始化各 Aggregation 阶段的 ObjectInspector
         *
         * @param mode       PARTIAL1:  原始数据到部分聚合，调用 iterate 和 terminatePartial -> map 阶段
         *                   PARTIAL2:  部分聚合到部分聚合，调用 merge 和 terminatePartial -> combine 阶段
         *                   FINAL:     部分聚合到完全聚合，调用 merge 和 terminate -> reduce 阶段
         *                   COMPLETE:  从原始数据直接到完全聚合，调用 iterate 和 terminate -> map 完直接聚合出最终结果，没有 combine 和 reduce 阶段
         * @param parameters UDAF 入参类型检查
         * @return
         * @throws HiveException
         */
        @Override
        public ObjectInspector init(Mode mode, ObjectInspector[] parameters) throws HiveException {
            super.init(mode, parameters);

            switch (mode) {
                case PARTIAL1:
                case PARTIAL2:
                case FINAL:
                case COMPLETE:
                default:
                    // 每个阶段的 outputOI 都是 String
                    return PrimitiveObjectInspectorFactory.javaStringObjectInspector;
            }

        }

        /**
         * 获取一个新的存放中间结果的缓存
         *
         * @return
         * @throws HiveException
         */
        @Override
        public AbstractAggregationBuffer getNewAggregationBuffer() throws HiveException {
            return new SortedFirstValueBuffer();
        }

        /**
         * Reset the aggregation. This is useful if we want to reuse the same aggregation.
         *
         * @param agg
         * @throws HiveException
         */
        @Override
        public void reset(AggregationBuffer agg) throws HiveException {
            ((SortedFirstValueBuffer) agg).container = null;
        }

        /**
         * 对应 map，将处理的一行数据写入 AggregationBuffer
         *
         * @param agg
         * @param parameters
         * @throws HiveException
         */
        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters) throws HiveException {
            Object value = parameters[0];
            Object orderBy = parameters[1];
            // parameters[2] 是 org.apache.hadoop.io.Text
            String order = parameters[2].toString();
            compareWithBuffer(value, orderBy, order, (SortedFirstValueBuffer) agg);
        }

        /**
         * 返回部分聚合数据 AggregationBuffer 中的持久化对象。
         * <p>
         * Note: 调用这个方法时，说明已经是 map 或者 combine 的结束了，必须将数据持久化以后交给 reduce 进行处理。
         * 持久化类型只支持 Java 原始数据类型及其封装类型、Hadoop Writable 类型、List、Map，不能返回自定义的类，
         * 即使实现了 Serializable 也不行，否则会出现问题或者错误的结果。
         *
         * @param agg
         * @return
         * @throws HiveException
         */
        @Override
        public Object terminatePartial(AggregationBuffer agg) throws HiveException {
            return ((SortedFirstValueBuffer) agg).container;
        }

        /**
         * 将 terminatePartial 返回的部分聚合数据进行合并，partial 是 List 类型
         *
         * @param agg
         * @param partial
         * @throws HiveException
         */
        @Override
        public void merge(AggregationBuffer agg, Object partial) throws HiveException {
            SortedFirstValueBuffer firstValueBuffer = (SortedFirstValueBuffer) agg;
            Row partialRow = Row.parseJson(partial.toString());
            compareWithBuffer(partialRow.value, partialRow.orderBy, partialRow.order, firstValueBuffer);
        }

        /**
         * 结束，生成最终结果
         *
         * @param agg
         * @return
         * @throws HiveException
         */
        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            return Row.parseJson(((SortedFirstValueBuffer) agg).container.toString()).value;
        }

        Comparator<Object> objectComparator = (o1, o2) -> {
            if (o1 == null && o2 == null) return 0;
            if (o1 == null) return -1; // NULL 算无限小
            if (o2 == null) return 1;

            if (o1 instanceof Number) {
                return Double.compare(((Number) o1).doubleValue(), ((Number) o2).doubleValue());
            } else {
                return o1.toString().compareTo(o2.toString());
            }
        };

        /**
         * @param value                  sorted_first_value 返回值的列
         * @param orderBy                sorted_first_value 用于排序的列
         * @param order                  asc｜desc 按指定列升序或降序后取 first_value
         * @param sortedFirstValueBuffer buffer 缓存
         */
        private void compareWithBuffer(Object value, Object orderBy, String order, SortedFirstValueBuffer sortedFirstValueBuffer) {
            Row row = new Row()
                    .setValue(value)
                    .setOrderBy(orderBy)
                    .setOrder(order);
            Object pCopy = ObjectInspectorUtils.copyToStandardObject(row.toJson(),
                    PrimitiveObjectInspectorFactory.javaStringObjectInspector);
            if (sortedFirstValueBuffer.container == null) {
                sortedFirstValueBuffer.container = pCopy;
            } else {
                Object elem = sortedFirstValueBuffer.container;
                Row bufferedRow = Row.parseJson(elem.toString());
                // compare > 0，代表当前值比 buffer 的值大
                int compare = Objects.compare(orderBy, bufferedRow.getOrderBy(), objectComparator);
                // 升序的 FirstValue 是取最小值，如果小于 buffer 就置换，反之是取最大值，如果大大于 buffer 就置换
                if (("ASC".equalsIgnoreCase(order) ? -compare : compare) > 0) {
                    sortedFirstValueBuffer.container = pCopy;
                }
            }
        }
    }

}