我的数据框的格式为：

+---+---+----+
|  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",10.),"b3",("a3",10.)],schema=StructType([StructField("A",StringType(),True),StructField("B",StructField("dist",FloatType(),True)]))

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

但是，由于Window函数使用OrderBy，因此该解决方案效率不高。

Pyspark还有另一种解决方案？例如，类似于用于RDD的.top（k，key =-）的方法？

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