PySpark中使用Jinja2模板生成列时序列化报错的解决方法

问题1：该错误是否与序列化无关？ #

这个错误完全和序列化相关，没有其他可能性。

PySpark执行UDF时，会将Driver端定义的依赖对象（包括你在外部创建的Template实例）序列化后分发到各个Worker节点。而Jinja2的Template类无法被Spark使用的cloudpickle（或其他序列化器）正确序列化和反序列化——报错中的TypeError: Template.__new__() missing 1 required positional argument: 'source'，本质是Worker节点反序列化时无法重建Template实例，因为序列化过程中丢失了创建实例必需的source参数。尝试更换序列化器无法解决问题，因为Template本身的结构不支持跨节点的序列化传递。

问题2：如何修复此问题，在PySpark中使用Jinja2？ #

核心思路是避免将Driver端创建的Template实例序列化到Worker节点，改为在Worker节点本地创建Template实例。以下是两种可行方案：

方案1：在UDF内部创建Template实例（简单易实现） #

将Template的创建逻辑移到UDF内部，让每个Worker在处理数据时自行创建实例，绕开序列化问题。同时注意修复原代码中udf_foo未返回渲染结果的问题：

from jinja2 import Template
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

# 定义模板字符串（字符串可正常序列化）
TEMPLATE = """
Hello {{ customize(name) }}!
"""

def customize(name):
    return name + "san"

def udf_foo(name):
    # 在UDF内部创建Template实例
    template = Template(source=TEMPLATE)
    template.globals["customize"] = customize
    # 返回渲染后的结果
    return template.render(name=name)

convertUDF = udf(udf_foo, StringType())

# 生成新列
df1 = df.select(col("name")).withColumn("new_name", convertUDF(col("name")))

方案2：使用广播变量+Worker端懒加载（性能更优） #

通过广播变量传递模板字符串（字符串序列化无问题），并在Worker节点第一次调用UDF时初始化Template实例，后续复用该实例，避免重复创建：

from jinja2 import Template
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType
from pyspark import SparkContext

# 广播模板字符串到所有Worker节点
broadcast_template = SparkContext.getOrCreate().broadcast(TEMPLATE)

def customize(name):
    return name + "san"

# 全局变量存储Worker端的Template实例，每个Worker仅初始化一次
worker_template = None

def udf_foo(name):
    global worker_template
    if worker_template is None:
        # 从广播变量获取模板字符串，创建Template实例
        worker_template = Template(source=broadcast_template.value)
        worker_template.globals["customize"] = customize
    return worker_template.render(name=name)

convertUDF = udf(udf_foo, StringType())

# 生成新列
df1 = df.select(col("name")).withColumn("new_name", convertUDF(col("name")))

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