PySpark中通过While Loop迭代关联实现层级描述补全

原始PySpark DataFrame定义及展示：

data = [('service 1','s1', 's2'),
       ('service 2','s2', 's4'),
       ('service 3','s3', 's5'),
       ('service 5','s5', 's6'),
       ('service 4','s4', 's3')]
sdf = spark.createDataFrame(data, schema = ['description', 'service', 'parent'])
sdf.show()

输出：

+-----------+-------+--------+
|description|service|parent  |
+-----------+-------+--------+
|  service 1|     s1|      s2|
|  service 2|     s2|      s4|
|  service 3|     s3|      s5|
|  service 5|     s5|      s6|
|  service 4|     s4|      s3|
|  service 6|     s6|    NULL|
+-----------+-------+--------+

需求 #

判断parent列的值是否存在于service列中，添加对应description的新列；最终删除service和原始parent列，期望结果如下：

+-----------+-----------+-----------+-----------+-----------+-----------+
|description|     parent|    parent1|    parent2|    parent3|    parent4|
+-----------+-----------+-----------+-----------+-----------+-----------+
|  service 1|  service 2|  service 4|  service 3|  service 5|  service 6|
|  service 2|  service 4|  service 3|  service 5|  service 6|       NULL|
|  service 3|  service 5|  service 6|       NULL|       NULL|       NULL|
|  service 5|  service 6|       NULL|       NULL|       NULL|       NULL|
|  service 4|  service 3|  service 5|  service 6|       NULL|       NULL|
+-----------+-----------+-----------+-----------+-----------+-----------+

当前问题 #

已实现基于service和parent列的循环关联，但生成的parent列仅保存service值，无法替换为对应的description。当前执行代码及最终输出如下：

执行代码：

i = 0
while(sdf.filter(F.col(f"parent{i if i>0 else ''}").isNotNull()).count() > 0):

  sdf = sdf.alias("a1").join(sdf.alias("a2").select("service", "parent"),
                            F.col(f"a1.parent{i if i>0 else ''}")==F.col("a2.service"), 
                            how="left") \
           .withColumn(f"parent{i+1}", F.col("a2.parent")) \
          .drop(F.col("a2.service")) \
          .drop(F.col("a2.parent")) 
  i += 1
  display(sdf)

最终输出：

+-----------+-------+------+-------+-------+-------+-------+-------+
|description|service|parent|parent1|parent2|parent3|parent4|parent5|
+-----------+-------+------+-------+-------+-------+-------+-------+
|  service 1|     s1|    s2|     s4|     s3|     s5|     s6|   null|
|  service 2|     s2|    s4|     s3|     s5|     s6|   null|   null|
|  service 3|     s3|    s5|     s6|   null|   null|   null|   null|
|  service 5|     s5|    s6|   null|   null|   null|   null|   null|
|  service 4|     s4|    s3|     s5|     s6|   null|   null|
+-----------+-------+------+-------+-------+-------+-------+-------+

方案1：循环关联+映射替换 #

核心思路是先建立service到description的映射，初始时将原始parent列替换为对应描述，后续每轮循环都通过当前parent的描述关联到对应记录，获取下一级parent的描述。

from pyspark.sql import functions as F

# 构建service与description的映射表
service_desc_map = sdf.select("service", "description").alias("map")

# 初始处理：替换原始parent为对应description，删除原始service列
sdf_processed = sdf.alias("main").join(
    service_desc_map,
    F.col("main.parent") == F.col("map.service"),
    "left"
).withColumnRenamed("map.description", "parent") \
 .drop("map.service", "service")

i = 0
while sdf_processed.filter(F.col(f"parent{i if i>0 else ''}").isNotNull()).count() > 0:
    current_col = f"parent{i if i>0 else ''}"
    next_col = f"parent{i+1}"
    
    # 关联原表，通过当前parent的description找到对应service记录，获取其parent的description
    sdf_processed = sdf_processed.alias("a1").join(
        sdf.alias("a2").select("description", "parent"),
        F.col(f"a1.{current_col}") == F.col("a2.description"),
        "left"
    ).join(
        service_desc_map,
        F.col("a2.parent") == F.col("map.service"),
        "left"
    ).withColumn(next_col, F.col("map.description")) \
     .drop("a2.description", "a2.parent", "map.service", "map.description")
    
    i += 1

# 展示最终结果
sdf_processed.show()

方案2：递归CTE构建层级链（更高效） #

利用Spark SQL的递归CTE直接构建每个节点向上的完整父级描述链，再将链展开为多列，避免循环关联的开销。

from pyspark.sql import functions as F

# 使用递归CTE生成每个节点的父级描述链
hierarchy_df = spark.sql("""
    WITH RECURSIVE node_hierarchy AS (
        -- 起始节点：所有有父节点的记录，初始父链为对应父节点的description
        SELECT 
            s.description AS root_desc,
            s.parent AS current_parent_service,
            ARRAY(COALESCE((SELECT description FROM sdf WHERE service = s.parent), NULL)) AS parent_descs
        FROM sdf s
        WHERE s.parent IS NOT NULL
        UNION ALL
        -- 递归遍历：向上查找父节点的父节点
        SELECT 
            nh.root_desc,
            s.parent AS current_parent_service,
            ARRAY_APPEND(nh.parent_descs, COALESCE((SELECT description FROM sdf WHERE service = s.parent), NULL)) AS parent_descs
        FROM node_hierarchy nh
        JOIN sdf s ON nh.current_parent_service = s.service
        WHERE s.parent IS NOT NULL
    )
    -- 合并无父节点的记录（父链为空数组）
    SELECT root_desc AS description, parent_descs FROM node_hierarchy
    UNION ALL
    SELECT description, ARRAY() AS parent_descs FROM sdf WHERE parent IS NULL
    -- 去重，保留每个节点的最长父链
    GROUP BY description, parent_descs
    ORDER BY description
""")

# 获取最大层级深度，用于展开数组为列
max_depth = hierarchy_df.select(F.size(F.col("parent_descs"))).agg(F.max("size(parent_descs)")).collect()[0][0]

# 将父链数组展开为多个parent列，不足长度补NULL
result_df = hierarchy_df.select(
    "description",
    *[F.coalesce(F.col("parent_descs")[i], F.lit(None)).alias(f"parent{i if i>0 else ''}") for i in range(max_depth)]
)

result_df.show()

方案说明 #

• 方案1适合层级较浅的场景，逻辑直观，容易调试；
• 方案2通过递归CTE一次性构建完整层级链，性能更优，适合层级较深的数据集。

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