python PySpark只选择第一个时间戳,该时间戳在同一组的另一个列时间戳之后

xjreopfe  于 2022-12-10  发布在  Python
关注(0)|答案(2)|浏览(124)

我有以下数据集:

ID Timestamp1             Timestamp2
 1  2022-01-01T12:00:00    2022-01-01T12:30:00
 1  2022-01-01T13:00:00    2022-01-01T12:30:00
 1  2022-01-01T14:00:00    2022-01-01T12:30:00
 2  2022-02-02T09:00:00    2022-02-02T10:15:00
 2  2022-02-02T10:00:00    2022-02-02T10:15:00
 2  2022-02-03T11:00:00    2022-02-02T10:15:00

每个唯一ID在Timestamp2中始终具有完全相同的值,而在Timestamp1中始终具有不同的值。目标是获取列Timestamp1中每个ID的timestamp2值之后的第一个时间戳。结果如下所示:

ID Timestamp1             Timestamp2
 1  2022-01-01T13:00:00    2022-01-01T12:30:00
 2  2022-02-03T11:00:00    2022-02-02T10:15:00

我不确定在这种情况下最好的解决方案是使用窗口函数还是直接使用groupby函数。

EDIT我确实找到了一个解决方案,但我认为它不是最好的,因此我仍在寻找更好的解决方案。我的想法是使用Timestamp2列和Timestamp1列之间的差异创建一个新列。然后我可以删除所有负值,最后对ID使用groupby并取每个Timestamp1的最小值。

xoefb8l8

xoefb8l81#

下面的简单group by怎么样?它只有一个shuffle操作(group by),所以我不认为窗口函数能真正帮助优化性能。

df \
.filter(col('Timestamp1') > col('Timestamp2')) \
.groupBy('Timestamp2') \
.agg(min('Timestamp1').alias('Timestamp1'))
c3frrgcw

c3frrgcw2#

像这样的东西?

from pyspark.sql import functions as F, types as T, Window

(
    df
    # Keep only rows where Timestamp1 is greater than Timestamp2
    .where(F.col('Timestamp1') > F.col('Timestamp2'))
    # add a column with the rank based on Timestamp1 for each ID
    .withColumn('rank', F.rank().over(Window.partitionBy('ID').orderBy('Timestamp1')))
    # Keep only rows with rank equal to one, i.e. the ones where Timestamp1 is the closest to Timestamp2
    .where(F.col('rank') == 1)
    .drop('rank')
)

其结果是:

+---+-------------------+-------------------+
|ID |Timestamp1         |Timestamp2         |
+---+-------------------+-------------------+
|1  |2022-01-01 13:00:00|2022-01-01 12:30:00|
|2  |2022-02-03 11:00:00|2022-02-02 10:15:00|
+---+-------------------+-------------------+

相关问题