Spark Dataframe - Sampling records randomly for each day

user_8229534 0
Viewed 338

I have some data in Hive wherein I have around 500k unique customerIds on a daily basis. The data spans for 2 months and is partitioned on date. Here is how the distribution looks like - 

+----------+---------------+
|      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|
+----------+---------------+

I would like to implement a function which takes N as input and gives me that many records per day in the output table.

For example, if I give N as 250k, then I want 250k randomly sampled unique customerIds for each day, for all the 60 days of data, so that I can maintain the consistency of audience size for each day in my output table.

So total records in output table would be 250k * 60. Here is how the distribution of my output table would look like -

+----------+---------------+
|      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|
+----------+---------------+

How can I achieve this using Spark?

apache-spark
1 Answers

I would just use window function partitionBy to partition by the date and order it by a random values. We add a 'rank' column using that window function, then filtering by ranks less than your 'n' value, and dropping the 'rank' column.

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