spark 函数 UDF和UDAF

spark函数功能丰富，主要有内置函数、UDF、UDAF等（UDAF用于聚合agg，同时对多行操作返回单个聚合值）

内置函数类型较多，见官网Built-in Functions - Spark 3.0.1 Documentation (apache.org)

UDF函数
使用步骤：定义UserDefinedFunction（重载call方法）、注册spark.udf().register(" ", userDefinedFunction)、使用（map或agg或withColunm）。

import org.apache.spark.sql.*;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.expressions.UserDefinedFunction;
import static org.apache.spark.sql.functions.udf;
import org.apache.spark.sql.types.DataTypes;

SparkSession spark = SparkSession
  .builder()
  .appName("Java Spark SQL UDF scalar example")
  .getOrCreate();

// Define and register a zero-argument non-deterministic UDF
// UDF is deterministic by default, i.e. produces the same result for the same input.
UserDefinedFunction random = udf(
  () -> Math.random(), DataTypes.DoubleType
);
random.asNondeterministic();
spark.udf().register("random", random);
spark.sql("SELECT random()").show();
// +-------+
// |UDF()  |
// +-------+
// |xxxxxxx|
// +-------+

// Define and register a one-argument UDF
spark.udf().register("plusOne", new UDF1<Integer, Integer>() {
  @Override
  public Integer call(Integer x) {
    return x + 1;
  }
}, DataTypes.IntegerType);
spark.sql("SELECT plusOne(5)").show();
// +----------+
// |plusOne(5)|
// +----------+
// |         6|
// +----------+

// Define and register a two-argument UDF
UserDefinedFunction strLen = udf(
  (String s, Integer x) -> s.length() + x, DataTypes.IntegerType
);
spark.udf().register("strLen", strLen);
spark.sql("SELECT strLen('test', 1)").show();
// +------------+
// |UDF(test, 1)|
// +------------+
// |           5|
// +------------+

// UDF in a WHERE clause
spark.udf().register("oneArgFilter", new UDF1<Long, Boolean>() {
  @Override
  public Boolean call(Long x) {
    return  x > 5;
  }
}, DataTypes.BooleanType);
spark.range(1, 10).createOrReplaceTempView("test");
spark.sql("SELECT * FROM test WHERE oneArgFilter(id)").show();
// +---+
// | id|
// +---+
// |  6|
// |  7|
// |  8|
// |  9|
// +---+

UDAF函数
Aggregator

import java.io.Serializable;

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.TypedColumn;
import org.apache.spark.sql.expressions.Aggregator;

public static class Employee implements Serializable {
  private String name;
  private long salary;

  // Constructors, getters, setters...

}

public static class Average implements Serializable  {
  private long sum;
  private long count;

  // Constructors, getters, setters...

}

public static class MyAverage extends Aggregator<Employee, Average, Double> {
  // A zero value for this aggregation. Should satisfy the property that any b + zero = b
  public Average zero() {
    return new Average(0L, 0L);
  }
  // Combine two values to produce a new value. For performance, the function may modify `buffer`
  // and return it instead of constructing a new object
  public Average reduce(Average buffer, Employee employee) {
    long newSum = buffer.getSum() + employee.getSalary();
    long newCount = buffer.getCount() + 1;
    buffer.setSum(newSum);
    buffer.setCount(newCount);
    return buffer;
  }
  // Merge two intermediate values
  public Average merge(Average b1, Average b2) {
    long mergedSum = b1.getSum() + b2.getSum();
    long mergedCount = b1.getCount() + b2.getCount();
    b1.setSum(mergedSum);
    b1.setCount(mergedCount);
    return b1;
  }
  // Transform the output of the reduction
  public Double finish(Average reduction) {
    return ((double) reduction.getSum()) / reduction.getCount();
  }
  // Specifies the Encoder for the intermediate value type
  public Encoder<Average> bufferEncoder() {
    return Encoders.bean(Average.class);
  }
  // Specifies the Encoder for the final output value type
  public Encoder<Double> outputEncoder() {
    return Encoders.DOUBLE();
  }
}

Encoder<Employee> employeeEncoder = Encoders.bean(Employee.class);
String path = "examples/src/main/resources/employees.json";
Dataset<Employee> ds = spark.read().json(path).as(employeeEncoder);
ds.show();
// +-------+------+
// |   name|salary|
// +-------+------+
// |Michael|  3000|
// |   Andy|  4500|
// | Justin|  3500|
// |  Berta|  4000|
// +-------+------+

MyAverage myAverage = new MyAverage();
// Convert the function to a `TypedColumn` and give it a name
TypedColumn<Employee, Double> averageSalary = myAverage.toColumn().name("average_salary");
Dataset<Double> result = ds.select(averageSalary);
result.show();
// +--------------+
// |average_salary|
// +--------------+
// |        3750.0|
// +--------------+


// Register the function to access it
spark.udf().register("myAverage", functions.udaf(new MyAverage(), Encoders.LONG()));
Dataset<Row> result = spark.sql("SELECT myAverage(salary) as average_salary FROM employees");

UDAF函数
UserDefinedAggregateFunction

import org.apache.spark.sql.Row;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructType;

public class MyAverage extends UserDefinedAggregateFunction {

    /**
     * 计算平均score，输入的应该是score这一列数据
     * StructField定义了列字段的名称score_column,字段的类型Double
     * @return StructType
     */
    @Override
    public StructType inputSchema() {
        return new StructType()
                .add("score_column", DataTypes.DoubleType);
    }

    /**
     * 缓存Schema，存储中间计算结果，
     * 比如计算平均score,需要计算score的总和和score的个数,然后average(score)=sum(score)/count(score)
     * 所以这里定义了StructType类型：两个StructField字段：sum和count
     * @return StructType
     */
    @Override
    public StructType bufferSchema() {
        return new StructType()
                .add("sum", DataTypes.DoubleType)
                .add("count", DataTypes.LongType);
    }

    /**
     * 自定义集成算子最终返回的数据类型
     * 也就是average(score)的类型，所以是Double
     * @return DataType 返回数据类型
     */
    @Override
    public DataType dataType() {
        return DataTypes.DoubleType;
    }


    /**
     * 数据一致性检验：对于同样的输入，输出是一样的
     * @return Boolean true 同样的输入，输出也一样
     */
    @Override
    public boolean deterministic() {
        return true;
    }


    /**
     * 初始化缓存sum和count
     * sum=0.0，count=0
     * @param mutableAggregationBuffer 中间数据
     */
    @Override
    public void initialize(MutableAggregationBuffer mutableAggregationBuffer) {
        mutableAggregationBuffer.update(0, 0.0);
        mutableAggregationBuffer.update(1, 0L);

    }

    /**
     * 每次计算更新缓存
     * @param mutableAggregationBuffer 缓存数据
     * @param row 输入数据score
     */
    @Override
    public void update(MutableAggregationBuffer mutableAggregationBuffer, Row row) {
        mutableAggregationBuffer.update(0, mutableAggregationBuffer.getDouble(0) + row.getDouble(0));
        mutableAggregationBuffer.update(1, mutableAggregationBuffer.getLong(1) + 1);
    }

    /**
     * 将更新后的buffer存储到缓存
     * @param mutableAggregationBuffer 缓存
     * @param row 更新后的buffer
     */
    @Override
    public void merge(MutableAggregationBuffer mutableAggregationBuffer, Row row) {
        mutableAggregationBuffer.update(0, mutableAggregationBuffer.getDouble(0) + row.getDouble(0));
        mutableAggregationBuffer.update(1, mutableAggregationBuffer.getLong(1) + row.getLong(1));

    }


    /**
     * 计算最终的结果：average(score)=sum(score)/count(score)
     * @param row
     * @return
     */
    @Override
    public Object evaluate(Row row) {
        return row.getDouble(0)/row.getLong(1);
    }
}



Encoder<Employee> employeeEncoder = Encoders.bean(Employee.class);
String path = "examples/src/main/resources/employees.json";
Dataset<Employee> ds = spark.read().json(path).as(employeeEncoder);
ds.show();
// +-------+------+
// |   name|salary|
// +-------+------+
// |Michael|  3000|
// |   Andy|  4500|
// | Justin|  3500|
// |  Berta|  4000|
// +-------+------+

MyAverage myAverage = new MyAverage();
Dataset<Row> result = ds.groupBy("name").agg(myAverage.apply(ds.col("salary")).as("score_column"));