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 显示如下警告： 警告记录 - 没有为窗口操作定义分区！将所有数据移动到单个分区，这可能会导致性能严重下降。

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

Answer 2

如果您的截止日期数量较少，您可以：

• 在dates_df数据帧上按截止日期添加一列，当DT_DATE在ref_date和截止日期之间时，值为1，否则为0
• 然后将每个截止日期列相加
• 最后转置结果数据帧得到你想要的数据帧

让我们一步一步来看看

按截止日期添加一栏：

from pyspark.sql import functions as F

deadline_rows = deadlines_df.collect()

dates_with_deadlines = dates_df
for row in deadline_rows:
    dates_with_deadlines = dates_with_deadlines.withColumn(
        str(row.DEADLINES),
        F.when(
          dates_df.DT_DATE.between(ref_date, row.DEADLINES), F.lit(1))
        .otherwise(
          F.lit(0)
        )
    )

通过您的示例，您会得到以下dates_with_deadlines 数据框：

+----------+----------+----------+----------+----------+----------+----------+----------+
|DT_DATE   |2023-07-15|2018-08-10|2022-03-28|2021-06-22|2021-12-18|2021-10-11|2021-11-13|
+----------+----------+----------+----------+----------+----------+----------+----------+
|2021-04-02|1         |0         |1         |1         |1         |1         |1         |
|2021-04-21|1         |0         |1         |1         |1         |1         |1         |
|2021-05-01|1         |0         |1         |1         |1         |1         |1         |
|2021-06-03|1         |0         |1         |1         |1         |1         |1         |
|2021-09-07|1         |0         |1         |0         |1         |1         |1         |
|2021-10-12|1         |0         |1         |0         |1         |0         |1         |
|2021-11-02|1         |0         |1         |0         |1         |0         |1         |
+----------+----------+----------+----------+----------+----------+----------+----------+

求和截止日期

aggregated_df = dates_with_deadlines.agg(*[F.sum(str(x.DEADLINES)).alias(str(x.DEADLINES)) for x in deadline_rows])

在这一步之后，你会得到以下aggregated_df 数据框：

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

转置数据帧

result_df = aggregated_df \
  .withColumn('merged', F.array(*[F.struct(F.lit(x.DEADLINES).alias('DEADLINES'), F.col(str(x.DEADLINES)).alias('dt_count')) for x in deadline_rows])) \
  .drop(*[str(x.DEADLINES) for x in deadline_rows]) \
  .withColumn('data', F.explode('merged')) \
  .drop('merged') \
  .withColumn('DEADLINES', F.col('data.DEADLINES')) \
  .withColumn('dt_count', F.col('data.dt_count')) \
  .drop('data')

你有你预期的result_df数据框：

+----------+--------+
|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       |
+----------+--------+

完整代码

from pyspark.sql import functions as F

deadline_rows = deadlines_df.collect()

dates_with_deadlines = dates_df
for row in deadline_rows:
    dates_with_deadlines = dates_with_deadlines.withColumn(
        str(row.DEADLINES),
        F.when(
          dates_df.DT_DATE.between(ref_date, row.DEADLINES), F.lit(1))
        .otherwise(
          F.lit(0)
        )
    )

aggregated_df = dates_with_deadlines.agg(*[F.sum(str(x.DEADLINES)).alias(str(x.DEADLINES)) for x in deadline_rows])

result_df = aggregated_df \
  .withColumn('merged', F.array(*[F.struct(F.lit(x.DEADLINES).alias('DEADLINES'), F.col(str(x.DEADLINES)).alias('dt_count')) for x in deadline_rows])) \
  .drop(*[str(x.DEADLINES) for x in deadline_rows]) \
  .withColumn('data', F.explode('merged')) \
  .drop('merged') \
  .withColumn('DEADLINES', F.col('data.DEADLINES')) \
  .withColumn('dt_count', F.col('data.dt_count')) \
  .drop('data')

此解决方案的优点和限制

使用此解决方案，使用分布式系统无法完成的唯一步骤是转置步骤。

此外，我们不是按顺序对每个截止日期列执行所有聚合，而不是您当前的解决方案。

但是，此解决方案仅在截止日期很少（数百个，可能是数千个截止日期）的情况下才有效，首先是因为我们在 Spark 驱动程序中使用.collect() 检索所有这些截止日期，其次是因为在第一步中我们创建了一个列每个截止日期，创建包含大量数据的行，最后因为最后一步也仅在一个执行器上执行。