StratifiedSplits.scala

class StratifiedSplits (data : DataFrame, labelCol : String, numSplits : Int) extends Serializable {

  private val labels : Seq[Double] = data.select(labelCol).distinct.map(row => row.getDouble(0)).collect

  private val dataFrames : Seq[Array[DataFrame]] = labels.map(label => {
    val newData = data.filter(data(labelCol) === label)

    val splits : Array[Double] = (0 until numSplits).map(k => 1 / numSplits.toDouble).toArray

    newData.randomSplit(splits)
  })

  dataFrames.foreach(e => e.foreach(df => df.persist))

  val dataSplits : Seq[DataFrame] = (0 until numSplits).map(k => {
    val datas : Seq[DataFrame] = dataFrames.map(e => e(k))

    datas.reduce((a, b) => a.unionAll(b))
  })
}