31. Train Test Split in Machine Learning

31.2 How train_test_split Works in sklearn

By Ravindra Bagale · Last updated: 24 October 2026

Why. sklearn gives one clear function so you do not hand-split and accidentally leak rows.

What. train_test_split from sklearn.model_selection randomly splits arrays or DataFrames into train and test parts.

How.

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame({
 "items": [2, 3, 5, 8, 1, 4, 7, 6, 2, 9],
 "distance_km": [1.2, 3.0, 4.5, 6.1, 0.8, 2.2, 5.0, 3.3, 1.0, 7.4],
 "was_late": [0, 0, 1, 1, 0, 0, 1, 0, 0, 1],
})
X = df[["items", "distance_km"]]
y = df["was_late"]

X_train, X_test, y_train, y_test = train_test_split(
 X, y, test_size=0.2, random_state=42
)
print("train rows", len(X_train), "test rows", len(X_test))
print(X_train.head())

What you should see. About 8 train rows and 2 test rows on this tiny set. On real data, keep more rows; this is only to learn the API.