opened
sethah
username_0
## What changes were proposed in this pull request?

This patch adds the ability to do stratified sampling in cross validation for ML pipelines. This is accomplished by modifying some of the methods in `StratifiedSamplingUtils` to support multiple _splits_ instead of a single subsample of the data. A method is added to `PairRDDFunctions` to support `randomSplitByKey`. Please see the detailed explanation below.


## How was this patch tested?

Unit tests were added to `PairRDDFunctionsSuite`, `MLUtilsSuite`, `CrossValidatorSuite`, and `TrainValidationSuite`.

## Algorithm changes

Currently, Spark implements a stratified sampling function on PairRDDs using the method `sampleByKeyExact` and `sampleByKey`. This method calls a stratified sampling routine that is implemented in `StratifiedSamplingUtils`. The underlying algorithm is described [here](http://jmlr.org/proceedings/papers/v28/meng13a.pdf) in the paper by Xiangrui Meng. When exact samples stratified samples are required, the algorithm makes an extra pass through the data. Each sample is mapped on to the interval [0, 1] (for sampling without replacement), and we expect that, say for a 50% sample, we will split the interval at 0.5 and accept the samples which fell below that threshold. Items near 0 are highly likely to be accepted, while items near 1 are highly unlikely to be accepted. Items near 0.5 are uncertain, and are added to a waitlist on the first pass. The items in the waitlist will be sorted and used to determine the exact split point which produces 50/50 sample. 

![image](https://cloud.githubusercontent.com/assets/7275795/17071720/2f90bd14-5018-11e6-91c2-2c1dc2191213.png)

This patch modifies the routine to produce multiple splits by generating multiple waitlists on the first pass. Each waitlist is sorted to determine the exact split points and then we can sample as normal. 

![image](https://cloud.githubusercontent.com/assets/7275795/17071744/4f3b6d76-5018-11e6-8591-857eaaf288f9.png)

One potential concern is that if this is used for a large number of splits, it may degrade to the point where sorting the entire dataset would be quicker, as the waitlists get closer and closer together. It could potentially cause OOM errors on the driver if there are too many waitlists collected. Still, before this patch there was not a way to actually take a single _split_ of the data, as `sampleByKey` does not return the complement of the sample. This patch fixes this as well.

## ML API

This patch also allows users to specify a stratified column in the `CrossValidator` and `TrainValidationSplit` estimators. This is done by converting the input dataframe to a PairRDD and calling the `randomSplitByKey` method. This is exposed via a `setStratifiedCol` parameter which, if set, will use _exact_ stratified splits for cross validation. 

## Future considerations

This can be implemented as a function on dataframes in the future, if there is interest. It is somewhat inconvient to convert the dataframe to a pair rdd, perform sampling, and then convert back to a dataframe.
[SPARK-8971][ML] Add stratified sampling to ML CrossValidator and TrainValidationSplit
issue
created
sethah
username_0
cc @MLnick @hhbyyh @mengxr I believe there is still interest in stratified sampling methods. Could you provide feedback/review on this patch? Thanks!
comment
created
pramitchoudhary
username_1
Has the progress on this initiative stalled for any reason. May be I could be of help. @username_0
comment
created
HyukjinKwon
username_2
I just happened to look at this PR. Is this still WIP or waiting more review comments? If it is simply that the author is not currently able to proceed this further, then, maybe it'd be better to close this for now.
comment
created
idlecool
username_3
Hi @username_0, any plans to work on it again?
comment