How to calculate the ApproxQuanitiles from list of Integers into Spark DataFrame column using scala

Viewed 74

I have a spark DataFrame with a column containing several arrays of Integers with varying lengths. I will need to create a new column to find the Quantiles for each of these. This is the input DataFrame :

+---------+------------------------+
|Comm     |List_Nb_total_operations|
+---------+------------------------+
|    comm1|         [1, 1, 2, 3, 4]|
|    comm4|                  [2, 2]|
|    comm3|                  [2, 2]|
|    comm0| [1, 1, 1, 2, 2, 2, 3,3]|
|    comm2|         [1, 1, 1, 2, 3]|
+---------+------------------------+

This is the desired result :

+---------+------------------------+----+----+
|Comm     |List_Nb_total_operations|QT25|QT75|
+---------+------------------------+----+----+
|    comm1|         [1, 1, 2, 3, 4]|   1|   3|
|    comm4|                  [2, 2]|   2|   2|
|    comm3|                  [2, 2]|   2|   2|
|    comm0| [1, 1, 1, 2, 2, 2, 3,3]|   1|   3|
|    comm2|         [1, 1, 1, 2, 3]|   1|   2|
+---------+------------------------+----+----+
2 Answers

The function you want to use is percentile_approx (since Spark 3.1):

val df = Seq(
  ("comm1", Seq(1,1,2,3,4)),
  ("comm4", Seq(2,2)),
  ("comm3", Seq(2,2)),
  ("comm0", Seq(1,1,1,2,2,2,3,3)),
  ("comm2", Seq(1,1,1,2,3))
).toDF("Comm", "ops")

val dfQ = df.select(
  col("Comm"),
  explode(col("ops")) as "ops")
  .groupBy("Comm")
  .agg(
    percentile_approx($"ops", lit(0.25), lit(100)) as "q25",
    percentile_approx($"ops", lit(0.75), lit(100)) as "q75"
  )

val dfWithQ = df.join(dfQ, Seq("Comm"))

The documentation has more information regarding tuning the parameters for accuracy.

Thank you for your help. I've found an other solution that works very well in my case:

import org.apache.spark.sql.functions._
import org.apache.spark.sql.Column
import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile

def percentile_approxx(col: Column, percentage: Column, accuracy: Column): Column = {
  val expr = new ApproximatePercentile(
      col.expr,  percentage.expr, accuracy.expr
  ).toAggregateExpression
  new Column(expr)
}
val perc_df = df.groupBy("Comm").agg(percentile_approxx(col("ops"), lit(0.75), lit(100)))
Related