如何在PySpark/数据农场中编写设备层级场景的优化递归SQL查询

场景说明 #

设备层级关系：设备1 → 设备3 → 设备53（设备53无父设备），Gold Asset表存储每条设备的device_id（设备编号）和parent_device_id（父设备编号）信息，需通过递归查询获取每个设备的顶层父设备。

PySpark SQL递归实现（推荐） #

Spark 3.0及以上支持WITH RECURSIVE语法，可直接编写递归CTE查询：

WITH RECURSIVE device_hierarchy AS (
    -- 基础节点：筛选顶层设备（无父设备），初始化顶层父ID为自身
    SELECT 
        device_id,
        parent_device_id,
        device_id AS top_parent_id,
        1 AS level
    FROM gold_asset
    WHERE parent_device_id IS NULL

    UNION ALL

    -- 递归关联：将子设备与父设备层级关联，传递顶层父ID
    SELECT 
        ga.device_id,
        ga.parent_device_id,
        dh.top_parent_id,
        dh.level + 1 AS level
    FROM gold_asset ga
    INNER JOIN device_hierarchy dh 
        ON ga.parent_device_id = dh.device_id
)
-- 提取最终结果：每个设备对应的顶层父设备
SELECT 
    device_id,
    top_parent_id
FROM device_hierarchy
ORDER BY device_id;

代码解释 #

• 基础CTE段：先定位所有无父设备的顶层节点，将自身ID设为顶层父ID
• 递归关联段：通过子设备的parent_device_id匹配父设备的device_id，把顶层父ID向下传递给子设备
• 最终查询：直接提取每个设备和对应的顶层父ID，完成层级穿透

PySpark DataFrame API实现 #

如果偏好DataFrame操作，可通过迭代递归方式实现：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit

# 初始化Spark会话
spark = SparkSession.builder.appName("DeviceTopParent").getOrCreate()

# 读取Gold Asset表（替换为你的表路径/名称）
gold_asset_df = spark.table("gold_asset")

# 初始化递归DataFrame：先加载所有顶层设备
hierarchy_df = gold_asset_df.filter(col("parent_device_id").isNull()) \
    .withColumn("top_parent_id", col("device_id")) \
    .withColumn("level", lit(1))

# 迭代递归，直到没有新节点需要处理
while True:
    # 找出未被处理的子设备，关联父设备层级获取顶层ID
    new_nodes = gold_asset_df.join(
        hierarchy_df,
        gold_asset_df.parent_device_id == hierarchy_df.device_id,
        "left_anti"
    ).join(
        hierarchy_df,
        gold_asset_df.parent_device_id == hierarchy_df.device_id
    ).select(
        gold_asset_df.device_id,
        gold_asset_df.parent_device_id,
        hierarchy_df.top_parent_id,
        (hierarchy_df.level + 1).alias("level")
    )

    # 无新节点则终止循环
    if new_nodes.count() == 0:
        break

    # 合并新节点到递归DataFrame
    hierarchy_df = hierarchy_df.union(new_nodes)

# 输出最终结果
result_df = hierarchy_df.select("device_id", "top_parent_id").orderBy("device_id")
result_df.show()

优化建议 #

• 索引优化：在device_id和parent_device_id字段上添加索引（如Databricks中的ZORDER排序），大幅提升递归关联速度
• 优先使用SQL CTE：Spark对递归CTE的执行计划优化更充分，性能优于DataFrame迭代方式
• 控制递归深度：若设备层级极深，可设置递归次数限制（Spark默认无限制，但过深层级建议拆分处理）

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