基于日期差在PySpark DataFrame中创建二元指示列的高效实现

高效实现方案 #

针对任意行数的场景，推荐使用窗口函数结合条件判断的方案，既适配需求又保证计算效率，以下提供两种可行写法：

写法一：窗口聚合+条件过滤（兼容PySpark全版本） #

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

# 定义窗口：按ID分区，date升序，窗口范围从当前行的下一行到分区末尾
window_spec = Window.partitionBy("ID").orderBy("date").rowsBetween(1, Window.unboundedFollowing)

# 计算D列：检查后续行是否存在符合条件的记录
df_result = df.withColumn(
    "D",
    F.max(
        F.when(
            (F.datediff(F.col("date"), F.lag("date").over(Window.partitionBy("ID").orderBy("date"))) < 14) & 
            (F.col("bool") == 1),
            1
        ).otherwise(0)
    ).over(window_spec)
).fillna(0, subset=["D"])  # 最后一行无后续行，填充0

df_result.show()

写法二：exists函数判断（PySpark 3.0+版本支持） #

逻辑更直观，直接判断窗口内是否存在符合条件的记录：

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

# 定义窗口：按ID分区，date升序，窗口范围从当前行的下一行到分区末尾
window_spec = Window.partitionBy("ID").orderBy("date").rowsBetween(1, Window.unboundedFollowing)

# 收集后续行的date和bool，用exists判断是否符合条件
df_result = df.withColumn(
    "D",
    F.when(
        F.exists(
            F.collect_list(F.struct("date", "bool")).over(window_spec),
            lambda x: (F.datediff(x.date, F.col("date")) < 14) & (x.bool == 1)
        ),
        1
    ).otherwise(0)
)

df_result.show()

方案优势 #

• 无需固定行数，自动适配两周内任意数量的后续行
• 仅在ID分区内计算，避免跨分区无效运算
• 窗口函数的计算逻辑由Spark优化引擎处理，性能优于多次调用lag的手动判断

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