Spark Java - 将多列收集到数组列中

Question

我有一个包含多列的数据框：

| a | b | c | d |
-----------------
| 0 | 4 | 3 | 6 |
| 1 | 7 | 0 | 4 |
| 2 | 4 | 3 | 6 |
| 3 | 9 | 5 | 9 |

我现在想合并[b,c,d]到一个列中。但是，我不知道列的列表有多大，否则我只能使用 UDF3 将三者结合起来。

所以想要的结果是：

| a | combined  |
-----------------
| 0 | [4, 3, 6] |
| 1 | [7, 0, 4] |
| 2 | [4, 3, 6] |
| 3 | [9, 5, 9] |

我怎样才能做到这一点？

非工作伪代码：

public static Dataset<Row> mergeColumns(Dataset<Row> ds, List<String> columns) {
   return ds.withColumn("combined", collectAsList(columns))
}

最坏情况的解决方法是输入列数的 switch 语句，然后为每个输入列编写一个 UDF，即 2-20 个输入列，如果提供更多输入列，则会抛出错误。

Answer 1

正如 Ramesh 在他的评论中提到的，你可以使用array函数。您只需要将列列表转换为Column数组。

public static Dataset<Row> mergeColumns(Dataset<Row> ds, List<String> columns) {
    return ds.withColumn("combined", functions.array(columns.stream().map(functions::col).toArray(Column[]::new)))
}