如何使用PySpark按3、6、12个月间隔获取V1数据生成衍生列

方案一：基于实际月份差的准确关联法 #

这种方法不受数据行顺序影响，通过日期计算匹配对应月份的数值，是更可靠的生产级实现方式：

from pyspark.sql import SparkSession
from pyspark.sql.functions import add_months, date_format, to_date, col

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

# 模拟输入数据
data = [
    ("202307", 10),
    ("202306", 20),
    ("202305", 30),
    ("202304", 40),
    ("202303", 50),
    ("202302", 60),
    ("202301", 70)
]
df = spark.createDataFrame(data, ["month", "V1"])

# 转换月份为日期格式，并计算3/6个月前的月份（转为YYYYMM格式）
df_with_target_months = df.withColumn("month_date", to_date(col("month"), "yyyyMM")) \
                          .withColumn("month_3ago", date_format(add_months(col("month_date"), -3), "yyyyMM")) \
                          .withColumn("month_6ago", date_format(add_months(col("month_date"), -6), "yyyyMM"))

# 关联自身获取对应月份的V1值，作为V2和V3
result_df = df_with_target_months.join(df.withColumnRenamed("V1", "V2"), 
                                       df_with_target_months["month_3ago"] == df["month"], 
                                       how="left") \
                                 .join(df.withColumnRenamed("V1", "V3"), 
                                       df_with_target_months["month_6ago"] == df["month"], 
                                       how="left") \
                                 .select(df_with_target_months["month"], 
                                         df_with_target_months["V1"], 
                                         "V2", "V3")

# 查看202307的结果
result_df.filter(col("month") == "202307").show()

执行后，202307的V2为40（对应202304的V1），V3为70（对应202301的V1），完全符合实际月份差的逻辑。

方案二：基于数据行顺序的滞后取值（匹配示例输出） #

如果你的需求是严格按照现有数据的降序行顺序，取当前行之后第3/6行的V1值（匹配你给出的示例结果），可以使用窗口函数结合降序排序：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import lead, col

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

# 模拟输入数据
data = [
    ("202307", 10),
    ("202306", 20),
    ("202305", 30),
    ("202304", 40),
    ("202303", 50),
    ("202302", 60),
    ("202301", 70)
]
df = spark.createDataFrame(data, ["month", "V1"])

# 定义窗口：保持数据的降序排列顺序
window_spec = Window.orderBy(col("month").desc())

# 使用lead函数获取当前行之后的指定行数值（示例中202307的V2对应第3行，需偏移2位；V3对应第6行，需偏移5位）
result_df = df.withColumn("V2", lead(col("V1"), 2).over(window_spec)) \
              .withColumn("V3", lead(col("V1"), 5).over(window_spec))

# 查看202307的结果
result_df.filter(col("month") == "202307").show()

执行后，202307的V2为30、V3为60，和你给出的示例输出完全一致。注意这种方法依赖固定的行顺序，若数据排序发生变化，结果会出错，仅适用于行顺序固定的场景。

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