PySpark实现基于列依赖关系生成runorder列的方案求助

问题场景 #

现有如下结构的PySpark DataFrame：

需求说明 #

新增runorder列，规则如下：

• 若questionkey字段值等于另一行的varid值，则被依赖的varid需优先执行，其runorder值更小
• 依赖项的runorder为依赖对象的runorder + 1

预期输出 #

遇到的问题 #

使用case when无法实现递归逻辑，普通Join仅能处理直接依赖，无法覆盖多层依赖场景（如id4依赖id2，id2又依赖id1）。已尝试的代码如下：

df_with_runorder = df.withColumn("runorder", lit(1)).withColumnRenamed("VarID","VarID_a") 
#renamed the columns to before joining
df_with_runorder=df_with_runorder.withColumnRenamed("QuestionKey","QuestionKey_a")\
 .withColumnRenamed("runorder","runorder_a") 

#here performed joined to check dependency of values in varid and questionkey and if match found(dependency found) then updated runorder by adding 1 to it

joined_df = df_with_runorder.alias("a").join(df.alias("b"), col("a.varid_a") == col("b.QuestionKey"), "left")\
 .withColumn("updated_runorder",when(col("VarID").isNotNull(),col("runorder_a")+1).otherwise(col("runorder_a"))) 
joined_df.show() 

解决方案：递归CTE实现多层依赖计算 #

Spark 2.1+支持递归CTE（Common Table Expression），可以遍历整个依赖链，从根节点开始逐层计算runorder，完美解决多层依赖问题。

实现步骤 #

1. 标记根节点：questionkey不在所有varid集合中的节点，初始runorder为1
2. 递归遍历：找到依赖当前节点的子节点，将子节点的runorder设为当前节点的runorder+1
3. 终止条件：无新节点可加入递归结果时停止

完整代码 #

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

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

# 创建原始DataFrame
data = [("id1", "qid1"), ("id2", "id1"), ("id3", "qid3"), ("id4", "id2")]
df = spark.createDataFrame(data, schema=["varid", "questionkey"])

# 生成varid集合字符串，用于SQL判断
varid_list = "'"+"','".join(df.select("varid").rdd.flatMap(lambda x: x).collect())+"'"

# 创建临时视图供SQL调用
df.createOrReplaceTempView("dependency_table")

# 递归CTE计算runorder
result_df = spark.sql(f"""
    WITH RECURSIVE dependency_cte AS (
        -- 基础查询：根节点，runorder初始为1
        SELECT 
            varid, 
            questionkey, 
            1 AS runorder
        FROM 
            dependency_table
        WHERE 
            questionkey NOT IN ({varid_list})
        
        UNION ALL
        
        -- 递归查询：关联父节点，计算子节点runorder
        SELECT 
            d.varid, 
            d.questionkey, 
            cte.runorder + 1 AS runorder
        FROM 
            dependency_table d
        JOIN 
            dependency_cte cte ON d.questionkey = cte.varid
        WHERE 
            d.varid NOT IN (SELECT varid FROM dependency_cte)
    )
    SELECT * FROM dependency_cte
""")

# 输出结果
result_df.show()

代码解释 #

• 根节点筛选：通过questionkey NOT IN ({varid_list})找出没有依赖的节点，这些节点是整个依赖链的起点，runorder设为1
• 递归关联：每次迭代将未处理的节点与已计算完成的父节点关联，子节点的runorder继承父节点值并加1
• 去重处理：通过d.varid NOT IN (SELECT varid FROM dependency_cte)避免重复计算同一节点

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