在PySpark中使用Scala类作为UDF的技术实现疑问

咱们先理清楚你遇到的场景：你想把部分计算逻辑从Python迁到Scala，用Java UDF接口实现带持久化变量的逻辑，但在PySpark调用时遇到了问题。下面我会一步步带你解决这个问题，同时点明代码里需要注意的关键细节。

问题核心先明确 #

你的Scala UDF里用了private var prop来维护状态，但要特别注意：Spark是分布式计算框架，每个Executor节点都会创建一个UDF类的独立实例。也就是说，每个实例的prop是各自独立的，不会在不同任务、不同节点间共享状态。如果你的需求是全局共享这个状态，那这种写法就不适用了，得改用广播变量、累加器或者外部存储；如果只是每个UDF实例内部（比如每个数据分区内）维护状态，那这个写法是可行的，只是要注意序列化和部署的问题。

解决步骤 #

1. 把Scala代码打包成JAR包 #

首先得把你的Scala类打包成Spark能识别的JAR包，这里以sbt为例（maven也可以，步骤类似）：

• 创建一个sbt项目，编写build.sbt配置文件，注意Scala版本要和你的Spark版本匹配（比如Spark 3.2+对应Scala 2.12）：

name := "SparkScalaUDF"
version := "1.0"
scalaVersion := "2.12.15"
libraryDependencies += "org.apache.spark" %% "spark-sql" % "3.3.0" % Provided

• 把你的SomeFun.scala放在src/main/scala/mwe/目录下
• 执行sbt package命令，生成的JAR包路径一般是target/scala-2.12/sparkscalaudf_2.12-1.0.jar

2. 在PySpark中引入JAR并注册UDF #

启动PySpark的时候，需要把刚才生成的JAR包添加到classpath，或者在代码里配置：

方式一：启动PySpark时指定JAR

pyspark --jars /path/to/your/sparkscalaudf_2.12-1.0.jar

方式二：在SparkSession中配置

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("CallScalaUDF") \
    .config("spark.jars", "/path/to/your/sparkscalaudf_2.12-1.0.jar") \
    .getOrCreate()

接下来注册你的Scala UDF：

# 用registerJavaFunction注册Java接口类型的UDF，参数分别是：UDF名称、Scala类的全路径、返回值类型
spark.udf.registerJavaFunction("some_fun", "mwe.SomeFun", "integer")

# 创建测试数据
test_df = spark.createDataFrame([(1,), (2,), (3,), (4,)], ["input_col"])

# 调用UDF计算
result_df = test_df.selectExpr("input_col", "some_fun(input_col) as output_col")
result_df.show()

3. 关键注意事项 #

• 状态的作用范围：如之前所说，每个Executor的UDF实例有自己的prop。比如如果数据分在2个分区，第一个分区的第一条数据会初始化该实例的prop，第二个分区的第一条数据会初始化另一个实例的prop，两者互不干扰。如果需要全局共享状态，建议用广播变量（只读场景）或者累加器（聚合场景），或者把状态存在Redis这类外部存储里。
• 序列化问题：UDF1接口本身继承了Serializable，所以你的SomeFun类是可序列化的，但如果后续要给prop换成复杂类型，一定要确保该类型也支持序列化，否则会出现序列化错误。
• 版本兼容性：Scala版本必须和Spark的Scala版本严格匹配，不然会出现类加载失败的问题。比如Spark 2.4.x对应Scala 2.11，Spark 3.x对应Scala 2.12。

测试结果示例 #

如果你的测试数据在同一个分区里，输出会是这样：

+---------+----------+
|input_col|output_col|
+---------+----------+
|        1|         2|
|        2|         3|
|        3|         4|
|        4|         5|
+---------+----------+

因为第一个输入1把prop初始化为1，后续计算都是1+输入值。如果数据分在两个分区，比如第一个分区是[(1,)]，第二个是[(2,), (3,), (4,)]，输出会变成：

+---------+----------+
|input_col|output_col|
+---------+----------+
|        1|         2|
|        2|         4|
|        3|         5|
|        4|         6|
+---------+----------+

第二个分区的UDF实例用2初始化了prop，所以后续计算都是2+输入值。

内容的提问来源于stack exchange，提问作者matz-e