Spark Scala中如何基于DataFrame列值累计和为每行分配分类标签

实现逻辑 #

你的数据已经按salary升序排列，我们可以通过自定义UDF维护累加状态的方式实现分组，全程不需要额外shuffle，性能优异。

完整代码 #

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
import spark.implicits._

// 1. 构造样例数据，你可以替换成自己的数据源
val rawDF = Seq(
  ("Emp1", 10),
  ("Emp2", 20),
  ("Emp3", 30),
  ("Emp4", 35),
  ("Emp5", 36),
  ("Emp6", 50),
  ("Emp7", 70)
).toDF("Employee", "salary")

// 2. 定义UDF：维护累加薪资状态，生成分组ID
val calcGroupId = udf(new Serializable {
  // 存储当前分组累计薪资
  private var currentTotal = 0
  // 存储当前分组ID
  private var groupId = 0
  def apply(salary: Int): Int = {
    // 累加后超过80则开启新分组
    if (currentTotal + salary > 80) {
      groupId += 1
      currentTotal = salary
    } else {
      currentTotal += salary
    }
    groupId
  }
}.apply _)

// 3. 定义UDF：将数字分组ID转换为A/B/C/...的字母标签
val groupIdToCategory = udf((id: Int) => {
  var label = ""
  var num = id
  while (num >= 0) {
    label = (('A' + num % 26).toChar) + label
    num = num / 26 - 1
  }
  label
})

// 4. 执行计算得到最终结果
val resultDF = rawDF
  // 保证行顺序和排序一致，你已经排好的话可以省略，建议加上确保顺序
  .orderBy("salary")
  .withColumn("group_id", calcGroupId(col("salary")))
  .withColumn("Category", groupIdToCategory(col("group_id")))
  .drop("group_id")

// 输出结果验证
resultDF.show()

输出结果 #

和预期完全一致：

+--------+------+--------+
|Employee|salary|Category|
+--------+------+--------+
|    Emp1|    10|       A|
|    Emp2|    20|       A|
|    Emp3|    30|       A|
|    Emp4|    35|       B|
|    Emp5|    36|       B|
|    Emp6|    50|       C|
|    Emp7|    70|       D|
+--------+------+--------+

注意事项 #

• 上述代码默认数据是全局排序且单分区运行，若有多分区全局排序的需求，可以改用Spark 3.x的transform配合有状态聚合，或者先按range分区后，每个分区单独计算分组ID再叠加偏移量即可。
• 字母标签转换支持超过26个分组的场景，第26个分组会返回AA，第27个返回AB，以此类推。

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