如何在PySpark中递推填充空的起止时间戳列？

核心思路 #

由于后续行的start_time依赖上一行的end_time，属于递推式填充，最适合用**递归CTE（Common Table Expression）**实现，配合行号确保数据顺序。

步骤实现 #

1. 准备数据并添加行号

首先给数据添加连续行号（如果已有天然排序键，比如业务ID，优先用该键）：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, TimestampType, LongType
from pyspark.sql.functions import monotonically_increasing_id

spark = SparkSession.builder.appName("TimestampFill").getOrCreate()

# 定义示例数据与Schema
data = [
    ("2024-01-01 00:00:00", "2024-01-01 00:01:00", 60),
    ("2024-01-01 00:01:00", "2024-01-01 00:03:00", 120),
    (None, None, 90),
    (None, None, 180),
    (None, None, 60)
]

schema = StructType([
    StructField("start_time", TimestampType(), True),
    StructField("end_time", TimestampType(), True),
    StructField("duration", LongType(), True)
])

df = spark.createDataFrame(data, schema)
# 添加行号确保顺序，若有业务排序键替换此处
df = df.withColumn("row_num", monotonically_increasing_id())
df.createOrReplaceTempView("temp_table")

2. 递归CTE填充空值

通过递归CTE分两部分处理：基准行（已有值的行）和递推行（空值行）：

filled_df = spark.sql("""
    WITH recursive cte AS (
        -- 基准部分：选取所有非空start_time的行
        SELECT row_num, start_time, end_time, duration
        FROM temp_table
        WHERE start_time IS NOT NULL
        UNION ALL
        -- 递归部分：关联上一行，递推计算当前行的时间戳
        SELECT 
            curr.row_num,
            prev.end_time AS start_time,
            -- 若duration是微秒，替换SECOND为MICROSECOND
            prev.end_time + INTERVAL curr.duration SECOND AS end_time,
            curr.duration
        FROM cte prev
        JOIN temp_table curr ON prev.row_num + 1 = curr.row_num
        WHERE curr.start_time IS NULL
    )
    SELECT start_time, end_time, duration FROM cte ORDER BY row_num
""")

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

关键注意事项 #

• 排序依据：monotonically_increasing_id()仅为临时方案，分布式环境下可能不严格连续，若业务数据有天然排序键（比如自增ID、原始时间列），必须用该键替换行号逻辑，保证递推顺序正确。
• 时间单位适配：如果duration是微秒，将SQL中的SECOND改为MICROSECOND即可。
• 递归深度：PySpark默认递归深度为100，若需处理超过100行的空值数据，需通过spark.sql.recursiveCTE.maxIterations参数调整最大迭代次数。

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