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

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

Answer 2

Как упоминал Рамеш в своем комментарии, вы можете использовать функцию 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)))
}