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

Posted 2023-04-15

【中文标题】使用 spark java 获取列的每个值的特定百分比数据

【英文标题】：Get certain percentage data on each values of a column using spark java

【发布时间】：2020-01-03 00:09:04

【问题描述】：

我想从更大的数据集中选择 100 万条记录。数据集有一个名为“城市”的列。 100 万条记录应包含来自“城市 1”的 15% 记录、来自“城市 2”的 30% 记录、来自“城市 3”的“55%”记录。我可以使用如下限制来选择。但这不会是随机选择。

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

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

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

有没有更好的方法来使用 spark java 获取数据的随机样本百分比？

【问题讨论】：

你有几个城市？

【参考方案1】：

你好像在找org.apache.spark.sql.DataFrameStatFunctions.sampleBy()

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

【讨论】：

虽然它没有给出确切的百分比，但调整种子值会得到更接近的结果。谢谢

哦，抱歉，我没有注意到您希望整个样本正好是 1M 记录。在这种情况下，您应该将采样率计算为ratio = count(city='city1') / (1M * 0.15) 等。

【参考方案2】：

您可以使用按城市分区并按随机列排序的窗口。然后你可以计算排名，然后只保留第一个排名。

假设您有一张地图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")