Цикл по столбцам данных для формирования вложенных данных - Spark

Question

У меня есть кадр данных, как показано ниже,

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] |
    +---+---+---+---+

Я хочу создать фрейм данных, создавая динамические имена столбцов, такие как A1, A2 .., циклически изменяя значения от 1 до 100 и делая выбор.

Как я могу это сделать?

Ура!

Answer 1

Просто соберите список столбцов для объединения в массив, преобразуйте имена столбцов в 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]|
// +---+----+---------+