如何在Spark SQL和PySpark中用CTE实现递归行生成？

问题背景 #

需要将trial表中每一行的started到ended区间，拆分为每行对应一个递增nr值的新行，类似SQL Server递归CTE的效果，但Spark SQL的递归CTE语法存在细微差异。

1. Spark SQL 实现方案 #

首先创建测试表并插入数据：

-- 创建trial表
CREATE TABLE IF NOT EXISTS trial (
    started INT,
    ended INT
);

-- 插入测试数据
INSERT INTO trial VALUES
(1, 3),
(5, 7),
(10, 12);

使用递归CTE实现行扩展：

WITH RECURSIVE trial_expanded AS (
    -- 基础查询：初始化nr为started的值
    SELECT started, ended, started AS nr
    FROM trial
    UNION ALL
    -- 递归查询：nr递增，直到不超过ended
    SELECT started, ended, nr + 1
    FROM trial_expanded
    WHERE nr < ended
)
-- 输出结果并排序
SELECT started, ended, nr
FROM trial_expanded
ORDER BY started, nr;

关键注意事项：

• Spark SQL必须显式使用WITH RECURSIVE关键字（SQL Server可省略RECURSIVE）
• 递归分支的列数、数据类型必须与基础分支完全匹配
• 必须添加明确的终止条件（WHERE nr < ended），避免无限递归

2. PySpark 实现方案 #

可以通过注册临时视图后执行Spark SQL语句，或者直接使用DataFrame API结合CTE：

方式一：使用Spark SQL语句

from pyspark.sql import SparkSession

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

# 创建测试DataFrame并注册为临时视图
trial_data = [(1, 3), (5, 7), (10, 12)]
df = spark.createDataFrame(trial_data, schema=["started", "ended"])
df.createOrReplaceTempView("trial")

# 执行递归CTE查询
expanded_df = spark.sql("""
WITH RECURSIVE trial_expanded AS (
    SELECT started, ended, started AS nr
    FROM trial
    UNION ALL
    SELECT started, ended, nr + 1
    FROM trial_expanded
    WHERE nr < ended
)
SELECT started, ended, nr
FROM trial_expanded
ORDER BY started, nr
""")

# 查看结果
expanded_df.show()

方式二：DataFrame API 替代方案（非CTE但更高效）

如果对CTE没有强制要求，使用sequence+explode的方式性能更优：

from pyspark.sql.functions import sequence, explode

expanded_df = df.withColumn("nr", explode(sequence("started", "ended")))
expanded_df.orderBy("started", "nr").show()

与SQL Server递归CTE的差异 #

• Spark SQL要求递归CTE必须声明RECURSIVE关键字
• Spark不允许递归分支直接引用外部表（除非显式关联），但本例中基础分支已包含ended字段，无需额外关联
• Spark对递归深度有默认限制（可通过spark.sql.recursiveCTE.maxIterations配置调整）

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