apache-spark - Spark Dataframe - 每天随机采样记录

Question

我在 Hive 中有一些数据，其中每天大约有 500k 个唯一的 customerId。数据跨度为 2 个月，并按日期进行分区。这是分布的样子 -

+----------+---------------+
|      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，那么我希望每天为所有 60 个客户随机采样 250k 的唯一 customerId天的数据，以便我可以在输出表中保持每天受众规模的一致性。

因此输出表中的总记录将为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按日期分区并按随机值排序。我们使用该窗口函数添加一个“排名”列，然后按小于“n”值的排名进行过滤，并删除“排名”列。

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")