[英]Sample a different number of random rows for every group in a dataframe in spark scala
目標是為每個組在數據幀中采樣(不替換)不同數量的行。 要為特定組采樣的行數在另一個數據幀中。
示例:idDF是要采樣的數據幀。 這些組由ID列表示。 數據幀planDF指定每個組的樣本行數,其中“datesToUse”表示行數,“ID”表示組。 “totalDates”是該組的總行數,可能有用也可能沒用。
最終結果應該具有從第一組(ID 1)采樣的3行,從第二組采樣的2行(ID 2)和從第三組采樣的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的新列,然后對該列進行過濾。
完全在入住Dataframes另一種選擇是計算概率使用planDF
,與加盟idDF
,追加隨機數的一列,然后進行篩選。 sql.functions
是, 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")
(你需要仔細檢查那不是整數除法。)
假設您的planDF足夠小以便collect
,您可以使用Scala的foldLeft
遍歷id
列表並累積每個id
的示例Dataframe:
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|
// +---+----------+
請注意,方法sample()
不一定生成方法參數中指定的確切數量的樣本。 這是一個相關的SO Q&A 。
如果你的planDF很大,你可能不得不考慮使用RDD的聚合 ,它具有以下簽名(跳過隱式參數):
def aggregate[U](zeroValue: U)(seqOp: (U, T) ⇒ U, combOp: (U, U) ⇒ U): U
它的工作方式有點像foldLeft
,除了它在一個分區中有一個累加運算符,另外一個用於匯總來自不同分區的結果。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.