如何从相对较大的 Spark 数据框中获取最高百分比并将其保存到文件中

Question

我正在开发一个 Spark 程序，它计算每个用户的概率，从而产生一个相对较大的数据帧（~137.5M 行）。我需要做的是取这些用户的前 10%（10 个是任意的，当然可以更改）并将它们保存到文件中。

一个最小化的例子是：

• 鉴于此数据框： hc.sparkContext.parallelize(Array(("uid1", "0.5"), ("uid2", "0.7"), ("uid3", "0.3"))).toDF("uuid", "prob")
• 并给定阈值 0.3
• 我希望输出为 ("uid2", "0.7") 并保存到文件 "output" 中，因为 "uid2" 的概率最高，我只需要从数据框中获取前 0.3% 的用户

所以我的问题是：使用相对较大的数据框来做到这一点的干净有效的方法是什么？

可以计算输入数据帧的 10% 有多少用户，然后使用 top 和大小。但是，我对此有两个担忧：

1. 使用 top 时 - 是否在洗牌前先减少数据，方法是从每个执行程序中取出前 10%，然后从洗牌数据中取出 10%？如果没有，是否有内置的方法来执行我的建议？或者我应该自己实现它？
2. Top 返回一个仍然很大的数组......我更愿意做的是将其保留为数据帧并保存其输出（也许在洗牌后重新分区数据）。有没有办法做到这一点而不将其转换为数组然后并行化它？

如果需要，我不介意使用 RDD 而不是数据帧

我目前使用的是 Spark 1.6.1

提前致谢

Answer 1

您可以使用窗口函数percent_rank()。但是，由于您不是按组进行排名，因此您将无法使用partitionBy(). 下面是一个例子pySpark：

from pyspark.sql.window import Window
from pyspark.sql.functions import percent_rank, col

window = Window.partitionBy().orderBy(df['prob'].desc())

df.select('*', percent_rank().over(window).alias('rank')) 
  .filter(col('rank') <= 0.3) # top 30% for example
  .show() 
+----+----+----+
|uuid|prob|rank|
+----+----+----+
|uid2| 0.7| 0.0|
+----+----+----+

数据：

df = sc.parallelize([("uid1", "0.5"), 
                     ("uid2", "0.7"), 
                     ("uid3", "0.3")]).toDF(["uuid", "prob"])