Spark：动态指定DataFrame列更新值并转为数组

问题背景 #

现有如下DataFrame定义：

import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
import org.apache.spark.sql.Row
import org.apache.spark.sql.functions.{col, array}

val schema = new StructType()
  .add(StructField("col1", IntegerType))
  .add(StructField("col2", IntegerType))
  .add(StructField("col3", IntegerType))
  .add(StructField("col4", IntegerType))

val data: RDD[Row] = spark.sparkContext.parallelize(Seq(
  (1, 2, 3, 4),
)).map(t => Row(t._1, t._2, t._3, t._4))

val sample = spark.createDataFrame(data, schema)

需要动态指定部分列（比如Seq("col1", "col4")、Seq("col3", "col4")等任意列组合），对这些列应用函数（例如值乘以3），并将处理后的结果存入一个新的Array类型列new_array。

当指定col1和col4时，预期输出为：

+----------+
 | new_array|
 +----------+
 |  [3, 12] |
 +----------+

解决方案 #

核心思路是根据动态指定的列名列表，生成对应的列处理表达式，再传入array()函数构建新列：

1. 定义需要处理的列列表（可替换为任意目标列组合）：

val columnsForHandle = Seq("col1", "col4")

2. 生成每个指定列的处理表达式：

// 这里以"乘以3"为例，可替换为任意自定义函数
val processedColumns = columnsForHandle.map(colName => col(colName) * 3)

3. 创建包含新数组列的DataFrame：

val resultDF = sample.withColumn("new_array", array(processedColumns: _*))

4. 查看结果：

resultDF.select("new_array").show()

说明 #

• processedColumns: _* 将序列转为可变参数，适配array()函数的参数要求，确保不管指定多少列都能动态生成数组。
• 处理函数可灵活替换，比如改为col(colName) + 5、upper(col(colName))（针对字符串列）等任意Spark支持的列操作。
• 测试不同列组合时，只需修改columnsForHandle的值即可，逻辑无需调整。

内容的提问来源于stack exchange，提问作者Jelly