Title: Added stratify option to train_test_split function.
Author: nandwalritik (CONTRIBUTOR)
State: closed
Created: 2022-05-12T08:00:31Z

This PR adds `stratify` option to `train_test_split` method. I took reference from scikit-learn's `StratifiedShuffleSplit` class for implementing stratified split and integrated the changes as were suggested by @lhoestq.

It fixes #3452.

@lhoestq Please review and let me know, if any changes are required.

Comments:
> Nice thank you ! This will be super useful :)
> 
> Could you also add some tests in test_arrow_dataset.py and add an example of usage in the `Example:` section of the `train_test_split` docstring ?

I will try to do it, is there any documentation for adding test cases? I have never done it before.

Thanks for the changes !

> I will try to do it, is there any documentation for adding test cases? I have never done it before.

You can just add a function `test_train_test_split_startify` in `test_arrow_dataset.py`.

In this function you can define a dataset and make sure that `train_test_split` with the `stratify` argument works as expected.

You can do `pytest tests/test_arrow_dataset.py::test_train_test_split_startify` to run your test.

Feel free to get some inspiration from other tests like `test_interleave_datasets` for example

I have added tests for stratified train_test_split in `test_arrow_dataset.py` file inside `test_train_test_split_startify` function. I have also added example usage with `stratify` arg in `Example:` section of the `train_test_split` docstring.
Results of tests:
```
(data) nandwalritik@hp:~/datasets$ pytest tests/test_arrow_dataset.py::test_train_test_split_startify -W ignore
============================================================================ test session starts ============================================================================
platform linux -- Python 3.9.5, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/nandwalritik/datasets
plugins: datadir-1.3.1, forked-1.4.0, xdist-2.5.0
collected 1 item                                                                                                                                                            

tests/test_arrow_dataset.py .                                                                                                                                         [100%]

============================================================================= 1 passed in 0.12s =============================================================================

```

Thanks a lot !

`utils/stratify.py` sounds good yes :)

Also feel free to merge `master` into your branch to fix the CI ;)

Added all the changes as were suggested and rebased with `main`.

_The documentation is not available anymore as the PR was closed or merged._

Hi, I encounter an error when I try to specify the stratify_by_column. However, I have a columns which specific the label of the row as a string. But an error showed when I try to do it. "ValueError: Stratifying by column is only supported for ClassLabel column, and column code is Value.".

Hi @Damon03 , you can change the type of your column to ClassLabel using
```python
ds = ds.class_encode_column(column_name)
```
then you'll be free to use `stratify` :)

Thank you so much. It worked.