通过将以逗号分隔的值列分配到行中来创建新的数据框

Question

我试图通过将列中的逗号分隔值分配给不同的行来使用给定的数据帧创建数据帧。 原始DF中的一列包含以逗号分隔的值对。这些对中的每对都需要添加到新数据框中的一行中。

这是我目前正在做的一个例子。

输入数据框

val df = spark.createDataFrame(Seq(
         (1, "11 111, 12 112"),
         (2, "21 211, 22 212, 23 213"))).toDF("id", "A")

+---+----------------------+
|id |A                     |
+---+----------------------+
|1  |11 111, 12 112        |
|2  |21 211, 22 212, 23 213|
+---+----------------------+

创建新数据框

val df_array = df.rdd.collect

val df_array2 = df_array.map(x => (x(0).toString.toInt, x(1).toString.trim.split(", ")))
  .flatMap(x => x._2.map(y => (y.take(y.indexOf(" ")).trim.toDouble, 
                               y.takeRight(y.length - y.indexOf(" ") - 1).trim.toDouble, x._1)))
  .zipWithIndex

val output = df_array2.toVector.toDF
                            .withColumn("new_id", $"_2" + 1)
                            .withColumn("id", $"_1._3")
                            .withColumn("AA", $"_1._1")
                            .withColumn("BB", $"_1._2")
                            .drop("_1", "_2")

输出

+------+---+----+-----+
|new_id| id|  AA|   BB|
+------+---+----+-----+
|     1|  1|11.0|111.0|
|     2|  1|12.0|112.0|
|     3|  2|21.0|211.0|
|     4|  2|22.0|212.0|
|     5|  2|23.0|213.0|
+------+---+----+-----+

此方法确实有效，我得到了想要的东西。但是，它不省时，并且如果输入很大，则代码会因java.lang.OutOfMemoryError：错误而崩溃。

有更好的方法吗？