How to test/train split by column value rather than by row in pyspark

user_9821273 0
Viewed 1086

I would like to generate a train and test set for machine learning. Let's say I have a dataframe with the following columns:

account_id | session_id | feature_1 | feature_2 | label

In this dataset, each row will have a unique session_id, but an account_id can show up multiple times. However, I want my train and test sets to have mutually exclusive account_ids. (Almost seems like the opposite of stratified sampling).

In pandas, this is simple enough. I have something like the following:

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

Now, my dataset is large enough that it cannot fit into memory, and I have to switch over from pandas to doing all of this in pyspark. How can I split a train and test set to have mutually exclusive account_ids in pyspark, without fitting everything into memory?

pyspark apache-spark-sql apache-spark machine-learning
1 Answers

You can use the rand() function from pyspark.sql.functions for generating a random number for each of the distinct account_id and create train and test dataframes based on this random number.

from psypark.sql import functions as F

TEST_FRACTION = 0.2

train_test_split = (df.select("account_id")
                      .distinct()  # removing duplicate account_ids
                      .withColumn("rand_val", F.rand())
                      .withColumn("data_type", F.when(F.col("rand_val") < TEST_FRACTION, "test")
                                                .otherwise("train")))

train_df = (train_test_split.filter(F.col("data_type") == "train")
                            .join(df, on="account_id"))  # inner join removes all rows other than train

test_df = (train_test_split.filter(F.col("data_type") == "test")
                           .join(df, on="account_id"))

Since an account_id cannot be both train and test at a time, train_df and test_df will have mutually exclusive account_ids.

user_37665680
Related