我想为机器学习生成一个训练和测试集。假设我有一个包含以下列的Dataframe：
account_id | session_id | feature_1 | feature_2 | label在这个数据集中，每一行都有一个唯一的session\u id，但是一个account\u id可以出现多次。但是，我希望我的训练集和测试集具有互斥的帐户ID(几乎与分层抽样相反）。
对Pandas来说，这很简单。我有如下几点：
def train_test_split(df, split_col, feature_cols, label_col, test_fraction=0.2):
"""
While sklearn train_test_split splits by each row in the dataset,
this function will split by a specific column. In that way, we can
separate account_id such that train and test sets will have mutually
exclusive accounts, to minimize cross-talk between train and test sets.
"""
split_values = df[split_col].drop_duplicates()
test_values = split_values.sample(frac=test_fraction, random_state=42)
df_test = df[df[split_col].isin(test_values)]
df_train = df[~df[split_col].isin(test_values)]
return df_test, df_train现在，我的数据集足够大，无法放入内存，我必须从pandas切换到pyspark中完成所有这些。如何在pyspark中拆分一个train和测试集，使其拥有互斥的account\u id，而不将所有内容都放入内存？
1条答案
按热度按时间4smxwvx51#
你可以用
rand()函数来自pyspark.sql.functions为每个不同的account_id创造train以及test基于此随机数的Dataframe。自从
account_id不能两者兼而有之train以及test一次，train_df以及test_df将具有互斥性account_ids。