Spark scala 从多列中获取字符串类型的数组

我正在使用 Spark 和 scala。

想象一下输入：

我想知道如何获得以下输出[请参阅下图中的列累加器]，它应该是字符串类型的数组 数组[字符串]

在我的真实数据框中，我有超过 3 列。我有几千个专栏。

我该如何继续以获得我想要的输出？

您可以使用array函数并映射一系列列：

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

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

where

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

如果列值不为零则采用列名，否则采用 null。

并使用 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]|
+-----+------+----+--------------+