在Synapse中使用PySpark UDF时遇无法序列化_thread.RLock对象错误

问题 #

在Synapse中使用PySpark创建DataFrame，通过自定义函数（UDF）添加计算列时，想在函数异常块中记录错误信息，但调用logger.error时触发了TypeError: cannot pickle '_thread.RLock' object错误。

代码片段 #

logger = sc._jvm.org.apache.log4j.LogManager.getLogger("com.test.test")
def log_test(arg):
    try:
        return 1/arg
    except Exception as e:
        logger.error("Divided by zero error occurred")

udf_log_test = udf(lambda col:log_test(col) if col != 0 else 0, IntegerType())
    
df = df.withColumn('TestColumn', udf_log_test(lit(0)))

错误栈信息 #

Traceback (most recent call last):
      File "/databricks/spark/python/pyspark/serializers.py", line 527, in dumps
        return cloudpickle.dumps(obj, pickle_protocol)
      File "/databricks/spark/python/pyspark/cloudpickle/cloudpickle_fast.py", line 73, in dumps
        cp.dump(obj)
      File "/databricks/spark/python/pyspark/cloudpickle/cloudpickle_fast.py", line 602, in dump
        return Pickler.dump(self, obj)
    TypeError: cannot pickle '_thread.RLock' object
    PicklingError: Could not serialize object: TypeError: cannot pickle '_thread.RLock' object

TypeError                                 Traceback (most recent call last)
File /databricks/spark/python/pyspark/serializers.py:527, in CloudPickleSerializer.dumps(self, obj)
    526 try:
--> 527     return cloudpickle.dumps(obj, pickle_protocol)
    528 except pickle.PickleError:

File /databricks/spark/python/pyspark/cloudpickle/cloudpickle_fast.py:73, in dumps(obj, protocol, buffer_callback)
     70 cp = CloudPickler(
     71     file, protocol=protocol, buffer_callback=buffer_callback
     72 )
--> 73 cp.dump(obj)
     74 return file.getvalue()

File /databricks/spark/python/pyspark/cloudpickle/cloudpickle_fast.py:602, in CloudPickler.dump(self, obj)
    601 try:
--> 602     return Pickler.dump(self, obj)
    603 except RuntimeError as e:

TypeError: cannot pickle '_thread.RLock' object

During handling of the above exception, another exception occurred:

PicklingError                             Traceback (most recent call last)
File <command-109665656057519>:1
----> 1 df = df.withColumn('TestColumn', udf_log_test(lit(0)))

File /databricks/spark/python/pyspark/sql/udf.py:322, in UserDefinedFunction._wrapped.<locals>.wrapper(*args)
    320 @functools.wraps(self.func, assigned=assignments)
    321 def wrapper(*args: "ColumnOrName") -> Column:
--> 322     return self(*args)

File /databricks/spark/python/pyspark/sql/udf.py:299, in UserDefinedFunction.__call__(self, *cols)
    297         sc.profiler_collector.add_profiler(id, memory_profiler)
    298 else:
--> 299     judf = self._judf
    300     jPythonUDF = judf.apply(_to_seq(sc, cols, _to_java_column))
    301 return Column(jPythonUDF)

File /databricks/spark/python/pyspark/sql/udf.py:219, in UserDefinedFunction._judf(self)
    212 @property
    213 def _judf(self) -> JavaObject:
    214     # It is possible that concurrent access, to newly created UDF,
    215     # will initialize multiple UserDefinedPythonFunctions.
    216     # This is unlikely, doesn't affect correctness,
    217     # and should have a minimal performance impact.
    218     if self._judf_placeholder is None:
--> 219         self._judf_placeholder = self._create_judf(self.func)
    220     return self._judf_placeholder

File /databricks/spark/python/pyspark/sql/udf.py:228, in UserDefinedFunction._create_judf(self, func)
    225 spark = SparkSession._getActiveSessionOrCreate()
    226 sc = spark.sparkContext
--> 228 wrapped_func = _wrap_function(sc, func, self.returnType)
    229 jdt = spark._jsparkSession.parseDataType(self.returnType.json())
    230 assert sc._jvm is not None

File /databricks/spark/python/pyspark/sql/udf.py:51, in _wrap_function(sc, func, returnType)
     47 def _wrap_function(
     48     sc: "SparkContext", func: Callable[..., Any], returnType: "DataTypeOrString"
     49 ) -> JavaObject:
     50     command = (func, returnType)
--> 51     pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command)
     52     assert sc._jvm is not None
     53     return sc._jvm.SimplePythonFunction(
     54         bytearray(pickled_command),
     55         env,
   (...)
     60         sc._javaAccumulator,
     61     )

File /databricks/spark/python/pyspark/rdd.py:5248, in _prepare_for_python_RDD(sc, command)
   5245 def _prepare_for_python_RDD(sc: "SparkContext", command: Any) -> Tuple[bytes, Any, Any, Any]:
   5246     # the serialized command will be compressed by broadcast
   5247     ser = CloudPickleSerializer()
-> 5248     pickled_command = ser.dumps(command)
   5249     assert sc._jvm is not None
   5250     if len(pickled_command) > sc._jvm.PythonUtils.getBroadcastThreshold(sc._jsc):  # Default 1M
   5251         # The broadcast will have same life cycle as created PythonRDD

File /databricks/spark/python/pyspark/serializers.py:537, in CloudPickleSerializer.dumps(self, obj)
    535     msg = "Could not serialize object: %s: %s" % (e.__class__.__name__, emsg)
    536 print_exec(sys.stderr)
-> 537 raise pickle.PicklingError(msg)

PicklingError: Could not serialize object: TypeError: cannot pickle '_thread.RLock' object

解决方案 #

问题原因 #

PySpark的UDF需要序列化后分发到集群的各个Executor节点执行，但你在Driver端创建的Log4j logger对象包含_thread.RLock（线程锁），这个对象无法被序列化，导致UDF在序列化时失败。

解决方法1：在UDF内部延迟初始化Logger #

不在Driver端提前创建logger，而是在UDF内部每次执行时初始化，这样每个Executor节点会自己创建独立的logger实例，避免序列化问题：

from pyspark.sql.functions import udf, lit
from pyspark.sql.types import IntegerType

def log_test(arg):
    try:
        return 1 / arg
    except Exception as e:
        # 延迟初始化logger，确保在Executor端创建
        logger = sc._jvm.org.apache.log4j.LogManager.getLogger("com.test.test")
        logger.error(f"Divided by zero error occurred: {str(e)}")
        return 0  # 必须返回对应类型的默认值，避免类型不匹配

udf_log_test = udf(log_test, IntegerType())

df = df.withColumn('TestColumn', udf_log_test(lit(0)))

解决方法2：使用Python标准logging模块 #

改用Python自带的logging模块，该模块的logger对象可以正常序列化，日志会输出到Executor的本地日志中：

import logging
from pyspark.sql.functions import udf, lit
from pyspark.sql.types import IntegerType

# 在Driver端配置日志级别
logging.basicConfig(level=logging.ERROR)
logger = logging.getLogger("com.test.test")

def log_test(arg):
    try:
        return 1 / arg
    except Exception as e:
        logger.error(f"Divided by zero error occurred: {str(e)}")
        return 0

udf_log_test = udf(log_test, IntegerType())

df = df.withColumn('TestColumn', udf_log_test(lit(0)))

注意事项 #

• 方法1的Log4j日志会集成到Spark的统一日志系统中，日志可以通过Spark的日志收集工具查看；
• 方法2的Python日志需要确保Executor节点的日志配置正确，默认会输出到Executor的本地日志文件。

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