Search code examples
apache-sparkapache-spark-sqlapache-spark-dataset

How to use approxQuantile by group?


Spark has SQL function percentile_approx(), and its Scala counterpart is df.stat.approxQuantile().

However, the Scala counterpart cannot be used on grouped datasets, something like df.groupby("foo").stat.approxQuantile(), as answered here: https://stackoverflow.com/a/51933027.

But it's possible to do both grouping and percentiles in SQL syntax. So I'm wondering, maybe I can define an UDF from SQL percentile_approx function and use it on my grouped dataset?


Solution

  • Spark >= 3.1

    Corresponding SQL functions have been added in Spark 3.1 - see SPARK-30569.

    Spark < 3.1

    While you cannot use approxQuantile in an UDF, and you there is no Scala wrapper for percentile_approx it is not hard to implement one yourself:

    import org.apache.spark.sql.functions._
    import org.apache.spark.sql.Column
    import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile
    
    
    object PercentileApprox {
      def percentile_approx(col: Column, percentage: Column, accuracy: Column): Column = {
        val expr = new ApproximatePercentile(
          col.expr,  percentage.expr, accuracy.expr
        ).toAggregateExpression
        new Column(expr)
      }
      def percentile_approx(col: Column, percentage: Column): Column = percentile_approx(
        col, percentage, lit(ApproximatePercentile.DEFAULT_PERCENTILE_ACCURACY)
      )
    }
    

    Example usage:

    import PercentileApprox._
    
    val df = (Seq.fill(100)("a") ++ Seq.fill(100)("b")).toDF("group").withColumn(
      "value", when($"group" === "a", randn(1) + 10).otherwise(randn(3))
    )
    
    df.groupBy($"group").agg(percentile_approx($"value", lit(0.5))).show
    
    +-----+------------------------------------+
    |group|percentile_approx(value, 0.5, 10000)|
    +-----+------------------------------------+
    |    b|                -0.06336346702250675|
    |    a|                   9.818985618591595|
    +-----+------------------------------------+
    
    df.groupBy($"group").agg(
      percentile_approx($"value", typedLit(Seq(0.1, 0.25, 0.75, 0.9)))
    ).show(false)
    
    +-----+----------------------------------------------------------------------------------+
    |group|percentile_approx(value, [0.1,0.25,0.75,0.9], 10000)                              |
    +-----+----------------------------------------------------------------------------------+
    |b    |[-1.2098351202406483, -0.6640768986666159, 0.6778253126144265, 1.3255676906697658]|
    |a    |[8.902067202468098, 9.290417382259626, 10.41767257153993, 11.067087075488068]     |
    +-----+----------------------------------------------------------------------------------+
    

    Once this is on the JVM classpath you can also add PySpark wrapper, using logic similar to built-in functions.