导入模块中UDF时SparkContext报错：仅能在Driver端创建访问

项目结构 #

dependencies/
└── spark.py
etl.py
shared/
└── tools.py

代码实现 #

dependencies/spark.py #

from pyspark.sql import SparkSession
def get_or_create_session(app_name, master="local[*]"):
    spark_builder = SparkSession.builder.master(master).appName(app_name)
    session = spark_builder.getOrCreate()
    return session

etl.py #

from dependencies.spark import get_or_create_session
from shared.tools import cleanup_pob_column

def main():
    spark = get_or_create_session(app_name="my_app")
    data = get_data(input_file)
    transformed_data = transform_data(data)
    transformed_data.printSchema()
    transformed_data.show(truncate=False)

def get_data(input_file):
    # 数据读取逻辑
    return data

def transform_data(data):
    return (
       data
       .transform(cleanup_pob_column)
)

if __name__ == "__main__":
    main()

shared/tools.py #

# 原代码缺失必要导入，需补充
from pyspark.sql import functions as F
import pandas as pd

def extract_iso(x):
    # 提取ISO字符串的逻辑
    return iso_string

@F.pandas_udf("string")
def cleanup_geo_column_udf(col: pd.Series) -> pd.Series:
    return col.apply(lambda x: extract_iso(x=x))

def cleanup_pob_column(df):
    return df.withColumn("pob_cln", cleanup_geo_column_udf(F.col("place_of_birth")))

报错循环 #

情况1：不在shared/tools.py添加会话创建代码 #

触发AssertionError，原因是SparkContext未初始化：

Caused by: org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/Users/gsimeone/PycharmProjects/assignment/shared/geographic_tools.py", line 39, in <module>
    def cleanup_geo_column_udf(col: pd.Series) -> pd.Series:
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/pandas/functions.py", line 450, in _create_pandas_udf
    return _create_udf(f, returnType, evalType)
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/udf.py", line 74, in _create_udf
    return udf_obj._wrapped()
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/udf.py", line 286, in _wrapped
    wrapper.returnType = self.returnType  # type: ignore[attr-defined]
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/udf.py", line 134, in returnType
    self._returnType_placeholder = _parse_datatype_string(self._returnType)
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/types.py", line 1010, in _parse_datatype_string
    assert sc is not None
AssertionError

情况2：在shared/tools.py添加会话创建代码 #

添加以下代码到shared/tools.py顶部：

from dependencies.spark import get_or_create_session
spark = get_or_create_session(app_name="my_app")

触发RuntimeError，违反SparkContext只能在driver端创建的规则：

Caused by: org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/Users/gsimeone/PycharmProjects/assignment/shared/geographic_tools.py", line 15, in <module>
    spark = get_or_create_session(app_name=config.get("app_name"))
  File "/Users/gsimeone/PycharmProjects/assignment/dependencies/spark.py", line 22, in get_or_create_session
    session = spark_builder.getOrCreate()
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/sql/session.py", line 277, in getOrCreate
    return session
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/context.py", line 485, in getOrCreate
    return SparkContext._active_spark_context
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/context.py", line 186, in __init__
    SparkContext._assert_on_driver()
  File "/Users/gsimeone/PycharmProjects/sayaritest/sayari_test/lib/python3.8/site-packages/pyspark/python/lib/pyspark.zip/pyspark/context.py", line 1533, in _assert_on_driver
    raise RuntimeError("SparkContext should only be created and accessed on the driver.")
RuntimeError: SparkContext should only be created and accessed on the driver.

临时可行方案 #

将shared/tools.py的全部内容粘贴到etl.py中，应用可正常运行。

根本解决方案及原因分析 #

问题根源 #

Pandas UDF的装饰器@F.pandas_udf在模块加载时会立即执行返回类型解析，这个过程依赖已存在的SparkContext（sc）：

• 若在SparkSession创建前导入shared.tools模块，此时sc未初始化，触发AssertionError。
• 若在shared.tools的顶级代码中创建SparkSession，当任务分发到worker节点时，worker会加载该模块并执行顶级代码，尝试在worker端创建SparkContext，违反Spark架构规则，触发RuntimeError。

正确解决方案 #

修改etl.py，将shared.tools的导入语句移到main()函数内部，确保SparkSession创建完成后再加载包含UDF的模块：

# 修改后的etl.py
from dependencies.spark import get_or_create_session

def main():
    # 先创建SparkSession
    spark = get_or_create_session(app_name="my_app")
    # 再导入包含UDF的模块
    from shared.tools import cleanup_pob_column
    
    data = get_data(input_file)
    transformed_data = transform_data(data)
    transformed_data.printSchema()
    transformed_data.show(truncate=False)

# 其余函数保持不变
def get_data(input_file):
    ... 
    return data

def transform_data(data):
    return (
       data
       .transform(cleanup_pob_column)
)

if __name__ == "__main__":
    main()

同时确保shared/tools.py补充必要的导入（如from pyspark.sql import functions as F和import pandas as pd）。

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