如何从Apache Spark中的数据框中选择相同大小的分层样本?

时间:2017-01-07 01:15:57

标签: apache-spark pyspark spark-dataframe

我在Spark 2中有一个数据框,如下所示,用户有50到数千个帖子。我想创建一个新的数据框,其中包含原始数据框中的所有用户,但每个用户只有5个随机抽样的帖子。

+--------+--------------+--------------------+
| user_id|       post_id|                text|
+--------+--------------+--------------------+
|67778705|44783131591473|some text...........|
|67778705|44783134580755|some text...........|
|67778705|44783136367108|some text...........|
|67778705|44783136970669|some text...........|
|67778705|44783138143396|some text...........|
|67778705|44783155162624|some text...........|
|67778705|44783688650554|some text...........|
|68950272|88655645825660|some text...........|
|68950272|88651393135293|some text...........|
|68950272|88652615409812|some text...........|
|68950272|88655744880460|some text...........|
|68950272|88658059871568|some text...........|
|68950272|88656994832475|some text...........|
+--------+--------------+--------------------+

posts.groupby('user_id').agg(sample('post_id'))之类的东西,但在pyspark中没有这样的功能。

有什么建议吗?

更新

这个问题与另一个密切相关的问题stratified-sampling-in-spark有两点不同:

  1. 它询问不成比例的分层抽样,而不是上述其他问题中的常用比例法。
  2. 它要求在Spark的Dataframe API而不是RDD中执行此操作。
  3. 我还更新了问题的标题以澄清这一点。

2 个答案:

答案 0 :(得分:10)

您可以对input http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.DataFrame.sampleBy

使用.sampleBy(...)方法

这是一个有效的例子:

DataFrames

答案 1 :(得分:4)

使用sampleBy将导致近似解决方案。这是一种替代方法,比上述方法更加晦涩,但总是产生完全相同的样本大小。

import org.apache.spark.sql.functions.row_number
import org.apache.spark.sql.expressions.Window

df.withColumn("row_num",row_number().over(Window.partitionBy($"user_id").orderBy($"something_random"))

如果您还没有随机ID,则可以使用org.apache.spark.sql.functions.rand创建一个具有随机值的列,以保证您的随机抽样。