python - 限制每个唯一的 pyspark 数据帧列值返回的行，而无需循环

Question

我有一个简短的问题，希望有人能回答。

是否可以在 pyspark 数据框中限制特定列中唯一值返回的结果数量？

例如，我有一个包含 4 列的数据框:

+-------+------------------+--------------+-------------+
|partner|productId         |    mediumtype|   prediction|
+-------+------------------+--------------+-------------+
|7005410|               544|             1| 0.0026476993|
|7005410|               549|             2|-2.6975607E-4|
|7005410|               626|             3| 2.0409889E-4|
|7005410|               840|             2| 3.6301462E-5|
|7005410|              1192|             3| 2.2148499E-5|
+-------+------------------+--------------+-------------+

合作伙伴列有一个数值(在本例中为 7005410)。在示例中，您看到 5 行具有唯一的合作伙伴 ID，但实际上有 7.000 多个唯一的合作伙伴 ID。这不仅适用于该合作伙伴，还适用于所有其他 4.000 多个合作伙伴，这意味着总共有 7.000 * 4.000 = 28.000.000 行。

我希望它只为每个唯一的合作伙伴 ID 获取前 5 个预测(限制(5))。到目前为止，我已经将数据帧放入一个带有过滤器的循环中，但由于 Spark 的延迟加载特性，这需要大量时间才能完成，我想知道是否只有一种方法可以将限制(5)应用于每个唯一的合作伙伴 ID。

Answer 1

尝试这个。拿 row_number() 超过 window 并过滤掉任何地方 rowNum 是 greater than 5 .考虑到数据已经排序，我们可以通过窗口中的文字常量进行排序。

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w=Window().partitionBy("partner").orderBy(F.lit(1))

df.withColumn("rowNum", F.row_number().over(w))\
  .filter('rowNum<=5').drop("rowNum").show()