Added stratify option to train_test_split function.
> 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._
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.
