How can I group and get the median num in Spark Dataframe

Viewed 109

I currently have a dataframe like this

+------------+----------+
|         A  |    B     |
+------------+----------+
|aaaaaaaaaaaa|11        |
|aaaaaaaaaaaa|44        |
|bbbbbbbbbbbb|22        |
|aaaaaaaaaaaa|33        |
+------------+----------+

I want to get the median of the value of column B in each column A.

+------------+----------+
|         A  |    B     |
+------------+----------+
|aaaaaaaaaaaa|33        |
|bbbbbbbbbbbb|22        |
+------------+----------+

How can I do it? Thanks for answer the question.

2 Answers

You can use UDF, groupBy and collect_list to achieve it. Here is code example in Scala:

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions._

def median[T: Numeric](xs: IndexedSeq[T]): Double = {
    if (xs.isEmpty) 0.0
    else {
      // There is faster algorithms (O(N), n-th order statistics) for finding percentiles,
      // but let's go with this one for simplicity - O(NlogN)
      val sorted = xs.sorted
      if (sorted.length % 2 == 1) implicitly[Numeric[T]].toDouble(sorted(sorted.length / 2))
      else {
        // [1, 2]
        val a = sorted(sorted.length / 2)
        val b = sorted(sorted.length / 2 - 1)
        implicitly[Numeric[T]].toDouble(implicitly[Numeric[T]].plus(a, b)) / 2
      }
    }
}

/// .... 

// This is important to make `toDF` visible!
import spark.sqlContext.implicits._

val medianUDF: UserDefinedFunction = udf[Double, IndexedSeq[Int]](median[Int])
val df: DataFrame = Seq(("aaaaaaaaaaaa", 11), ("aaaaaaaaaaaa", 44), ("bbbbbbbbbbbb", 22), ("aaaaaaaaaaaa", 33))
  .toDF("A", "B")
df.show()
//  +------------+---+
//  |           A|  B|
//  +------------+---+
//  |aaaaaaaaaaaa| 11|
//  |aaaaaaaaaaaa| 44|
//  |bbbbbbbbbbbb| 22|
//  |aaaaaaaaaaaa| 33|
//  +------------+---+

// Using UDF as aggregation function. Input for this UDF is indexed sequence - result from collect_list
df.groupBy(col("A"))
  .agg(medianUDF(collect_list(col("B"))).as("median"))
  .show()
//    +------------+------+
//    |           A|median|
//    +------------+------+
//    |bbbbbbbbbbbb|  22.0|
//    |aaaaaaaaaaaa|  33.0|
//    +------------+------+

Here is a way to compute all the medians in parallel (approximately of course).

Let's start by collecting all the possible keys:

// generating data (BTW, please provide that code next time you ask a question)
val df = Seq(("aaaaaaaaaaaa", 11), ("aaaaaaaaaaaa", 44),
             ("bbbbbbbbbbbb", 22), ("aaaaaaaaaaaa", 33))
    .toDF("A", "B")

val cols = df.select("A").distinct.collect.map(_.getAs[String](0))

Let's say cols has size N. One way to go at it would be to iterate over cols and compute the median via N separate jobs. Another answer provide the code for that.

Yet, it is possible to compute all the medians in parallel (and therefore just one job) using a pivot.

val precision = 1e-3
val medians = df
    // the index is artificial, it is just meant to treat each line individually
    .withColumn("index", monotonicallyIncreasingId)
    .groupBy("index")
    .pivot("A").agg(first('B))
    .stat.approxQuantile(cols, Array(0.5), precision)

val result = cols.indices.map(i => cols(i) -> medians(i)(0)).toMap

It may not be worth it with only a few keys be could be interesting if you have more.

EDIT The first solution keeps one line per line in the original dataframe and will work, I think, all the time. In case you have many keys, it could be interesting to compute the index with a window like this.

.withColumn("index", row_number() over Window.partitionBy("A").orderBy("B"))

But do not use windows if you have millions of rows per key. It could be very slow or even crash your job.

Related