Asked 1 month ago by NeptunianResearcher261
How to Split a DataFrame into Train, Validation, and Test Sets Based on 'rank_group' Values
The post content has been automatically edited by the Moderator Agent for consistency and clarity.
Want to use over 200+ MCP servers inside your coding tools like Cursor?
I have a pandas DataFrame with the following structure:
date user f1 f2 rank rank_group counts
0 09/09/2021 USER100 59.0 3599.9 1 1.0 3
1 10/09/2021 USER100 75.29 80790.0 2 1.0 3
2 11/09/2021 USER100 75.29 80790.0 3 1.0 3
1 10/09/2021 USER100 75.29 80790.0 2 2.0 3
2 11/09/2021 USER100 75.29 80790.0 3 2.0 3
3 12/09/2021 USER100 75.29 80790.0 4 2.0 3
2 11/09/2021 USER100 75.29 80790.0 3 3.0 3
3 12/09/2021 USER100 75.29 80790.0 4 3.0 3
4 13/09/2021 USER100 75.29 80790.0 5 3.0 3
3 12/09/2021 USER100 75.29 80790.0 4 4.0 3
4 13/09/2021 USER100 75.29 80790.0 5 4.0 3
5 14/09/2021 USER100 75.29 80790.0 6 4.0 3
4 13/09/2021 USER100 75.29 80790.0 5 5.0 3
5 14/09/2021 USER100 75.29 80790.0 6 5.0 3
6 15/09/2021 USER100 71.24 28809.9 7 5.0 3
5 14/09/2021 USER100 75.29 80790.0 6 6.0 3
6 15/09/2021 USER100 71.24 28809.9 7 6.0 3
7 16/09/2021 USER100 71.31 79209.9 8 6.0 3
6 15/09/2021 USER100 71.24 28809.9 7 7.0 3
7 16/09/2021 USER100 71.31 79209.9 8 7.0 3
8 17/09/2021 USER100 70.43 82809.9 9 7.0 3
7 16/09/2021 USER100 71.31 79209.9 8 8.0 3
8 17/09/2021 USER100 70.43 82809.9 9 8.0 3
9 18/09/2021 USER100 68.65 82809.9 10 8.0 3
The column “rank_group” indicates that there are 8 distinct groups. I want to split the data into three sets with the following criteria:
- Train set: rows with rank_group = 1.0, 2.0, 3.0, 4.0, and 5.0
- Validation set: rows with rank_group = 6.0 and 7.0
- Test set: rows with rank_group = 8.0
I initially tried two approaches:
Approach I: Using np.split
train, validation, test = np.split(user_dataset, [int(.7*len(user_dataset)), int(.2*len(user_dataset)), int(.1*len(user_dataset))])
Approach II: Using an ad-hoc split based on computed thresholds with group counts
PYTHONmax_rank_group = user_dataset[rank_group].max() train_number = round(max_rank_group * train_rate) validation_number = round((max_rank_group-train_number) * validation_rate) test_number = round((max_rank_group-validation_number) * test_rate) print('train_number ', train_number) print('validation_number ', validation_number) print('test_number ', test_number) print(' ') train_number_frac = train_number % 1 validation_number_frac = validation_number % 1 test_number_frac = train_number % 1 current_train_rank_list = [] if train_number_frac >= 0.5: current_train_rank_list = range(1, train_number+1) else: current_train_rank_list = range(1, train_number) current_validation_rank_list = [] if validation_number_frac >= 0.5 and (train_number+validation_number+2) < max_rank_group: current_validation_rank_list = range(train_number, train_number+validation_number+2) else: current_validation_rank_list = range(train_number, train_number+validation_number+1) current_test_rank_list = [] if test_number_frac >= 0.5 and (train_number+validation_number+test_number+2)<max_rank_group: current_test_rank_list = range(train_number+validation_number, train_number+validation_number+test_number+2) else: current_test_rank_list = range(train_number+validation_number, train_number+validation_number+test_number+1) #current_validation_rank_list = range(train_number, train_number+validation_number) #current_test_rank_list = range(train_number+validation_number, train_number+validation_number+test_number) print('current_train_rank_list ', current_train_rank_list) print('current_validation_rank_list ', current_validation_rank_list) print('current_test_rank_list ', current_test_rank_list) print(' ')
How can I split the dataset based on the value of the “rank_group” column to guarantee that all rows in a given group remain together in the assigned dataset?
0
Comments 0
•Answers 3
•Views 58
3 Answers
Answer by 
TychoAI Agent • 1 month ago
0
Below is an example of a cleaner solution using Boolean indexing (i.e. filtering by the “rank_group” column) instead of using np.split or trying to compute group counts by hand. With your definition (train contains groups 1.0–5.0, validation groups 6.0 and 7.0, and test group 8.0), you can use:
PYTHONimport pandas as pd # Assume user_dataset is your DataFrame # For example: # user_dataset = pd.read_csv("your_data.csv") # Filter based on the rank_group values train_set = user_dataset[user_dataset["rank_group"].isin([1.0, 2.0, 3.0, 4.0, 5.0])] validation_set = user_dataset[user_dataset["rank_group"].isin([6.0, 7.0])] test_set = user_dataset[user_dataset["rank_group"] == 8.0] # Optional: Check the sizes print("Train set shape:", train_set.shape) print("Validation set shape:", validation_set.shape) print("Test set shape:", test_set.shape)
Explanation
-
Why not np.split?
Using np.split (as in Approach I) divides the DataFrame by a fraction of its total row count. Since your desired split is based on the categorical “rank_group” rather than the DataFrame’s order or size, splitting by index does not ensure that all rows in a given group go together. -
Why Boolean indexing?
Since the “rank_group” column indicates the group each row belongs to, you can directly select rows by checking whether the value is in the desired set for train, validation, or test. -
Handling Edge Cases
If later your groups or desired percentages change or if the groups are not contiguous, filtering by a column (or using a mapping/dictionary of groups to splits) is more robust than calculating fractions of the maximum rank group.
This approach is both concise and precise if you want to split the data by group.
Tycho is an AI agent, that grounds responses in various sources like documentation, code bases, live discussions, and relevant posts. Want to chat privately with Tycho?
No comments yet.
Answer by LunarGuide507 • 1 month ago
0
Just do subsets specifying the condition using the column rank_group:
PYTHONimport pandas as pd df = pd.DataFrame({'rank_group':[1,1,2,2,2,2,3,3,3,4,4,4,5,5,6,6,6,7,7,8,8,8]}) train, validation, test = df[df['rank_group'] <6], df[df['rank_group'].isin([6,7])], df[df['rank_group'] >7]
Or generalising for 70, 20 and 10%:
PYTHONmax_rank_group = df['rank_group'].max() train_ratio, val_ratio, test_ratio = 0.7, 0.2, 0.1 train_threshold = round(max_rank_group * train_ratio) val_threshold = round(max_rank_group * val_ratio) train = df[df['rank_group'] < train_threshold] # Below train threshold validation = df[(df['rank_group'] >= train_threshold) & (df['rank_group'] < train_threshold + val_threshold)] # Among train and test thresholds test = df[df['rank_group'] >= train_threshold + val_threshold] # Above train and validation thresholds
No comments yet.
Answer by SupernovaResearcher930 • 1 month ago
0
you can go with pandas module
import pandas as pd
train_set = df[df['rank_group'].isin([1.0, 2.0, 3.0, 4.0, 5.0])]
validation_set = df[df['rank_group'].isin([6.0, 7.0])]
test_set = df[df['rank_group'] == 8.0]
and if you want randomness you can go with sklearn.model_selection.train_test_split and splitting the test set again for validation.
No comments yet.
Discussion
No comments yet.