PySpark：如何实现动态时间间隔的数值聚合

现有一个按5分钟间隔采集数值的PySpark DataFrame，数据结构如下：

TIMESTAMP_COL               interval_length     value
2019-08-03T07:05:00.000Z    5                   5.34
2019-08-03T07:10:00.000Z    5                   4.5
2019-08-03T07:15:00.000Z    5                   4.74
2019-08-03T08:20:00.000Z    5                   5.64
2019-08-03T08:25:00.000Z    5                   4.5
2019-08-03T08:30:00.000Z    5                   4.56
2019-08-03T08:35:00.000Z    5                   4.56
2019-08-03T09:40:00.000Z    5                   4.56
2019-08-03T09:45:00.000Z    5                   4.68

需要将这些数值按可变时间间隔（如15、20、30、60分钟，要求为60的因数，范围5-60分钟）进行聚合。已实现固定15分钟间隔的聚合逻辑，现在要把该逻辑改造成支持传入agg_freq参数的动态函数，函数定义为def agg_input(df, agg_freq):。

原固定15分钟逻辑：

import pyspark.sql.functions as F

df = df.withColumn('time_utc_flat',
                         F.when(F.col("interval_length") < 15,
                                F.when(((F.minute(F.col(TIMESTAMP_COL)) > 0) &
                                       (F.minute(F.col(TIMESTAMP_COL)) <= 15)),
                                       F.date_trunc('hour', F.col(TIMESTAMP_COL)) +
                                       F.expr('INTERVAL 15 MINUTES'))
                                 .when(((F.minute(F.col(TIMESTAMP_COL)) > 15) &
                                       (F.minute(F.col(TIMESTAMP_COL)) <= 30)),
                                       F.date_trunc('hour', F.col(TIMESTAMP_COL)) +
                                       F.expr('INTERVAL 30 MINUTES'))
                                 .when(((F.minute(F.col(TIMESTAMP_COL)) > 30) &
                                       (F.minute(F.col(TIMESTAMP_COL)) <= 45)),
                                       F.date_trunc('hour', F.col(TIMESTAMP_COL)) +
                                       F.expr('INTERVAL 45 MINUTES'))
                                 .when(((F.minute(F.col(TIMESTAMP_COL)) > 45) &
                                       (F.minute(F.col(TIMESTAMP_COL)) <= 59) &
                                       (F.second(F.col(TIMESTAMP_COL)) <= 59)),
                                       F.date_trunc('hour', F.col(TIMESTAMP_COL)) +
                                       F.expr('INTERVAL 60 MINUTES'))
                                 .otherwise(F.date_trunc('hour', F.col(TIMESTAMP_COL))))
                          .otherwise(F.col(TIMESTAMP_COL)))

原聚合操作：

df_agg = df \
.groupBy("time_utc_flat") \
.agg(F.sum(value).alias("value")) \
.withColumnRenamed("time_utc_flat", TIMESTAMP_COL)

动态聚合函数实现 #

核心是根据传入的agg_freq自动生成时间区间判断逻辑，不用硬编码每个区间的条件。完整代码如下：

import pyspark.sql.functions as F

def agg_input(df, agg_freq):
    # 验证参数合法性
    if not isinstance(agg_freq, int) or agg_freq < 5 or agg_freq > 60 or 60 % agg_freq != 0:
        raise ValueError("agg_freq必须是5-60之间的整数，且为60的因数")
    
    TIMESTAMP_COL = "TIMESTAMP_COL"  # 若列名不固定，可改为函数参数
    
    # 计算每小时内的区间数量
    intervals_per_hour = 60 // agg_freq
    
    # 初始化条件链
    interval_conditions = None
    
    for i in range(1, intervals_per_hour + 1):
        start_min = (i - 1) * agg_freq
        end_min = i * agg_freq
        
        # 构建当前区间的判断条件
        if i < intervals_per_hour:
            condition = (F.minute(F.col(TIMESTAMP_COL)) > start_min) & (F.minute(F.col(TIMESTAMP_COL)) <= end_min)
        else:
            # 最后一个区间包含59分59秒的情况
            condition = (F.minute(F.col(TIMESTAMP_COL)) > start_min) & \
                        (F.minute(F.col(TIMESTAMP_COL)) <= 59) & \
                        (F.second(F.col(TIMESTAMP_COL)) <= 59)
        
        # 计算当前区间对应的截断时间
        trunc_time = F.date_trunc('hour', F.col(TIMESTAMP_COL)) + F.expr(f'INTERVAL {end_min} MINUTES')
        
        # 拼接条件链
        if interval_conditions is None:
            interval_conditions = F.when(condition, trunc_time)
        else:
            interval_conditions = interval_conditions.when(condition, trunc_time)
    
    # 处理分钟为0的默认情况
    interval_conditions = interval_conditions.otherwise(F.date_trunc('hour', F.col(TIMESTAMP_COL)))
    
    # 添加聚合分组用的时间列
    df = df.withColumn('time_utc_flat',
                       F.when(F.col("interval_length") < agg_freq, interval_conditions)
                       .otherwise(F.col(TIMESTAMP_COL)))
    
    # 执行聚合操作
    df_agg = df.groupBy("time_utc_flat") \
               .agg(F.sum("value").alias("value")) \
               .withColumnRenamed("time_utc_flat", TIMESTAMP_COL)
    
    return df_agg

代码说明 #

• 参数校验：确保agg_freq符合要求，避免非法输入引发错误
• 动态条件生成：根据聚合间隔自动计算每小时内的区间，循环生成对应判断条件，无需手动编写每个区间
• 逻辑兼容：保留原逻辑中对interval_length的判断——仅当原始间隔小于聚合间隔时才做区间映射，否则直接保留原时间戳
• 聚合复用：沿用原有的分组求和逻辑，最后将分组列重命名回原时间戳列名

使用示例 #

# 按15分钟间隔聚合
df_15min = agg_input(df, 15)

# 按30分钟间隔聚合
df_30min = agg_input(df, 30)

# 按60分钟（1小时）间隔聚合
df_hourly = agg_input(df, 60)

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