如何让PySpark UDF在Python类中访问实例变量？

在Python类中使用PySpark UDF时，静态方法形式的UDF无法直接访问类的实例变量（如self.increase），直接引用会触发NameError。已知可通过在调用方法内部定义UDF解决，以下是另外三种可行的替代方案：

方案1：使用广播变量传递实例变量 #

将需要访问的实例变量封装为Spark广播变量，在UDF中引用广播变量的值。这种方式适合在分布式环境中共享只读变量，避免重复传输大对象，提升性能。

import numpy as np
from pyspark.sql.types import StringType, IntegerType, StructType, StructField
from pyspark.sql.functions import udf, col

class Example():
    def __init__(self):
        self.students = [[f'student_{i}', np.random.randint(80)] for i in range(3)]
        self.increase = 10
        # 将实例变量转为广播变量
        self.broadcast_increase = spark.sparkContext.broadcast(self.increase)
    
    def create_spark_df(self):
        cSchema = StructType([StructField("Name", StringType()),
                             StructField("Marks", IntegerType())])
        return spark.createDataFrame(self.students, schema=cSchema)

    @staticmethod
    @udf(returnType=IntegerType())
    def add_increase_marks(marks, broadcast_val):
        return marks + broadcast_val.value
    
    def calculate_new_marks(self):
        df = self.create_spark_df()
        df = df.withColumn("New Marks", self.add_increase_marks(col("Marks"), self.broadcast_increase))
        return df

c = Example()
c.calculate_new_marks().show()

方案2：使用可序列化的实例方法作为UDF #

直接使用实例方法作为UDF，但需要确保类的所有实例变量都支持序列化（Python中大部分基础类型默认支持）。Spark会将类实例序列化后分发到各个工作节点，实例方法即可访问实例变量。

import numpy as np
from pyspark.sql.types import StringType, IntegerType, StructType, StructField
from pyspark.sql.functions import udf, col

class Example():
    def __init__(self):
        self.students = [[f'student_{i}', np.random.randint(80)] for i in range(3)]
        self.increase = 10
    
    def create_spark_df(self):
        cSchema = StructType([StructField("Name", StringType()),
                             StructField("Marks", IntegerType())])
        return spark.createDataFrame(self.students, schema=cSchema)

    # 定义实例方法，无需静态装饰器
    def add_increase_marks(self, marks):
        return marks + self.increase
    
    def calculate_new_marks(self):
        df = self.create_spark_df()
        # 将实例方法包装为UDF
        add_marks_udf = udf(self.add_increase_marks, IntegerType())
        df = df.withColumn("New Marks", add_marks_udf(col("Marks")))
        return df

c = Example()
c.calculate_new_marks().show()

注意：如果类包含不可序列化的变量，需要重写__getstate__和__setstate__方法，手动处理序列化逻辑。

方案3：将实例变量作为常量列加入DataFrame #

把实例变量作为常量列添加到DataFrame中，在UDF中同时引用原数据列和这个常量列。这种方式简单直接，适合变量值较小的场景。

import numpy as np
from pyspark.sql.types import StringType, IntegerType, StructType, StructField
from pyspark.sql.functions import udf, col, lit

class Example():
    def __init__(self):
        self.students = [[f'student_{i}', np.random.randint(80)] for i in range(3)]
        self.increase = 10
    
    def create_spark_df(self):
        cSchema = StructType([StructField("Name", StringType()),
                             StructField("Marks", IntegerType())])
        return spark.createDataFrame(self.students, schema=cSchema)

    @staticmethod
    @udf(returnType=IntegerType())
    def add_increase_marks(marks, increase_val):
        return marks + increase_val
    
    def calculate_new_marks(self):
        df = self.create_spark_df()
        # 添加存储实例变量的常量列
        df = df.withColumn("Increase", lit(self.increase))
        df = df.withColumn("New Marks", self.add_increase_marks(col("Marks"), col("Increase")))
        # 可选：删除临时常量列
        df = df.drop("Increase")
        return df

c = Example()
c.calculate_new_marks().show()

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