Spark Dataframe 中多列的每行排名

Question

我正在使用 spark 和 Scala 来转换 Dataframe ，我想在其中计算一个新变量，该变量计算许多变量中每行一个变量的排名。

例子-

Input DF-

+---+---+---+
|c_0|c_1|c_2|
+---+---+---+
| 11| 11| 35|
| 22| 12| 66|
| 44| 22| 12|
+---+---+---+

Expected DF-

+---+---+---+--------+--------+--------+
|c_0|c_1|c_2|c_0_rank|c_1_rank|c_2_rank|
+---+---+---+--------+--------+--------+
| 11| 11| 35|        2|        3|        1|
| 22| 12| 66|       2|       3|       1|
| 44| 22| 12|       1|       2|       3|
+---+---+---+--------+--------+--------+

这已经用 R 回答了 - Rank per row over multiple columns in R，

但我需要在 spark-sql 中使用 scala 做同样的事情。谢谢您的帮助！

编辑- 4/1 。遇到一种情况，如果值相同，则排名应该不同。编辑第一行以复制情况。

Answer 1

如果我理解正确，您希望在每一行中获得每列的排名。

让我们首先定义数据，以及要“排名”的列。

val df = Seq((11,  21,  35),(22,  12, 66),(44, 22 , 12))
    .toDF("c_0", "c_1", "c_2")
val cols = df.columns

然后我们定义一个UDF来查找数组中元素的索引。

val pos = udf((a : Seq[Int], elt : Int) => a.indexOf(elt)+1)

我们最终创建了一个排序数组（按降序排列）并使用 UDF 来查找每列的排名。

val ranks = cols.map(c => pos(col("array"), col(c)).as(c+"_rank"))
df.withColumn("array", sort_array(array(cols.map(col) : _*), false))
  .select((cols.map(col)++ranks) :_*).show 
+---+---+---+--------+--------+--------+
|c_0|c_1|c_2|c_0_rank|c_1_rank|c_2_rank|
+---+---+---+--------+--------+--------+
| 11| 12| 35|       3|       2|       1|
| 22| 12| 66|       2|       3|       1|
| 44| 22| 12|       1|       2|       3|
+---+---+---+--------+--------+--------+

编辑： 从 Spark 2.4 开始，我定义的 pos UDF 可以替换为内置函数 array_position(column: Column, value: Any)，其工作方式完全相同（第一个索引为 1）。这样可以避免使用效率稍低的 UDF。

编辑2： 如果您有重复的键，上面的代码将生成重复的索引。如果你想避免它，你可以创建数组，压缩它以记住哪一列是哪一列，对其进行排序并再次压缩它以获得最终排名。它看起来像这样：

val colMap = df.columns.zipWithIndex.map(_.swap).toMap
val zip = udf((s: Seq[Int]) => s
    .zipWithIndex
    .sortBy(-_._1)
    .map(_._2)
    .zipWithIndex
    .toMap
    .mapValues(_+1))
val ranks = (0 until cols.size)
    .map(i => 'zip.getItem(i) as colMap(i) + "_rank")
val result = df
    .withColumn("zip", zip(array(cols.map(col) : _*)))
    .select(cols.map(col) ++ ranks :_*)

Answer 2

解决此问题的一种方法是使用 windows。

val df = Seq((11,  21,  35),(22,  12, 66),(44, 22 , 12))
    .toDF("c_0", "c_1", "c_2")
(0 to 2)
    .map("c_"+_)
    .foldLeft(df)((d, column) => 
          d.withColumn(column+"_rank", rank() over Window.orderBy(desc(column))))
    .show
+---+---+---+--------+--------+--------+                                        
|c_0|c_1|c_2|c_0_rank|c_1_rank|c_2_rank|
+---+---+---+--------+--------+--------+
| 22| 12| 66|       2|       3|       1|
| 11| 21| 35|       3|       2|       2|
| 44| 22| 12|       1|       1|       3|
+---+---+---+--------+--------+--------+

但这不是一个好主意。所有的数据最终都会放在一个分区中，如果所有数据都不能放入一个执行器中，则会导致 OOM 错误。

另一种方法需要对数据框进行 3 次排序，但至少可以扩展到任何大小的数据。

让我们定义一个函数来压缩具有连续索引的数据帧（它存在于 RDD，但不存在于数据帧）

def zipWithIndex(df : DataFrame, name : String) : DataFrame = {
    val rdd = df.rdd.zipWithIndex
      .map{ case (row, i) => Row.fromSeq(row.toSeq :+ (i+1)) }
    val newSchema = df.schema.add(StructField(name, LongType, false))
    df.sparkSession.createDataFrame(rdd, newSchema)
}

让我们在同一个数据帧上使用它df：

(0 to 2)
    .map("c_"+_)
    .foldLeft(df)((d, column) => 
        zipWithIndex(d.orderBy(desc(column)), column+"_rank"))
    .show

提供与上述完全相同的结果。

Answer 3

您可能可以创建一个窗口函数。请注意，如果您有太多数据，这很容易受到 OOM 的影响。但是，我只是想在这里介绍一下窗口函数的概念。

inputDF.createOrReplaceTempView("my_df")
val expectedDF =  spark.sql("""
    select 
        c_0
        , c_1
        , c_2
        , rank(c_0) over (order by c_0 desc) c_0_rank
        , rank(c_1) over (order by c_1 desc) c_1_rank
        , rank(c_2) over (order by c_2 desc) c_2_rank 
    from my_df""")
expectedDF.show()

+---+---+---+--------+--------+--------+
|c_0|c_1|c_2|c_0_rank|c_1_rank|c_2_rank|
+---+---+---+--------+--------+--------+
| 44| 22| 12|       3|       3|       1|
| 11| 21| 35|       1|       2|       2|
| 22| 12| 66|       2|       1|       3|
+---+---+---+--------+--------+--------+