Question

我是Spark的初学者，我正在尝试创建一个包含每个键的前3个值的RDD（不仅仅是前3个值）。我当前的RDD包含以下格式的数千个条目：

(key, String, value)

所以想象一下我的RDD内容是这样的：

[("K1", "aaa", 6), ("K1", "bbb", 3), ("K1", "ccc", 2), ("K1", "ddd", 9),
("B1", "qwe", 4), ("B1", "rty", 7), ("B1", "iop", 8), ("B1", "zxc", 1)]

我当前可以在RDD中显示前三个值，如下所示：

("K1", "ddd", 9)
("B1", "iop", 8)
("B1", "rty", 7)

使用：

top3RDD = rdd.takeOrdered(3, key = lambda x: x[2])

相反，我想要的是收集RDD中每个键的前三个值，所以我想返回它：

("K1", "ddd", 9)
("K1", "aaa", 6)
("K1", "bbb", 3)
("B1", "iop", 8)
("B1", "rty", 7)
("B1", "qwe", 4)

Answer 1

您需要将key分组，然后您可以使用heapq.nlargest从每个组中获取前3个值：

from heapq import nlargest
rdd.groupBy(
    lambda x: x[0]
).flatMap(
    lambda g: nlargest(3, g[1], key=lambda x: x[2])
).collect()

[('B1', 'iop', 8), 
 ('B1', 'rty', 7), 
 ('B1', 'qwe', 4), 
 ('K1', 'ddd', 9), 
 ('K1', 'aaa', 6), 
 ('K1', 'bbb', 3)]

Answer 2

如果您愿意将rdd转换为DataFrame，则可以按key定义要分区的窗口，然后按value降序排序。使用此窗口计算行号，并选择行号小于或等于3的行。

import pyspark.sql.functions as f
import pyspark.sql.Window

w = Window.partitionBy("key").orderBy(f.col("value").desc())

rdd.toDF(["key", "String", "value"])\
    .select("*", f.row_number().over(w).alias("rowNum"))\
    .where(f.col("rowNum") <= 3)\
    .drop("rowNum")
    .show()

获取Spark中RDD中每个键的前三个值

2 个答案: