PySpark调用无Spark依赖Python模块的机制与优化问询

现有项目代码 #

Project A #

class FeatureBuilder:
    def __init__(self):
        self.artifact = read_artifacts_from_s3()

    def create_features(self, pdf):
        # 基于artifact和输入pdf生成特征
        pass

Project B #

from pyspark.sql import DataFrame
from builder import FeatureBuilder

def pandas_udf(df: DataFrame):
    feature_builder = FeatureBuilder()
    def create_features(pdf):
        # 注：原代码变量名笔误，应为feature_builder
        feature_vector = feature_builder.create_features(pdf)
        return feature_vector

    # 注：原代码缺少返回schema参数，实际使用需补充
    return df.groupby("id").applyInPandas(create_features, schema=...)

问题解答 #

1. 集群中的每台机器都会从S3读取该文件吗？ #

是的，每台执行Task的Worker节点都会重复读取S3文件，甚至同一Worker上的多个Task都会各自触发一次读取。原因是：

• FeatureBuilder的初始化逻辑被包含在Pandas UDF的代码路径中，PySpark会将UDF相关代码序列化后分发到各个Worker节点；
• 每个Task执行时都会反序列化代码并初始化FeatureBuilder，进而触发read_artifacts_from_s3()调用。

2. 若可修改Project A（但不能添加Spark代码），如何优化？ #

核心优化思路是将S3文件的读取限制在Driver端仅执行一次，再通过广播将数据分发到所有Worker节点，避免重复读取。以下是两种可行方案：

方案一：修改FeatureBuilder支持传入预加载的artifact

修改Project A的代码，允许在初始化时传入已加载的artifact，而非在构造函数中硬编码读取逻辑（兼容原有使用方式）：

class FeatureBuilder:
    def __init__(self, artifact=None):
        self.artifact = artifact if artifact is not None else read_artifacts_from_s3()

    def create_features(self, pdf):
        # 基于artifact和输入pdf生成特征
        pass

然后在Project B的Driver端读取一次artifact并广播，在UDF中使用广播的artifact初始化FeatureBuilder：

from pyspark.sql import DataFrame
from builder import FeatureBuilder

def process_data(df: DataFrame):
    # Driver端仅读取一次S3文件
    artifact = read_artifacts_from_s3()
    # 将artifact广播到所有Worker节点
    broadcast_artifact = spark.sparkContext.broadcast(artifact)

    def pandas_udf_func(pdf):
        # 使用广播的artifact初始化FeatureBuilder，避免重复读取S3
        fb = FeatureBuilder(artifact=broadcast_artifact.value)
        return fb.create_features(pdf)

    return df.groupby("id").applyInPandas(pandas_udf_func, schema=...)

方案二：广播已初始化的FeatureBuilder对象

如果FeatureBuilder实例是可序列化的（即artifact是可序列化类型），可以直接在Driver端初始化实例并广播，进一步减少Worker端的初始化操作：

from pyspark.sql import DataFrame
from builder import FeatureBuilder

def process_data(df: DataFrame):
    # Driver端初始化FeatureBuilder，仅读取一次S3
    fb = FeatureBuilder()
    # 广播FeatureBuilder实例到所有Worker
    broadcast_fb = spark.sparkContext.broadcast(fb)

    def pandas_udf_func(pdf):
        # 获取广播的实例并调用方法
        return broadcast_fb.value.create_features(pdf)

    return df.groupby("id").applyInPandas(pandas_udf_func, schema=...)

注意：若artifact体积较大，优先选择方案一广播原始数据，避免广播整个对象带来的额外内存开销。

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