PySpark如何为每个cost_type筛选距离ds最近的日期对应行

核心思路 #

• 先统一日期字段的时间戳类型，方便区间判断和差值计算
• 构造优先级排序字段：优先按「ds是否落在[min_ts, max_ts]区间」倒序排列，再按「ds与区间的最小时间差值」正序排列
• 用窗口函数按cost_type分组取排序第一的行即可

完整实现代码 #

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

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

# 1. 样例数据构造（实际使用时替换为你自己的DataFrame即可）
data = [
    (100000, "reorder", 1.0, "2020-10-16 10:16:31", "2020-11-21 18:50:27", "2021-05-29"),
    (250000, "overusage", 1.0, "2020-11-21 18:48:02", "2021-02-09 20:07:28", "2021-05-29"),
    (100000, "overusage", 1.0, "2019-05-12 16:00:40", "2020-11-21 18:44:04", "2021-05-29"),
    (200000, "reorder", 1.0, "2020-11-21 19:00:09", "2021-05-29 23:56:25", "2021-05-29")
]
df = spark.createDataFrame(data, schema=["amount", "cost_type", "place2", "min_ts", "max_ts", "ds"])

# 2. 预处理字段类型，新增计算字段
df = df.withColumn("ds_ts", F.col("ds").cast("timestamp")) \
       .withColumn("min_ts", F.col("min_ts").cast("timestamp")) \
       .withColumn("max_ts", F.col("max_ts").cast("timestamp")) \
       # 判断是否在区间内
       .withColumn("in_range", (F.col("ds_ts") >= F.col("min_ts")) & (F.col("ds_ts") <= F.col("max_ts"))) \
       # 计算ds和区间的最小时间差（单位为秒）
       .withColumn("time_diff", F.when(F.col("in_range"), 0) \
                   .otherwise(F.least(F.abs(F.unix_timestamp("ds_ts") - F.unix_timestamp("min_ts")),
                                      F.abs(F.unix_timestamp("ds_ts") - F.unix_timestamp("max_ts")))))

# 3. 窗口函数按cost_type分组排序取第一行
window_spec = Window.partitionBy("cost_type").orderBy(F.desc("in_range"), F.asc("time_diff"))
result_df = df.withColumn("rn", F.row_number().over(window_spec)) \
              .filter(F.col("rn") == 1) \
              .drop("ds_ts", "in_range", "time_diff", "rn")

# 查看结果
result_df.show()

效果验证 #

• 当ds='2021-05-29'时，返回cost_type为overusage的第2行、reorder的第4行，符合要求
• 当ds='2020-05-01'时，所有行都不在区间内，按最小时间差取reorder的第1行、overusage的第3行，符合要求

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