[SPARK-12913] [SQL] Improve performance of stat functions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala

@@ -347,14 +347,10 @@ object HiveTypeCoercion {
 
       case Sum(e @ StringType()) => Sum(Cast(e, DoubleType))
       case Average(e @ StringType()) => Average(Cast(e, DoubleType))
-      case StddevPop(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
-        StddevPop(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
-      case StddevSamp(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
-        StddevSamp(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
-      case VariancePop(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
-        VariancePop(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
-      case VarianceSamp(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
-        VarianceSamp(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
+      case StddevPop(e @ StringType()) => StddevPop(Cast(e, DoubleType))
+      case StddevSamp(e @ StringType()) => StddevSamp(Cast(e, DoubleType))
+      case VariancePop(e @ StringType()) => VariancePop(Cast(e, DoubleType))
+      case VarianceSamp(e @ StringType()) => VarianceSamp(Cast(e, DoubleType))
       case Skewness(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
         Skewness(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
       case Kurtosis(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala

@@ -57,7 +57,7 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
 
   override def dataType: DataType = DoubleType
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
+  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType)
 
   override def checkInputDataTypes(): TypeCheckResult =
     TypeUtils.checkForNumericExpr(child.dataType, s"function $prettyName")
@@ -109,7 +109,7 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
    * Update the central moments buffer.
    */
   override def update(buffer: MutableRow, input: InternalRow): Unit = {
-    val v = Cast(child, DoubleType).eval(input)
+    val v = child.eval(input)
     if (v != null) {
       val updateValue = v match {
         case d: Double => d

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala

@@ -17,65 +17,136 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
 
-case class StddevSamp(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0)
-  extends CentralMomentAgg(child) {
+// Compute standard deviation based on online algorithm specified here:
+// http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
 
-  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+  override def children: Seq[Expression] = child :: Nil
 
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+  override def nullable: Boolean = true
 
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
+  override def dataType: DataType = resultType
 
-  override def prettyName: String = "stddev_samp"
+  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType)
+
+  protected val resultType = DoubleType
+
+  protected val count = AttributeReference("count", resultType, nullable = false)()
+  protected val avg = AttributeReference("avg", resultType, nullable = false)()
+  protected val mk = AttributeReference("mk", resultType, nullable = false)()
+
+  override val aggBufferAttributes = count :: avg :: mk :: Nil
 
-  override protected val momentOrder = 2
+  override val initialValues: Seq[Expression] = Seq(
+    /* count = */ Literal(0.0),
+    /* avg = */ Literal(0.0),
+    /* mk = */ Literal(0.0)
+  )
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
+  override val updateExpressions: Seq[Expression] = {
+    val newCount = count + Literal(1.0)
 
-    if (n == 0.0) {
-      null
-    } else if (n == 1.0) {
-      Double.NaN
+    // update average
+    // avg = avg + (value - avg)/count
+    val newAvg = avg + (child - avg) / newCount
+
+    // update sum ofference from mean
+    // Mk = Mk + (value - preAvg) * (value - updatedAvg)
+    val newMk = mk + (child - avg) * (child - newAvg)
+
+    if (child.nullable) {
+      Seq(
+        /* count = */ If(IsNull(child), count, newCount),
+        /* avg = */ If(IsNull(child), avg, newAvg),
+        /* mk = */ If(IsNull(child), mk, newMk)
+      )
     } else {
-      math.sqrt(moments(2) / (n - 1.0))
+      Seq(
+        /* count = */ newCount,
+        /* avg = */ newAvg,
+        /* mk = */ newMk
+      )
     }
   }
-}
 
-case class StddevPop(
-    child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0)
-  extends CentralMomentAgg(child) {
+  override val mergeExpressions: Seq[Expression] = {
 
-  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+    // count merge
+    val newCount = count.left + count.right
 
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+    // average merge
+    val newAvg = ((avg.left * count.left) + (avg.right * count.right)) / newCount
 
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
+    // update sum of square differences
+    val newMk = {
+      val avgDelta = avg.right - avg.left
+      val mkDelta = (avgDelta * avgDelta) * (count.left * count.right) / newCount
+      mk.left + mk.right + mkDelta
+    }
+
+    Seq(
+      /* count = */ newCount,
+      /* avg = */ newAvg,
+      /* mk = */ newMk
+    )
+  }
+}
+
+// Compute the population standard deviation of a column
+case class StddevPop(child: Expression) extends StddevAgg(child) {
+
+  override val evaluateExpression: Expression = {
+    // when count == 0, return null
+    // when count >0, sqrt (mk/count)
+    If(EqualTo(count, Literal(0.0)), Literal.create(null, resultType),
+      Sqrt(mk / count))
+  }
 
   override def prettyName: String = "stddev_pop"
+}
 
-  override protected val momentOrder = 2
+// Compute the sample standard deviation of a column
+case class StddevSamp(child: Expression) extends StddevAgg(child) {
+  override val evaluateExpression: Expression = {
+    // when count == 0, return null
+    // when count == 1, return Na
+    // when count >1, sqrt(mk/(count -1))
+    If(EqualTo(count, Literal(0.0)), Literal.create(null, resultType),
+      If(EqualTo(count, Literal(1.0)), Literal(Double.NaN),
+        Sqrt(mk / (count - Literal(1.0)))))
+  }
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
+  override def prettyName: String = "stddev_samp"
+}
 
-    if (n == 0.0) {
-      null
-    } else {
-      math.sqrt(moments(2) / n)
-    }
+// Compute the population variance of a column
+case class VariancePop(child: Expression) extends StddevAgg(child) {
+
+  override val evaluateExpression: Expression = {
+    // when count == 0, return null
+    // when count >1, sqrt (mk/count)
+
+    If(EqualTo(count, Literal(0.0)), Literal.create(null, resultType),
+      mk / count)
+  }
+
+  override def prettyName: String = "var_pop"
+}
+
+// Compute the sample variance of a column
+case class VarianceSamp(child: Expression) extends StddevAgg(child) {
+  override val evaluateExpression: Expression = {
+    // when count == 0, return null
+    // when count == 1, return Na
+    // when count >1N, sqrt (mk/(count -1))
+    If(EqualTo(count, Literal(0.0)), Literal.create(null, resultType),
+      If(EqualTo(count, Literal(1.0)), Literal(Double.NaN),
+        mk / (count - Literal(1.0))))
   }
+
+  override def prettyName: String = "var_samp"
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegen.scala

@@ -190,7 +190,7 @@ case class WholeStageCodegen(plan: CodegenSupport, children: Seq[SparkPlan])
      """
     // try to compile, helpful for debug
     // println(s"${CodeFormatter.format(source)}")
-    CodeGenerator.compile(source)
+    // CodeGenerator.compile(source)
 
     rdd.mapPartitions { iter =>
       val clazz = CodeGenerator.compile(source)

sql/core/src/test/scala/org/apache/spark/sql/execution/BenchmarkWholeStageCodegen.scala

@@ -27,10 +27,11 @@ import org.apache.spark.util.Benchmark
   *  build/sbt "sql/test-only *BenchmarkWholeStageCodegen"
   */
 class BenchmarkWholeStageCodegen extends SparkFunSuite {
+  lazy val conf = new SparkConf().setMaster("local[1]").setAppName("benchmark")
+  lazy val sc = SparkContext.getOrCreate(conf)
+  lazy val sqlContext = SQLContext.getOrCreate(sc)
+
   def testWholeStage(values: Int): Unit = {
-    val conf = new SparkConf().setMaster("local[1]").setAppName("benchmark")
-    val sc = SparkContext.getOrCreate(conf)
-    val sqlContext = SQLContext.getOrCreate(sc)
 
     val benchmark = new Benchmark("Single Int Column Scan", values)
 
@@ -54,7 +55,42 @@ class BenchmarkWholeStageCodegen extends SparkFunSuite {
     benchmark.run()
   }
 
-  ignore("benchmark") {
-    testWholeStage(1024 * 1024 * 200)
+  def testStddev(values: Int): Unit = {
+
+    val benchmark = new Benchmark("stddev", values)
+
+    benchmark.addCase("stddev w/o codegen") { iter =>
+      sqlContext.setConf("spark.sql.codegen.wholeStage", "false")
+      sqlContext.range(values).groupBy().agg("id" -> "stddev").collect()
+    }
+
+    benchmark.addCase("stddev w codegen") { iter =>
+      sqlContext.setConf("spark.sql.codegen.wholeStage", "true")
+      sqlContext.range(values).groupBy().agg("id" -> "stddev").collect()
+    }
+
+    /**
+    Using ImperativeAggregate:
+
+      Intel(R) Core(TM) i7-4558U CPU @ 2.80GHz
+      stddev:                Avg Time(ms)    Avg Rate(M/s)  Relative Rate
+      -------------------------------------------------------------------
+      stddev w/o codegen         10157.82            10.32         1.00 X
+      stddev w codegen           10528.03             9.96         0.96 X
+
+      Using DeclarativeAggregate:
+
+      Intel(R) Core(TM) i7-4558U CPU @ 2.80GHz
+      stddev:                Avg Time(ms)    Avg Rate(M/s)  Relative Rate
+      -------------------------------------------------------------------
+      stddev w/o codegen          4128.44            25.40         1.00 X
+      stddev w codegen            1400.25            74.88         2.95 X
+      */
+    benchmark.run()
+  }
+
+  test("benchmark") {
+    // testWholeStage(200 << 20)
+    testStddev(100 << 20)
   }
 }