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:
There are a few things here:

_Is a hash based split a good idea?_

From what I read, I'm not convinced at all that any of the issues "solved" by this approach, are doing the right thing. For instance, if you're worried about "model seeing the validation / test set over time" in cases where you do a "warm started" kinda learning, you probably really should be a lot more careful about how you're separating that validation set instead of relying on any kind of random split. You might want to make sure data related to the same customer, or the same time period, for instance, is not included in both train and test.

On the other point of "needing to have a very stable model performance", that seems like a very bad idea to me. If your model is very sensitive to the way data is split, then you really should be reporting a confidence interval of your model performance, instead of a single performance value, and you're probably over estimating how good the model is.

But it seems like enough people like doing this, so I wouldn't mind having a hash based splitter in sklearn.

_Allowing `train_test_split` to use an alternative splitter_

I don't mind allowing to pass a cv object to `train_test_split`, and then passing the hash based splitter to it. It would also help with other cases of passing a group fold, or a time series splitter. To me that's a better design than hard coding this in `train_test_split`.

cc @ogrisel since he's been investigating splitters.

Labels: New Feature