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:
Thanks for the issue, it seems like this would make sense in the following two cases:
- your dataset order is changed while working on your problem. This seems like a bit of a corner case if you ask me ...
- you gather more data and you strongly insist that the samples that were previously (before more data gathering) in the train set still go in the train set and same for the test set.

I have to admit I am not 100% sure those corner cases are worth adding the functionality but we can leave this open and see whether there is some interest. I was not able to find an existing issue about this in our github issue tracker.

By the way, your last snippet doesn't run, `unique_ids` is not defined ...

Labels: New Feature