Scala：如何基于可变数量的列创建DataFrame数组列？

方法1：已知最大活动编号 #

如果已经明确最大的activityNum值，可直接生成目标列数组：

import org.apache.spark.sql.functions.{array, col}

// 替换为实际的最大活动编号
val maxActivityNum = 3

// 生成包含所有目标列的列表
val activityColumns = (1 to maxActivityNum).map(num => col(s"activity_${num}_id")).toList

// 创建数组表达式
val array_activityId = array(activityColumns: _*)

// 添加新列到DataFrame
val wellnessActivity1 = wellnessActivity.withColumn("activityIds", array_activityId)

方法2：自动识别现有活动列 #

若不确定最大编号，可从DataFrame现有列中自动提取符合规则的列：

import org.apache.spark.sql.functions.{array, col}
import scala.util.matching.Regex

// 定义列名匹配正则
val activityColPattern: Regex = "activity_(\\d+)_id".r

// 从现有列中提取所有活动编号
val existingActivityNums = wellnessActivity.columns.collect {
  case activityColPattern(num) => num.toInt
}

// 获取最大编号（若无匹配列需额外处理空值，此处假设至少存在一列）
val maxActivityNum = existingActivityNums.max

// 生成列数组并添加新列
val activityColumns = (1 to maxActivityNum).map(num => col(s"activity_${num}_id")).toList
val array_activityId = array(activityColumns: _*)
val wellnessActivity1 = wellnessActivity.withColumn("activityIds", array_activityId)

关键说明 #

• 通过(1 to maxActivityNum).map遍历编号，动态生成对应col对象列表
• array(activityColumns: _*)中的: _*用于将列表转换为可变参数传递给array函数
• 两种方式均避免了手动逐个列写col("activity_X_id")，实现动态生成数组列

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