如何高效递归遍历PySpark DataFrame？性能优化方案求助

问题背景 #

原始PySpark DataFrame结构如下：

需求：针对每个Material，递归展开其所有底层组件，新增对应行，并按公式RequiredQuantity / BatchSize * Parent RequiredQuantity计算新行的RequiredQuantity，最终得到目标结构的DataFrame（示例见问题描述）。

现有递归实现性能极差：单个Material处理约5分钟，面对5000种Material、每种约10个Component的场景，总耗时无法接受。原始代码如下：

def recurse_components(df, material):
    if df.isEmpty():
        return df

    filtered_material = df.where(F.col("Material") == material)
    batch_size = filtered_material.select("BatchSize").first()["BatchSize"]

    component_list = (
        filtered_material.select("Component").rdd.flatMap(lambda x: x).collect()
    )

    for component in component_list:
        component_table = df.where(F.col("Material") == component)
        if not component_table.isEmpty():
            required_quantity = (
                filtered_material.where(F.col("Component") == component)
                .select("RequiredQuantity")
                .first()["RequiredQuantity"]
            )
            recursive_call = recurse_components(df, component).withColumns(
                {
                    "Material": F.lit(material),
                    "RequiredQuantity": F.col("RequiredQuantity")
                    * required_quantity
                    / F.col("BatchSize"),
                    "BatchSize": F.lit(batch_size),
                }
            )
            filtered_material = filtered_material.union(recursive_call)

    return filtered_material

material_list = df.select("Material").distinct().rdd.flatMap(lambda x: x).collect()

extended_df = spark.createDataFrame([], df.schema)
for material in material_list:
    extended_df = extended_df.union(recurse_components(df, material))

性能瓶颈分析 #

• 频繁Driver-Job交互：多次调用collect()、first()，将分布式数据拉取到Driver端，产生大量IO开销，破坏Spark分布式计算特性。
• 串行处理+多次Union：遍历每个Material串行处理，递归中反复union导致DataFrame lineage冗长，Spark优化器无法有效优化。
• 递归逻辑局限性：Python递归无法利用Spark并行计算能力，单线程处理完全浪费集群资源。

高效实现方案：迭代式Join（无需额外依赖） #

核心思路 #

1. 为每个根Material保留基准BatchSize，初始化**乘数因子multiplier**为1.0（代表从根节点到当前节点的累积计算系数）。
2. 迭代关联原始DataFrame，展开所有可关联组件，更新multiplier为父multiplier * 父RequiredQuantity / 父BatchSize。
3. 直到无新组件可展开，最终计算RequiredQuantity = 原始RequiredQuantity * multiplier，统一使用根Material的BatchSize。

代码实现 #

from pyspark.sql import functions as F

# 1. 预处理：获取每个Material的基准BatchSize
material_batch_map = df.select("Material", "BatchSize").distinct() \
    .withColumnRenamed("BatchSize", "RootBatchSize")

# 初始化DataFrame：保留根Material信息，初始化multiplier为1.0
initial_df = df.join(material_batch_map, on="Material", how="inner") \
    .withColumn("RootMaterial", F.col("Material")) \
    .withColumn("multiplier", F.lit(1.0))

current_df = initial_df

# 迭代展开组件，直到没有新数据
while True:
    # 关联原始数据，找到可展开的组件（当前Component是某个Material）
    expandable = current_df.join(
        df,
        current_df["Component"] == df["Material"],
        how="inner"
    ).select(
        current_df["RootMaterial"],
        current_df["RootBatchSize"],
        df["Component"].alias("ChildComponent"),
        df["RequiredQuantity"].alias("ChildRequired"),
        df["BatchSize"].alias("ChildBatch"),
        # 计算新的乘数因子
        current_df["multiplier"] * current_df["RequiredQuantity"] / current_df["BatchSize"]
    ).withColumnRenamed("(multiplier * RequiredQuantity) / BatchSize", "new_multiplier")

    if expandable.isEmpty():
        break

    # 将展开的行转换为目标格式
    expanded_rows = expandable.select(
        F.col("RootMaterial").alias("Material"),
        F.col("ChildComponent").alias("Component"),
        F.col("RootBatchSize").alias("BatchSize"),
        (F.col("ChildRequired") * F.col("new_multiplier")).alias("RequiredQuantity"),
        F.col("RootMaterial"),
        F.col("RootBatchSize"),
        F.col("new_multiplier").alias("multiplier")
    )

    # 合并数据并去重（避免重复处理，防止潜在循环依赖）
    current_df = current_df.union(expanded_rows).dropDuplicates(["Material", "Component"])

# 整理最终结果，只保留目标列
result_df = current_df.select("Material", "Component", "BatchSize", "RequiredQuantity") \
    .orderBy("Material", "Component")

result_df.show(truncate=False)

方案优势 #

• 全分布式计算：无Driver端数据拉取，充分利用集群资源。
• 迭代替代递归：避免Python递归单线程瓶颈，Spark优化器高效处理Join操作。
• 批量处理所有Material：无需串行遍历，一次性处理所有根节点，大幅提升速度。

额外优化建议 #

• 处理循环依赖：若业务存在组件循环引用（如A包含B，B包含A），需在迭代中加入判断逻辑避免无限循环。
• 优化Join性能：为Material和Component列建立索引，或调整spark.sql.shuffle.partitions适配集群资源。
• 内存优化：数据量极大时，可将中间结果写入磁盘（如Parquet），避免内存溢出。

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