Added stratify option to train_test_split function.
CONTRIBUTOR
> 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.
2022-05-12T11:30:53Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1124879187
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DDEtT
https://api.github.com/repos/huggingface/datasets/issues/comments/1124879187/reactions
2022-05-12T11:30:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1124879187
https://avatars.githubusercontent.com/u/48522685?v=4
https://api.github.com/users/nandwalritik/events{/privacy}
https://api.github.com/users/nandwalritik/followers
https://api.github.com/users/nandwalritik/following{/other_user}
https://api.github.com/users/nandwalritik/gists{/gist_id}
https://github.com/nandwalritik
nandwalritik
MDQ6VXNlcjQ4NTIyNjg1
https://api.github.com/users/nandwalritik/orgs
https://api.github.com/users/nandwalritik/received_events
https://api.github.com/users/nandwalritik/repos
https://api.github.com/users/nandwalritik/starred{/owner}{/repo}
https://api.github.com/users/nandwalritik/subscriptions
User
https://api.github.com/users/nandwalritik
public
MEMBER
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
2022-05-12T11:49:53Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1124895298
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DDIpC
https://api.github.com/repos/huggingface/datasets/issues/comments/1124895298/reactions
2022-05-12T11:49:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1124895298
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
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 =============================================================================

```
2022-05-13T04:40:13Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1125650655
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DGBDf
https://api.github.com/repos/huggingface/datasets/issues/comments/1125650655/reactions
2022-05-13T04:40:13Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1125650655
https://avatars.githubusercontent.com/u/48522685?v=4
https://api.github.com/users/nandwalritik/events{/privacy}
https://api.github.com/users/nandwalritik/followers
https://api.github.com/users/nandwalritik/following{/other_user}
https://api.github.com/users/nandwalritik/gists{/gist_id}
https://github.com/nandwalritik
nandwalritik
MDQ6VXNlcjQ4NTIyNjg1
https://api.github.com/users/nandwalritik/orgs
https://api.github.com/users/nandwalritik/received_events
https://api.github.com/users/nandwalritik/repos
https://api.github.com/users/nandwalritik/starred{/owner}{/repo}
https://api.github.com/users/nandwalritik/subscriptions
User
https://api.github.com/users/nandwalritik
public
MEMBER
Thanks a lot !

`utils/stratify.py` sounds good yes :)

Also feel free to merge `master` into your branch to fix the CI ;)
2022-05-16T16:55:42Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1127909475
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DOohj
https://api.github.com/repos/huggingface/datasets/issues/comments/1127909475/reactions
2022-05-16T16:56:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1127909475
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
Added all the changes as were suggested and rebased with `main`.
2022-05-17T05:13:00Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1128419268
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DQk_E
https://api.github.com/repos/huggingface/datasets/issues/comments/1128419268/reactions
2022-05-17T05:13:00Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1128419268
https://avatars.githubusercontent.com/u/48522685?v=4
https://api.github.com/users/nandwalritik/events{/privacy}
https://api.github.com/users/nandwalritik/followers
https://api.github.com/users/nandwalritik/following{/other_user}
https://api.github.com/users/nandwalritik/gists{/gist_id}
https://github.com/nandwalritik
nandwalritik
MDQ6VXNlcjQ4NTIyNjg1
https://api.github.com/users/nandwalritik/orgs
https://api.github.com/users/nandwalritik/received_events
https://api.github.com/users/nandwalritik/repos
https://api.github.com/users/nandwalritik/starred{/owner}{/repo}
https://api.github.com/users/nandwalritik/subscriptions
User
https://api.github.com/users/nandwalritik
public
NONE
_The documentation is not available anymore as the PR was closed or merged._
2022-05-18T12:16:07Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1129933249
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5DWWnB
https://api.github.com/repos/huggingface/datasets/issues/comments/1129933249/reactions
2022-05-25T20:52:10Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1129933249
https://avatars.githubusercontent.com/u/99929124?v=4
https://api.github.com/users/HuggingFaceDocBuilderDev/events{/privacy}
https://api.github.com/users/HuggingFaceDocBuilderDev/followers
https://api.github.com/users/HuggingFaceDocBuilderDev/following{/other_user}
https://api.github.com/users/HuggingFaceDocBuilderDev/gists{/gist_id}
https://github.com/HuggingFaceDocBuilderDev
HuggingFaceDocBuilderDev
U_kgDOBfTMJA
https://api.github.com/users/HuggingFaceDocBuilderDev/orgs
https://api.github.com/users/HuggingFaceDocBuilderDev/received_events
https://api.github.com/users/HuggingFaceDocBuilderDev/repos
https://api.github.com/users/HuggingFaceDocBuilderDev/starred{/owner}{/repo}
https://api.github.com/users/HuggingFaceDocBuilderDev/subscriptions
User
https://api.github.com/users/HuggingFaceDocBuilderDev
public
NONE
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.".
2022-11-22T12:42:00Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1323621095
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5O5Nrn
https://api.github.com/repos/huggingface/datasets/issues/comments/1323621095/reactions
2022-11-22T12:42:00Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1323621095
https://avatars.githubusercontent.com/u/25030871?v=4
https://api.github.com/users/Damon03/events{/privacy}
https://api.github.com/users/Damon03/followers
https://api.github.com/users/Damon03/following{/other_user}
https://api.github.com/users/Damon03/gists{/gist_id}
https://github.com/Damon03
Damon03
MDQ6VXNlcjI1MDMwODcx
https://api.github.com/users/Damon03/orgs
https://api.github.com/users/Damon03/received_events
https://api.github.com/users/Damon03/repos
https://api.github.com/users/Damon03/starred{/owner}{/repo}
https://api.github.com/users/Damon03/subscriptions
User
https://api.github.com/users/Damon03
public
MEMBER
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` :)
2022-11-22T13:14:25Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1323661613
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5O5Xkt
https://api.github.com/repos/huggingface/datasets/issues/comments/1323661613/reactions
2022-11-22T13:14:25Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1323661613
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Thank you so much. It worked.
2022-11-22T14:53:55Z
https://github.com/huggingface/datasets/pull/4322#issuecomment-1323804641
https://api.github.com/repos/huggingface/datasets/issues/4322
IC_kwDODunzps5O56fh
https://api.github.com/repos/huggingface/datasets/issues/comments/1323804641/reactions
2022-11-22T14:53:55Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1323804641
https://avatars.githubusercontent.com/u/25030871?v=4
https://api.github.com/users/Damon03/events{/privacy}
https://api.github.com/users/Damon03/followers
https://api.github.com/users/Damon03/following{/other_user}
https://api.github.com/users/Damon03/gists{/gist_id}
https://github.com/Damon03
Damon03
MDQ6VXNlcjI1MDMwODcx
https://api.github.com/users/Damon03/orgs
https://api.github.com/users/Damon03/received_events
https://api.github.com/users/Damon03/repos
https://api.github.com/users/Damon03/starred{/owner}{/repo}
https://api.github.com/users/Damon03/subscriptions
User
https://api.github.com/users/Damon03
public
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.
