Scala Spark中无需逐个输入列名使用.collect_list()的方法

问题背景 #

我有一个规模极大的数据集（550亿行、2300列），需要按指定列分组后，对其余所有列应用collect_list函数。但手动逐个输入2300个列名完全不现实，求Scala Spark中的可行实现方案，且方案必须适配大数据场景。

示例数据集 #

import org.apache.spark.sql.functions._

val Enormous_df = Seq(("Apples", "a", "b"),
                      ("Apples", "a", "d"),
                      ("Banana", "e", "f"))
                      .toDF("Grouping_by_this_Column", "One_of_Thousands_of_Columns", "Second_of_Thousands_of_Columns")

Enormous_df.show()

输出：

+-----------------------+---------------------------+------------------------------+
|Grouping_by_this_Column|One_of_Thousands_of_Columns|Second_of_Thousands_of_Columns|
+-----------------------+---------------------------+------------------------------+
|                 Apples|                          a|                             b|
|                 Apples|                          a|                             d|
|                 Banana|                          e|                             f|
+-----------------------+---------------------------+------------------------------+

错误尝试及问题 #

之前尝试直接传入列名序列，代码报错：

// 创建目标列名序列：排除分组列
val All_Thousands_of_Column_Titles = Seq(Enormous_df.drop("Grouping_by_this_Column").columns)

// 此代码无法运行
val The_df_I_Want = Enormous_df.groupBy("Grouping_by_this_Column").agg(collect_list(All_Thousands_of_Column_Titles))

报错信息：

command-33353371:11: error: overloaded method value collect_list with alternatives:
  (columnName: String)org.apache.spark.sql.Column <and>
  (e: org.apache.spark.sql.Column)org.apache.spark.sql.Column
 cannot be applied to (Seq[Array[String]])
val The_df_I_Want = Enormous_df.groupBy("Grouping_by_this_Column").agg(collect_list(All_Thousands_of_Column_Titles))

问题根源：collect_list仅接受单个列名或Column对象，无法直接处理序列/数组类型的批量输入。

解决方案（适配大数据场景） #

方案1：动态生成聚合表达式列表（推荐） #

通过列名列表动态生成每个列对应的collect_list表达式，再批量传入agg方法，完全适配分布式大数据处理：

import org.apache.spark.sql.functions._

val groupCol = "Grouping_by_this_Column"
// 获取所有需要应用collect_list的列：排除分组列
val targetCols = Enormous_df.columns.filter(_ != groupCol)

// 为每个列生成collect_list表达式，并保留原列名
val aggExprs = targetCols.map(colName => collect_list(col(colName)).alias(colName))

// 执行分组聚合：将表达式序列转成可变参数传入agg
val The_df_I_Want = Enormous_df.groupBy(groupCol).agg(aggExprs.head, aggExprs.tail:_*)

The_df_I_Want.show()

方案2：Struct打包聚合（可选） #

如果需要对所有列做统一打包处理，可先将目标列封装为Struct，聚合后再展开：

import org.apache.spark.sql.functions._

val groupCol = "Grouping_by_this_Column"
val targetCols = Enormous_df.columns.filter(_ != groupCol)

// 先打包目标列为Struct，聚合后再展开每个字段
val The_df_I_Want = Enormous_df.groupBy(groupCol)
  .agg(collect_list(struct(targetCols.map(col):_*)).alias("collected_struct"))
  .select(
    col(groupCol),
    targetCols.map(colName => col(s"collected_struct.$colName").alias(colName)):_*
  )

The_df_I_Want.show()

期望输出结果 #

两种方案均可得到如下目标DataFrame：

+-----------------------+---------------------------+------------------------------+
|Grouping_by_this_Column|One_of_Thousands_of_Columns|Second_of_Thousands_of_Columns|
+-----------------------+---------------------------+------------------------------+
|                 Apples|                     [a, a]|                        [b, d]|
|                 Banana|                        [e]|                           [f]|
+-----------------------+---------------------------+------------------------------+

大数据场景注意事项 #

• 避免使用collect/toLocalIterator等将数据拉到Driver端的操作，上述方案均为Executor端分布式执行。
• 针对2300列的大表，建议调整spark.sql.shuffle.partitions配置，避免Shuffle阶段出现性能瓶颈。
• 方案1内存开销更低，推荐优先使用，无需额外的Struct打包/解包操作。

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