How to split data into train and test sets with scikit-learn

A held-out test set keeps model evaluation separate from fitting data. In scikit-learn, train_test_split() splits aligned feature, label, and row-id arrays in one call, so the inputs stay matched while a fixed share is reserved for final checks.

train_test_split() accepts arrays, sparse matrices, and dataframes with the same row count. Set test_size for the holdout share, pass a fixed random_state when an example or experiment needs repeatable rows, and use stratify=y for classification labels whose class mix should stay close between train and test sets.

A simple random split fits rows that can be shuffled independently. Grouped subjects, repeated measurements, or time-ordered data need a group-aware or time-series splitter because a random row split can leak information from training into evaluation.

1. Create split_train_test.py with the dataset, row IDs, and split call.

   split_train_test.py

   from collections import Counter
    
   import numpy as np
   import sklearn
   from sklearn.datasets import load_breast_cancer
   from sklearn.model_selection import train_test_split
    
    
   dataset = load_breast_cancer()
   X = dataset.data
   y = dataset.target
   row_ids = np.arange(X.shape[0])
    
   X_train, X_test, y_train, y_test, train_ids, test_ids = train_test_split(
       X,
       y,
       row_ids,
       test_size=0.25,
       stratify=y,
       random_state=42,
   )
    
   _, _, _, _, _, repeat_test_ids = train_test_split(
       X,
       y,
       row_ids,
       test_size=0.25,
       stratify=y,
       random_state=42,
   )
    
    
   def class_counts(labels):
       counts = Counter(labels)
       return {
           str(dataset.target_names[int(label)]): int(count)
           for label, count in sorted(counts.items())
       }
    
    
   print(f"scikit-learn {sklearn.__version__}")
   print(f"total rows: {X.shape[0]}")
   print(f"train shape: {X_train.shape}")
   print(f"test shape: {X_test.shape}")
   print(f"train class counts: {class_counts(y_train)}")
   print(f"test class counts: {class_counts(y_test)}")
   print(f"test share: {len(test_ids) / len(row_ids):.3f}")
   print(f"overlap rows: {len(np.intersect1d(train_ids, test_ids))}")
   print(f"repeat split matches: {np.array_equal(test_ids, repeat_test_ids)}")
   print(f"first five test row ids: {test_ids[:5].tolist()}")

   The optional row_ids array is included to prove that the training and test sets do not share rows.

2. Run the script.

   $ python split_train_test.py
   scikit-learn 1.9.0
   total rows: 569
   train shape: (426, 30)
   test shape: (143, 30)
   train class counts: {'malignant': 159, 'benign': 267}
   test class counts: {'malignant': 53, 'benign': 90}
   test share: 0.251
   overlap rows: 0
   repeat split matches: True
   first five test row ids: [519, 408, 291, 518, 385]

3. Confirm that test_size=0.25 produces the expected holdout size.

   The breast cancer dataset has 569 rows. A 25 percent split gives 426 training rows and 143 test rows because the requested share is rounded to whole samples.

4. Confirm that stratify=y keeps both labels represented in each split.

   The train and test class counts both include malignant and benign labels, with similar label proportions in each split.

5. Confirm that random_state=42 makes the split repeatable and that no row ID appears in both sets.

   overlap rows: 0 proves the held-out rows are separate from the training rows. repeat split matches: True shows the same test rows are selected when the same inputs and random state are used again.