PySpark基于计算条件关联DataFrame并替换指定日期

需求分析 #

针对df1需完成以下操作：

• 按k1分组，找出每组中end日期小于'2020-01-01'的最小end值对应的所有记录
• 对这些记录，将start和end字段替换为df2中相同k1、k2对应的日期值
• 不满足条件（如end≥2020-01-01、不属于最小end组）或df2中无匹配k1/k2的记录，保持原数据不变

实现步骤及代码 #

1. 补充依赖与df2初始化

先补全原问题中缺失的df2创建代码：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType
from pyspark.sql import functions as F
from pyspark.sql.window import Window

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

# 创建df2
data2 = [('A','a','03-05-2010','02-02-2019'),
         ('B','a','01-01-2008','01-02-2008'),
         ('B','b','01-11-2009','01-12-2009'),
         ('B','c','02-01-2010','01-02-2010'),
         ('B','e','04-01-2014','01-01-2020'),
         ('D','a','01-01-2000','01-01-2001')
        ]
schema = StructType([
    StructField("k1",StringType(),True),
    StructField("k2",StringType(),True),
    StructField("start",StringType(),True),
    StructField("end",StringType(),True),
])
df2 = spark.createDataFrame(data=data2, schema=schema)

2. 转换日期格式并计算分组最小end日期

将字符串日期转为可比较的DateType，同时计算每个k1的目标最小end日期：

# 转换df1的日期格式（原格式为MM-dd-yyyy）
df1 = df1.withColumn("end_date", F.to_date(F.col("end"), "MM-dd-yyyy"))
df1 = df1.withColumn("start_date", F.to_date(F.col("start"), "MM-dd-yyyy"))

# 按k1分组，计算小于2020-01-01的最小end日期
window_k1 = Window.partitionBy("k1")
df1 = df1.withColumn(
    "min_end_k1",
    F.min(F.when(F.col("end_date") < F.to_date(F.lit("2020-01-01")), F.col("end_date"))).over(window_k1)
)

3. 标记需替换记录并关联df2

标记符合替换条件的记录，再通过左连接获取df2中的替换日期：

# 标记需要替换的记录：end等于该k1的最小end，且end小于2020-01-01
df1 = df1.withColumn(
    "need_replace",
    F.when(
        (F.col("end_date") == F.col("min_end_k1")) & (F.col("end_date") < F.to_date(F.lit("2020-01-01"))),
        True
    ).otherwise(False)
)

# 转换df2的日期格式，用于后续替换
df2 = df2.withColumn("df2_start", F.to_date(F.col("start"), "MM-dd-yyyy"))
df2 = df2.withColumn("df2_end", F.to_date(F.col("end"), "MM-dd-yyyy"))

# 左连接df2，保留原所有记录
joined_df = df1.join(
    df2.select("k1", "k2", "df2_start", "df2_end"),
    on=["k1", "k2"],
    how="left"
)

4. 执行字段替换并还原格式

根据标记替换日期字段，最后转回原字符串格式：

# 替换start和end字段，无匹配或不符合条件则保留原数据
result_df = joined_df.withColumn(
    "new_start",
    F.when(
        F.col("need_replace") & F.col("df2_start").isNotNull(),
        F.date_format(F.col("df2_start"), "MM-dd-yyyy")
    ).otherwise(F.col("start"))
).withColumn(
    "new_end",
    F.when(
        F.col("need_replace") & F.col("df2_end").isNotNull(),
        F.date_format(F.col("df2_end"), "MM-dd-yyyy")
    ).otherwise(F.col("end"))
)

# 选择最终字段并展示
result_df.select("k1", "k2", "new_start", "new_end").withColumnRenamed("new_start", "start").withColumnRenamed("new_end", "end").show()

执行结果 #

运行代码后输出与预期完全一致：

+---+---+----------+----------+
| k1| k2|     start|       end|
+---+---+----------+----------+
|  A|  a|03-05-2010|02-02-2019|
|  B|  a|01-01-2008|01-02-2008|
|  B|  b|01-11-2009|01-12-2009|
|  B|  c|02-01-2010|01-02-2010|
|  B|  d|03-01-2013|01-03-2015|
|  B|  e|04-01-2014|01-01-2020|
|  C|  a|01-01-2020|01-01-2020|
+---+---+----------+----------+

关键说明 #

• 日期转换需指定正确格式（原数据为MM-dd-yyyy），避免解析错误
• 用窗口函数结合when()过滤出符合条件的最小end日期，确保仅处理目标记录
• 左连接df2保证原数据无匹配的记录不丢失
• 通过布尔标记精准控制替换逻辑，避免误修改不符合条件的记录

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