dataframe - 如何按特定字段对 Spark DataFrame 中的结构类型数组进行排序？

Question

给出以下代码:

import java.sql.Date
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

object SortQuestion extends App{

  val spark = SparkSession.builder().appName("local").master("local[*]").getOrCreate()
  import spark.implicits._
  case class ABC(a: Int, b: Int, c: Int)

  val first = Seq(
    ABC(1, 2, 3),
    ABC(1, 3, 4),
    ABC(2, 4, 5),
    ABC(2, 5, 6)
  ).toDF("a", "b", "c")

  val second = Seq(
    (1, 2, (Date.valueOf("2018-01-02"), 30)),
    (1, 3, (Date.valueOf("2018-01-01"), 20)),
    (2, 4, (Date.valueOf("2018-01-02"), 50)),
    (2, 5, (Date.valueOf("2018-01-01"), 60))
  ).toDF("a", "b", "c")

  first.join(second.withColumnRenamed("c", "c2"), Seq("a", "b")).groupBy("a").agg(sort_array(collect_list("c2")))
    .show(false)

}

Spark 产生以下结果:

+---+----------------------------------+
|a  |sort_array(collect_list(c2), true)|
+---+----------------------------------+
|1  |[[2018-01-01,20], [2018-01-02,30]]|
|2  |[[2018-01-01,60], [2018-01-02,50]]|
+---+----------------------------------+

这意味着 Spark 正在按日期对数组进行排序(因为它是第一个字段)，但我想指示 Spark 按该嵌套结构中的特定字段进行排序。

我知道我可以将数组 reshape 为 (value, date) 但这似乎不方便，我想要一个通用的解决方案(想象我有一个大的嵌套结构，5层深，我想排序该结构由特定列组成)。

有办法做到这一点吗？我错过了什么吗？

Answer 1

对于 Spark 3+，您可以将自定义比较器函数传递给 array_sort :

The comparator will take two arguments representing two elements ofthe array. It returns -1, 0, or 1 as the first element is less than,equal to, or greater than the second element. If the comparatorfunction returns other values (including null), the function will failand raise an error.

val df = first
  .join(second.withColumnRenamed("c", "c2"), Seq("a", "b"))
  .groupBy("a")
  .agg(collect_list("c2").alias("list"))

val df2 = df.withColumn(
  "list",
  expr(
    "array_sort(list, (left, right) -> case when left._2 < right._2 then -1 when left._2 > right._2 then 1 else 0 end)"
  )
)

df2.show(false)
//+---+------------------------------------+
//|a  |list                                |
//+---+------------------------------------+
//|1  |[[2018-01-01, 20], [2018-01-02, 30]]|
//|2  |[[2018-01-02, 50], [2018-01-01, 60]]|
//+---+------------------------------------+

其中 _2 是您想要用于排序的结构体字段的名称