如何用PySpark实现多列1-12期滞后列生成（适配海量列场景）

原Pandas代码核心逻辑 #

• 读取日度数据，剔除空值后重采样为月度均值
• 扩展数据索引至原结束日期后13个月（补充空行）
• 计算每个指标的12个月同比变化率：((当前值 - 12个月前值)/12个月前值)*100
• 对每个同比指标生成滞后1至12个月的衍生列

适配数百列场景的PySpark实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

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

# --------------------------
# 1. 读取并预处理日度数据（替换为你的实际数据源读取逻辑）
# --------------------------
# 示例：模拟日度数据，实际场景可替换为读取CSV/数据库/对象存储等
df_daily = spark.createDataFrame(
    [
        ("2016-01-01", 37.05), ("2016-01-02", 36.78),
        # 省略更多日度数据
        ("2022-12-30", 79.56), ("2022-12-31", 80.23)
    ],
    ["date", "dcoilwtico"]
).withColumn("date", F.to_date("date"))

# 剔除空值，按月份聚合计算均值
df_monthly = df_daily.filter(F.col("dcoilwtico").isNotNull()) \
    .withColumn("month", F.date_trunc("month", "date")) \
    .groupBy("month") \
    .agg(F.mean("dcoilwtico").alias("dcoilwtico")) \
    .orderBy("month")

# --------------------------
# 2. 扩展数据范围（新增后13个月的空行）
# --------------------------
# 获取原始数据最后一个月份
last_month = df_monthly.select(F.max("month")).first()[0]
# 生成后续13个月的时间序列
future_months = spark.createDataFrame(
    [(last_month + F.expr(f"interval {i} months")) for i in range(1, 14)],
    ["month"]
)
# 合并原始数据与未来空行，空值自动填充为null
df_extended = df_monthly.join(future_months, on="month", how="full_outer") \
    .orderBy("month")

# --------------------------
# 3. 计算12个月同比变化率
# --------------------------
# 定义全局时间窗口（按月份排序）
window_spec = Window.orderBy("month")
# 批量处理所有指标列（排除month列）
metric_cols = [col for col in df_extended.columns if col != "month"]
for col in metric_cols:
    # 获取12个月前的数值
    lag_12_val = F.lag(col, 12).over(window_spec)
    # 计算同比变化率，空值场景保留null
    df_yoy = df_extended.withColumn(
        f"{col}_yoy",
        F.when(lag_12_val.isNotNull(), ((F.col(col) - lag_12_val)/lag_12_val)*100)
        .otherwise(None)
    )

# --------------------------
# 4. 生成每个同比指标的1-12个月滞后列
# --------------------------
# 筛选所有同比列
yoy_cols = [col for col in df_yoy.columns if "_yoy" in col]
for col in yoy_cols:
    base_col_name = col.replace("_yoy", "")
    # 生成1到12个月的滞后列
    for lag_num in range(1, 13):
        df_yoy = df_yoy.withColumn(
            f"{base_col_name}_{lag_num}",
            F.lag(col, lag_num).over(window_spec)
        )

# 查看结果（倒序展示最近18行）
df_yoy.orderBy("month", ascending=False).show(18)

关键适配说明 #

• 重采样逻辑：用date_trunc提取月份维度，结合groupBy+agg实现月度均值聚合，替代Pandas的resample
• 批量多列处理：通过遍历列名自动处理数百个指标，无需手动编写重复逻辑
• 空行扩展：通过生成未来月份序列+全外连接实现空行补充，适配后续滞后计算需求
• 窗口函数：用Window.orderBy定义全局时间窗口，lag函数对应Pandas的shift功能

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