[SPARK-4500] Improve exact stratified sampling implementation

• Type: Improvement
• Status: Resolved
• Priority: Major
• Resolution: Incomplete
• Affects Version/s: 1.2.0
• Component/s: MLlib
• Labels: bulk-closed
• Assignee: Unassigned
• Reporter: Joseph K. Bradley
• Created: 19/Nov/14 21:59
• Updated: 21/May/19 05:37
• Resolved: 21/May/19 05:37

Description

The current implementation for exact stratified sampling (sampleByKeyExact) could be more efficient. Proposed algorithm sketch:

• Sampling is done separately for each stratum. Here, all counts are w.r.t. a fixed stratum.
• Let n_partition = number of elements in a given partition
• This method uses 2-tiered sampling:
  • Tier 1 (on driver): Select the sample sizes {n_partition}

    for each partition.

    • Without replacement, this uses a multivariate hypergeometric distribution.
    • With replacement, this should use a multinomial distribution.
  • Tier 2 (in parallel): Select a sample of size n_partition.

If anyone is interested in implementing this, I have a rough draft which works without replacement, but it needs to be cleaned up and augmented to do sampling with replacement too.