Scala 将数组转换为 DataFrame 列

Question

我正在尝试将值数组作为新列添加到 DataFrame。

例如：让我们假设有一个 Array(4,5,10) 和一个 dataframe

+----------+-----+
|   name   | age |
+----------+-----+
|   John   | 32  |
| Elizabeth| 28  |
|   Eric   | 41  |
+----------+-----+

我的要求是将上述数组作为新列添加到数据框中。我的预期输出如下：

+----------+-----+------+
|   name   | age | rank |
+----------+-----+------+
|   John   | 32  | 4    | 
| Elizabeth| 28  | 5    |
|   Eric   | 41  | 10   |
+----------+-----+------+

我正在尝试是否可以使用 rdd 和 zipWithIndex 实现此目的。

df.rdd.zipWithIndex.map(_.swap).join(array_rdd.zipWithIndex.map(_.swap))

这导致了这种情况。

(0,([John, 32],4))

我想将上面的 RDD 转换回所需的数据帧。让我知道如何实现这一目标。

除了使用 rdd 和 zipWithIndex 之外，是否还有其他方法可以达到预期的效果？最好的方法是什么？

PS:

更好理解的上下文：

我正在使用 Xpress 优化套件来解决一个数学问题。 Xpress 接受数组的输入，并在数组中输出结果。我将输入作为 DataFrame 获取，并将列提取为数组（使用收集）并传递给 Xpress。 Xpress 输出 Array[Double] 作为解决方案。我想将此解决方案作为一列添加回数据框，解决方案数组中的每个值都对应于索引处数据框的行，即输出数组的索引 'n' 处的值对应于 'n'数据框的第行

Answer 1

加入后只需将结果映射到您要查找的内容即可。您可以在加入 RDD 后将其转换回数据框。


val originalDF = Seq(("John", 32), ("Elizabeth", 28), ("Eric", 41)).toDF("name", "age")

val rank = Array(4, 5, 10)

// convert to Seq first      
val rankDF = rank.toSeq.toDF("rank")

val joined = originalDF.rdd.zipWithIndex.map(_.swap).join(rankDF.rdd.zipWithIndex.map(_.swap))

val finalRDD = joined.map{ case (k,v) => (k, v._1.getString(0), v._1.getInt(1), v._2.getInt(0)) }

val finalDF = finalRDD.toDF("id", "name", "age", "rank")

finalDF.show()
/*
+---+---------+---+----+
| id|     name|age|rank|
+---+---------+---+----+
|  0|     John| 32|   4|
|  1|Elizabeth| 28|   5|
|  2|     Eric| 41|  10|
+---+---------+---+----+
*/

我能想到的唯一替代方法是使用 org.apache.spark.sql.functions.row_number() window 函数。这基本上通过向数据帧添加一个递增的连续行号来实现相同的目的。

这样做的缺点是将大量数据混洗到一个分区中，因为我们需要为数据帧中的所有行提供不重复的行号。如果您的数据非常大，这可能会导致内存不足问题。（注意：这可能不适用于您的情况，因为您提到您正在对数据进行收集并且没有在其中提及任何内存问题）。

转换为 rdd 并使用 zipWithIndex 的方法是一个可以接受的解决方案，但通常不推荐从数据帧转换为 rdd，因为使用 RDD 而不是数据框。

Scala 将数组转换为 DataFrame 列

Scala convert Array to DataFrame Column

dataframe

scala

apache-spark

apache-spark-sql

rdd