PySpark实现多父层级节点扁平化生成node_map方案问询

针对你提到的需要生成包含节点及其所有父节点的node_map需求，这里有个非常简洁的实现方式，不需要繁琐的多次自连接，而是利用**Spark SQL的递归CTE（Common Table Expression）**来完成层级遍历，刚好适配你最多9级父节点的场景。

核心思路 #

1. 先定位叶子节点：也就是那些从未作为其他节点父节点的节点——这些是我们最终需要生成node_map的目标节点。
2. 递归遍历层级：从每个叶子节点出发，向上递归收集所有祖先节点（包括自身），直到根节点（parent_node为NULL）。
3. 聚合路径为列表：将每个叶子节点的所有路径节点聚合为一个列表，得到期望的node_map。

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list

# 初始化SparkSession
spark = SparkSession.builder.appName("NodeMapGenerator").getOrCreate()

# 构造输入DataFrame（替换成你的实际数据）
data = [
    ("pl3182020828", None),
    ("pl3182020830", "pl3182020828"),
    ("pl3182020829", "pl3182020830"),
    ("l66434", "pl3182020829"),
    ("l66435", "pl3182020829"),
    ("pl3182020832", None),
    ("pl3182020831", "pl3182020832"),
    ("l66900", "pl3182020831"),
    ("pl3182020833", "pl3182020832"),
    ("l66916", "pl3182020833")
]
df = spark.createDataFrame(data, ["node", "parent_node"])
# 将DataFrame注册为临时表，方便SQL操作
df.createOrReplaceTempView("df")

# 第一步：识别所有叶子节点（不在parent_node列中的node）
spark.sql("CREATE OR REPLACE TEMP VIEW leaf_nodes AS SELECT node FROM df EXCEPT ALL SELECT parent_node FROM df WHERE parent_node IS NOT NULL")

# 第二步：用递归CTE遍历并收集路径
result_df = spark.sql("""
    WITH RECURSIVE node_path AS (
        -- 初始层：叶子节点，路径初始化为自身
        SELECT 
            node AS original_node,
            node AS path_node,
            parent_node,
            1 AS depth
        FROM df
        WHERE node IN (SELECT node FROM leaf_nodes)
        
        UNION ALL
        
        -- 递归层：向上连接父节点，直到parent_node为NULL
        SELECT 
            np.original_node,
            df.node AS path_node,
            df.parent_node,
            np.depth + 1 AS depth
        FROM node_path np
        JOIN df ON np.parent_node = df.node
    )
    -- 聚合路径为列表，顺序是自身→父节点→根节点
    SELECT 
        original_node AS node,
        collect_list(path_node) AS node_map
    FROM node_path
    GROUP BY original_node
    ORDER BY original_node
""")

# 查看结果
result_df.show(truncate=False)

代码说明 #

• 递归CTE工作原理：
  • 初始阶段从叶子节点开始，把节点自身加入路径，记录当前深度为1。
  • 递归阶段每次用当前记录的parent_node关联原表，获取父节点信息，深度加1；当parent_node为NULL时，因为无法匹配到原表的节点，递归自动终止，完美适配你的层级上限。
• 聚合结果：collect_list会按照递归遍历的顺序收集节点（自身→直接父→...→根），刚好符合你期望的输出格式。

替代方案（GraphFrames） #

如果你的环境允许安装graphframes库，也可以用图遍历实现，但需要额外依赖：

from graphframes import GraphFrame

# 构建图结构
vertices = df.select("node").distinct()
edges = df.filter(df.parent_node.isNotNull()).select("parent_node", "node").withColumnRenamed("parent_node", "src").withColumnRenamed("node", "dst")
g = GraphFrame(vertices, edges)

# 从叶子节点出发，向上遍历到根节点
bfs_result = g.bfs("node IN (SELECT node FROM leaf_nodes)", "parent_node IS NULL")
# 后续可通过提取路径列并转换为列表得到node_map，不过递归CTE的方式更直接轻量

递归CTE的方式不需要额外依赖，代码简洁且性能稳定，完全满足你的需求。

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