使用 spark java 获取列的每个值的特定百分比数据

Get certain percentage data on each values of a column using spark java

我想 select 从更大的数据集中提取 100 万条记录。数据集有一列名为 "city"。 100 万条记录应包含 "city 1" 中的 15% 记录,"city 2" 中的 30% 记录,"city 3" 中的“55%”。我可以使用如下限制 select。但它不会是随机的 selection.

dataset.filter(col("city").equals("city1")).limit(.15* 1 million)

我可以使用示例函数。但我不想过滤每个城市并对其进行采样。

dataset.filter(col("city").equals("city1")).sample(false,percentage);

是否有更好的方法使用 spark java 获取数据的随机样本百分比?

您可以使用 window,按城市分区并按随机列排序。然后您可以计算排名,然后只保留第一名。

假设您有一张地图 cityMap,它将每个城市与您要保留的样本数相关联。你可以从数据中计算出来,或者提前知道你想要什么。

我在 scala 中编写解决方案以确保它有效,但如果您熟悉 spark Java API,那么转换为 Java 应该很简单。我尝试使用 java 之类的语法 ;)

// creating the window
val win = Window.partitionBy("city").orderBy("random")

// defining a UDF that decides what records to sample
val isSampled = udf((r : Long, city : String) => r <= cityMap(city))

val sampledData = df
    .withColumn("random", rand())
    .withColumn("rank", rank().over(win))
    .where(isSampled(col("rank"), col("city")))
    .drop("rank", "random")

您似乎在寻找org.apache.spark.sql.DataFrameStatFunctions.sampleBy()

dataset.stat().sampleBy("city", ImmutableMap.of("city1", 0.15, "city2", 0.3, "city3", 0.55)