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

Posted 2023-04-15

【中文标题】Spark Java - 将多列收集到数组列中

【英文标题】：Spark Java - Collect multiple columns into array column

【发布时间】：2018-07-05 06:50:46

【问题描述】：

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

| 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，如果提供更多输入列，则会引发错误。

【问题讨论】：

使用数组内置函数

【参考方案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)))

【参考方案2】：

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()

【讨论】：

那会产生一列StringType，而不是一个数组，或者？