Pyspark 從每個組的列中獲取第一個值

Question

我在 pyspark 中有一個數據框，看起來像這樣

|Id1| id2  |row  |grp    |
|12 | 1234 |1    | 1     |
|23 | 1123 |2    | 1     |
|45 | 2343 |3    | 2     |
|65 | 2345 |1    | 2     |
|67 | 3456 |2    | 2     |```

I need to retrieve value for id2 corresponding to row = 1 and update all id2 values within a grp to that value.
This should be the final result

|Id1 | id2  |row |grp|
|12  |1234  |1   |1  |
|23  |1234  |2   |1  |
|45  |2345  |3   |2  |
|65  |2345  |1   |2  |
|67  |2345  |2   |2  |

我嘗試做類似 df.groupby('grp').sort('row').first('id2') 但顯然 sort 和 orderby 在 pyspark 中不適用於 groupby。

知道如何解決這個問題嗎？

Answer 1

與.rowsBetween的回答非常相似，不使用.rowsBetween

您基本上為每個grp創建一個Window ，然后按row對行進行排序並為每個grp選擇第一個id2 。

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

w = Window.partitionBy('grp').orderBy('row')

df = df.withColumn('id2', F.first('id2').over(w))

df.show()

+---+----+---+---+
|Id1| id2|row|grp|
+---+----+---+---+
| 12|1234|  1|  1|
| 23|1234|  2|  1|
| 65|2345|  1|  2|
| 67|2345|  2|  2|
| 45|2345|  3|  2|
+---+----+---+---+

Answer 2

試試這個 ：

from pyspark.sql import functions as F, Window as W


df.withColumn(
    "id2",
    F.first("id2").over(
        W.partitionBy("grp")
        .orderBy("row")
        .rowsBetween(W.unboundedPreceding, W.currentRow)
    ),
).show()

+---+----+---+---+                                                              
|id1| id2|row|grp|
+---+----+---+---+
| 12|1234|  1|  1|
| 23|1234|  2|  1|
| 65|2345|  1|  2|
| 45|2345|  2|  2|
| 45|2345|  3|  2|
+---+----+---+---+