Spark Dataset按指定列分组时如何保留全部原有列？

为什么当前代码拿不到全列 #

Spark的DataFrame执行groupBy操作的核心机制是：仅保留两类列——你指定的分组键列、agg方法里传入的聚合计算列。哪怕你提前select("*")把全列查出来，只要进入groupBy的聚合流程，所有非分组键、非聚合逻辑涉及的列都会被优化器直接裁剪掉，这是模型设计的固有逻辑，不是代码写法错误。
你之前尝试的join方案性能差，本质是引入了两次全量shuffle（一次分组聚合shuffle、一次和原表关联的join shuffle），很容易触发数据倾斜、小文件过多等问题，完全可以通过无join的写法实现同等需求。

方案1：窗口函数实现全列保留+分组聚合计算（推荐） #

如果你的需求是保留原数据集的每一行、不做行合并，同时给每一行带上对应分组的聚合统计值，这是性能最优的无join实现：

1. 用和groupBy完全一致的分区规则定义窗口
2. 将原来的聚合计算改为在窗口上执行
3. 提前把所有原表列放在select列表中，再叠加窗口计算的聚合值

示例代码：

import org.apache.spark.sql.expressions.Window;
import static org.apache.spark.sql.functions.count;

// 定义和原groupBy逻辑完全对齐的窗口分区规则
WindowSpec groupPartitionWindow = Window
    .partitionBy(groupByAttributes.toArray(new Column[0]));

// 直接输出全原列 + 对应分组的聚合值
Dataset<Row> groupedRows = initDataset
    .select(
        initDataset.col("*"), // 完整保留所有原有列
        count("*").over(groupPartitionWindow).alias("group_row_count")
    );

这个方案的优势是：全程只扫描一次原数据，shuffle逻辑和原groupBy完全一致，没有额外join带来的数据膨胀、倾斜问题，输出的每一行都100%保留原表的所有字段，不会出现列丢失。

方案2：聚合阶段收集全列值（输出为一行一组的聚合结果） #

如果你的需求是保持groupBy“一组输出一行”的聚合形态，同时在结果中带上每个分组下的所有原始列值，可以直接在agg阶段把非分组键的列收集为集合，不需要和原表join：

1. 遍历原表的所有列，跳过已经是分组键的列
2. 对剩余列用collect_list收集为对应字段的明细数组
3. 把收集表达式和你原本需要的聚合值一起传入agg方法

示例代码：

import static org.apache.spark.sql.functions.collect_list;

// 构造聚合表达式：把所有非分组键的原列收集为列表
List<Column> aggColumns = Arrays.stream(initDataset.columns())
    .filter(colName -> !groupByAttributes.contains(colName))
    .map(colName -> collect_list(initDataset.col(colName)).alias(colName + "_in_group"))
    .collect(Collectors.toList());

// 补充原逻辑需要的count(*)聚合值
aggColumns.add(count("*").alias("group_row_count"));

Dataset<Row> groupedRows = initDataset
    .groupBy(groupByAttributes.toArray(new Column[0]))
    .agg(aggColumns.stream().toArray(Column[]::new));

这个方案输出的每个分组对应一行，分组键列完整保留，其他原表列以数组形式存储，数组内的元素就是该分组下所有原始行的对应字段值，所有计算在一次shuffle聚合中完成，没有额外join开销。

常见避坑点 #

• 不要为了保留列在agg里随便加first(col)这类无意义聚合，这种写法拿到的只是分组内随机某一行的字段值，不是全量原有列，和需求偏差很大。
• 动态创建groupBy的Column数组时，不需要硬算数组长度，直接传new Column[0]即可，Spark会自动根据传入的元素数量适配数组长度，不会触发类型异常。
• 窗口函数的分区规则必须和你原本要groupBy的列完全一致，否则计算出来的分组聚合值会出现逻辑错误。

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