Spark中的分层采样

Question

我有包含用户和购买数据的数据集。这是一个示例，其中第一个元素是 userId，第二个是 productId，第三个表示布尔值。

(2147481832,23355149,1)
(2147481832,973010692,1)
(2147481832,2134870842,1)
(2147481832,541023347,1)
(2147481832,1682206630,1)
(2147481832,1138211459,1)
(2147481832,852202566,1)
(2147481832,201375938,1)
(2147481832,486538879,1)
(2147481832,919187908,1)
... 

我想确保我只获取每个用户数据的 80% 并构建一个 RDD，而将其余 20% 的数据构建另一个 RDD。让我们调用训练和测试。我想远离使用 groupBy 开始，因为它会产生内存问题，因为数据集很大。最好的方法是什么？

我可以进行关注，但这不会为每个用户提供 80% 的用户。

val percentData = data.map(x => ((math.random * 100).toInt, x._1. x._2, x._3)
val train = percentData.filter(x => x._1 < 80).values.repartition(10).cache()

Answer 1

Holden 的回答中有一个可能的解决方案，这里有一些其他的解决方案：

使用 RDD：

您可以使用 PairRDDFunctions 类中的 sampleByKeyExact 转换。

sampleByKeyExact（boolean withReplacement，scala.collection.Map 分数，长种子） 返回此 RDD 的一个子集，该子集按键（通过分层采样）采样，准确包含每个层（具有相同键的对组）的 math.ceil(numItems * samplingRate)。

这就是我会做的：

考虑以下列表：

val seq = Seq(
                (2147481832,23355149,1),(2147481832,973010692,1),(2147481832,2134870842,1),(2147481832,541023347,1),
                (2147481832,1682206630,1),(2147481832,1138211459,1),(2147481832,852202566,1),(2147481832,201375938,1),
                (2147481832,486538879,1),(2147481832,919187908,1),(214748183,919187908,1),(214748183,91187908,1)
           )

我会创建一个RDD Pair，将所有用户映射为键：

val data = sc.parallelize(seq).map(x => (x._1,(x._2,x._3)))

然后我将为每个键设置fractions，如下所示，因为 sampleByKeyExact 为每个键获取一个分数映射：

val fractions = data.map(_._1).distinct.map(x => (x,0.8)).collectAsMap

我在这里所做的是映射键以找到不同的键，然后将每个键关联到等于0.8 的分数。我将整体收集为地图。

现在去采样：

import org.apache.spark.rdd.PairRDDFunctions
val sampleData = data.sampleByKeyExact(false, fractions, 2L)

或

val sampleData = data.sampleByKeyExact(withReplacement = false, fractions = fractions,seed = 2L)

您可以检查您的键或数据或数据样本的计数：

scala > data.count
// [...]
// res10: Long = 12

scala > sampleData.count
// [...]
// res11: Long = 10

使用数据帧：

让我们考虑与上一节相同的数据 (seq)。

val df = seq.toDF("keyColumn","value1","value2")
df.show
// +----------+----------+------+
// | keyColumn|    value1|value2|
// +----------+----------+------+
// |2147481832|  23355149|     1|
// |2147481832| 973010692|     1|
// |2147481832|2134870842|     1|
// |2147481832| 541023347|     1|
// |2147481832|1682206630|     1|
// |2147481832|1138211459|     1|
// |2147481832| 852202566|     1|
// |2147481832| 201375938|     1|
// |2147481832| 486538879|     1|
// |2147481832| 919187908|     1|
// | 214748183| 919187908|     1|
// | 214748183|  91187908|     1|
// +----------+----------+------+

我们将需要底层的RDD 来执行我们创建此RDD 中元素的元组的操作，方法是将我们的键定义为第一列：

val data: RDD[(Int, Row)] = df.rdd.keyBy(_.getInt(0))
val fractions: Map[Int, Double] = data.map(_._1)
                                      .distinct
                                      .map(x => (x, 0.8))
                                      .collectAsMap

val sampleData: RDD[Row] = data.sampleByKeyExact(withReplacement = false, fractions, 2L)
                               .values

val sampleDataDF: DataFrame = spark.createDataFrame(sampleData, df.schema) // you can use sqlContext.createDataFrame(...) instead for spark 1.6)

您现在可以检查您的密钥或df 或数据样本的计数：

scala > df.count
// [...]
// res9: Long = 12

scala > sampleDataDF.count
// [...]
// res10: Long = 10

由于Spark 1.5.0你可以使用DataFrameStatFunctions.sampleBy方法：

df.stat.sampleBy("keyColumn", fractions, seed)

Answer 2

像这样的东西可能非常适合像“Blink DB”这样的东西，但让我们看看这个问题。有两种方法可以解释您所问的一种是：

1) 您想要 80% 的用户，并且想要他们的所有数据。 2) 你想要每个用户数据的 80%

对于#1，您可以做一个映射来获取用户 ID，调用 distinct，然后对其中的 80% 进行采样（您可能需要查看 MLUtils 或 BernoulliCellSampler 中的 kFold）。然后，您可以将输入数据过滤为您想要的一组 ID。

对于 #2，您可以查看 BernoulliCellSampler 并直接应用它。