为spark scala

时间:2018-03-01 19:25:13

标签: scala apache-spark apache-spark-sql spark-dataframe

目标是为每个组在数据框中采样(不替换)不同数量的行。要为特定组采样的行数在另一个数据帧中。

示例: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的新列,然后对该列进行过滤。

2 个答案:

答案 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 类似,不同之处在于它在一个分区中有一个累加运算符,另一个用于汇总来自不同分区的结果。