Как сгруппировать по несколько столбцов в список и поддерживать порядок?

Question

Я пытаюсь сгруппировать по двум разным столбцам, но сортирую сгруппированные элементы по отметке времени и сталкиваюсь с некоторыми проблемами.

Исходный кадр данных:

| user_id|group_id|item_id|value|timestamp|
+--------+--------+-------+-----+---------+
|  A     |1.0     |1      |5    |100      | 
|  A     |1.0     |2      |6    |90       | 
|  A     |0.0.    |3      |7    |80       |
|  A     |0.0.    |4      |8    |80       |
+---+----+--------+-------+-----+----------

Что я хочу:

| user_id|group_id|grouped_item_value|timestamps|
+--------+--------+------------------+----------+
|  A     |1.0     |[[2,6],[1,5]]     |[90,100]  |
|  A     |0.0     |[[4,7],[3,8]]     |[80,80]   |

Столбец grouped_item_value сортируется в порядке убывания на основе второго значения в кортеже.

Любая помощь очень ценится. Я бился головой о стену в течение 9 часов.

Answer 1

Вы можете сделать что-то вроде этого:

import pyspark.sql.functions as F

df = (
    df.withColumn("item_value", F.array(F.col("item"), F.col("value")))
    .orderBy("timestamp")
    .groupby(["user_id", "group_id"])
    .agg(F.collect_list(F.col("item_value")), F.collect_list("timestamp"))
)

1. Сначала вы создаете массив из пар item и value.
2. Затем сортируйте фрейм данных по timestamp
3. Затем вы группируете user_id и group_id до collect_list() для item_value и timestamp