Pyspark:采用平衡类的n个样本

Question

我有一个相当大的数据集，大约有5bil。记录。我想从其中随机抽取1mio的样本。问题是标签不平衡。

+---------+----------+
|    label|     count|
+---------+----------+
|A        | 768866802|
|B        |4241039902|
|C        | 584150833|
+---------+----------+

标签B比其他标签有更多的数据。我知道有一个向上和向下采样的概念，但是考虑到大量的数据，我可能不需要这样做，因为我可以很容易地从每个标签中找到1mio记录。

我想知道我怎样才能有效地使用~1mio。随机样本(没有替换)，这样我就有一个偶数的所有标签~ 333k在每个标签。使用PySpark

我的一个想法是将数据集分成3个不同的df。从其中随机抽取30万个样本，并将它们缝合在一起。但也许有更有效的方法来做到这一点。

Answer 1

您可以创建一个包含随机值的列，并使用row_number为每个标签筛选1M个随机样本：

from pyspark.sql.types import *
from pyspark.sql import functions as F
from pyspark.sql.functions import *
from pyspark.sql.window import Window

n = 333333 # number of samples
df = df.withColumn('rand_col', F.rand())
sample_df1 = df.withColumn("row_num",row_number().over(Window.partitionBy("label")\
                .orderBy("rand_col"))).filter(col("row_num")<=3)\
                .drop("rand_col", "row_num")

sample_df1.groupBy("label").count().show()

这将始终为每个标签提供100万个样本。

另一种方法是使用spark的stat.sampleBy进行分层采样

n = 333333
seed = 12345

# Creating a dictionary of fractions for eacch label
fractions = df.groupBy("label").count().withColumn("required_n", n/col("count"))\
                .drop("count").rdd.collectAsMap()

sample_df2 = df.stat.sampleBy("label", fractions, seed)

sample_df2.groupBy("label").count().show()

然而，sampleBy根据运行情况得出一个近似的解决方案，并且不能保证每个标签的确切记录数量。

示例数据帧：

schema = StructType([StructField("id", IntegerType()), StructField("label", IntegerType())])
data = [[1, 2], [1, 2], [1, 3], [2, 3], [1, 2],[1, 1], [1, 2], [1, 3], [2, 2], [1, 1],[1, 2], [1, 2], [1, 3], [2, 3], [1, 1]]
df = spark.createDataFrame(data,schema=schema)

df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
|    1|    3|
|    2|    7|
|    3|    5|
+-----+-----+

方法1:

# Sampling 3 records from each label
n = 3 

# Assign a column with random values
df = df.withColumn('rand_col', F.rand())
sample_df1 = df.withColumn("row_num",row_number().over(Window.partitionBy("label")\
                .orderBy("rand_col"))).filter(col("row_num")<=3)\
                .drop("rand_col", "row_num")

sample_df1.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
|    1|    3|
|    2|    3|
|    3|    3|
+-----+-----+

方法2:

# Sampling 3 records from each label
n = 3
seed = 12

fractions = df.groupBy("label").count().withColumn("required_n", n/col("count"))\
                .drop("count").rdd.collectAsMap()

sample_df2 = df.stat.sampleBy("label", fractions, seed)
sample_df2.groupBy("label").count().show()

+-----+-----+
|label|count|
+-----+-----+
|    1|    3|
|    2|    3|
|    3|    4|
+-----+-----+

正如您所看到的，sampleBy倾向于为您提供大致相等的分布。但不完全是。我更喜欢方法1来解决你的问题。