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

Question

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

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

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

В моей реальной dataframe у меня есть более чем в 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]| 
+-----+------+----+--------------+