PySpark动态扁平化层级数据表（Databricks无递归CTE场景）

问题背景 #

现有PySpark DataFrame结构如下：

HierarchyNode       ParentNode
 U1.1               U1
 U1.1.1             U1.1
 24.0.1             24
 24.1.1             24.0.1
 24.1.1.1           24.1.1

这是一张层级关联表，需要将其扁平化处理，预期输出如下：

ParentWBSElementExternalID  LN-1            LN-2              LN-3
 U1                         NULL            NULL              NULL
 U1                         U1.1            NULL              NULL
 U1                         U1.1            U1.1.1            NULL
 24                         NULL            NULL              NULL
 24                         24.0.1          NULL              NULL
 24                         24.0.1          24.1.1            NULL
 24                         24.0.1          24.1.1        24.1.1.1

需满足的限制：

• Databricks环境不支持递归CTE
• 层级深度为动态（可能是2层、3层或更多）

实现思路 #

核心通过迭代遍历层级关系构建完整路径，再拆分路径生成多列，最后展开所有中间层级记录：

1. 识别所有根节点（即未出现在HierarchyNode中的ParentNode）
2. 迭代拼接每个节点的完整路径（从根节点到当前节点）
3. 拆分路径为对应层级的列
4. 展开所有中间层级的记录，补全NULL值

代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, array, explode_outer, size, expr

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

# 创建示例数据
data = [
    ("U1.1", "U1"),
    ("U1.1.1", "U1.1"),
    ("24.0.1", "24"),
    ("24.1.1", "24.0.1"),
    ("24.1.1.1", "24.1.1")
]
df = spark.createDataFrame(data, ["HierarchyNode", "ParentNode"])

# 步骤1：识别根节点
root_nodes = df.select("ParentNode").subtract(df.select("HierarchyNode")).withColumnRenamed("ParentNode", "HierarchyNode")

# 步骤2：迭代构建完整路径
path_df = root_nodes.withColumn("path", array(col("HierarchyNode")))

while True:
    # 关联原表获取子节点，拼接新路径
    new_path_df = path_df.join(df, path_df["HierarchyNode"] == df["ParentNode"], "inner") \
        .withColumn("new_path", expr("array_union(path, array(HierarchyNode))")) \
        .select(df["HierarchyNode"].alias("HierarchyNode"), "new_path") \
        .withColumnRenamed("new_path", "path")
    
    if new_path_df.count() == 0:
        break
    
    # 合并路径数据并去重
    path_df = path_df.union(new_path_df).dropDuplicates(["HierarchyNode"])

# 步骤3：生成所有中间层级的路径
expanded_df = path_df.withColumn("path_length", size(col("path"))) \
    .withColumn("index", explode_outer(expr("sequence(1, path_length)"))) \
    .withColumn("partial_path", expr("slice(path, 1, index)"))

# 步骤4：动态生成层级列，补全NULL值
max_level = expanded_df.selectExpr("max(size(partial_path))").first()[0]
final_columns = []

# 生成根节点列
final_columns.append(expr("partial_path[0]").alias("ParentWBSElementExternalID"))

# 生成LN-n列
for i in range(1, max_level):
    col_name = f"LN-{i}"
    final_columns.append(expr(f"if(size(partial_path) > {i}, partial_path[{i}], null)").alias(col_name))

# 构建最终结果
final_df = expanded_df.select(*final_columns)

# 排序后展示
final_df.orderBy("ParentWBSElementExternalID", *[f"LN-{i}" for i in range(1, max_level)]).show(truncate=False)

代码说明 #

• 根节点识别：通过集合差集筛选出没有父节点的顶层节点
• 路径迭代：循环关联原表，逐步拼接每个节点的完整路径，直到无新节点可加入
• 中间层级展开：利用sequence和slice生成所有从根到当前节点的中间路径，保证每个层级都有对应记录
• 动态列适配：根据最大层级数自动生成对应LN列，适配任意深度的层级结构
• NULL值补全：对短路径的后续列设置NULL，完全匹配预期输出格式

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