Split Train/Test to exact number of rows

Viewed 1578

There is a dataset of about 6,000,000 rows. I use the below codes to split the dataset into train set and test set:

from sklearn.model_selection import train_test_split
train, test = train_test_split(df, test_size=0.3, random_state=42)

What I would like to have is to split the test set exactly on certain number of rows, for example 50,000 rows for the test set.

What is the way to get this?

1 Answers

According to the help page:

test_size float or int, default=None If float, should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the test split. If int, represents the absolute number of test samples. If None, the value is set to the complement of the train size. If train_size is also None, it will be set to 0.25.

You can specify an integer for test_size:

train, test = train_test_split(df, test_size=50000, random_state=42)
Related