Scala Dataframe 获取特定行的最大值

Question

给定一个带有索引列（“Z”）的数据框：

val tmp= Seq(("D",0.1,0.3, 0.4), ("E",0.3, 0.1, 0.4), ("F",0.2, 0.2, 0.5)).toDF("Z", "a", "b", "c")

+---+---+---+---+
| Z |  a|  b|  c|
 ---+---+---+---+
| "D"|0.1|0.3|0.4|
| "E"|0.3|0.1|0.4|
| "F"|0.2|0.2|0.5|
+---+---+---+---+

说我对 Z = "D" 的第一行感兴趣：

tmp.filter(col("Z")=== "D")
+---+---+---+---+
| Z |  a|  b|  c|
+---+---+---+---+
|"D"|0.1|0.3|0.4|
+---+---+---+---+

如何在保留索引列的同时获取该 Dataframe 行的最小值和最大值及其对应的列名？

如果我想要前 2 个最大值，则需要输出

+---+---+---
| Z |  b|c  |
+---+---+--+
| D |0.3|0.4|
+---+---+---

如果我想要 min 所需的输出

+---+---+
| Z |  a|
+---+---+
| D |0.1|
+---+---+

我尝试了什么：

// first convert that DF to an array
val tmp = df.collect.map(_.toSeq).flatten
// returns 
tmp: Array[Any] = Array(0.1, 0.3, 0.4) <---dont know why Any is returned


//take top values of array
val n = 1
tmp.zipWithIndex.sortBy(-_._1).take(n).map(_._2)

但出现错误：

   No implicit Ordering defined for Any.

有什么方法可以直接从数据框而不是数组中完成？

Answer 1

你可以这样做

tmp
  .where($"a" === 0.1)
  .take(1)
  .map { row =>
      Seq(row.getDouble(0), row.getDouble(1), row.getDouble(2))
  }
  .head
  .sortBy(d => -d)
  .take(2)

或者，如果您有大量字段，您可以将架构和模式匹配行字段与这样的架构数据类型进行对比

import org.apache.spark.sql.types._

val schemaWithIndex = tmp.schema.zipWithIndex

tmp
.where($"a" === 0.1)
.take(1)
.map { row =>
    for {
        tuple <- schemaWithIndex
    } yield {
        val field = tuple._1
        val index = tuple._2
        field.dataType match {
            case DoubleType => row.getDouble(index)
        }
    }
}
.head
.sortBy(d => -d)
.take(2)

也许有更简单的方法来做到这一点。

Answer 2

绝对不是最快的方式，而是直接来自数据框

更通用的解决方案：

// somewhere in codebase
import spark.implicits._
import org.apache.spark.sql.functions._

def transform[T, R : Encoder](ds: DataFrame, colsToSelect: Seq[String])(func: Map[String, T] => Map[String, R])
                            (implicit encoder: Encoder[Map[String, R]]): DataFrame = {
    ds.map(row => func(row.getValuesMap(colsToSelect)))
      .toDF()
      .select(explode(col("value")))
      .withColumn("idx", lit(1))
      .groupBy(col("idx")).pivot(col("key")).agg(first(col("value")))
      .drop("idx")
  }

现在是关于使用 Map，其中地图 key 是 field name，地图 value 是 field value。

def fuzzyStuff(values: Map[String, Any]): Map[String, String] = {
  val valueForA = values("a").asInstanceOf[Double]
  //Do whatever you want to do
  // ...
  //use map as a return type where key is a column name and value is whatever yo want to
  Map("x" -> (s"fuzzyA-$valueForA"))
}


def maxN(n: Int)(values: Map[String, Double]): Map[String, Double] = {
 println(values)
 values.toSeq.sorted.reverse.take(n).toMap
}

用法：

val tmp = Seq((0.1,0.3, 0.4), (0.3, 0.1, 0.4), (0.2, 0.2, 0.5)).toDF("a", "b", "c")
val filtered = tmp.filter(col("a") === 0.1)

transform(filtered, colsToSelect = Seq("a", "b", "c"))(maxN(2))
   .show()

+---+---+
|  b|  c|
+---+---+
|0.3|0.4|
+---+---+

transform(filtered, colsToSelect = Seq("a", "b", "c"))(fuzzyStuff)
   .show()

+----------+
|         x|
+----------+
|fuzzyA-0.1|
+----------+

---

1. 定义max 和min 函数

  def maxN(values: Map[String, Double], n: Int): Map[String, Double] = {
    values.toSeq.sorted.reverse.take(n).toMap
  }

  def min(values: Map[String, Double]): Map[String, Double] = {
    Map(values.toSeq.min)
  }

1. 创建数据集

val tmp= Seq((0.1,0.3, 0.4), (0.3, 0.1, 0.4), (0.2, 0.2, 0.5)).toDF("a", "b", "c")
val filtered = tmp.filter(col("a") === 0.1)

1. Exple 和 Pivot 地图类型

val df = filtered.map(row => maxN(row.getValuesMap(Seq("a", "b", "c")), 2)).toDF()

val exploded = df.select(explode($"value"))
+---+-----+
|key|value|
+---+-----+
|  a|  0.1|
|  b|  0.3|
+---+-----+

//Then pivot
exploded.withColumn("idx", lit(1))
      .groupBy($"idx").pivot($"key").agg(first($"value"))
      .drop("idx")
      .show()

+---+---+
|  b|  c|
+---+---+
|0.3|0.4|
+---+---+