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:
Hey, thanks for the consideration. This issue is more prevalent than it appears. One reason is that people need to actively look out for this mistake, which is often caused by [splitting after a filter](https://scikit-learn.org/stable/common_pitfalls.html#how-to-avoid-data-leakage), or else it flies under the radar.

A secondary point is that sometimes this advice cannot be followed. For example, if you are getting your dataset from a big database, you'll have to choose a query condition.

A lot of best practices are not absolutely necessary for correctness, but they do help immensely by bringing down the mental load of juggling many things, and in my opinion, this falls under that category. There are plenty of cases where models under development go through a change of scope of data, multiple queries, multiple people working on the same project. Then a lot of care is taken to make sure no leakage happens. Having UID-stable splits would provide stronger guarantees against data leakage, and can be easily checked at different stages of development.

( I fixed the variable name in my examples)

Labels: New Feature