Introduction

Creating training and test samples from one DataFrame is a standard step in machine-learning workflows because you need one subset to fit the model and another to evaluate it on unseen data. The most reliable approach is usually train_test_split from scikit-learn, with pandas DataFrames passed in directly.

Use train_test_split for the Usual Case

The standard workflow looks like this:

1import pandas as pd
2from sklearn.model_selection import train_test_split
3
4_df = pd.DataFrame(
5    {
6        "age": [21, 35, 42, 19, 28, 50],
7        "income": [35000, 72000, 68000, 22000, 54000, 91000],
8        "bought": [0, 1, 1, 0, 0, 1],
9    }
10)
11
12X = _df[["age", "income"]]
13y = _df["bought"]
14
15X_train, X_test, y_train, y_test = train_test_split(
16    X,
17    y,
18    test_size=0.33,
19    random_state=42,
20    shuffle=True,
21)

This keeps rows aligned correctly between features and target, which is why it is better than sampling X and y separately.

If You Want Train and Test as DataFrames

Sometimes you want full DataFrames instead of separate feature and target objects. You can split the original frame directly.

1train_df, test_df = train_test_split(
2    _df,
3    test_size=0.2,
4    random_state=42,
5    shuffle=True,
6)

That is useful when the next stage of the pipeline still expects a single DataFrame and you are not ready to separate predictors from labels yet.

Use stratify for Classification Targets

If the target classes are imbalanced, use stratification so the train and test sets keep roughly the same class proportions.

1train_df, test_df = train_test_split(
2    _df,
3    test_size=0.2,
4    random_state=42,
5    stratify=_df["bought"],
6)

Without stratification, a small or imbalanced dataset can produce a misleading split where one subset underrepresents an important class.

Why Pure pandas.sample Is Usually Not Enough

You can create a quick split with pandas alone, but you must do it carefully to avoid overlap and leakage.

train_df = _df.sample(frac=0.8, random_state=42)
test_df = _df.drop(train_df.index)

This works for simple situations, but once you need stratification, reproducibility conventions, or aligned X and y handling, train_test_split is usually clearer and safer.

Split Before Fitting Preprocessing

One important rule is to split the data before fitting scalers, imputers, encoders, or feature selectors. If you fit preprocessing on the entire DataFrame first, information from the future test set leaks into training and makes evaluation too optimistic.

That is one of the most common mistakes in beginner machine-learning workflows.

Keep the Index Behavior in Mind

After splitting, the resulting DataFrames keep their original row indexes unless you reset them. That is often useful for traceability, but it can surprise people who expect the train and test sets to start at zero. If the downstream code assumes contiguous indexes, call reset_index(drop=True) explicitly instead of assuming the split changed them for you.

Common Pitfalls

• Sampling feature rows and target rows separately and breaking row alignment.
• Fitting preprocessing on the full dataset before the split.
• Forgetting to set random_state when reproducibility matters.
• Ignoring class imbalance and skipping stratify in classification problems.
• Using pandas-only sampling when the project already depends on scikit-learn’s safer split utilities.

Summary

• train_test_split is the standard way to create train and test samples from a pandas DataFrame.
• You can split the full DataFrame or split features and target together.
• Use stratify when class balance matters.
• Pure pandas sampling can work, but it is easier to make mistakes with it.
• Split before fitting preprocessing to avoid data leakage.
• Index handling after the split should be explicit, not accidental.