PYTHON / MACHINE LEARNING WITH PYTHON
Loading data and splitting train from test
Load a dataset into an X/y pair and carve out an honest hold-out test set with train_test_split, using random_state, stratify and shuffle correctly.
What you will learn
- Load data into X of shape (n_samples, n_features) and y of length n_samples
- Split with train_test_split and pin random_state so runs are reproducible
- Use stratify=y to keep class proportions identical in train and test
- Pass shuffle=False for time-ordered data so the test set is the future
Understanding Loading data and splitting train from test
Every scikit-learn estimator expects the same two objects: X, a two-dimensional container of shape (n_samples, n_features), and y, a one-dimensional container of length n_samples. Row i of X must describe the same observation as element i of y, which is why you split features from labels once, up front, and never sort or filter one without the other. Bundled datasets hand you both directly with load_iris(return_X_y=True); a CSV becomes a DataFrame via pandas.read_csv, and you select the feature columns into X and the target column into y.
A model that has seen a row can simply remember it, so scoring on training rows measures memory, not generalisation. The test set exists to imitate data that arrives after the model is finished, and it only imitates that faithfully while it stays untouched: no fitting, no inspecting distributions, no picking hyperparameters against it. This is why the split happens immediately after loading, before anything else in the pipeline gets to look at the numbers.
train_test_split works by generating a random permutation of the row indices and slicing it into two blocks, so the split is a shuffle followed by a cut. random_state fixes that permutation, which matters because otherwise a model change and a split change are indistinguishable in your results. The default shuffling is essential for files stored in label order — iris is sorted by species, so an unshuffled 20% test set would be nothing but species 2 — while stratify=y goes further and permutes within each class so the class proportions survive the cut. Time series invert the rule: there you pass shuffle=False, because training on later rows to predict earlier ones lets the model peek at the future.
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import numpy as np
X, y = load_iris(return_X_y=True)
print("X shape:", X.shape)
print("first 5 labels:", y[:5])
print("last 5 labels:", y[-5:])
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=0, stratify=y
)
print("train/test sizes:", X_train.shape[0], X_test.shape[0])
print("train label counts:", np.bincount(y_train))
print("test label counts:", np.bincount(y_test))The test set is a stand-in for data the model has never seen, so it must be separated from the training data before any decision is influenced by it.
Worked examples
Splitting a DataFrame read from CSV
Shows that a pandas DataFrame survives the split as a DataFrame, keeping column names and the original row index.
import io
import pandas as pd
from sklearn.model_selection import train_test_split
csv = """age,income,clicked
22,31000,0
25,42000,0
31,58000,1
35,61000,1
41,52000,0
44,88000,1
28,39000,0
52,95000,1
"""
df = pd.read_csv(io.StringIO(csv))
X = df[["age", "income"]]
y = df["clicked"]
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.25, random_state=42, stratify=y
)
print(type(X_train).__name__, X_train.shape, X_test.shape)
print(list(X_train.columns))
print(sorted(y_test.tolist()))Example explained
Line 1df[["age", "income"]] uses a list of columns, so X stays two-dimensional; df["clicked"] with a single string gives a 1D Series for y.
Line 2train_test_split returns the same types it was given, so X_train is still a DataFrame with usable column names.
Line 3test_size=0.25 of 8 rows is 2 rows, and stratify=y forces one clicked=0 and one clicked=1 into that pair.
Line 4The returned Series keep their original index labels, so you can join predictions back to the source rows.
Time-ordered data: no shuffling
Demonstrates shuffle=False, which makes the test set the final rows of the file instead of a random sample.
import numpy as np
from sklearn.model_selection import train_test_split
day = np.arange(1, 11).reshape(-1, 1)
sales = np.array([10, 12, 13, 15, 14, 18, 19, 22, 21, 25])
X_train, X_test, y_train, y_test = train_test_split(
day, sales, test_size=0.3, shuffle=False
)
print("train days:", X_train.ravel())
print("test days:", X_test.ravel())
print("test sales:", y_test)Example explained
Line 1reshape(-1, 1) turns the 10 day numbers into shape (10, 1), because X must be two-dimensional even with one feature.
Line 2shuffle=False makes the split a plain cut: the first 7 rows train, the last 3 test, in file order.
Line 3random_state is pointless here and would be ignored, since nothing is randomised.
Line 4Adding stratify=y with shuffle=False raises a ValueError, as stratification requires permuting rows.
Important notes
test_size can be a fraction or an integer count; with a fraction, scikit-learn rounds the test set up, so 0.2 of 150 gives exactly 30 and 0.2 of 7 gives 2.
Tuning against the same test set repeatedly leaks it slowly through your own decisions; hold out a separate validation split for that work.
Common mistakes
Fitting anything on the full X before splitting: the transformation absorbs information from the test rows, and the test score comes out higher than the model deserves.
Leaving random_state unset while comparing models: each run draws a different split, so the score difference you see may be the split moving rather than the model improving.
Passing shuffle=False on a dataset sorted by label: the test set ends up holding one class only, and accuracy on it says nothing about the other classes.
Slicing X and y separately, for example X[:120] with a re-sorted y: rows and labels stop lining up and the model trains on nonsense with no error raised.
Try it yourself
Change, predict, then run
Load load_wine(return_X_y=True), split it 70/30 with random_state=1 and stratify=y, then print np.bincount for both label arrays and check that each class keeps roughly the same proportion it had in the full dataset.
Open the Python workspaceCheck your understanding
You have 1000 emails of which 30 are spam, and you split off 20% as test data with a plain random split. What is the concrete risk?
- The test set may contain only a handful of spam emails, so any spam-related score computed on it swings wildly between splits
- The training set will be biased toward spam because the rare class is oversampled during shuffling
- The split will fail with an error, since train_test_split requires balanced classes
- The test set will be too small at 200 rows for the model to learn the spam pattern
Show answer
With 30 positives, a random 20% test set contains about 6 spam emails on average and can easily land on 2 or 11, so one misclassification moves the spam recall by a large fraction; stratify=y pins it at 6 every time. The last option confuses the two sets: the model never learns from the test set, that is the whole point of holding it out.