如何在Presto/Trino自定义UDF执行结束时运行清理代码

Trino（原PrestoSQL）原生支持UDF生命周期回调能力，完全可以实现单split处理完成后的close类清理逻辑，无需修改内核源码。

实现方案 #

标量UDF的实例与Worker节点上的单个split处理任务一一绑定：每处理一个split，Trino都会创建独立的UDF实例，split处理完成后实例会被销毁，你只需要遵循Trino的方法约定即可实现自动回调：

• 在你的UDF类中实现java.io.Closeable接口，或者直接定义无参的public close()方法
• 给close方法添加@UsedByGeneratedCode注解，避免被Trino的代码生成器当作无用代码优化
• 所有split级的清理逻辑、IPC资源释放逻辑都可以写在这个close方法中，会在当前split的最后一行数据处理完成后自动触发

代码示例 #

import io.trino.spi.function.Description;
import io.trino.spi.function.ScalarFunction;
import io.trino.spi.function.SqlType;
import io.trino.spi.function.UsedByGeneratedCode;
import java.io.Closeable;
import java.io.IOException;

@ScalarFunction("custom_udf_demo")
@Description("自定义标量UDF示例")
public class CustomDemoUdf implements Closeable
{
    // 声明split级别的资源，比如IPC客户端、本地缓存等
    private IpcConnection ipcConnection;

    public CustomDemoUdf() {
        // 构造方法会在split处理前调用，适合做资源初始化
        this.ipcConnection = new IpcConnection("your-ipc-address");
    }

    @UsedByGeneratedCode
    @SqlType("VARCHAR")
    public String evaluate(@SqlType("VARCHAR") String input) {
        // 每行数据的处理逻辑，每一行都会调用一次
        return ipcConnection.invoke(input);
    }

    @Override
    @UsedByGeneratedCode
    public void close() throws IOException
    {
        // split处理完成后自动触发，这里写资源释放逻辑
        if (ipcConnection != null) {
            ipcConnection.close();
        }
    }
}

注意事项 #

• 不要用静态变量存储split级资源，静态变量是Worker进程级别的，会被多个split的UDF实例共享，清理时会出现并发冲突
• close方法仅在split正常处理完成时触发，如果split执行失败抛出异常，不会自动调用close，异常场景的兜底清理可以在evaluate方法的异常捕获逻辑中补充
• 如果你用的是Trino更名前的PrestoSQL版本（0.2xx系列），逻辑完全一致，仅需要把代码中的io.trino包名替换为com.facebook.presto即可

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