UID-based Stable Train-Test Split

### Describe the workflow you want to enable

During model development, it's common to perform train-test splits multiple times on a dataset. This may occur during development iterations, when the dataset evolves over time, or when working with different data subsets. However, the current `sklearn.model_selection.train_test_split` function has a subtle failure mode: even with a fixed random seed, it doesn't guarantee consistent splits in these scenarios. A simple modification like adding a single row or reordering the dataset can result in completely different splits, making debugging particularly challenging.

This limitation is well-recognized in the data science community. This is documented in a blog post [1], and popular book like Aurélien Géron's "Hands-On Machine Learning with Scikit-Learn and Tensorflow" includes custom implementations to address this problem [2-4].

To resolve this, we propose implementing a stable splitting mechanism based on unique identifiers. This approach would ensure that specific entries consistently remain in the same split, regardless of dataset modifications. 

```python
import pandas as pd
from sklearn.model_selection import train_test_split
df = pd.DataFrame({
    'id': [101, 102, 103, 104, 105, 106, 107, 108, 109, 110],
    'feature': [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1],
    'target': [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
})

train, test= train_test_split( df, test_size=0.3, random_state=42  )
print("initial split:\n",
    test)

df = pd.concat([df, pd.DataFrame(
                {'id': [111], 
                'feature': [11.1],
                'target': [1]}
                )],
                ignore_index=True)
train, test = train_test_split( df, test_size=0.3, random_state=42)
print("split after adding a new row:\n",
    test)

# initial split:
#      id  feature  target
# 8  109      9.9       0
# 1  102      2.2       1
# 5  106      6.6       1
# split after adding a new row:
#       id  feature  target
# 5   106      6.6       1
# 0   101      1.1       0
# 9   110     10.1       1
# 10  111     11.1       1
```

[1] [Reproducible ML: Maybe you shouldn't be using train_test_split](https://github.com/julian-west/e4ds-snippets/blob/master/best-practices/repeatable-splitting/Reproducible%20ML%20-%20Maybe%20you%20shouldn't%20be%20using%20sklearn's%20train_test_split.ipynb)

[2]  `split_data_with_id_hash()`: [Hands-on Machine Learning Chapter 2](https://github.com/ageron/handson-ml3/blob/main/02_end_to_end_machine_learning_project.ipynb)

[3]  [Data Science Stack Exchange: Splitting train/test sets by an identifier?](https://datascience.stackexchange.com/questions/51348/splitting-train-test-sets-by-an-identifier)

[4] [Data Science Stack Exchange: I dont understand this way of having a stable train/test split even after updating the dataset](https://datascience.stackexchange.com/questions/126648/i-dont-understand-this-way-of-having-a-stable-train-test-split-even-after-updati)

### Describe your proposed solution


Currently, `train_test_split` uses `BaseShuffleSplit` class to generate randomized splits. The proposed solution should implement a deterministic splitting mechanism that:

1. Uses a computationally efficient hash function to map unique IDs to consistent numeric values, assumed to be the first column of input
2. Ensures consistent test-size proportions to the best extent possible
3. Maintains compatibility with sklearn's `random_state` parameter

While this method sacrifices exact split ratios, the trade-off is worth it for the benefit of knowing that specific entries are consistently in the same split, preventing potential data leakage.

Here's a sketch implementation: 

```python
import pandas as pd
import numpy as np
from itertools import chain
from sklearn.utils import indexable, _safe_indexing, murmurhash3_32

class StableIDSplit:
    """A stable train-test splitter that uses unique IDs to ensure consistency."""
    
    def __init__(self, test_size=0.25, random_state=None):
        self.test_size = test_size
        self.random_state = random_state

    def split(self, X, y=None):
        hashed_values = np.array([
            murmurhash3_32(str(uid).encode(), seed=self.random_state, positive=True) 
            for uid in X
        ])
        
        test_mask = hashed_values < self.test_size * 2**32
        
        train_indices = np.where(~test_mask)[0]
        test_indices = np.where(test_mask)[0]

        return train_indices, test_indices

def stable_train_test_split(
    *arrays,
    test_size=0.25,
    random_state=None,
):
    """A modified train-test split function that ensures stable splits using unique IDs."""
    arrays = indexable(*arrays)

    stable_splitter = StableIDSplit(test_size=test_size, random_state=random_state)
    train_indices, test_indices = stable_splitter.split(arrays[0])

    return  list(
        chain.from_iterable(
            (_safe_indexing(a, train_indices), _safe_indexing(a, test_indices))
            for a in arrays
        )
    )


if __name__=="__main__":
    df = pd.DataFrame({
        'id': [101, 102, 103, 104, 105, 106, 107, 108, 109, 110],
        'feature': [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1],
        'target': [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
    })

    train_ids, test_ids, train, test= stable_train_test_split( df["id"], df, test_size=0.3, random_state=42  )
    print("initial split:\n",
        test)

    # Add a new row to the dataframe
    df = pd.concat([df, pd.DataFrame(
                    {'id': [111], 
                    'feature': [11.1],
                    'target': [1]}
                    )],
                    ignore_index=True)
    train_ids, test_ids, train, test = stable_train_test_split( df["id"], df, test_size=0.3, random_state=42)
    print("split after adding a new row:\n",
        test)

    # initial split:
    #      id  feature  target
    # 1  102      2.2       1
    # 6  107      7.7       0
    # 9  110     10.1       1
    # split after adding a new row:
    #      id  feature  target
    # 1  102      2.2       1
    # 6  107      7.7       0
    # 9  110     10.1       1
```


### Describe alternatives you've considered, if relevant

- Géron's function used `crc32()`. We switched to `murmurhash3_32()` as the function is built into scikit-learn
- This functionality could be built into `train_test_split`, possibly with a new flag (e.g., `strategy='uid'`)

### Additional context

Code sample and proposal edited with the help of GPT-4o and claude

COMMENT:
It is true that when the practitioner does things correctly, data leakage should never happen. This is unfortunately a common pitfall and I think it is harder to guard against, because it often results from lengthy development timeline 

Another thing I want to point out is that this is meant to solve a different problem than data splitting strategies. For example, even if I split based on customer id, I may still ask myself if this customer has been in the training set in previous rounds of training. Perhaps I need to store train-test split of customer id, and append to the list when I have a fresh set of data. Yes this would work, but a hash based method would avoid the diffing and appending. Even if I have only a single new row, I can immediately know whether it should belong in train or test.

To the second point about "stable model performance", I agree with you. I don't think the author of that notebook made a good point there

Labels: New Feature