如何在 Spark 中找到分组数据的精确中位数

Question

我需要使用 Scala 计算 Spark 中 Double 数据类型的分组数据集的精确中位数。

它与类似查询不同：在 Spark SQL 中查找多个双数据类型列的中位数。这个问题是关于分组数据的查找数据，而另一个问题是关于在 RDD 级别上查找中位数。

这是我的示例数据

scala> sqlContext.sql("select * from test").show()

+---+---+
| id|num|
+---+---+
|  A|0.0|
|  A|1.0|
|  A|1.0|
|  A|1.0|
|  A|0.0|
|  A|1.0|
|  B|0.0|
|  B|1.0|
|  B|1.0|
+---+---+

预期答案：

+--------+
| Median |
+--------+
|   1    |
|   1    |
+--------+

我尝试了以下选项，但没有运气：

1）Hive函数百分位数，仅适用于BigInt。

2) Hive 函数percentile_approx，但它没有按预期工作（返回 0.25 vs 1）。

scala> sqlContext.sql("select percentile_approx(num, 0.5) from test group by id").show()

+----+
| _c0|
+----+
|0.25|
|0.25|
+----+

Answer 1

最简单的方法（需要 Spark 2.0.1+ 并且不是精确的中位数）

正如参考第一个问题Findmedian in Spark SQL for double datatype columns 的评论中所述，我们可以使用它percentile_approx来计算 Spark 2.0.1+ 的中位数。要将其应用于 Apache Spark 中的分组数据，查询将如下所示：

val df = Seq(("A", 0.0), ("A", 0.0), ("A", 1.0), ("A", 1.0), ("A", 1.0), ("A", 1.0), ("B", 0.0), ("B", 1.0), ("B", 1.0)).toDF("id", "num")
df.createOrReplaceTempView("df")
spark.sql("select id, percentile_approx(num, 0.5) as median from df group by id order by id").show()

输出为：

+---+------+
| id|median|
+---+------+
|  A|   1.0|
|  B|   1.0|
+---+------+

话虽如此，这是一个近似值（而不是每个问题的精确中位数）。

计算分组数据的精确中位数

有多种方法，所以我确信 SO 中的其他人可以提供更好或更有效的示例。但这里有一个计算 Spark 中分组数据中位数的代码片段（在 Spark 1.6 和 Spark 2.1 中验证）：

import org.apache.spark.SparkContext._

val rdd: RDD[(String, Double)] = sc.parallelize(Seq(("A", 1.0), ("A", 0.0), ("A", 1.0), ("A", 1.0), ("A", 0.0), ("A", 1.0), ("B", 0.0), ("B", 1.0), ("B", 1.0)))

// Scala median function
def median(inputList: List[Double]): Double = {
  val count = inputList.size
  if (count % 2 == 0) {
    val l = count / 2 - 1
    val r = l + 1
    (inputList(l) + inputList(r)).toDouble / 2
  } else
    inputList(count / 2).toDouble
}

// Sort the values
val setRDD = rdd.groupByKey()
val sortedListRDD = setRDD.mapValues(_.toList.sorted)

// Output DataFrame of id and median
sortedListRDD.map(m => {
  (m._1, median(m._2))
}).toDF("id", "median_of_num").show()

输出为：

+---+-------------+
| id|median_of_num|
+---+-------------+
|  A|          1.0|
|  B|          1.0|
+---+-------------+

我应该指出一些警告，因为这可能不是最有效的实现：

• 它目前使用的groupByKey性能不是很好。您可能想将其更改为 a reduceByKey（更多信息请参阅避免 GroupByKey）
• 使用 Scala 函数计算median.

此方法对于较小量的数据应该可以正常工作，但如果每个键有数百万行，建议使用 Spark 2.0.1+ 并使用该percentile_approx方法。