如何将数据框中的列添加到spark中的列表中

时间:2017-07-24 17:05:17

标签: scala apache-spark

我的DataFrame有80列,我需要将其中的12列放入集合中,ArrayList都可以。我做了一点谷歌,发现了这个:

dataFrame.select("YOUR_COLUMN_NAME").rdd.map(r => r(0)).collect()

问题是,这适用于一列。如果我df.select(col1,col2,col3...).rdd.map.collect(),那么就给我这样的内容:Array[[col1,col2,col3]]

我想要的是Array[[col1],[col2],[col3]]。有没有办法在Spark中做到这一点?

提前致谢。

更新

例如,我有一个数据框:

----------
A   B   C
----------
1   2   3
4   5   6

我需要将列格式化为这种格式:

Array[[1,4],[2,5],[3,6]]

希望这一点更加明确......对不起因为混乱

1 个答案:

答案 0 :(得分:2)

您可以通过执行以下操作获得Array[Array[Any]]

scala> df.select("col1", "col2", "col3", "col4").rdd.map(row => (Array(row(0)), Array(row(1)), Array(row(2)), Array(row(3))))
res6: org.apache.spark.rdd.RDD[(Array[Any], Array[Any], Array[Any], Array[Any])] = MapPartitionsRDD[34] at map at <console>:32

RDD就像一个Array,所以你需要的数组在上面。如果你想要RDD[Array[Array[Any]]],那么你可以做

scala> df.select("col1", "col2", "col3", "col4").rdd.map(row => Array(Array(row(0)), Array(row(1)), Array(row(2)), Array(row(3))))
res7: org.apache.spark.rdd.RDD[Array[Array[Any]]] = MapPartitionsRDD[39] at map at <console>:32

您可以按照相同的方式进行十二列

<强>更新

您更新的问题更加明确。因此,您可以在转换为collect_list之前使用rdd函数,并像以前一样继续使用。

scala> import org.apache.spark.sql.functions._
import org.apache.spark.sql.functions._

scala> val rdd = df.select(collect_list("col1"), collect_list("col2"), collect_list("col3"), collect_list("col4")).rdd.map(row => Array(row(0), row(1), row(2), row(3)))
rdd: org.apache.spark.rdd.RDD[Array[Any]] = MapPartitionsRDD[41] at map at <console>:36

scala> rdd.map(array => array.map(element => println(element))).collect
[Stage 11:>                                                         (0 + 0) / 2]WrappedArray(1, 1)
WrappedArray(2, 2)
WrappedArray(3, 3)
WrappedArray(4, 4)
res8: Array[Array[Unit]] = Array(Array((), (), (), ())) 

仅限数据框

您可以在数据框本身中执行所有这些操作,而无需转换为rdd

假设您将数据框设为

scala> df.show(false)
+----+----+----+----+----+----+
|col1|col2|col3|col4|col5|col6|
+----+----+----+----+----+----+
|1   |2   |3   |4   |5   |6   |
|1   |2   |3   |4   |5   |6   |
+----+----+----+----+----+----+

您可以简单地执行以下操作

scala> import org.apache.spark.sql.functions._
import org.apache.spark.sql.functions._

scala> df.select(array(collect_list("col1"), collect_list("col2"), collect_list("col3"), collect_list("col4")).as("collectedArray")).show(false)
+--------------------------------------------------------------------------------+
|collectedArray                                                                  |
+--------------------------------------------------------------------------------+
|[WrappedArray(1, 1), WrappedArray(2, 2), WrappedArray(3, 3), WrappedArray(4, 4)]|
+--------------------------------------------------------------------------------+