Spark Scala: разбить фрейм данных на несколько фреймов данных после x строк

Question

Я ищу способ разбить фрейм данных на несколько фреймов данных после x строк. Поскольку я хочу использовать метод несколько раз, было бы здорово его параметризовать. Например, у меня есть такой фрейм данных с несколькими тысячами строк:

<table>
  <tr>
    <th>Id</th>
    <th>Name</th>
  </tr>
  <tr>
    <td>1</td>
    <td>A</td>
  </tr>
  <tr>
    <td>2</td>
    <td>B</td>
  </tr>
  <tr>
    <td>3</td>
    <td>C</td>
  </tr>
  <tr>
    <td>4 .. </td>
    <td>..</td>
  </tr>
</table>

После x строк я хочу разделить его на два, три или x фреймов данных. Чтобы сделать его более понятным, я попытаюсь объяснить логику c метода, о котором я думаю.

def divideDF(df: DataFrame, delimiter: Integer): Seq[DataFrame] = {
  val num = df.count
  val start = 0
  val end = fn.round(num/delimiter) // this is the number of dfs i want to receive
  val i = 0

  while(i <= end){
    // split df in multiple data frames
  }  
}

Я очень ценю вашу помощь, если вам нужна дополнительная информация, я предоставлю ее: )

Answer 1

Попробуйте функцию randomSplit:

import org.apache.spark.sql.SparkSession

object SampleFoo extends App {
  val spark = SparkSession
    .builder()
    .master("local[2]")
    .getOrCreate()

  spark.sparkContext.setLogLevel("WARN")

  import spark.implicits._

  val splits = (1 to 100)
    .toDS
    .randomSplit(Array(.25, .25, .25, .25), 1)

  println(splits.length)

  splits.head.printSchema()

  splits.foreach(s => {
    s.show(40)
  })
}