具有条件重置的 PySpark Windows 函数

Question

我有一个这样的数据框

|  user_id  | acivity_date |
|  -------- | ------------ |
| 49630701  | 1/1/2019     |
| 49630701  | 1/10/2019    |
| 49630701  | 1/28/2019    |
| 49630701  | 2/5/2019     |
| 49630701  | 3/10/2019    |
| 49630701  | 3/21/2019    |
| 49630701  | 5/25/2019    |
| 49630701  | 5/28/2019    |
| 49630701  | 9/10/2019    |
| 49630701  | 1/1/2020     |
| 49630701  | 1/10/2020    |
| 49630701  | 1/28/2020    |
| 49630701  | 2/10/2020    |
| 49630701  | 3/10/2020    |

我需要创建的是“组”列，逻辑是对于每个用户，我们需要保留组 # 直到累积日期差异小于 30 天，只要累积日期差异大于 30 天那么我们需要增加组#并将累积日期差异重置为零

|  user_id  | acivity_date | Group |
|  -------- | ------------ | ----- | 
| 49630701  | 1/1/2019     |  1    |
| 49630701  | 1/10/2019    |  1    |
| 49630701  | 1/28/2019    |  1    | 
| 49630701  | 2/5/2019     |  2    | <- Cumulative date diff till here is 35, which is greater than 30, so increment the Group by 1 and reset the cumulative diff to 0 
| 49630701  | 3/10/2019    |  3    |
| 49630701  | 3/21/2019    |  3    |
| 49630701  | 5/25/2019    |  4    |
| 49630701  | 5/28/2019    |  4    |
| 49630701  | 9/10/2019    |  5    |
| 49630701  | 1/1/2020     |  6    |
| 49630701  | 1/10/2020    |  6    |
| 49630701  | 1/28/2020    |  6    |
| 49630701  | 2/10/2020    |  7    |
| 49630701  | 3/10/2020    |  7    |

我尝试使用带有循环的以下代码，但效率不高，它运行了几个小时。有没有更好的方法来实现这一目标？任何帮助将不胜感激

df= spark.read.table('excel_file)
df1 = df.select(col("user_id"), col("activity_date")).distinct()
partitionWindow = Window.partitionBy("user_id").orderBy(col("activity_date").asc())
lagTest = lag(col("activity_date"), 1, "0000-00-00 00:00:00").over(partitionWindow)
df1 = df1.select(col("*"), (datediff(col("activity_date"),lagTest)).cast("int").alias("diff_val_with_previous"))
df1 = df1.withColumn('diff_val_with_previous', when(col('diff_val_with_previous').isNull(), lit(0)).otherwise(col('diff_val_with_previous')))
distinctUser = [i['user_id'] for i in df1.select(col("user_id")).distinct().collect()]
rankTest = rank().over(partitionWindow)
df2 = df1.select(col("*"), rankTest.alias("rank"))

interimSessionThreshold = 30
totalSessionTimeThreshold = 30
rowList = []

for x in distinctUser:
  tempDf = df2.filter(col("user_id") == x).orderBy(col('activity_date'))
  cumulDiff = 0
  group = 1
  startBatch = True
  len_df = tempDf.count()
  dp = 0
  for i in range(1, len_df+1):
    r = tempDf.filter(col("rank") == i)
    dp = r.select("diff_val_with_previous").first()[0]
    cumulDiff += dp
    if ((dp <= interimSessionThreshold) & (cumulDiff <= totalSessionTimeThreshold)):
      startBatch=False
      rowList.append([r.select("user_id").first()[0], r.select("activity_date").first()[0], group])
    else:
      group += 1
      cumulDiff = 0
      startBatch = True
      dp = 0
      rowList.append([r.select("user_id").first()[0], r.select("activity_date").first()[0], group])

ddf = spark.createDataFrame(rowList, ['user_id', 'activity_date', 'group'])

Answer 1

我能想到两种解决方案，但没有一个完全符合您的要求：

from pyspark.sql import functions as F, Window

df.withColumn(
    "idx", F.monotonically_increasing_id()
).withColumn(
    "date_as_num", F.unix_timestamp("activity_date")
).withColumn(
    "group", F.min("idx").over(Window.partitionBy('user_id').orderBy("date_as_num").rangeBetween(- 60 * 60 * 24 * 30, 0))
).withColumn(
    "group", F.dense_rank().over(Window.partitionBy("user_id").orderBy("group"))
).show()

+--------+-------------+----------+-----------+-----+                           
| user_id|activity_date|       idx|date_as_num|group|
+--------+-------------+----------+-----------+-----+
|49630701|   2019-01-01|         0| 1546300800|    1|
|49630701|   2019-01-10|         1| 1547078400|    1|
|49630701|   2019-01-28|         2| 1548633600|    1|
|49630701|   2019-02-05|         3| 1549324800|    2|
|49630701|   2019-03-10|         4| 1552176000|    3|
|49630701|   2019-03-21|         5| 1553126400|    3|
|49630701|   2019-05-25|         6| 1558742400|    4|
|49630701|   2019-05-28|8589934592| 1559001600|    4|
|49630701|   2019-09-10|8589934593| 1568073600|    5|
|49630701|   2020-01-01|8589934594| 1577836800|    6|
|49630701|   2020-01-10|8589934595| 1578614400|    6|
|49630701|   2020-01-28|8589934596| 1580169600|    6|
|49630701|   2020-02-10|8589934597| 1581292800|    7|
|49630701|   2020-03-10|8589934598| 1583798400|    8|
+--------+-------------+----------+-----------+-----+

或

df.withColumn(
    "group",
    F.datediff(
        F.col("activity_date"),
        F.lag("activity_date").over(
            Window.partitionBy("user_id").orderBy("activity_date")
        ),
    ),
).withColumn(
    "group", F.sum("group").over(Window.partitionBy("user_id").orderBy("activity_date"))
).withColumn(
    "group", F.floor(F.coalesce(F.col("group"), F.lit(0)) / 30)
).withColumn(
    "group", F.dense_rank().over(Window.partitionBy("user_id").orderBy("group"))
).show()

+--------+-------------+-----+                                                  
| user_id|activity_date|group|
+--------+-------------+-----+
|49630701|   2019-01-01|    1|
|49630701|   2019-01-10|    1|
|49630701|   2019-01-28|    1|
|49630701|   2019-02-05|    2|
|49630701|   2019-03-10|    3|
|49630701|   2019-03-21|    3|
|49630701|   2019-05-25|    4|
|49630701|   2019-05-28|    4|
|49630701|   2019-09-10|    5|
|49630701|   2020-01-01|    6|
|49630701|   2020-01-10|    6|
|49630701|   2020-01-28|    7|
|49630701|   2020-02-10|    7|
|49630701|   2020-03-10|    8|
+--------+-------------+-----+