Учитывая, что вы ищете верхние n 'y' элементов, которые принадлежат каждой группе 'x'.
from pyspark.sql import Window
from pyspark.sql import functions as F
import sys
rdd = sc.parallelize([(1, 10), (3, 11), (1, 8), (1, 12), (3, 7), (3, 9)])
df = spark.createDataFrame(rdd, ['x', 'y'])
df.show()
df_g = df.groupBy('x').agg(F.collect_list('y').alias('y'))
df_g = df_g.withColumn('y_sorted', F.sort_array('y', asc = False))
df_g.withColumn('y_slice', F.slice(df_g.y_sorted, 1, 2)).show()
Вывод
+---+-----------+-----------+--------+
| x| y| y_sorted| y_slice|
+---+-----------+-----------+--------+
| 1|[10, 8, 12]|[12, 10, 8]|[12, 10]|
| 3| [11, 7, 9]| [11, 9, 7]| [11, 9]|
+---+-----------+-----------+--------+