如何在Scala DataFrame中按条件计算滚动求和

我有如下Scala DataFrame：

+---+------+----------+-------+-----------+
|id |amount|date      |ndxdays|ndxdaysDate|
+---+------+----------+-------+-----------+
|1  |77    |2023-10-01|2      |2023-09-29 |
|1  |14    |2023-10-02|2      |2023-09-30 |
|1  |70    |2023-10-04|2      |2023-10-02 |
|1  |70    |2023-10-05|2      |2023-10-03 |
|1  |14    |2023-10-07|2      |2023-10-05 |
|2  |65    |2023-10-01|3      |2023-09-28 |
|2  |32    |2023-10-02|3      |2023-09-29 |
|2  |21    |2023-10-03|3      |2023-09-30 |
|2  |78    |2023-10-07|3      |2023-10-04 |
|3  |80    |2023-10-01|2      |2023-09-29 |
|3  |98    |2023-10-03|2      |2023-10-01 |
|3  |36    |2023-10-06|2      |2023-10-04 |
+---+------+----------+-------+-----------+

需求 #

为每一行计算滚动求和，仅纳入当前行ndxDaysDate小于该行date的amount值，即在遵守ndxdaysDate < date条件的前提下，计算同id下符合条件的amount累计和。

预期输出 #

+---+------+----------+-------+-----------+-----------+
|id |amount|date      |ndxdays|ndxdaysDate|rollingsum |
+---+------+----------+-------+-----------+-----------+
|1  |77    |2023-10-01|2      |2023-09-29 |       77  |<- 仅77符合：2023-10-01(date) > 2023-09-29(ndxdaysDate)（id=1）
|1  |14    |2023-10-02|2      |2023-09-30 |       91  |<- 14+77：2023-10-02和2023-10-01的date均大于2023-09-30（id=1）
|1  |70    |2023-10-04|2      |2023-10-02 |       70  |<- 仅70符合：2023-10-04(date) > 2023-10-02(ndxdaysDate)（id=1）
|1  |70    |2023-10-05|2      |2023-10-03 |       140 |<- 70+70：2023-10-05和2023-10-04的date均大于2023-10-03（id=1）
|1  |14    |2023-10-07|2      |2023-10-05 |       14  |<- 仅14符合：2023-10-07(date) > 2023-10-05(ndxdaysDate)（id=1）
|2  |65    |2023-10-01|3      |2023-09-28 |       65  |<- 仅65符合：2023-10-01(date) > 2023-09-28(ndxdaysDate)（id=2）
|2  |32    |2023-10-02|3      |2023-09-29 |       97  |<- 65+32：2023-10-02和2023-10-01的date均大于2023-09-29（id=2）
|2  |21    |2023-10-03|3      |2023-09-30 |       118 |<- 65+32+21：2023-10-03、2023-10-02和2023-10-01的date均大于2023-09-30（id=2）
|2  |78    |2023-10-07|3      |2023-10-04 |       78  |<- 仅78符合：2023-10-07(date) > 2023-10-04(ndxdaysDate)（id=2）
|3  |80    |2023-10-01|2      |2023-09-29 |       80  |<- 仅80符合：2023-10-01(date) > 2023-09-29(ndxdaysDate)（id=3）
|3  |98    |2023-10-03|2      |2023-10-01 |       98  |<- 仅98符合：2023-10-03(date) > 2023-10-01(ndxdaysDate)（id=3）
|3  |36    |2023-10-06|2      |2023-10-04 |       36  |<- 仅36符合：2023-10-06(date) > 2023-10-04(ndxdaysDate)（id=3）
+---+------+----------+-------+-----------+-----------+

解决方案 #

方法1：使用内置聚合函数（高效推荐）

先确保日期列是Date类型，再通过collect_list收集同id的所有日期与金额，过滤符合条件的记录后求和：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.DateType

// 转换日期列（若原始为字符串类型）
val formattedDF = originalDF
  .withColumn("date", to_date(col("date"), "yyyy-MM-dd"))
  .withColumn("ndxdaysDate", to_date(col("ndxdaysDate"), "yyyy-MM-dd"))

// 计算条件滚动求和
val resultDF = formattedDF.withColumn(
  "rollingsum",
  expr("""
    aggregate(
      filter(
        collect_list(struct(date, amount)), 
        x -> x.date > ndxdaysDate
      ), 
      0, 
      (acc, x) -> acc + x.amount
    )
  """)
)

resultDF.show()

方法2：自连接分组求和（逻辑直观）

通过自连接关联同id的记录，过滤符合条件的行后分组求和：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.DateType

// 转换日期列（若原始为字符串类型）
val formattedDF = originalDF
  .withColumn("date", to_date(col("date"), "yyyy-MM-dd"))
  .withColumn("ndxdaysDate", to_date(col("ndxdaysDate"), "yyyy-MM-dd"))

// 自连接+过滤+分组求和
val resultDF = formattedDF.alias("a")
  .join(formattedDF.alias("b"), col("a.id") === col("b.id"), "left")
  .filter(col("b.date") > col("a.ndxdaysDate"))
  .groupBy(
    col("a.id"), col("a.amount"), col("a.date"), 
    col("a.ndxdays"), col("a.ndxdaysDate")
  )
  .agg(sum("b.amount").alias("rollingsum"))
  .orderBy(col("a.id"), col("a.date"))

resultDF.show()

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