我有一个包含多列的数据框:
| 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语句,然后为每个输入列(即2-20个输入列)编写一个UDF,如果提供了更多的输入列,则会抛出错误。
答案 0 :(得分: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)))
}
答案 1 :(得分:0)
val newDF = df.drop("a")
//create a new dataframe to get the combination of columns except the first column
val df2 = df.withColumn("combined", concat_ws(",",newDF.columns.map(collist =>
col(collist)): _*))
//new dataframe to select only the first and combined columns
val columnNames = Seq("a", "combined")
df2.select(columnNames.head, columnNames.tail:_*).show()