Spark DataFrame分组聚合保留原有列并计算分组内credit/debit总额

现有代码问题 #

你当前编写的代码存在三点问题：

• 提前过滤仅保留credit类型行，直接丢失了所有debit类型的原始数据
• 使用普通groupBy聚合仅返回分组键和聚合结果列，自然会丢失其他原始列
• 分组键设置错误，没有使用你指定的id、txnId、date作为分组维度

解决方案 #

要实现保留所有原始列的分组统计，使用Spark窗口函数即可：窗口函数可以在不减少原始行数、保留所有原有列的前提下，完成指定分组内的聚合计算。

依赖导入 #

import org.apache.spark.sql.expressions.Window;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import static org.apache.spark.sql.functions.*;

核心实现代码 #

// 定义窗口分区规则，按你指定的id、txnId、date三个字段分组
Window spec = Window.partitionBy("id", "txnId", "date");

Dataset<Row> resultDf = df2
    // 统计分组内所有credit类型的金额总和，非credit类型统计值为0
    .withColumn("totalcredit", sum(when(col("type").equalTo("credit"), col("amount")).otherwise(0)).over(spec))
    // 统计分组内所有debit类型的金额总和，非debit类型统计值为0
    .withColumn("totaldebit", sum(when(col("type").equalTo("debit"), col("amount")).otherwise(0)).over(spec));

逻辑说明 #

• Window.partitionBy 对应你要求的分组规则，所有分区内的行共享同一个统计结果
• when().otherwise() 条件判断实现仅对指定交易类型的金额做求和
• over(spec) 标识该聚合是窗口聚合，计算结果会追加到分组内的每一行，不会丢失任何原始列和数据行

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