如何惰性地 GroupShuffleSplit parquet 数据框？

Question

我有一个看起来像这样的镶木地板数据集（我使用的是极坐标，但任何数据框库都可以）：

df = pl.DataFrame(
    {
        "match_id": [
            1, 1, 1,
            2, 2, 2, 2,
            3, 3, 3, 3,
        ],
        "team_id": [
            1, 2, 2,
            1, 1, 2, 2,
            1, 2, 2, 3,
        ],
        "player_name": [
            "kevin", "james", "kelly",
            "john", "jenny", "jim", "josh",
            "alice", "kevin", "lilly", "erica",
        ],
    }
)

我想按 match_id 进行分组并测试训练集，以便 80% 的匹配在训练集中，其余的在测试集中。所以像这样：

group_df = df.group_by(["match_id"])
train, test = group_split(group_df, test_size=0.20)

我需要一个Python解决方案，最好带有dask、pandas或其他数据框架库。目前 pandas 不支持延迟计算，因为数据集非常大。所以使用 pandas 似乎是不可能的。另一方面，Dask 不支持任何

sklearn.model_selection

拆分器，因为它没有基于整数的索引支持。

理想情况下，我只需要一个简单的

GroupShuffleSplit

使用 dask 即可。还有其他库支持这个吗？如果是这样，我该如何以懒惰的方式用镶木地板做到这一点？

Answer 1

不幸的是，Dask 目前没有提供内置的拆分器，用于基于组的训练测试拆分，就像 scikit-learn 中的

GroupShuffleSplit

一样。但是，您可以使用

dask_ml.model_selection.GroupShuffleSplit

库中的

dask_ml

类来实现此功能。

以下是如何使用 Dask 和 Parquet 数据集执行基于组的训练测试分割的示例：

import dask.dataframe as dd
from dask_ml.model_selection import GroupShuffleSplit

# Load the Parquet dataset into a Dask DataFrame
df = dd.read_parquet("path/to/your/data.parquet")

# Define the grouping column
group_column = "match_id"

# Create the GroupShuffleSplit object
gss = GroupShuffleSplit(n_splits=1, train_size=0.8, random_state=42)

# Perform the train-test split
train_indices, test_indices = next(gss.split(df, groups=df[group_column]))

# Split the DataFrame based on the indices
train = df.loc[train_indices].compute()
test = df.loc[test_indices].compute()

在此示例中，我们首先使用

dd.read_parquet()

将 Parquet 数据集加载到 Dask DataFrame 中。然后，我们创建一个具有所需列车大小 (80%) 和随机状态的

GroupShuffleSplit

对象。我们使用

split()

的

GroupShuffleSplit

方法根据分组列生成训练和测试索引。最后，我们使用

loc

访问器从 DataFrame 中选择相应的行，并使用

.compute()

计算它们，将它们转换为 pandas DataFrame。

请注意，此方法会将整个数据集加载到内存中以进行训练和测试拆分，因为 Dask 不支持此特定操作的延迟评估。如果内存是一个问题，您可能需要考虑替代方法，例如对数据集的子集进行采样以进行训练测试 sp