PySpark按ID分组求各属性对应最小day_time的实现问询

问题分析 #

当前代码存在两个核心问题：

• 按id+country+activities+is_stream合并分组时，同一个country会因activities或is_stream的差异生成多条记录，后续collect_set会保留该country的多个时间值，不符合仅保留最小时间的需求。
• activities是数组类型，未拆分为单个元素处理，无法为每个元素关联对应的最小时间。

解决方案 #

需要针对每个属性维度单独计算最小时间，再按id聚合结果，具体实现步骤如下：

1. 导入必要依赖 #

补充处理数组所需的explode函数：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, ArrayType, TimestampType
from pyspark.sql.functions import collect_set, min, col, struct, explode
import datetime

2. 初始化SparkSession与测试数据 #

spark = SparkSession.builder \
    .appName("Data Transformation") \
    .getOrCreate()

schema = StructType([
    StructField("id", StringType(), True),
    StructField("day_time", TimestampType(), True),
    StructField("country", StringType(), True),
    StructField("activities", ArrayType(StringType()), True),
    StructField("is_stream", StringType(), True)
])

data = [("1", datetime.strptime("2024-01-08T13:14:41", '%Y-%m-%dT%H:%M:%S') , "US", ["fun"], "FALSE"),
        ("1", datetime.strptime("2024-01-08T13:14:42", '%Y-%m-%dT%H:%M:%S') , "US", ["action","adventure"], "TRUE"),
        ("1", datetime.strptime("2024-01-08T13:14:44", '%Y-%m-%dT%H:%M:%S') , "US", ["action","adventure"], "TRUE")
]

df = spark.createDataFrame(data, schema)

3. 分维度计算最小时间 #

处理Country维度

直接按id+country分组，取该组合的最小时间：

country_df = df.groupBy("id", "country") \
    .agg(min("day_time").alias("timestamp")) \
    .groupBy("id") \
    .agg(collect_set(struct(col("country").alias("key"), col("timestamp"))).alias("country_details"))

处理Activities维度

先通过explode拆分数组为单行元素，再按id+activity分组取最小时间：

activities_df = df.withColumn("activity", explode(col("activities"))) \
    .groupBy("id", "activity") \
    .agg(min("day_time").alias("timestamp")) \
    .groupBy("id") \
    .agg(collect_set(struct(col("activity").alias("key"), col("timestamp"))).alias("activities_details"))

处理is_stream维度

同Country维度逻辑，按id+is_stream分组取最小时间：

is_stream_df = df.groupBy("id", "is_stream") \
    .agg(min("day_time").alias("timestamp")) \
    .groupBy("id") \
    .agg(collect_set(struct(col("is_stream").alias("key"), col("timestamp"))).alias("is_stream_details"))

4. 合并所有结果 #

将三个维度的结果按id关联得到最终输出：

final_df = country_df.join(activities_df, on="id", how="inner") \
    .join(is_stream_df, on="id", how="inner")

final_df.show(truncate=False)

最终输出结果 #

+---+------------------------------+------------------------------------------------------------------------------------+--------------------------------------------+
|id |country_details               |activities_details                                                                  |is_stream_details                           |
+---+------------------------------+------------------------------------------------------------------------------------+--------------------------------------------+
|1  |[{US, 2024-01-08 13:14:41}]   |[{fun, 2024-01-08 13:14:41}, {adventure, 2024-01-08 13:14:42}, {action, 2024-01-08 13:14:42}]|[{FALSE, 2024-01-08 13:14:41}, {TRUE, 2024-01-08 13:14:42}]|
+---+------------------------------+------------------------------------------------------------------------------------+--------------------------------------------+

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