Spark中用户自定义聚合函数能否与Dataset.observe结合使用？

问题重现 #

Spark的observe方法是实用的作业监控手段，但要实现原生不支持的复杂聚合（比如用sketch对列采样）时，使用自定义Aggregator（Typed UDAF）会触发序列化异常。

已尝试两种调用方式：

方式一：

session.udf().register("myAgg", functions.udaf(new MyAggregator(), Encoders.STRING()));
Column aggCol = functions.expr("myAgg(inputCol)");

方式二：

Column aggCol = functions.udaf(new MyAggregator(), Encoders.STRING()).apply(ds.col("inputCol"));

执行ds = ds.observe("foo", aggCol);后抛出异常：

org.apache.spark.SparkException: Job aborted due to stage failure: task 0.0 in stage 4.0 (TID 2) had a not serializable result: org.apache.spark.sql.catalyst.expressions.GeneratedClass$SpecificUnsafeProjection
Serialization stack:
    - object not serializable (class: org.apache.spark.sql.catalyst.expressions.GeneratedClass$SpecificUnsafeProjection, value: <function1>)
    - field (class: org.apache.spark.sql.execution.aggregate.ScalaAggregator, name: inputProjection, type: class org.apache.spark.sql.catalyst.expressions.UnsafeProjection)
    - object (class org.apache.spark.sql.execution.aggregate.ScalaAggregator, myAgg(input[5, string, true]))
    - element of array (index: 0)
    - array (class [Lorg.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate;, size 1)
    - field (class: org.apache.spark.sql.execution.AggregatingAccumulator, name: typedImperatives, type: class [Lorg.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate;)
    - object (class org.apache.spark.sql.execution.AggregatingAccumulator, AggregatingAccumulator(id: 111, name: Some(Collected metrics), value: [empty row]))
    - writeExternal data
    - externalizable object (class org.apache.spark.scheduler.DirectTaskResult, org.apache.spark.scheduler.DirectTaskResult@7c17b6a8)

异常原因 #

observe机制依赖AggregatingAccumulator收集聚合结果，而Typed UDAF（基于Aggregator实现）对应的ScalaAggregator内部会生成UnsafeProjection实例处理输入数据。这个运行时生成的投影类未实现Serializable接口，当任务结束后AggregatingAccumulator需要序列化回Driver时，就会触发序列化失败。

解决方案 #

方案1：改用传统UDAF（基于UserDefinedAggregateFunction实现） #

传统UDAF的实现模式不会引入不可序列化的UnsafeProjection对象，能兼容observe方法。示例实现如下：

import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.*;
import org.apache.spark.sql.Row;

public class MySketchUDAF extends UserDefinedAggregateFunction {
    // 输入列的schema，这里定义为字符串类型
    @Override
    public StructType inputSchema() {
        return new StructType().add("inputCol", DataTypes.StringType);
    }

    // 聚合缓冲区的schema，根据你的sketch逻辑定义结构
    @Override
    public StructType bufferSchema() {
        // 示例：用字符串存储sketch的中间状态（实际根据你的sketch实现调整）
        return new StructType().add("sketchState", DataTypes.StringType);
    }

    // 输出结果的数据类型
    @Override
    public DataType dataType() {
        return DataTypes.StringType;
    }

    // 是否为确定性聚合（相同输入返回相同结果）
    @Override
    public boolean deterministic() {
        return true;
    }

    // 初始化缓冲区
    @Override
    public void initialize(MutableAggregationBuffer buffer) {
        // 根据你的sketch逻辑初始化，比如空字符串
        buffer.update(0, "");
    }

    // 单条输入数据更新缓冲区
    @Override
    public void update(MutableAggregationBuffer buffer, Row input) {
        String currentState = buffer.getString(0);
        String newInput = input.getString(0);
        // 替换为你的sketch更新逻辑，比如将新输入加入sketch
        String newState = currentState + "," + newInput;
        buffer.update(0, newState);
    }

    // 合并两个缓冲区
    @Override
    public void merge(MutableAggregationBuffer buffer1, Row buffer2) {
        String state1 = buffer1.getString(0);
        String state2 = buffer2.getString(0);
        // 替换为你的sketch合并逻辑
        buffer1.update(0, state1 + state2);
    }

    // 根据缓冲区计算最终结果
    @Override
    public Object evaluate(Row buffer) {
        String finalState = buffer.getString(0);
        // 替换为你的sketch最终结果生成逻辑，比如返回采样结果
        return finalState.replaceFirst("^,", "");
    }
}

调用方式：

// 注册UDAF
spark.udf().register("mySketchAgg", new MySketchUDAF());
// 使用observe执行聚合监控
Dataset<Row> ds = ...; // 你的数据集
ds = ds.observe("sketch_metrics", expr("mySketchAgg(inputCol)"));

方案2：调整Typed UDAF实现（仅限特殊场景） #

如果必须使用Aggregator，需要确保MyAggregator的所有成员变量都是可序列化的，并且避免依赖Spark内部生成的不可序列化对象。但这种方式复杂度较高，且不一定能完全规避ScalaAggregator内部的序列化问题，因此优先推荐方案1。

内容的提问来源于stack exchange，提问作者Alex Field