具有特定条件的时间戳的唯一ID

Question

我有一个由这两列组成的Spark数据帧，我想根据timestamp列上的一个简单条件为每一行提供一个ID。 对于每个用户，如果下一个时间戳和上一个时间戳之间的差异小于10秒，我会继续追加相同的ID，否则我会更新ID并继续，直到我为每一行分配了一个ID。
因为我使用的是PySpark，所以所有内容都应该用python编码。
为了让事情更容易理解，这里有一个例子：

启动DF

    +------------------+
    |   User| timestamp|
    +------------------+
    |  user0|     100  |
    |  user1|     102  |
    |  user0|     109  |
    |  user2|     103  |
    |  user1|     108  |
    |  user0|     119  |
    |  user0|     140  |
    |  user0|     142  |
    +------------------+

理想的DF是这样的

    +----------------------+
    |   User| timestamp| ID|
    +----------------------+
    |  user0|     100  |  1|
    |  user1|     102  |  2|
    |  user0|     109  |  1|
    |  user2|     103  |  3|
    |  user1|     108  |  2|
    |  user0|     119  |  1|
    |  user0|     140  |  4|
    |  user0|     142  |  4|
    +----------------------+

或者，如果算法在给定用户的ID之前分配，则可能是类似的情况。我真的不在乎，也没关系

    +----------------------+
    |   User| timestamp| ID|
    +----------------------+
    |  user0|     100  |  1|
    |  user1|     102  |  3|
    |  user0|     109  |  1|
    |  user2|     103  |  4|
    |  user1|     108  |  3|
    |  user0|     119  |  1|
    |  user0|     140  |  2|
    |  user0|     142  |  2|
    +----------------------+

如您所见，时间戳为140的用户0具有不同的ID（2），因为与前一个时间戳的差异大于10

如果我可以循环并动态分配每个单元格，这将非常容易，但这违背了使用spark数据帧的目的，而且我认为这是不可能的，因为它们是不可变的

在Spark中，最有效的方法是什么

Answer 1

您可以先为每个用户生成一个ID，然后在不同的用户之间组合它们，如下所示

import pyspark.sql.functions as F
from pyspark.sql.window import Window

w = Window.partitionBy('User').orderBy('timestamp')

df2 = df.withColumn(
    'begin',
    F.coalesce(
        F.col('timestamp') - F.lag('timestamp').over(w) > 10, 
        F.lit(True)
    ).cast('int')
).withColumn(
    'userid',
    F.sum('begin').over(w.rowsBetween(Window.unboundedPreceding, 0))
).withColumn(
    'ID',
    F.dense_rank().over(Window.orderBy('userid', 'User'))
)

# If you just want to keep your columns, do:
# df2 = df2.select('User', 'timestamp', 'ID')

df2.show()
+  -+    -+  -+   + -+
| User|timestamp|begin|userid| ID|
+  -+    -+  -+   + -+
|user0|      100|    1|     1|  1|
|user0|      109|    0|     1|  1|
|user0|      119|    0|     1|  1|
|user1|      102|    1|     1|  2|
|user1|      108|    0|     1|  2|
|user2|      103|    1|     1|  3|
|user0|      140|    1|     2|  4|
|user0|      142|    0|     2|  4|
+  -+    -+  -+   + -+