Train-Test Split Done Right: Stratified, Time-Based, and Group Splits

Contents

Figure 1: The opening decides the game — the split decides whether your metrics mean anything

Stratified, time-based, and group splits, each fixing a different way the default lies

train_test_split(X, y) is one line, and it's the line most likely to make your results meaningless. The default works for tidy, independent, balanced data — real data is rarely any of those.

The three fixes (stratified, time-based, and group splits) were tested on data where the truth was known, and then measured what happens when you tune on the wrong split. One of the results contradicted expectation, so it stayed in.

What the Default Actually Does

train_test_split shuffles your rows randomly and cuts them. Scikit-learn's docs list the defaults: shuffle=True, random_state=None, and if you set neither size, the test set is 25% of the data. The docs add one rule: if shuffle=False, then stratify must be None.

Random shuffling assumes every row is independent and every class is common. Break either assumption and the default misleads you.

Stratified Splits: Protect the Rare Class

Imagine fraud, disease, or churn: a class that's 2% of your data. Built: 500 rows with 10 positives and a 20% test set of 100 rows, then 1,000 splits each way:

Split	Positives in test (min–max)	Splits with 0 positives	Splits with 0 or 1
Random	0–6	10.7%	37.3%
Stratified	2–2	0%	0%

A random split left the test set with no positives at all one time in nine. Stratifying gave exactly 2 every time, matching the 2% base rate:

Xtr, Xte, ytr, yte = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42,
)

The effect on a real metric is visible — logistic model test AUC across 300 splits:

Split	AUC mean	AUC sd	Splits where AUC can't be computed
Random	0.959	0.064	36
Stratified	0.966	0.058	0

In 36 of 300 random splits, the test set had no positives and AUC doesn't exist. The averages differ only slightly, but you don't want to learn about missing positives from a crash. The spread (sd 0.058 to 0.064) is a reminder that 10 positives is a thin basis for any estimate.

FYI: the docs don't spell out the error for very rare classes. A class with only one member makes stratification fail outright, so check your smallest class before you split.

One Split Is One Draw

Even with stratification, a single split is one random sample of your data. On the breast cancer dataset (569 rows), the same logistic regression across 200 different random splits:

Test size	Mean accuracy	Spread (sd)	Min	Max
10%	0.979	0.018	0.912	1.000
20%	0.978	0.013	0.939	1.000
30%	0.976	0.011	0.936	1.000

The same model scored anywhere from 0.912 to 1.000, depending on the seed. A bigger test set shrinks the noise, but it also leaves less for training. Repeated 5-fold cross-validation (5 repeats) did much better: across 20 different seeds, its mean accuracy had a spread of only 0.0019 — the kind of stability the cross-validation master class measures in depth.

If your dataset is small, trust cross-validation over a single split, and always report a spread alongside your score. Don't celebrate a lucky seed.

Time-Based Splits: Respect the Arrow of Time

For time-ordered data, shuffling lets the model peek at the future — the same leak the time-series feature guide covers from the feature side. TimeSeriesSplit makes every test window come after its training window. Confirmed against the docs' example: with n_splits=3, test_size=2, and gap=2 on 12 rows, fold 0 trains on rows 0–3 and tests on rows 6–7.

The three parameters that matter:

• gap: rows excluded between train and test. The docs say it's useful when targets depend on recent observations.
• test_size: the length of each test window.
• max_train_size: caps the training window. Without it, training sets only grow.

tscv = TimeSeriesSplit(n_splits=5, test_size=30, gap=7)
for tr, te in tscv.split(X):
    ...  # every te index follows every tr index, minus the gap

The last one raises a real question: should old data stay in? Tested on a relationship that drifts, with a linear model on one feature, scored as mean walk-forward R² across 10 folds:

Window	Fast drift (slope flips every ~200 rows)	Slow drift (slope rises steadily)
Expanding	−0.140	0.628
Sliding, last 100 rows	−0.755	0.847
Sliding, last 300 rows	−0.579	0.824

There's no universal winner. With slow drift, forgetting old rows helped a lot (0.628 to 0.847). With fast drift, the sliding windows did worse than the expanding one, because a short window learns a noisy slope that is already out of date when the test starts. I expected sliding to win in both — it doesn't. Choose the window by how fast your world changes, and test both.

Group Splits: Keep Entities Together

When rows share an owner (patients, customers, devices), a random split puts the same entity on both sides. The leakage experiments showed AUC falling from 0.796 to 0.495 when that was fixed — here the focus is the splitters:

• GroupKFold guarantees a group never appears in both train and test, per the docs.
• StratifiedGroupKFold adds class balance on top: it tries to keep class percentages "as much as possible given the constraint of non-overlapping groups."
• GroupShuffleSplit gives a single holdout split by group.

from sklearn.model_selection import GroupKFold, StratifiedGroupKFold

gkf = StratifiedGroupKFold(n_splits=5)
for tr, te in gkf.split(X, y, groups=patient_ids):
    assert set(patient_ids[tr]).isdisjoint(patient_ids[te])

Made: 300 patients with between 1 and 40 visits each and about 15% positive rows, averaged over 20 datasets:

Splitter	Worst fold's positive rate off by	Largest fold-size deviation
GroupKFold	0.048	0.1%
StratifiedGroupKFold	0.005	1.5%

GroupKFold kept fold sizes nearly equal but let positive rates drift by almost 5 percentage points. StratifiedGroupKFold kept class rates tight and gave up a little size balance. A GroupShuffleSplit with test_size=0.2 put exactly 20% of the patients in test, which came to 20.8% of the rows — group sizes vary, so "20%" can mean different things.

One more caveat from the docs: with a small number of huge groups, stratification may be impossible, and StratifiedGroupKFold then behaves like GroupKFold.

A real-world version of the entity trap: a session-level fraud model looked strong in validation, then collapsed when deployed — the split had been row-level, so the same high-risk accounts had sessions in both train and test. The model had memorized accounts, not behavior. After splitting by account_id, the score dropped to an honest number and the model finally told the truth about unseen accounts. The fix took three lines; finding it took one question: can the same owner appear on both sides?

Don't Tune on Your Test Set

The last mistake is the quietest: using the test set to choose settings. 40 random forest configurations trained on weak-signal data, picking the best one two ways, averaged over 20 repeats. "Fresh rows" are 5,000 new samples that stand in for the real world:

Selection	Score reported	Score on fresh rows	Optimism
Chosen on the test set	0.682	0.650	+0.032
Chosen on a separate validation set	0.651	0.649	+0.003

Picking the winner on the test set inflated the reported AUC by 0.032, because the "best of 40" is partly the luckiest of 40. A separate validation set made the final test score honest (0.651 against 0.649 on fresh rows). The fix is a three-way split: train, validation (or cross-validation) for choices, and a test set you touch once — the same "untouched test set" rule from the preprocessing guide.

Which Split Should You Use?

• Independent rows, balanced classes: plain shuffled split is fine.
• Rare or imbalanced classes: stratify=y.
• Small dataset: repeated cross-validation instead of a single split.
• Time-ordered data: TimeSeriesSplit, with a gap if targets look ahead, and test sliding versus expanding windows.
• Repeated entities: GroupKFold, StratifiedGroupKFold, or GroupShuffleSplit.
• Any tuning: use validation data or cross-validation, and keep the test set for the end.

Common Pitfalls

• Letting rare classes fall out of the test set — 10.7% of random splits had zero positives here.
• Reporting one split's score as "the" score — the same model ranged 0.912 to 1.000 across seeds.
• Shuffling time-ordered data — every test row gets future neighbors in training.
• Row-level splits on entity data — the same patient, customer, or device on both sides.
• Assuming 20% of groups equals 20% of rows — group sizes vary.
• Choosing sliding vs. expanding windows by habit — the winner flips with drift speed.
• Stratifying a class with one member — it fails outright; check your smallest class first.
• Tuning, iterating, or "just checking" against the test set — the best-of-N is partly the luckiest-of-N.
• Forgetting random_state when comparing runs — you can't debug a split you can't reproduce.

Unlock AI That Actually Works

Get lifetime access to GPT-6 Astra, Claude Fable 5.1, Gemini 3.5, Grok 4.5, and more — all in one platform. Build websites, apps, videos, content, and digital products from a single command. No monthly fees. No tool-hopping.

Click here to get GPTAstra Max now — one-time payment, lifetime access.

Frequently Asked Questions

When should you use stratify in train_test_split?

Whenever a class is rare enough that a random draw can miss it. With 500 rows and 10 positives, random splits produced test sets with zero positives 10.7% of the time and zero-or-one positive 37.3% of the time — in 36 of 300 random splits the test AUC didn't exist at all, because AUC needs both classes. Passing stratify=y gave exactly 2 positives (the 2% base rate) in every one of 1,000 splits, and AUC was computable every time. Check your smallest class first: a class with a single member makes stratification fail outright.

Should you use a sliding or expanding window for time-series splits?

It depends on how fast your data drifts — test both. On slow drift (slope rising steadily), a sliding window of the last 100 rows lifted mean walk-forward R² from 0.628 (expanding) to 0.847, because forgetting old rows helps when old rows are misleading. On fast drift (slope flipping every ~200 rows), the expanding window won: −0.140 versus −0.755 for sliding, because a short window learns a noisy slope that's already out of date when the test window starts. There's no universal winner — choose by how fast your world changes, and set gap when targets depend on recent observations.

How do you split data with repeated entities like patients or customers?

Use a group-aware splitter so no entity appears on both sides: GroupKFold guarantees non-overlapping groups (nearly equal fold sizes — largest deviation 0.1% — but class rates drifted by up to 0.048), StratifiedGroupKFold adds class balance on top (positive rates within 0.005, at a cost of 1.5% fold-size deviation), and GroupShuffleSplit gives a single holdout by group. Beware that "20% of groups" isn't "20% of rows": with 300 patients holding 1–40 visits each, test_size=0.2 put exactly 20% of patients in test, which came to 20.8% of the rows. With few huge groups, stratification may be impossible and StratifiedGroupKFold degrades to GroupKFold.

Why shouldn't you tune hyperparameters on the test set?

Because the best of many candidates is partly the luckiest. Picking the best of 40 random forest configurations on the test set reported 0.682 while the same model scored 0.650 on 5,000 fresh rows — optimism of +0.032, averaged over 20 repeats. Choosing the winner on a separate validation set gave 0.651 reported versus 0.649 fresh: honest. Use a three-way split — train, validation (or cross-validation) for every choice, and a test set you touch exactly once at the end.

Wrapping This Up

A random split produced test sets with no positives, spreads from 0.912 to 1.000 on one model, and optimistic scores when tuning happened on the test data. Stratifying, grouping, and respecting time each fixed a specific failure — though the best time window depended on how fast the data drifted.

Pick the splitter that matches how your model will really be used, and write that reason in one line above your split code.

Related Articles

• Data Leakage in Machine Learning: How to Spot and Prevent It
• Scikit-Learn Cross-Validation Master Class
• Feature Engineering for Dates and Times in Python
• Data Preprocessing for Machine Learning: Complete Step-by-Step Guide