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来回答这个问题，

但是我需要使用scala在spark-sql中执行相同的操作。谢谢您的帮助！

编辑-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。

EDIT2：如果键重复，上面的代码将生成重复的索引。如果要避免这种情况，可以创建数组，将其压缩以记住是哪一列，对其进行排序并再次压缩以得到最终排名。看起来像这样：

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错误。

另一种方法将需要对数据帧进行三次排序，但至少可以缩放到任意大小的数据。

让我们定义一个函数，该函数压缩具有连续索引的数据框（它存在于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|
+---+---+---+--------+--------+--------+

Spark Dataframe中多列的每行排名

3 个答案: