能否在Spark UDF中创建/共享Spark Session？代码报错求助

问题背景 #

需要基于主表数据抽取子表数据，尝试两种方案均报错：

1. 在Spark UDF中直接使用Driver端的Spark Session，触发序列化错误（Spark Session不可序列化，无法传递到Executor节点）
2. 在UDF内创建Spark Session，触发异常：Exception: SparkContext should only be created and accessed on the driver.

原始UDF代码 #

def select(entity):
    query = f"SELECT * FROM `{database.value}`.`{table.value}`  WHERE id='{entity}'"
    records = spark.sql(query)
    # store records on S3 in CSVformat, filename table_name.csv
    return records
    
ingestion = F.udf(select, ArrayType(StringType()))

try:
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    query = f"SELECT * FROM `{database}`.`{entity_table}`"        
    entities = spark.sql(query)
    # store records on S3 in CSV format, filename entities.csv
    
    tables_with_relations = ["entity.some_child_table", "another_child_table"]
    for child_table in tables_with_relations:
        table = spark.sparkContext.broadcast(child_table)
        response = entities.withColumn("response", ingestion("id"))
        response.show()
            
except Exception as e:
    raise e

修改后UDF代码及报错 #

def select(entity):
    query = f"SELECT * FROM `{database.value}`.`{table.value}` WHERE id='{entity}'"
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    records = spark.sql(query)
    return records

报错信息：

Exception: SparkContext should only be created and accessed on the driver.

错误核心原因 #

Spark的UDF运行在Executor节点，而Spark Session/Spark Context是Driver端专属对象：

• 它们无法被序列化，不能传递到Executor
• 也不允许在Executor端创建，违反Spark架构设计
• 你试图在UDF里执行Spark SQL，本质是在Executor端发起新的Spark任务，这完全不符合Spark的分布式运行逻辑

可行解决方案 #

方案1：使用Spark Join（推荐，最符合Spark分布式特性） #

直接通过表关联获取子表数据，这是Spark处理此类场景的标准做法，性能最优。

try:
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    # 读取主表数据（仅保留id用于关联）
    entities = spark.sql(f"SELECT id FROM `{database}`.`{entity_table}`")
    # 保存主表到S3
    entities.write.csv("s3://your-target-path/entities.csv", header=True)
    
    tables_with_relations = ["entity.some_child_table", "another_child_table"]
    for child_table in tables_with_relations:
        # 读取子表全量数据
        child_df = spark.sql(f"SELECT * FROM `{database}`.`{child_table}`")
        # 关联主表，仅保留主表中存在的id对应的子表数据
        joined_df = entities.join(child_df, on="id", how="inner")
        # 保存子表结果到S3，文件名替换特殊字符避免问题
        output_path = f"s3://your-target-path/{child_table.replace('.', '_')}.csv"
        joined_df.write.csv(output_path, header=True)
        # 查看结果
        joined_df.show()
            
except Exception as e:
    raise e

方案2：批量收集主表ID后查询子表 #

如果子表数据量极大，直接Join内存压力大，可先收集主表ID到Driver端，再批量查询子表（注意：主表数据量不能过大，否则会导致Driver内存溢出）。

try:
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    entities = spark.sql(f"SELECT id FROM `{database}`.`{entity_table}`")
    entities.write.csv("s3://your-target-path/entities.csv", header=True)
    
    # 收集主表id到Driver端
    entity_ids = [row.id for row in entities.collect()]
    # 转换为SQL IN子句支持的格式
    ids_str = ",".join([f"'{id}'" for id in entity_ids])
    
    tables_with_relations = ["entity.some_child_table", "another_child_table"]
    for child_table in tables_with_relations:
        # 批量查询子表中符合条件的数据
        child_df = spark.sql(f"SELECT * FROM `{database}`.`{child_table}` WHERE id IN ({ids_str})")
        output_path = f"s3://your-target-path/{child_table.replace('.', '_')}.csv"
        child_df.write.csv(output_path, header=True)
        child_df.show()
            
except Exception as e:
    raise e

方案3：广播子表后用UDF过滤（仅特殊场景使用，不推荐） #

如果业务逻辑必须在UDF内处理，可先将子表全量广播到Executor，再在UDF中根据ID过滤数据（注意：子表数据量不能过大，否则广播会占用过多Executor内存）。

from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType
import json

try:
    glueContext = GlueContext(SparkContext.getOrCreate())
    spark = glueContext.spark_session
    entities = spark.sql(f"SELECT * FROM `{database}`.`{entity_table}`")
    entities.write.csv("s3://your-target-path/entities.csv", header=True)
    
    tables_with_relations = ["entity.some_child_table", "another_child_table"]
    for child_table in tables_with_relations:
        # 读取子表并转为字典（id为key，整行数据转为JSON字符串）
        child_df = spark.sql(f"SELECT * FROM `{database}`.`{child_table}`")
        child_dict = {row.id: json.dumps(row.asDict()) for row in child_df.collect()}
        # 广播子表字典到所有Executor
        broadcast_child = spark.sparkContext.broadcast(child_dict)
        
        # 定义UDF：根据id从广播字典中获取对应数据
        def get_child_data(entity_id):
            return broadcast_child.value.get(entity_id, None)
        
        ingestion_udf = udf(get_child_data, StringType())
        # 关联数据
        response = entities.withColumn("response", ingestion_udf(col("id")))
        response.show()
        
        # 保存结果到S3
        output_path = f"s3://your-target-path/{child_table.replace('.', '_')}_with_response.csv"
        response.write.csv(output_path, header=True)
            
except Exception as e:
    raise e

关键注意事项 #

• 永远不要在UDF中创建或使用Spark Session/Spark Context，这是Spark架构的基本规则
• 优先使用Join操作，能充分利用集群的分布式计算能力，性能最优
• 批量收集ID和广播子表的方式仅适合对应表数据量较小的场景，避免内存溢出问题

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