Sparkscala将列名映射为值true的输出数组类型列

Question

如何创建列colMap的ArrayType[StringType]哪个值是数组与元素是字符串匹配的列的名称哪些值为真？

我有这样的输入DataFrame：

+-----+-----+-----+
|col1 |col2 |col3 |
+-----+-----+-----+
|true |false|true |
|false|false|false|
|false|false|true |
+-----+-----+-----+

我想创建这样的输出DataFrame：

+-----+-----+-----+------------+
|col1 |col2 |col3 |colMap      |
+-----+-----+-----+------------+
|true |false|true |[col1, col3]|
|false|false|false|[]          |
|false|false|true |[col3]      |
+-----+-----+-----+------------+

编辑：我发现了这个重复的问题：

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

但想知道是否有更好的方法来实现产出？

Answer 1

您可以使用内置的高阶函数过滤器，而不是使用UDF过滤数组中的值。

val df = Seq((true, false, true),
    (false, false, false),
    (false, false, true)).toDF("col1", "col2", "col3")


df.withColumn("colMap", array(df.columns.map(c=> when(col(c) === "true", c)):_*))
  .withColumn("colMap", expr("filter(colMap, c-> c is not null)"))
  .show(false)

+-----+-----+-----+------------+
|col1 |col2 |col3 |colMap      |
+-----+-----+-----+------------+
|true |false|true |[col1, col3]|
|false|false|false|[]          |
|false|false|true |[col3]      |
+-----+-----+-----+------------+