Scala DataFrame获取每条记录最大值与次大值列名

解决方案代码 #

Spark 3.x+ 版本（支持自定义排序逻辑） #

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

// 构造示例DataFrame
val dfIn = spark.createDataFrame(Seq(
  ("r0", 0, 1, 2, "a0"),
  ("r1", 1, 2, 0, "a1"),
  ("r2", 2, 0, 1, "a2")
)).toDF("prev_column", "c0", "c1", "c2", "post_column")

// 指定需要分析的数值列
val targetCols = List("c0", "c1", "c2")

// 生成包含列名与对应数值的结构体数组
val colsArray = array(
  targetCols.map(colName => struct(lit(colName).alias("name"), col(colName).alias("value"))): _*
)

// 排序后提取最大、第二大的列名
val dfResult = dfIn
  .withColumn("sorted_cols", array_sort(colsArray, (a, b) => b.getField("value").compareTo(a.getField("value"))))
  .withColumn("first", col("sorted_cols")(0).getField("name"))
  .withColumn("second", col("sorted_cols")(1).getField("name"))
  .drop("sorted_cols")

// 查看结果
dfResult.show()

Spark 2.x 版本兼容写法 #

如果你的Spark版本低于3.0，无法使用带自定义比较器的array_sort，可以改用sort_array配合结构体顺序实现降序：

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

val dfIn = spark.createDataFrame(Seq(
  ("r0", 0, 1, 2, "a0"),
  ("r1", 1, 2, 0, "a1"),
  ("r2", 2, 0, 1, "a2")
)).toDF("prev_column", "c0", "c1", "c2", "post_column")

val targetCols = List("c0", "c1", "c2")

// 将数值放在结构体首位，方便按数值降序排序
val colsArray = array(
  targetCols.map(colName => struct(col(colName).alias("value"), lit(colName).alias("name"))): _*
)

val dfResult = dfIn
  .withColumn("sorted_cols", sort_array(colsArray, asc = false))
  .withColumn("first", col("sorted_cols")(0).getField("name"))
  .withColumn("second", col("sorted_cols")(1).getField("name"))
  .drop("sorted_cols")

dfResult.show()

代码说明 #

1. 构造结构体数组：将每个目标列的名称（通过lit生成常量字符串）和对应数值打包成struct，再组合成数组，这样每行都会得到一个包含所有目标列信息的数组。
2. 排序数组：
   • Spark 3.x+ 使用array_sort并传入自定义比较器，实现按数值从大到小排序。
   • Spark 2.x 则通过调整结构体字段顺序，用sort_array的asc=false参数实现降序。
3. 提取列名：从排序后的数组中取第1个元素（索引0）的name字段作为最大数值的列名，第2个元素（索引1）的name字段作为第二大数值的列名，最后删除中间辅助列即可得到目标DataFrame。

结果验证 #

执行代码后输出的DataFrame与期望结构完全一致：

+------------+---+---+---+-----------+-----+------+
|prev_column |c0 |c1 |c2 |post_column|first|second|
+------------+---+---+---+-----------+-----+------+
|r0          |0  |1  |2  |a0         |c2   |c1    |
|r1          |1  |2  |0  |a1         |c1   |c0    |
|r2          |2  |0  |1  |a2         |c0   |c2    |
+------------+---+---+---+-----------+-----+------+

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