Groupby 列並為 pyspark 中的另一列值創建列表

Question

我有一個數據框如下：

dummy = pd.DataFrame([[1047,2021,0.38],[1056,2021,0.19]],columns=['reco','user','score'])
dummy

reco    user    score
0   1047    2021    0.38
1   1056    2021    0.19

我希望 output 看起來像這樣：

user    score   reco
2021    [0.38, 0.19]    [1047, 1056]

我想按用戶分組，然后應按分數按降序創建列表，並且記錄應與其分數值相對應。

我嘗試了 collect_list 但順序發生了變化。 我想保持相同的順序。

Answer 1

您可以通過在 window function 上應用collect_list來保留排序。 在這種情況下， window 由user分區並按score descending排序。

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql import Window as W

dummy = pd.DataFrame([[1047,2021,0.38],[1056,2021,0.19]],columns=['reco','user','score'])

df = spark.createDataFrame(dummy)

window_spec = W.partitionBy("user").orderBy(F.desc("score"))
ranged_spec = window_spec.rowsBetween(W.unboundedPreceding, W.unboundedFollowing)

df.withColumn("reco", F.collect_list("reco").over(window_spec))\
  .withColumn("score", F.collect_list("score").over(window_spec))\
  .withColumn("rn", F.row_number().over(window_spec))\
  .where("rn == 1")\
  .drop("rn").show()

Output

+------------+----+------------+
|        reco|user|       score|
+------------+----+------------+
|[1047, 1056]|2021|[0.38, 0.19]|
+------------+----+------------+