时间戳

首页 > 解决方案 > 有效地计算 PySpark GroupedData 上的前 k 个元素（不是 scala）

pyspark - 有效地计算 PySpark GroupedData 上的前 k 个元素（不是 scala）

问题描述

我有一个形式的数据框：

+---+---+----+
|  A|  B|dist|
+---+---+----+
| a1| b1| 1.0|
| a1| b2| 2.0|
| a2| b1|10.0|
| a2| b2|10.0|
| a2| b3| 2.0|
| a3| b1|10.0|
+---+---+----+

并且，固定max_rank=2，我想得到以下一个

+---+---+----+----+
|  A|  B|dist|rank|
+---+---+----+----+
| a3| b1|10.0|   1|
| a2| b3| 2.0|   1|
| a2| b1|10.0|   2|
| a2| b2|10.0|   2|
| a1| b1| 1.0|   1|
| a1| b2| 2.0|   2|
+---+---+----+----+

执行此操作的经典方法如下

df = sqlContext.createDataFrame([("a1", "b1", 1.), ("a1", "b2", 2.), ("a2", "b1", 10.), ("a2", "b2", 10.), ("a2", "b3", 2.), ("a3", "b1", 10.)], schema=StructType([StructField("A", StringType(), True), StructField("B", StringType(), True),StructField("dist", FloatType(), True)]))

win = Window().partitionBy(df['A']).orderBy(df['dist'])
out = df.withColumn('rank', rank().over(win))
out = out.filter('rank<=2')

但是，由于使用 OrderBy 的 Window 函数，此解决方案效率低下。

Pyspark 还有另一种解决方案吗？例如类似于 .top(k, key=--) 的 RDD 方法？

我在这里找到了类似的答案，但使用 scala 而不是 python。

标签： pysparkpyspark-dataframes

解决方案

推荐阅读