PySpark中如何高效获取数组最大时间戳对应元素的指定字段？

针对你需要从数组列idInfo中，找到lastUpdated.sourceSystemTimestamp最新的元素并返回其individualNumber的需求，推荐两种PySpark实现方式，优先选择性能更优的高阶函数方案：

方案一：使用aggregate高阶函数（高效推荐） #

aggregate是PySpark内置的数组高阶函数，无需展开数组即可完成遍历和比较，适合处理大数组场景，性能更优。

代码示例 #

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

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

# 构造示例数据
data = [
    (1, [
        {"accountId": "A1", "lastUpdated": {"sourceSystemTimestamp": "2023-01-01T10:00:00"}, "individualNumber": "IN1"},
        {"accountId": "A2", "lastUpdated": {"sourceSystemTimestamp": "2023-02-01T12:00:00"}, "individualNumber": "IN2"},
        {"accountId": "A3", "lastUpdated": {"sourceSystemTimestamp": "2023-03-01T09:00:00"}, "individualNumber": "IN3"}
    ]),
    (2, [
        {"accountId": "B1", "lastUpdated": {"sourceSystemTimestamp": "2022-12-01T08:00:00"}, "individualNumber": "IN4"},
        {"accountId": "B2", "lastUpdated": {"sourceSystemTimestamp": "2023-04-01T15:00:00"}, "individualNumber": "IN5"}
    ])
]

df = spark.createDataFrame(data, ["id", "idInfo"])

# 使用aggregate遍历数组，保留最新时间戳对应的individualNumber
df = df.withColumn(
    "latest_individualNumber",
    F.aggregate(
        # 待处理的数组列
        F.col("idInfo"),
        # 初始化累加器：存储时间戳和对应的individualNumber，初始值为null
        F.struct(F.lit(None).alias("ts"), F.lit(None).alias("num")),
        # 迭代逻辑：比较当前元素时间戳与累加器中的时间戳，更新为更大的那个
        lambda acc, elem: F.when(
            elem["lastUpdated"]["sourceSystemTimestamp"] > acc["ts"],
            F.struct(
                elem["lastUpdated"]["sourceSystemTimestamp"].alias("ts"),
                elem["individualNumber"].alias("num")
            )
        ).otherwise(acc)
    )["num"]  # 最后取出存储的individualNumber
)

df.show(truncate=False)

逻辑说明 #

1. 初始化一个结构体累加器，用于临时存储当前找到的最新时间戳和对应的individualNumber
2. 遍历数组中的每个元素，比较元素的时间戳与累加器中的时间戳
3. 若当前元素时间戳更新，则替换累加器中的值；否则保留原累加器内容
4. 遍历结束后，从累加器中提取目标字段individualNumber

方案二：展开数组+窗口函数（逻辑简单，适合小数组） #

如果数组元素数量较少，也可以通过explode展开数组，结合窗口函数筛选最新时间戳对应的字段，但该方法会增加数据行数，大数组场景下性能不如方案一。

代码示例 #

from pyspark.sql.window import Window

# 展开数组为单行记录
df_exploded = df.select(
    "id",
    F.explode("idInfo").alias("info")
).withColumn(
    "source_ts", F.col("info.lastUpdated.sourceSystemTimestamp")
)

# 窗口函数按id分组，按时间戳降序排序
window_spec = Window.partitionBy("id").orderBy(F.col("source_ts").desc())

# 取每组第一条的individualNumber
df_result = df_exploded.withColumn(
    "latest_individualNumber",
    F.first("info.individualNumber").over(window_spec)
).select("id", "latest_individualNumber").distinct()

df_result.show(truncate=False)

逻辑说明 #

1. 使用explode将数组列拆分为多行，每行对应一个数组元素
2. 提取时间戳字段，通过窗口函数按id分组并按时间戳降序排列
3. 取每组的第一条记录，即为时间戳最新的元素，提取其individualNumber后去重

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