Scala - Spark In Dataframe 检索具有最大值的行、列名

Question

我有一个数据框：

name     column1  column2  column3  column4
first    2        1        2.1      5.4
test     1.5      0.5      0.9      3.7
choose   7        2.9      9.1      2.5

我想要一个包含包含列的新数据框，列名的行具有最大值：

| name   | max_column |
|--------|------------|
| first  | column4    |
| test   | column4    |
| choose | column3    |

非常感谢您的支持。

Answer 1

可能有更好的方式来编写 UDF。但这可能是可行的解决方案

val spark: SparkSession = SparkSession.builder.master("local").getOrCreate

//implicits for magic functions like .toDf
import spark.implicits._

import org.apache.spark.sql.functions.udf

//We have hard code number of params as UDF don't support variable number of args
val maxval = udf((c1: Double, c2: Double, c3: Double, c4: Double) =>
  if(c1 >= c2 && c1 >= c3 && c1 >= c4)
    "column1"
  else if(c2 >= c1 && c2 >= c3 && c2 >= c4)
    "column2"
  else if(c3 >= c1 && c3 >= c2 && c3 >= c4)
    "column3"
  else
    "column4"
)

//create schema class
case class Record(name: String, 
                    column1: Double, 
                    column2: Double, 
                    column3: Double, 
                    column4: Double)

val df = Seq(
  Record("first", 2.0, 1, 2.1, 5.4),
  Record("test", 1.5, 0.5, 0.9, 3.7),
  Record("choose", 7, 2.9, 9.1, 2.5)
).toDF();

df.withColumn("max_column", maxval($"column1", $"column2", $"column3", $"column4"))
  .select("name", "max_column").show

输出

+------+----------+
|  name|max_column|
+------+----------+
| first|   column4|
|  test|   column4|
|choose|   column3|
+------+----------+

Answer 2

您可以绕道到 RDD 并使用“getValuesMap”完成工作。

val dfIn = Seq(
  ("first", 2.0, 1., 2.1, 5.4),
  ("test", 1.5, 0.5, 0.9, 3.7),
  ("choose", 7., 2.9, 9.1, 2.5)
).toDF("name","column1","column2","column3","column4")

简单的解决办法是

val dfOut = dfIn.rdd
  .map(r => (
       r.getString(0),
       r.getValuesMap[Double](r.schema.fieldNames.filter(_!="name"))
     ))
  .map{case (n,m) => (n,m.maxBy(_._2)._1)}
  .toDF("name","max_column")

但是，如果您想从原始数据框中取回所有列（例如在 Scala/Spark dataframes: find the column name corresponding to the max 中），则必须在合并行和扩展架构方面进行一些尝试

import org.apache.spark.sql.types.{StructType,StructField,StringType}
import org.apache.spark.sql.Row
val dfOut = sqlContext.createDataFrame(
  dfIn.rdd
    .map(r => (r, r.getValuesMap[Double](r.schema.fieldNames.drop(1))))
    .map{case (r,m) => Row.merge(r,(Row(m.maxBy(_._2)._1)))},
  dfIn.schema.add(StructField("max_column",StringType))
)

Answer 3

我想发布我的最终解决方案：

val finalDf = originalDf.withColumn("name", maxValAsMap(keys, values)).select("cookie_id", "max_column")

val maxValAsMap = udf((keys: Seq[String], values: Seq[Any]) => {

    val valueMap:Map[String,Double] = (keys zip values).filter( _._2.isInstanceOf[Double] ).map{
      case (x,y) => (x, y.asInstanceOf[Double])
    }.toMap

    if (valueMap.isEmpty) "not computed" else valueMap.maxBy(_._2)._1
  })

它的工作速度非常快。