apache-spark - Pyspark --- 添加新列，其中包含每组的值

Question

假设我有以下数据集:

a | b   
1 | 0.4 
1 | 0.8 
1 | 0.5 
2 | 0.4
2 | 0.1

我想添加一个名为“label”的新列，其中的值是在本地为 a 中的每组值确定的。组 a 中 b 的最高值标记为 1，所有其他值标记为 0。

输出如下:

a | b   | label
1 | 0.4 | 0
1 | 0.8 | 1
1 | 0.5 | 0
2 | 0.4 | 1
2 | 0.1 | 0

如何使用 PySpark 高效地完成此操作？

Answer 1

您可以使用窗口函数来做到这一点。首先，您需要一些导入:

from pyspark.sql.functions import desc, row_number, when
from pyspark.sql.window import Window

和窗口定义:

w = Window().partitionBy("a").orderBy(desc("b"))

最后你使用这些:

df.withColumn("label", when(row_number().over(w) == 1, 1).otherwise(0))

例如数据:

df = sc.parallelize([
    (1, 0.4), (1, 0.8), (1, 0.5), (2, 0.4), (2, 0.1)
]).toDF(["a", "b"])

结果是:

+---+---+-----+
|  a|  b|label|
+---+---+-----+
|  1|0.8|    1|
|  1|0.5|    0|
|  1|0.4|    0|
|  2|0.4|    1|
|  2|0.1|    0|
+---+---+-----+