Train-Test Split — Module 5.1

How This Lesson Fits the Module

You understand train, validation, and test conceptually. Train-test split is the sklearn workhorse that implements those partitions: train_test_split and friends turn one dataframe into reproducible subsets.

Splitting is not a one-liner to forget—it encodes business reality (time, groups, balance) and is the first line of defense against the leakage patterns from Volume 04.

Learning Objectives

By the end of this lesson, students should be able to:

Use train_test_split for train/test and nested train/val splits.
Apply stratify, random_state, and test_size correctly.
Choose shuffle vs temporal splitting for the problem domain.
Split with GroupShuffleSplit when rows share entities.
Verify split integrity with pandas index alignment checks.
Document split parameters in experiment configs.

train_test_split Basics

sklearn.model_selection.train_test_split randomly partitions arrays or DataFrames. Default shuffle is fine for i.i.d. tabular data; time-series and grouped data need specialized splitters.

Parameter	Effect	Typical value
`test_size`	Fraction or count for holdout	`0.2`
`stratify`	Preserve class ratios	`y` for classification
`random_state`	Reproducible shuffle	Fixed integer
`shuffle`	Randomize before split	`False` for ordered time

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    stratify=y,
    random_state=42,
)

# Verify alignment
assert X_train.index.equals(y_train.index)
print("Train positives:", y_train.mean())
print("Test positives:", y_test.mean())

Non-Random Splits

When rows are not independent, random splits lie. Use group-aware or time-aware strategies from the same model_selection module.

from sklearn.model_selection import GroupShuffleSplit

groups = df["user_id"]
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
train_idx, test_idx = next(gss.split(X, y, groups=groups))

X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]

Critical Mistake — Splitting After Preprocessing

Fit scalers or imputers on the full dataset, then split—test statistics leaked into training. Split first (or use pipelines inside CV) as Volume 04 taught.

Split Strategy by Problem

Data type	Recommended split
i.i.d. tabular	`train_test_split` + stratify
Multiple rows per user	`GroupShuffleSplit` / `GroupKFold`
Time-ordered events	Cutoff date or `TimeSeriesSplit`
Spatial clusters	Group by region or site

Engineering Habit — Split Manifest

Save train_idx.parquet and test_idx.parquet with split version, seed, and SQL filter. Retrains must reuse or consciously bump the version.

Knowledge Check

Short Answer: What does test_size=0.2 mean? Answer: 20% of rows go to the holdout set.
True/False: stratify=y works for regression targets. Answer: False—for classification-style discrete labels.
Multiple Choice: Same user in train and test causes: (a) group leakage, (b) faster training, (c) better calibration. Answer: (a).
Short Answer: Why set random_state? Answer: Reproducible splits across runs and teammates.
Short Answer: When set shuffle=False? Answer: Preserve temporal order for time-series splits.
True/False: Preprocessing before split is leakage-safe. Answer: False.
Multiple Choice: GroupShuffleSplit needs: (a) groups array, (b) GPU, (c) one-hot labels only. Answer: (a).
Short Answer: Why verify index alignment? Answer: Ensures each y row matches the correct X row.
Short Answer: Nested splits create what overall ratio from 80% dev with 25% val? Answer: 60% train, 20% val, 20% test of original.
Multiple Choice: Daily sales forecast test window should be: (a) random rows, (b) most recent dates, (c) earliest dates. Answer: (b).

Key Takeaways

train_test_split is the default tool; parameters encode reproducibility and balance.
Match split strategy to independence of rows (time, groups).
Split before fitting transforms; persist indices for audit.
Next: Cross Validation—robust estimates with limited data.

Trainer’s Guide

Hands-on idea: Same dataset, three splits: random, group by user, time cutoff. Compare validation AUC and discuss which is deployment-honest.

Discussion prompt: Marketing wants to include last week’s data in training for a model scoring tomorrow. Any leakage concern?

Recap: Train-test split partitions data so learning and evaluation stay honest. Continue with Cross Validation.