如何从 PySpark 中的另一个分类值（同一列内）获取下一个最早日期？

Question

我不确定我的问题标题是否有意义，所以我将在下面举例说明。

我有一个数据框（在使用 PySpark 的 Databricks 中），如下所示：

自 2020 年 11 月起，table 所有 Stack Exchange 网站现在都支持 Markdown。

Date	Category
2021-01-02	A
2021-01-06	A
2021-01-07	A
2021-01-10	B
2021-01-15	A
2021-01-16	A
2021-01-18	A
2021-01-23	B
2021-01-25	A

我正在尝试获取类别 A 和下一个类别 B 之间的日期差异（以天为单位），示例如下：

Date	Category	Days Diff
2021-01-02	A	8
2021-01-06	A	4
2021-01-07	A	3
2021-01-10	B
2021-01-15	A	8
2021-01-16	A	7
2021-01-18	A	5
2021-01-23	B
2021-01-25	A	...

有没有人知道执行此操作的最佳方法？谢谢！

Answer 1

您可以使用 window 函数来实现。

首先，创建一个仅包含类别为 B 的日期的列。

df = df.withColumn('BDate', when(df.Category == 'B', df.Date))

+----------+--------+----------+
|      Date|Category|     BDate|
+----------+--------+----------+
|2021-01-02|       A|      null|
|2021-01-06|       A|      null|
|2021-01-07|       A|      null|
|2021-01-10|       B|2021-01-10|
|2021-01-15|       A|      null|
...

然后，创建一个 window，它可以查找当前行的任何后续行并找到 window 中的第一个 BDate。取当前 Date 和找到的 BDate 之间的日期差异。

w = Window.orderBy('Date').rowsBetween(0, Window.unboundedFollowing)
df = df.withColumn('Days Diff', F.datediff(F.first(df.BDate, ignorenulls=True).over(w), df.Date))

结果

+----------+--------+----------+---------+
|      Date|Category|     BDate|Days Diff|
+----------+--------+----------+---------+
|2021-01-02|       A|      null|        8|
|2021-01-06|       A|      null|        4|
|2021-01-07|       A|      null|        3|
|2021-01-10|       B|2021-01-10|        0|
|2021-01-15|       A|      null|        8|
|2021-01-16|       A|      null|        7|
|2021-01-18|       A|      null|        5|
|2021-01-23|       B|2021-01-23|        0|
|2021-01-25|       A|      null|     null|
+----------+--------+----------+---------+

如何从 PySpark 中的另一个分类值（同一列内）获取下一个最早日期？

How to get next earliest date from another categorical value (within the same column) in PySpark?

python

date

dataframe

pyspark

databricks