python - Pyspark - 计算两个数据帧之间的日期

Question

我有两个数据框，每个数据框都有一个日期列。即:

+-----------+
|  DEADLINES|
+-----------+
| 2023-07-15|
| 2018-08-10|
| 2022-03-28|
| 2021-06-22|
| 2021-12-18|
| 2021-10-11|
| 2021-11-13|
+-----------+

+----------+
|   DT_DATE|
+----------+
|2021-04-02|
|2021-04-21|
|2021-05-01|
|2021-06-03|
|2021-09-07|
|2021-10-12|
|2021-11-02|
+----------+

我需要计算在给定引用日期和每个 DEADLINES 日期之间有多少 DT_DATE 日期。

例如:使用 2021-03-31 作为引用日期应该给出以下结果集。

+-----------+------------+
|  DEADLINES|    dt_count|
+-----------+------------+
| 2023-07-15|           7|
| 2018-08-10|           0|
| 2022-03-28|           7|
| 2021-06-22|           4|
| 2021-12-18|           7|
| 2021-10-11|           5|
| 2021-11-13|           7|
+-----------+------------+

我设法让它迭代每一行截止日期数据框，但是对于更大的数据集，性能变得非常差。

有没有人有更好的解决方案？

编辑:这是我目前的解决方案:

def count_days(deadlines_df, dates_df, ref_date):
    for row in deadlines_df.collect():
        qtt = dates_df.filter(dates_df.DT_DATE.between(ref_date, row.DEADLINES)).count()
        yield row.DEADLINES, qtt


new_df = spark.createDataFrame(count_days(deadlines_df, dates_df, "2021-03-31"), ["DEADLINES", "dt_count"])

Answer 1

两个数据帧可以结合不同的权重，并使用范围从开始到当前行的窗口函数(Scala):

val deadlines = Seq(
  ("2023-07-15"),
  ("2018-08-10"),
  ("2022-03-28"),
  ("2021-06-22"),
  ("2021-12-18"),
  ("2021-10-11"),
  ("2021-11-13")
).toDF("DEADLINES")

val dates = Seq(
  ("2021-04-02"),
  ("2021-04-21"),
  ("2021-05-01"),
  ("2021-06-03"),
  ("2021-09-07"),
  ("2021-10-12"),
  ("2021-11-02")
).toDF("DT_DATE")

val referenceDate = "2021-03-31"
val united = deadlines.withColumn("weight", lit(0))
  .unionAll(
    dates
      .where($"DT_DATE" >= referenceDate)
      .withColumn("weight", lit(1))
  )

val fromStartToCurrentRowWindow = Window.orderBy("DEADLINES").rangeBetween(Window.unboundedPreceding, Window.currentRow)

val result = united
  .withColumn("dt_count", sum("weight").over(fromStartToCurrentRowWindow))
  .where($"weight" === lit(0))
  .drop("weight")

输出:

+----------+--------+
|DEADLINES |dt_count|
+----------+--------+
|2018-08-10|0       |
|2021-06-22|4       |
|2021-10-11|5       |
|2021-11-13|7       |
|2021-12-18|7       |
|2022-03-28|7       |
|2023-07-15|7       |
+----------+--------+

注意:计算将在一个分区中执行，Spark 显示这样的警告:WARN 日志记录 - 没有为 Window 操作定义分区!将所有数据移动到单个分区，这会导致严重的性能下降。

还有其他可能的解决方案，按范围连接两个数据帧，这导致笛卡尔连接。