如何让PySpark中的Java UDF获取数据库读取的映射数据？

我的驱动脚本基于PySpark编写，需要对列执行复杂转换，出于性能考虑，这些转换用Java UDF实现。

PySpark中注册Java UDF的代码：

session.udf.registerJavaFunction("my_udf", "com.example.MyUDF", StringType())

执行转换的代码：

output_df = input_df.withColumn(f"transformed", F.expr("my_udf(col1, col2)"))

Java UDF的定义：

import org.apache.spark.sql.api.java.UDF2;

public class MyUDF implements UDF2<Integer, Integer, String> {

  @Override
  public String call(Integer x, Integer y) {
       // 原有逻辑
  }
}

现在需要给这个UDF传入一个从数据库读取的映射表myMap，用于计算映射值的总和：

@Override
  public String call(Integer x, Integer y) {
       // 伪代码
       return myMap.get(x) + myMap.get(y);
  }

核心问题：从数据库读取myMap并提供给Java UDF的最优方案是什么？已知广播变量适合这类场景，但不知道如何在Java UDF中访问PySpark创建的广播变量，或者有没有更优方案？

方案一：使用广播变量（推荐） #

广播变量是Spark分发只读大数据集的最优方式，能避免每个任务重复加载数据。具体步骤如下：

1. PySpark端读取并广播映射表 #

从数据库读取数据转成字典后创建广播变量：

# 从数据库读取映射表（以JDBC为例）
map_df = session.read.jdbc(url="jdbc:mysql://host/db", table="mapping_table", properties={"user": "xxx", "password": "xxx"})
# 转为Python字典
my_map = {row.key: row.value for row in map_df.collect()}
# 创建广播变量
broadcast_map = session.sparkContext.broadcast(my_map)

2. 修改Java UDF接收广播变量 #

Java UDF无法直接读取PySpark广播变量，需要通过构造函数注入的方式传递，同时调整PySpark的UDF注册逻辑：

调整Java UDF代码

让UDF持有Spark广播对象，在call方法中获取映射表：

import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.sql.api.java.UDF2;

import java.util.Map;

public class MyUDF implements UDF2<Integer, Integer, String> {

    private final Broadcast<Map<Integer, Integer>> broadcastMap;

    // 构造函数注入广播变量
    public MyUDF(Broadcast<Map<Integer, Integer>> broadcastMap) {
        this.broadcastMap = broadcastMap;
    }

    @Override
    public String call(Integer x, Integer y) {
        Map<Integer, Integer> myMap = broadcastMap.value();
        return String.valueOf(myMap.get(x) + myMap.get(y));
    }
}

PySpark端注册带参数的Java UDF

原有的registerJavaFunction不支持传递构造参数，改用udf方法包装UDF实例：

# 将Python广播变量转为Java端的Broadcast对象
java_broadcast = broadcast_map._jbroadcast

# 创建Java UDF实例并传入广播变量
udf_instance = session._jvm.com.example.MyUDF(java_broadcast)

# 注册UDF并指定类型
my_udf = F.udf(udf_instance, StringType())

# 使用UDF执行转换
output_df = input_df.withColumn("transformed", my_udf(F.col("col1"), F.col("col2")))

方案二：预加载映射表到UDF（适合小数据集） #

如果映射表数据量极小，可以在UDF初始化时直接从数据库读取。但这种方式每个Executor节点都会加载一次数据，仅适合数据量非常小的场景：

import org.apache.spark.sql.api.java.UDF2;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.util.HashMap;
import java.util.Map;

public class MyUDF implements UDF2<Integer, Integer, String> {

    private final Map<Integer, Integer> myMap;

    public MyUDF() {
        // 构造函数中从数据库加载映射表
        myMap = new HashMap<>();
        try (Connection conn = DriverManager.getConnection("jdbc:mysql://host/db", "user", "password")) {
            ResultSet rs = conn.createStatement().executeQuery("SELECT key, value FROM mapping_table");
            while (rs.next()) {
                myMap.put(rs.getInt("key"), rs.getInt("value"));
            }
        } catch (Exception e) {
            throw new RuntimeException("加载映射表失败", e);
        }
    }

    @Override
    public String call(Integer x, Integer y) {
        return String.valueOf(myMap.get(x) + myMap.get(y));
    }
}

这种方式不需要修改PySpark代码，直接用原来的registerJavaFunction注册即可，但要注意数据库连接的重复开销。

方案三：用Spark原生操作替代UDF（性能最优） #

如果你的逻辑可以用Spark原生SQL操作实现，优先选这种方式——Spark对原生操作的优化远优于UDF，尤其是大数据量场景：

# 读取映射表
map_df = session.read.jdbc(url="jdbc:mysql://host/db", table="mapping_table", properties={"user": "xxx", "password": "xxx"})

# 关联输入表与映射表，获取col1对应的映射值
map_col1 = map_df.withColumnRenamed("key", "col1_key").withColumnRenamed("value", "col1_value")
input_df = input_df.join(map_col1, input_df.col1 == map_col1.col1_key, "left")

# 再次关联映射表，获取col2对应的映射值
map_col2 = map_df.withColumnRenamed("key", "col2_key").withColumnRenamed("value", "col2_value")
input_df = input_df.join(map_col2, input_df.col2 == map_col2.col2_key, "left")

# 计算转换列并转为字符串类型
output_df = input_df.withColumn("transformed", (F.col("col1_value") + F.col("col2_value")).cast(StringType()))

# 清理中间列
output_df = output_df.drop("col1_key", "col1_value", "col2_key", "col2_value")

这种方式完全依赖Spark原生操作，Spark会自动优化执行计划（比如自动触发广播join），性能通常比UDF方案更好。

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