pyspark - 在 pyspark 中创建训练集和测试集时，如何跨不同组进行分层抽样？

Question

我正在寻找一种解决方案，将我的数据拆分为测试集和训练集，但我希望在测试和训练中拥有所有级别的分类变量。我的变量有 200 个级别，数据是 1800 万条记录。我尝试sampleBy了带有分数（0.8）的函数，并且可以获得训练集，但由于 Spark 中没有索引，即使创建键，使用左连接或减法获取测试集的速度也非常慢！

我想groupBy根据我的分类变量做一个并随机抽样每个类别，如果该类别只有一个观察值，则将其放入训练集中。

是否有默认函数或库来帮助执行此操作？

标签： pysparkpyspark-sqlsamplingapache-spark-ml

Answer 1

一个相当难的问题。

我不知道可以帮助您获得此功能的内置功能。使用sampleBy然后这样的减法减法会起作用，但正如你所说 - 会很慢。

或者，想知道你是否可以试试这个*：

1. 使用窗口函数，添加第 num 行并将所有具有 rownum=1 的内容删除到一个单独的数据框中，您将在最后将其添加到您的训练中。

2. 在剩余的数据上，使用randomSplit（一个数据框函数）分为training和test

3. 将步骤 1 中分离的数据添加到训练中。

这应该工作得更快。

*（我以前没试过！如果你能分享一下最后的效果就太好了！）