python - pyspark 2.4.0 中带有 sampleBy 的样本大小

Question

我df在 pyspark 2.4.0 中有一个数据框，我想使用sampleBy. 它包含一个列category，我有一个这样的字典来采样：

dict = {
"cat_A" : 0.1,
"cat_B" : 0.5, 
"cat_C" : 0.4
}

我希望我的样本有 5 000 行，如果我使用这样的函数sampleBy：df_sampled = df.sampleBy(col = 'category', fractions = fraction, seed = 1)生成的数据框具有每个类别的正确比例，但似乎该函数没有选择样本大小的参数。

你知道怎么做吗？

我的初始数据框中大约有 100 000 行，还有大约 10 列其他列，因此我可以直接使用 Pandas 或 scikit-learn 切换到 Python，但我也不知道那里是否有解决方案。

标签： pythondataframeapache-sparkpysparksample

Answer 1

我认为您可以根据所需的行数缩放分数，然后使用以下命令设置数据框的最终大小limit：

fraction = {
"cat_A" : 0.1,
"cat_B" : 0.5,
"cat_C" : 0.4
}

scale = 5000 / df.count() / sum(fraction.values())
fraction = {k, v * scale for (k, v) in fraction.items()}

df_sampled = df.sampleBy(col = 'category', fractions = fraction, seed = 1).limit(5000)

只是为了确保您获得 5000 行（而不是更少），您可以在该行中使用稍大的数字scale，例如 5100。