title: Added stratify option to train_test_split function.
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 =============================================================================
  
  ```
body: 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.
  