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

我正在寻找一种解决方案，将我的数据拆分为测试集和训练集，但我希望在测试和训练中都有分类变量的所有级别。 我的变量有200个级别，数据是1800万条记录。我尝试了带分数（0.8）的

sampleBy

函数，可以得到训练集，但很难得到测试集，因为Spark中没有索引，即使创建了一个键，使用left join或subtract也很难得到测试集

我想根据我的分类变量做一个

groupBy

，随机抽取每个类别的样本，如果该类别只有一个观察值，则将其放入训练集中

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

这是一个相当困难的问题

我不知道有什么内置函数可以帮助您实现这一点。使用

sampleBy

，然后进行减法运算，但正如你所说的那样，速度会非常慢

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

使用窗口函数，添加row num，并将rownum=1的所有内容删除到一个单独的数据框中，您将在最后添加到培训中

对于剩余的数据，使用

randomspilt

（数据帧函数）将其分为

训练

和

测试

将步骤1中分离的数据添加到培训中

这应该更快

---

*（我以前从未尝试过！如果你能分享最后的成功经验，那就太好了！）

是的，我就是这么做的！就在第一步，我做了一个频率测试，并通过一次观察将所有级别都放在一边。randomSplit非常快，但我不确定训练和测试中每个级别的百分比是多少。我必须检查，我希望他们至少接近0.7，0.3！我想这就是对随机性的信任！如果维护发行版非常重要的话——也许值得花时间为每个类别做一次分组和抽样，然后坚持下去。