如何用PySpark按Job ID选取最接近指定时间戳的数据？

我需要按小时频率提取时间戳数据，但有时不存在精确整点的记录，因此要为每个job_id选取与指定目标时间最接近的那条记录。

原始DataFrame #

需求示例 #

比如指定目标时间为2022-02-15T23:00:00.000+0000，但该整点记录不存在，此时要为每个job_id选取最接近这个时间的23:10:00的记录，期望输出如下：

PySpark实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, unix_timestamp, lit, abs
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

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

# 创建测试DataFrame（如果你的数据已经是DataFrame可以跳过这步）
data = [
    ("job1", "2022-02-15T23:40:00.000+0000", [0.2,3.4,13.2]),
    ("job1", "2022-02-15T23:10:00.000+0000", [0.1,2.9,11.2]),
    ("job2", "2022-02-15T23:40:00.000+0000", [1.2,3.1,16.0]),
    ("job1", "2022-02-15T23:20:00.000+0000", [0.4,0.4,16.2]),
    ("job2", "2022-02-15T23:10:00.000+0000", [0.7,8.4,11.2]),
    ("job2", "2022-02-15T23:20:00.000+0000", [0.3,1.5,19.1]),
    ("job3", "2022-02-15T23:20:00.000+0000", [0.7,7.4,13.2]),
    ("job3", "2022-02-15T23:40:00.000+0000", [0.7,7.4,13.2]),
    ("job3", "2022-02-15T23:10:00.000+0000", [0.7,7.4,13.2])
]

df = spark.createDataFrame(data, ["job_id", "timestamp", "item_values"])

# 1. 将字符串类型的timestamp转换为TimestampType（如果原始数据已经是时间类型可跳过）
df = df.withColumn("timestamp", col("timestamp").cast("timestamp"))

# 2. 定义目标时间
target_time = "2022-02-15T23:00:00.000+0000"

# 3. 计算每条记录与目标时间的时间差绝对值（单位：秒）
df_with_diff = df.withColumn(
    "time_diff",
    abs(unix_timestamp(col("timestamp")) - unix_timestamp(lit(target_time)))
)

# 4. 按job_id分组，为每组按时间差升序排序，取第一条（最接近的记录）
window_spec = Window.partitionBy("job_id").orderBy("time_diff")
result_df = df_with_diff.withColumn(
    "row_num",
    row_number().over(window_spec)
).filter(col("row_num") == 1).drop("time_diff", "row_num")

# 显示结果
result_df.show(truncate=False)

代码说明 #

1. 时间类型转换：确保timestamp列是Spark的TimestampType，否则无法正确计算时间差。
2. 时间差计算：用unix_timestamp将时间转换为秒级时间戳，计算与目标时间的差值绝对值，差值越小说明时间越接近。
3. 窗口函数筛选：通过Window.partitionBy("job_id")按任务分组，再按时间差排序，用row_number()标记每组内的顺序，取row_num=1的记录就是每组最接近目标时间的那条。

如果Spark版本在3.0及以上，也可以用timestamp_diff函数替代unix_timestamp，写法更直观：

from pyspark.sql.functions import timestamp_diff, abs

df_with_diff = df.withColumn(
    "time_diff",
    abs(timestamp_diff(col("timestamp"), lit(target_time).cast("timestamp"), "second"))
)

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