遍历数据框列以形成嵌套数据框-Spark

时间:2019-10-31 01:18:24

标签: scala dataframe apache-spark

我有一个如下数据框,

val x = Seq(("A", "B", "C", "D")).toDF("DOC", "A1", "A2", "A3")

+---+---+---+---+
|DOC| A1| A2| A3|
+---+---+---+---+
|  A|  B|  C|  D|
+---+---+---+---+

这里的A可以一直到100,所以我想循环获取所有的A并将它们嵌套在如下所示的通用结构下,

   +---+---+---+----+
    |DOC|A LIST     |
    +---+---+---+---+
    |  A| [B, C, D] |
    +---+---+---+---+

我想通过从1到100的循环创建动态列名称(如A1,A2 ..)来创建数据框。

我该怎么做?

干杯!

1 个答案:

答案 0 :(得分:2)

只需将要组合的列列表组合成一个数组,即可通过Columns将列名转换为col并将方法array应用于结果列表:

val df = Seq(
  (1, "a", "b", "c", 10.0),
  (2, "d", "e", "f", 20.0)
).toDF("id", "a1", "a2", "a3", "b")

val selectedCols = df.columns.filter(_.startsWith("a")).map(col)
val otherCols = df.columns.map(col) diff selectedCols

df.select((otherCols :+ array(selectedCols: _*).as("a_list")): _*).show
// +---+----+---------+
// | id|   b|   a_list|
// +---+----+---------+
// |  1|10.0|[a, b, c]|
// |  2|20.0|[d, e, f]|
// +---+----+---------+