Scala spark在窗口分区中找到中位数答案

【问题标题】：Scala spark find median in a window partitionScala spark在窗口分区中找到中位数
【发布时间】：2017-05-17 13:51:13
【问题描述】：

我有一个这样的数据框：

df = 
--------------
|col1 | col2 |
--------------
| A   | 1    |
| A   | 5    |
| B   | 0    |
| A   | 2    |
| B   | 6    |
| B   | 8    |
--------------

我想按 col1 进行分区，在每个分区中找到 col2 的中位数，并将结果追加形成一个新列。结果应如下所示：

result = 
---------------------
|col1 | col2 | col3 |
---------------------
| A   | 1    | 2    |
| A   | 5    | 2    |
| B   | 0    | 6    |
| A   | 2    | 2    |
| B   | 6    | 6    |
| B   | 8    | 8    |
---------------------

目前，我正在使用以下代码：

val df2 = df
.withColumn("tmp", percent_rank over Window.partition('col1).orderBy('col2))
.where("tmp <= 0.5")
.groupBy("col1").agg(max(col2) as "col3")

val result = df.join(df2, df("col1") === df2("col1")).drop(df2("col1"))

但是当数据帧很大时，这需要太多的时间和空间资源来运行。请帮助我找到一种更有效地完成上述工作的方法！非常感谢任何帮助！

【问题讨论】：

标签： sql scala apache-spark dataframe bigdata

【解决方案1】：

使用您拥有的数据，您可以使用 percentile_approx 执行 Spark DataFrame groupBy 语句来执行计算。

// Creating the `df` dataset
val df = Seq(("A", 1), ("A", 5), ("B", 0), ("A", 2), ("B", 6), ("B", 8)).toDF("col1", "col2")
df.createOrReplaceTempView("df")

使用percentile_approx 和groupBy 进行中位数计算：

val df2 = spark.sql("select col1, percentile_approx(col2, 0.5) as median from df group by col1 order by col1")
df2.show()

df2 的输出为：

+----+------+
|col1|median|
+----+------+
|   A|   2.0|
|   B|   6.0|
+----+------+

现在运行join 重新创建最终结果：

val result = df.join(df2, df("col1") === df2("col1"))
result.show()

//// output
+----+----+----+------+
|col1|col2|col1|median|
+----+----+----+------+
|   A|   1|   A|   2.0|
|   A|   5|   A|   2.0|
|   B|   0|   B|   6.0|
|   A|   2|   A|   2.0|
|   B|   6|   B|   6.0|
|   B|   8|   B|   6.0|
+----+----+----+------+

【讨论】：