目标是为每个组在数据框中采样(不替换)不同数量的行。要为特定组采样的行数在另一个数据帧中。
示例:idDF是要采样的数据帧。这些组由ID列表示。数据框planDF指定每个组的采样行数,其中" datesToUse"表示行数," ID"表示该组。 " totalDates"是该组的总行数,可能有用也可能没用。
最终结果应该从第一组(ID 1)采样3行,从第二组采样2行(ID 2),从第3组采样1行(ID 3)。
val idDF = Seq(
(1, "2017-10-03"),
(1, "2017-10-22"),
(1, "2017-11-01"),
(1, "2017-10-02"),
(1, "2017-10-09"),
(1, "2017-12-24"),
(1, "2017-10-20"),
(2, "2017-11-17"),
(2, "2017-11-12"),
(2, "2017-12-02"),
(2, "2017-10-03"),
(3, "2017-12-18"),
(3, "2017-11-21"),
(3, "2017-12-13"),
(3, "2017-10-08"),
(3, "2017-10-16"),
(3, "2017-12-04")
).toDF("ID", "date")
val planDF = Seq(
(1, 3, 7),
(2, 2, 4),
(3, 1, 6)
).toDF("ID", "datesToUse", "totalDates")
这是结果数据框应该是什么样子的一个例子:
+---+----------+
| ID| date|
+---+----------+
| 1|2017-10-22|
| 1|2017-11-01|
| 1|2017-10-20|
| 2|2017-11-12|
| 2|2017-10-03|
| 3|2017-10-16|
+---+----------+
到目前为止,我尝试使用DataFrame的示例方法:https://spark.apache.org/docs/1.5.0/api/java/org/apache/spark/sql/DataFrame.html 这是一个适用于整个数据框的示例。
def sampleDF(DF: DataFrame, datesToUse: Int, totalDates: Int): DataFrame = {
val fraction = datesToUse/totalDates.toFloat.toDouble
DF.sample(false, fraction)
}
我无法弄清楚如何为每个群体使用这样的东西。我尝试将planDF表加入idDF表并使用窗口分区。
我的另一个想法是以某种方式创建一个随机标记为True / false的新列,然后对该列进行过滤。
答案 0 :(得分:1)
完全停留在Dataframes中的另一个选择是使用planDF
计算概率,使用idDF
加入,附加一列随机数然后过滤。有用的是,sql.functions
具有rand
功能。
import org.apache.spark.sql.functions._
import spark.implicits._
val probabilities = planDF.withColumn("prob", $"datesToUse" / $"totalDates")
val dfWithProbs = idDF.join(probabilities, Seq("ID"))
.withColumn("rand", rand())
.where($"rand" < $"prob")
(你想要仔细检查那不是整数除法。)
答案 1 :(得分:0)
假设您的planDF足够小,可以collect
,您可以使用Scala的foldLeft
遍历id
列表并累积每个{的示例数据帧{1}}:
id
请注意,方法import org.apache.spark.sql.{Row, DataFrame}
def sampleByIdDF(DF: DataFrame, id: Int, datesToUse: Int, totalDates: Int): DataFrame = {
val fraction = datesToUse.toDouble / totalDates
DF.where($"id" === id ).sample(false, fraction)
}
val emptyDF = Seq.empty[(Int, String)].toDF("ID", "date")
val planList = planDF.rdd.collect.map{ case Row(x: Int, y: Int, z: Int) => (x, y, z) }
// planList: Array[(Int, Int, Int)] = Array((1,3,7), (2,2,4), (3,1,6))
planList.foldLeft( emptyDF ){
case (accDF: DataFrame, (id: Int, num: Int, total: Int)) =>
accDF union sampleByIdDF(idDF, id, num, total)
}
// res1: org.apache.spark.sql.DataFrame = [ID: int, date: string]
// res1.show
// +---+----------+
// | ID| date|
// +---+----------+
// | 1|2017-10-03|
// | 1|2017-11-01|
// | 1|2017-10-02|
// | 1|2017-12-24|
// | 1|2017-10-20|
// | 2|2017-11-17|
// | 2|2017-11-12|
// | 2|2017-12-02|
// | 3|2017-11-21|
// | 3|2017-12-13|
// +---+----------+
不一定会生成方法参数中指定的确切数量的样本。这是一个相关的SO Q&A。
如果你的planDF很大,你可能不得不考虑使用RDD&#39; aggregate,它具有以下签名(跳过隐式参数):
sample()
它的工作方式与def aggregate[U](zeroValue: U)(seqOp: (U, T) ⇒ U, combOp: (U, U) ⇒ U): U
类似,不同之处在于它在一个分区中有一个累加运算符,另一个用于汇总来自不同分区的结果。