如何为用户事件提供索引并使用 PySpark 保持顺序？

Question

我有以下数据框：

+------------+------------------+--------------------+
|id.         |install_time_first|           timestamp|
+------------+------------------+--------------------+
|           2|        2022-02-02|2022-02-01 10:03:...|
|           3|        2022-02-01|2022-02-01 10:00:...|
|           2|        2022-02-02|                null|
|           3|        2022-02-01|2022-02-03 11:35:...|
|           1|        2022-02-01|                null|
|           2|        2022-02-02|2022-02-02 10:05:...|
|           3|        2022-02-01|2022-02-01 10:05:...|
|           4|        2022-02-02|                null|
|           1|        2022-02-01|2022-02-01 10:05:...|
|           2|        2022-02-02|2022-02-02 10:05:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           1|        2022-02-01|                null|
|           1|        2022-02-01|2022-02-01 10:03:...|
|           1|        2022-02-01|2022-02-01 10:05:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           2|        2022-02-02|2022-02-02 11:00:...|
|           4|        2022-02-02|2022-02-03 11:35:...|
|           3|        2022-02-01|2022-02-04 11:35:...|
|           1|        2022-02-01|2022-02-01 10:00:...|
+------------+------------------+--------------------+

我想按 install_time_first 对数据框进行排序，并为每个用户（他的所有事件）添加一个索引并保留顺序。例如：

+------------+------------------+--------------------+-----+
|id.         |install_time_first|           timestamp|index|
+------------+------------------+--------------------+-----+
|           1|        2022-02-01|                null|    1|
|           1|        2022-02-01|                null|    1|
|           1|        2022-02-01|2022-02-01 10:00:...|    1|
|           1|        2022-02-01|2022-02-01 10:03:...|    1|
|           1|        2022-02-01|2022-02-01 10:05:...|    1|
|           1|        2022-02-01|2022-02-01 10:05:...|    1|
|           3|        2022-02-01|2022-02-01 10:00:...|    2|
|           3|        2022-02-01|2022-02-01 10:05:...|    2|
|           3|        2022-02-01|2022-02-03 11:35:...|    2|
|           3|        2022-02-01|2022-02-04 11:35:...|    2|
|           2|        2022-02-02|                null|    3|
|           2|        2022-02-02|2022-02-01 10:03:...|    3|
|           2|        2022-02-02|2022-02-02 10:05:...|    3|
|           2|        2022-02-02|2022-02-02 10:05:...|    3|
|           2|        2022-02-02|2022-02-02 11:00:...|    3|
|           4|        2022-02-02|                null|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
|           4|        2022-02-02|2022-02-03 11:35:...|    4|
+------------+------------------+--------------------+-----+

我怎样才能做到这一点？我做不到并保持排序

Answer 1

这里的关键观察是“index”列具有相同的“id”列值，但按“install_time_first”排序，一种查看方式是在 (install_time_first, id) 上分区/排序并为每对分配一个 unqiue 索引，我做了 2 个解决方案，第一个使用连接，第二个使用带有一些技巧的窗口，我更喜欢第一个解决方案，因为第二个解决方案可能性能很重：

准备数据：

spark = SparkSession.builder.master("local[*]").getOrCreate()
df = spark.createDataFrame([
    (2, "2022-02-02", "2022-02-01 10:03"),
    (3, "2022-02-01", "2022-02-01 10:00"),
    (2, "2022-02-02", None),
    (3, "2022-02-01", "2022-02-03 11:35"),
    (1, "2022-02-01", None),
    (2, "2022-02-02", "2022-02-02 10:05"),
    (3, "2022-02-01", "2022-02-01 10:05"),
    (4, "2022-02-02", None),
    (1, "2022-02-01", "2022-02-01 10:05"),
    (2, "2022-02-02", "2022-02-02 10:05"),
    (4, "2022-02-02", "2022-02-03 11:35"),
    (1, "2022-02-01", None),
    (1, "2022-02-01", "2022-02-01 10:03"),
    (1, "2022-02-01", "2022-02-01 10:05"),
    (4, "2022-02-02", "2022-02-03 11:35"),
    (2, "2022-02-02", "2022-02-02 11:00"),
    (4, "2022-02-02", "2022-02-03 11:35"),
    (3, "2022-02-01", "2022-02-04 11:35"),
    (1, "2022-02-01", "2022-02-01 10:00"),
], ("id", "install_time_first", "timestamp"))

解决方案 1：

df_with_index = df.select("id", "install_time_first").distinct().orderBy("install_time_first", "id")
    .withColumn("index",monotonically_increasing_id() + 1)
    .withColumnRenamed("id", "id2").withColumnRenamed("install_time_first", "install_time_first2")
df.join(df_with_index, (df.id == df_with_index.id2) & (df.install_time_first == df_with_index.install_time_first2),
        "left").orderBy("install_time_first", "id").drop("id2", "install_time_first2").show()

解决方案 2：

w = Window.partitionBy(col("id")).orderBy(col("install_time_first"))
w2 = Window.orderBy(col("install_time_first"))
df = df.withColumn("prev_id", lag("id", 1, None).over(w))
df.withColumn("index", when(df.prev_id.isNull() | (df.prev_id != df.id), 1).otherwise(0))
    .withColumn("index", sum("index").over(w2.rowsBetween(Window.unboundedPreceding, Window.currentRow)))
    .orderBy("install_time_first", "id").drop("prev_id").show()

两者给出相同的结果：

+---+------------------+----------------+-----+
| id|install_time_first|       timestamp|index|
+---+------------------+----------------+-----+
|  1|        2022-02-01|2022-02-01 10:05|    1|
|  1|        2022-02-01|2022-02-01 10:00|    1|
|  1|        2022-02-01|            null|    1|
|  1|        2022-02-01|            null|    1|
|  1|        2022-02-01|2022-02-01 10:03|    1|
|  1|        2022-02-01|2022-02-01 10:05|    1|
|  3|        2022-02-01|2022-02-03 11:35|    2|
|  3|        2022-02-01|2022-02-01 10:00|    2|
|  3|        2022-02-01|2022-02-04 11:35|    2|
|  3|        2022-02-01|2022-02-01 10:05|    2|
|  2|        2022-02-02|            null|    3|
|  2|        2022-02-02|2022-02-02 10:05|    3|
|  2|        2022-02-02|2022-02-02 10:05|    3|
|  2|        2022-02-02|2022-02-02 11:00|    3|
|  2|        2022-02-02|2022-02-01 10:03|    3|
|  4|        2022-02-02|            null|    4|
|  4|        2022-02-02|2022-02-03 11:35|    4|
|  4|        2022-02-02|2022-02-03 11:35|    4|
|  4|        2022-02-02|2022-02-03 11:35|    4|
+---+------------------+----------------+-----+