Spark Dataframe -每天随机采样记录

Question

我在蜂巢中有一些数据，其中我每天都有大约500k唯一的customerIds。数据跨度为2个月，并在date上分区。下面是这个发行版的样子-

+----------+---------------+
|      date|unique_visitors|
+----------+---------------+
|2019-01-01|        500,000|
|2019-01-02|        500,000|
|2019-01-03|        500,000|
|2019-01-04|        500,000|
|2019-01-05|        500,000|
|2019-01-06|        500,000|
|..........|         ..... |
|2019-02-27|        500,000|
|2019-02-28|        500,000|
+----------+---------------+

我想实现一个函数，它接受N作为输入，每天在输出表中给我那么多记录。

例如，如果我将N设为250k，那么我希望250k每天为所有60天的数据随机采样唯一的customerIds，这样我就可以在输出表中保持每天观众人数的一致性。

所以输出表中的总记录应该是250k * 60。下面是我的输出表的分布情况：

+----------+---------------+
|      date|unique_visitors|
+----------+---------------+
|2019-01-01|        250,000|
|2019-01-02|        250,000|
|2019-01-03|        250,000|
|2019-01-04|        250,000|
|2019-01-05|        250,000|
|2019-01-06|        250,000|
|..........|         ..... |
|2019-02-27|        250,000|
|2019-02-28|        250,000|
+----------+---------------+

我如何使用Spark来实现这一点？

Answer 1

我只会使用窗口函数partitionBy对日期进行分区，并按随机值进行排序。我们使用该窗口函数添加一个‘ranks’列，然后按小于'n‘值的排名进行过滤，并删除’ranks‘列。

import org.apahce.spark.sql.functions._
import org.apache.spark.sql.expressions.Window
import spark.implicits._

val n = 250000
val w = Window.partitionBy($"date").orderBy(rand())
val res = df.withColumn("rank", rank().over(w)).filter($"rank" <= n).drop("rank")