我的数据框的格式为：
+---+---+----+
| 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')
但是，由于Window函数使用OrderBy，因此该解决方案效率不高。
Pyspark还有另一种解决方案？例如，类似于用于RDD的.top（k，key =-）的方法？
我找到了类似的答案here，但是使用了scala而不是python。