如何在Apache Spark上进行非随机数据集拆分?

时间:2016-12-02 14:50:32

标签: apache-spark apache-spark-sql apache-spark-dataset apache-spark-2.0

我知道我可以使用randomSplit方法随机拆分:

val splittedData: Array[Dataset[Row]] = 
        preparedData.randomSplit(Array(0.5, 0.3, 0.2))

我可以使用一些'nonRandomSplit方法'将数据拆分成连续的部分吗?

Apache Spark 2.0.1。 提前谢谢。

UPD:数据顺序非常重要,我将使用“较小ID”训练我的模型数据并使用“较大ID”对数据进行测试。所以我想将数据拆分成连续的部分而不需要改组。

e.g。

my dataset = (0,1,2,3,4,5,6,7,8,9)
desired splitting = (0.8, 0.2)
splitting = (0,1,2,3,4,5,6,7), (8,9)

我能想到的唯一解决方案是使用计数限制,但可能会有更好的解决方案。

1 个答案:

答案 0 :(得分:7)

这是我实施的解决方案:数据集 - > Rdd - >数据集。

我不确定这是否是最有效的方式,所以我很乐意接受更好的解决方案。

val count = allData.count()
val trainRatio = 0.6  
val trainSize = math.round(count * trainRatio).toInt
val dataSchema = allData.schema

// Zipping with indices and skipping rows with indices > trainSize.
// Could have possibly used .limit(n) here
val trainingRdd =
  allData
    .rdd
    .zipWithIndex()
    .filter { case (_, index) => index < trainSize }
    .map { case (row, _) => row }

// Can't use .limit() :(
val testRdd =
allData
  .rdd
  .zipWithIndex()
  .filter { case (_, index) => index >= trainSize }
  .map { case (row, _) => row }

val training = MySession.createDataFrame(trainingRdd, dataSchema)
val test = MySession.createDataFrame(testRdd, dataSchema)