Spark Scala получить массив типа строки из нескольких столбцов

Question

Я использую искру со скалой.

Представьте себе вход:

Я хотел бы знать, как получить следующий вывод [см. Аккумулятор столбца на следующем изображении], который должен быть массивом типа String Array[String]

В моем реальном фрейме данных у меня более 3 столбцов. У меня есть несколько тысяч столбцов.

Как я могу продолжить, чтобы получить желаемый результат?

Answer 1

Вы можете использовать array Функция и карта последовательности столбцов:

import org.apache.spark.sql.functions.{array, col, udf} 

val tmp = array(df.columns.map(c => when(col(c) =!= 0, c)):_*)

где

when(col(c) =!= 0, c)

принимает имя столбца, если значение столбца отличается от нуля, и ноль в противном случае.

и использовать UDF для фильтрации нулей:

val dropNulls = udf((xs: Seq[String]) => xs.flatMap(Option(_)))
df.withColumn("accumulator", dropNulls(tmp))

Итак, с примером данных:

val df = Seq((1, 0, 1), (0, 1, 1), (1, 0, 0)).toDF("apple", "orange", "kiwi")

вы сначала получите:

+-----+------+----+--------------------+
|apple|orange|kiwi|                 tmp|
+-----+------+----+--------------------+
|    1|     0|   1| [apple, null, kiwi]|
|    0|     1|   1|[null, orange, kiwi]|
|    1|     0|   0| [apple, null, null]|
+-----+------+----+--------------------+

и наконец:

+-----+------+----+--------------+
|apple|orange|kiwi|   accumulator|
+-----+------+----+--------------+
|    1|     0|   1| [apple, kiwi]|
|    0|     1|   1|[orange, kiwi]|
|    1|     0|   0|       [apple]|
+-----+------+----+--------------+