Databricks中不同类别下两个日期区间数值求和对比问题

原始DataFrame结构 #

+----------+---+-----+
|      date|cat|value|
+----------+---+-----+
|2022-08-11|  a|    1|
|2022-08-12|  a|    1|
|2022-08-13|  a|    1|
|2022-08-14|  a|    1|
|2022-08-15|  a|    1|
|2022-08-16|  a|    1|
|2022-08-17|  a|    2|
|2022-08-18|  a|    2|
|2022-08-19|  a|    2|
|2022-08-20|  a|    2|
|2022-08-21|  a|    2|
|2022-08-22|  a|    2|
|2022-08-11|  b|    1|
|2022-08-12|  b|    1|
|2022-08-13|  b|    1|
|2022-08-14|  b|    1|
|2022-08-15|  b|    1|
|2022-08-16|  b|    1|
|2022-08-17|  b|    3|
|2022-08-18|  b|    3|
|2022-08-19|  b|    3|
|2022-08-20|  b|    3|
|2022-08-21|  b|    3|
|2022-08-22|  b|    3|
+----------+---+-----+

需求说明 #

计算两个预定义日期区间内各cat分组的value总和：

• week1：2022-08-17 至 2022-08-22
• week2：2022-08-11 至 2022-08-16

将这两个总和（命名为w1、w2）附加到原DataFrame的每一行，并计算差值diff（w2 - w1），最终得到如下结构的结果：

+----------+---+-----+----+------+--------+
|      date|cat|value|  w1|    w2|    diff|
+----------+---+-----+----+------+--------+
|2022-08-11|  a|    1|   6|    12|       6|
|2022-08-12|  a|    1|   6|    12|       6|
|2022-08-13|  a|    1|   6|    12|       6|
|2022-08-14|  a|    1|   6|    12|       6|
|2022-08-15|  a|    1|   6|    12|       6|
|2022-08-16|  a|    1|   6|    12|       6|
|2022-08-17|  a|    2|   6|    12|       6|
|2022-08-18|  a|    2|   6|    12|       6|
|2022-08-19|  a|    2|   6|    12|       6|
|2022-08-20|  a|    2|   6|    12|       6|
|2022-08-21|  a|    2|   6|    12|       6|
|2022-08-22|  a|    2|   6|    12|       6|
|2022-08-11|  b|    3|  18|    30|      12|
|2022-08-12|  b|    3|  18|    30|      12|
|2022-08-13|  b|    3|  18|    30|      12|
|2022-08-14|  b|    3|  18|    30|      12|
|2022-08-15|  b|    3|  18|    30|      12|
|2022-08-16|  b|    3|  18|    30|      12|
|2022-08-17|  b|    5|  18|    30|      12|
|2022-08-18|  b|    5|  18|    30|      12|
|2022-08-19|  b|    5|  18|    30|      12|
|2022-08-20|  b|    5|  18|    30|      12|
|2022-08-21|  b|    5|  18|    30|      12|
|2022-08-22|  b|    5|  18|    30|      12|
+----------+---+-----+----+------+--------+

原代码问题分析 #

你尝试的窗口函数代码存在两个核心问题：

1. sum('value').over(w)是对整个cat分组的所有value求和，而非仅指定日期区间内的数值
2. 外层的when只会在符合日期条件的行显示求和结果，其他行返回null，无法实现每行都显示分组总和的需求

import pyspark.sql.functions as f

w = (Window.partitionBy('cat'))

df = (df
    .withColumn('date', f.to_date('date', 'yyyy-MM-dd'))
    .withColumn('value_week_1', 
               f.when(
                 (f.col('date') >= '2022-08-17') & 
                 (f.col('date') <= '2022-08-22'),
                 f.sum('value').over(w)
               )
              )
    .withColumn('value_week_2', 
                f.when(
                 (f.col('date') >= '2022-08-11') & 
                 (f.col('date') <= '2022-08-16'),
                 f.sum('value').over(w)
               )
              )
   )

解决方案 #

方法1：窗口函数+条件求和（简洁直观） #

直接在窗口函数内部使用when筛选日期区间，对分组内符合条件的value求和，确保每行都能获取对应分组的两个区间总和：

import pyspark.sql.functions as f

# 定义日期常量，便于维护
WEEK1_START = '2022-08-17'
WEEK1_END = '2022-08-22'
WEEK2_START = '2022-08-11'
WEEK2_END = '2022-08-16'

# 处理逻辑
df_result = df.withColumn('date', f.to_date('date', 'yyyy-MM-dd')) \
    .withColumn('w1', 
                f.sum(f.when(
                    (f.col('date') >= WEEK1_START) & (f.col('date') <= WEEK1_END),
                    f.col('value')
                )).over(f.Window.partitionBy('cat'))) \
    .withColumn('w2', 
                f.sum(f.when(
                    (f.col('date') >= WEEK2_START) & (f.col('date') <= WEEK2_END),
                    f.col('value')
                )).over(f.Window.partitionBy('cat'))) \
    .withColumn('diff', f.col('w2') - f.col('w1'))  # 差值方向可根据需求调整

方法2：先聚合再关联（适合大数据量） #

先对分组进行聚合计算两个区间的总和，再通过join将结果关联回原DataFrame，减少窗口函数的计算开销：

import pyspark.sql.functions as f

# 定义日期常量
WEEK1_START = '2022-08-17'
WEEK1_END = '2022-08-22'
WEEK2_START = '2022-08-11'
WEEK2_END = '2022-08-16'

# 1. 聚合计算各分组的w1、w2总和及差值
agg_df = df.withColumn('date', f.to_date('date', 'yyyy-MM-dd')) \
    .groupBy('cat') \
    .agg(
        f.sum(f.when(
            (f.col('date') >= WEEK1_START) & (f.col('date') <= WEEK1_END),
            f.col('value')
        )).alias('w1'),
        f.sum(f.when(
            (f.col('date') >= WEEK2_START) & (f.col('date') <= WEEK2_END),
            f.col('value')
        )).alias('w2')
    ) \
    .withColumn('diff', f.col('w2') - f.col('w1'))

# 2. 关联原DataFrame，将聚合结果附加到每行
df_result = df.withColumn('date', f.to_date('date', 'yyyy-MM-dd')) \
    .join(agg_df, on='cat', how='left')

两种方法均能生成符合需求的结果，可根据数据规模选择合适的方案。

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