根据类标签和标识符拆分数据框

Question

我正在处理一个非常大的数据集~50GB，我正在尝试对其进行采样以减小其大小。采样参数应该是动态的，例如初始数据集的20％。

我通常使用sklearn.model_selection包和train_test_split函数。此函数允许我根据类标签获取数据集的分层样本（在此特定情况下为二进制）。

对于我正在考虑的数据集，我需要根据标识符进行一次拆分。如果可以多次出现的标识符在训练集中出现一次或多次，则无法出现在测试集中。

总结我需要的是：

1. 减少数据集大小（始终保持类标签频率）

2. 进行火车测试分裂：

   2.1关于班级标签的分层

   2.2，在

之前定义了标识符约束

数据集结构的一个例子是：

ID    F1    LABEL
1     2     1
2     2     1
1     2     1
4     2     1
6     2     1
1     2     1
4     2     0
6     2     0
4     2     0
6     2     0

将尺寸缩小50％：

ID    F1    LABEL
1     2     1
2     2     1
1     2     1
4     2     0
6     2     0

可能的火车＆amp;测试分裂：

      Train
ID    F1    LABEL
1     2     1
1     2     1
4     2     0

      Test
ID    F1    LABEL
2     2     1
4     2     0

请注意，ID拆分不是有序的，所以我不能使用任何带有pandas的掩码，因为它在许多回答拆分问题（Pandas split DataFrame by column value）时提出了建议。 目前我的管道有以下步骤：

减少数据集维度：

 X_tr,Y_te,l_tr,l_te = sklearn.model_selection.train_test_split(test_size = 0.5)
 (join on test removing train)

在ID上选择分层样本：

(same as before, join both on train and test sets to reconstruct the two datasets)
The ID selection is performed getting the set of all the IDs and randomly choose the ID for training and testing.

正如您可能注意到的，这个解决方案非常慢，我想知道如何减少计算时间（现在超过3个小时）。