Practice Now
sklearn
data splitting
machine learning
Python
data preprocessing

How to split data based on a column value in sklearn

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Machine LearningData SciencePython

Introduction

Scikit-learn's train_test_split splits data randomly, but sometimes you need to split based on a specific column's values — for example, keeping all records from certain users in the same split, or separating data by date. Scikit-learn provides GroupShuffleSplit and GroupKFold for group-aware splitting. For simple value-based filtering, pandas boolean indexing is the most direct approach.

Simple Split by Column Value (pandas)

python
1import pandas as pd
2
3df = pd.DataFrame({
4    'feature1': [1, 2, 3, 4, 5, 6],
5    'feature2': [10, 20, 30, 40, 50, 60],
6    'label': ['A', 'B', 'A', 'B', 'A', 'B'],
7    'group': ['train', 'train', 'train', 'test', 'test', 'test']
8})
9
10# Split by a column value
11train_df = df[df['group'] == 'train']
12test_df = df[df['group'] == 'test']
13
14print(train_df)
15#    feature1  feature2 label  group
16# 0         1        10     A  train
17# 1         2        20     B  train
18# 2         3        30     A  train

Split by Category Values

python
1import pandas as pd
2from sklearn.model_selection import train_test_split
3
4df = pd.DataFrame({
5    'user_id': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5],
6    'action': ['click', 'buy', 'click', 'view', 'buy', 'click', 'view', 'buy', 'click', 'view'],
7    'value': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
8})
9
10# Split keeping all records for each user together
11unique_users = df['user_id'].unique()
12train_users, test_users = train_test_split(unique_users, test_size=0.4, random_state=42)
13
14train_df = df[df['user_id'].isin(train_users)]
15test_df = df[df['user_id'].isin(test_users)]
16
17print(f"Train users: {train_users}, Test users: {test_users}")
18print(f"Train rows: {len(train_df)}, Test rows: {len(test_df)}")

This ensures that all records from one user appear in the same split, preventing data leakage.

GroupShuffleSplit (sklearn)

GroupShuffleSplit splits data while keeping groups intact:

python
1from sklearn.model_selection import GroupShuffleSplit
2import numpy as np
3
4X = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
5y = np.array([0, 0, 1, 1, 0, 1])
6groups = np.array(['A', 'A', 'B', 'B', 'C', 'C'])
7
8gss = GroupShuffleSplit(n_splits=1, test_size=0.33, random_state=42)
9
10for train_idx, test_idx in gss.split(X, y, groups):
11    X_train, X_test = X[train_idx], X[test_idx]
12    y_train, y_test = y[train_idx], y[test_idx]
13    print(f"Train groups: {groups[train_idx]}")
14    print(f"Test groups: {groups[test_idx]}")

All samples with the same group value end up in the same split.

GroupKFold for Cross-Validation

python
1from sklearn.model_selection import GroupKFold
2import numpy as np
3
4X = np.random.randn(12, 2)
5y = np.array([0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1])
6groups = np.array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])
7
8gkf = GroupKFold(n_splits=4)
9
10for fold, (train_idx, test_idx) in enumerate(gkf.split(X, y, groups)):
11    print(f"Fold {fold}: train groups {np.unique(groups[train_idx])}, "
12          f"test group {np.unique(groups[test_idx])}")
13
14# Fold 0: train groups [2 3 4], test group [1]
15# Fold 1: train groups [1 3 4], test group [2]
16# Fold 2: train groups [1 2 4], test group [3]
17# Fold 3: train groups [1 2 3], test group [4]

Each fold holds out one group entirely, preventing data leakage in cross-validation.

Time-Based Split

python
1import pandas as pd
2
3df = pd.DataFrame({
4    'date': pd.date_range('2025-01-01', periods=100, freq='D'),
5    'feature': range(100),
6    'target': [i % 2 for i in range(100)]
7})
8
9# Split by date threshold
10cutoff = '2025-03-01'
11train_df = df[df['date'] < cutoff]
12test_df = df[df['date'] >= cutoff]
13
14print(f"Train: {len(train_df)} rows (before {cutoff})")
15print(f"Test: {len(test_df)} rows (from {cutoff})")

For time series, always split chronologically — never randomly — to prevent future data leaking into training.

TimeSeriesSplit (sklearn)

python
1from sklearn.model_selection import TimeSeriesSplit
2import numpy as np
3
4X = np.random.randn(100, 3)
5y = np.random.randn(100)
6
7tscv = TimeSeriesSplit(n_splits=5)
8
9for fold, (train_idx, test_idx) in enumerate(tscv.split(X)):
10    print(f"Fold {fold}: train [{train_idx[0]}:{train_idx[-1]}], "
11          f"test [{test_idx[0]}:{test_idx[-1]}]")
12
13# Fold 0: train [0:16], test [17:33]
14# Fold 1: train [0:33], test [34:49]
15# ...

TimeSeriesSplit ensures training always uses earlier data than testing.

Stratified Split by Column

When you want balanced class distribution in both splits:

python
1from sklearn.model_selection import train_test_split
2import pandas as pd
3
4df = pd.DataFrame({
5    'feature': range(100),
6    'category': ['A'] * 70 + ['B'] * 30
7})
8
9# Stratified split ensures same class proportions in train and test
10X = df[['feature']]
11y = df['category']
12
13X_train, X_test, y_train, y_test = train_test_split(
14    X, y, test_size=0.2, stratify=y, random_state=42
15)
16
17print(f"Train class distribution: {y_train.value_counts().to_dict()}")
18print(f"Test class distribution: {y_test.value_counts().to_dict()}")
19# Train: {'A': 56, 'B': 24}
20# Test: {'A': 14, 'B': 6}

Combining Group and Stratified Splitting

python
1from sklearn.model_selection import StratifiedGroupKFold
2import numpy as np
3
4X = np.random.randn(100, 5)
5y = np.array([0] * 50 + [1] * 50)
6groups = np.array([i // 5 for i in range(100)])  # 20 groups of 5
7
8sgkf = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=42)
9
10for fold, (train_idx, test_idx) in enumerate(sgkf.split(X, y, groups)):
11    print(f"Fold {fold}: {len(train_idx)} train, {len(test_idx)} test, "
12          f"class balance: {y[test_idx].mean():.2f}")

StratifiedGroupKFold (sklearn 1.0+) keeps groups together while maintaining class balance.

Common Pitfalls

  • Data leakage with group splits: If the same user/session/entity appears in both train and test, the model learns from future data about that entity. Always use group-aware splitting when records are not independent.
  • Uneven group sizes: GroupShuffleSplit splits by groups, not by rows. If one group has 1000 rows and another has 10, the row counts in train/test may not match the specified test_size ratio.
  • Forgetting random_state: Without it, splits are different every time. Set random_state=42 (or any fixed value) for reproducibility.
  • Time series random split: Never use train_test_split on time series data — it mixes future and past data. Use TimeSeriesSplit or date-based filtering.
  • Dropping the split column from features: After splitting by a column, drop that column from the feature set if it should not be a model input: X_train = train_df.drop(columns=['group']).

Summary

  • Use pandas boolean indexing (df[df['col'] == value]) for simple value-based splits
  • Use GroupShuffleSplit to split while keeping groups (users, sessions) intact
  • Use GroupKFold for group-aware cross-validation
  • Use TimeSeriesSplit or date filtering for temporal data
  • Use stratify=y in train_test_split to maintain class balance
  • Always use group-aware splitting when records within a group are not independent
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.