Как преобразовать в список из многих столбцов данных в Spark?

Question

у меня есть этот пример dataframe:

id |  A  |  B  |  C  |  D 
1  |NULL | 1   | 1   |NULL
2  | 1   | 1   | 1   | 1
3  | 1   |NULL |NULL |NULL

и я хочу изменить на этот формат:

id | newColumn
1  | {"B", "C"}
2  | {"A","B","C","D"}
3  | {"A"}

Другими словами, я хочу создать новый столбец со списком, содержащим имена столбцов, в которых значения строк не равны нулю.

Как я могу сделать это в Spark, используя Scala?

Answer 1

Сначала получите имена столбцов там, где есть фактическое значение, а не ноль. Это можно сделать с помощью такой функции, как:

val notNullColNames = Seq("A", "B", "C", "D").map(c => when(col(c).isNotNull, c))

Обычно для создания массива значений используется array, однако, это все равно будет возвращать ноль, когда входное значение равно нулю. Вместо этого одним из решений является использование concat_ws и split для удаления любых нулевых значений:

df.select($"id", split(concat_ws(",", notNullColNames:_*), ",").as("newColumn"))

Для примера ввода это выведет:

+---+------------+
| id|   newColumn|
+---+------------+
|  1|      [B, C]|
|  2|[A, B, C, D]|
|  3|         [A]|
+---+------------+