Make Aggregate2Sort work with both algebraic AggregateFunctions and non-algebraic AggregateFunctions.

Author: yhuai

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

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.expressions.aggregate2.{Complete, AggregateExpression2, AggregateFunction2}
 import org.apache.spark.sql.catalyst.{SimpleCatalystConf, CatalystConf}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical._
@@ -483,7 +484,10 @@ class Analyzer(
         q transformExpressions {
           case u @ UnresolvedFunction(name, children) =>
             withPosition(u) {
-              registry.lookupFunction(name, children)
+              registry.lookupFunction(name, children) match {
+                case agg2: AggregateFunction2 => AggregateExpression2(agg2, Complete, false)
+                case other => other
+              }
             }
         }
     }
@@ -501,6 +505,7 @@ class Analyzer(
     def containsAggregates(exprs: Seq[Expression]): Boolean = {
       exprs.foreach(_.foreach {
         case agg: AggregateExpression => return true
+        case agg2: AggregateExpression2 => return true
         case _ =>
       })
       false

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
+import org.apache.spark.sql.catalyst.expressions.aggregate2.MyDoubleSum
+
 import scala.reflect.ClassTag
 import scala.util.{Failure, Success, Try}
 
@@ -143,6 +145,7 @@ object FunctionRegistry {
     expression[Max]("max"),
     expression[Min]("min"),
     expression[Sum]("sum"),
+    expression[MyDoubleSum]("mydoublesum"),
 
     // string functions
     expression[Ascii]("ascii"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate2/aggregates.scala

@@ -95,10 +95,56 @@ abstract class AggregateFunction2
   override def eval(buffer: InternalRow = null): Any
 }
 
+case class MyDoubleSum(child: Expression) extends AggregateFunction2 {
+  override val bufferSchema: StructType =
+    StructType(StructField("currentSum", DoubleType, true) :: Nil)
+
+  override val bufferAttributes: Seq[Attribute] = bufferSchema.toAttributes
+
+  override def initialize(buffer: MutableRow): Unit = {
+    buffer.update(bufferOffset, null)
+  }
+
+  override def update(buffer: MutableRow, input: InternalRow): Unit = {
+    val inputValue = child.eval(input)
+    if (inputValue != null) {
+      if (buffer.isNullAt(bufferOffset) == null) {
+        buffer.setDouble(bufferOffset, inputValue.asInstanceOf[Double])
+      } else {
+        val currentSum = buffer.getDouble(bufferOffset)
+        buffer.setDouble(bufferOffset, currentSum + inputValue.asInstanceOf[Double])
+      }
+    }
+  }
+
+  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
+    if (!buffer2.isNullAt(bufferOffset)) {
+      if (buffer1.isNullAt(bufferOffset)) {
+        buffer1.setDouble(bufferOffset, buffer2.getDouble(bufferOffset))
+      } else {
+        val currentSum = buffer1.getDouble(bufferOffset)
+        buffer1.setDouble(bufferOffset, currentSum + buffer2.getDouble(bufferOffset))
+      }
+    }
+  }
+
+  override def eval(buffer: InternalRow = null): Any = {
+    if (buffer.isNullAt(bufferOffset)) {
+      null
+    } else {
+      buffer.getDouble(bufferOffset)
+    }
+  }
+
+  override def nullable: Boolean = true
+  override def dataType: DataType = DoubleType
+  override def children: Seq[Expression] = child :: Nil
+}
+
 /**
  * A helper class for aggregate functions that can be implemented in terms of catalyst expressions.
  */
-abstract class AlgebraicAggregate extends AggregateFunction2 with Serializable{
+abstract class AlgebraicAggregate extends AggregateFunction2 with Serializable {
   self: Product =>
 
   val initialValues: Seq[Expression]
@@ -109,6 +155,11 @@ abstract class AlgebraicAggregate extends AggregateFunction2 with Serializable{
   /** Must be filled in by the executors */
   var inputSchema: Seq[Attribute] = _
 
+  override def withBufferOffset(newBufferOffset: Int): AlgebraicAggregate = {
+    bufferOffset = newBufferOffset
+    this
+  }
+
   def offsetExpressions: Seq[Attribute] = Seq.fill(bufferOffset)(AttributeReference("offset", NullType)())
 
   lazy val rightBufferSchema = bufferAttributes.map(_.newInstance())
@@ -182,7 +233,7 @@ case class Average(child: Expression) extends AlgebraicAggregate {
 
   val evaluateExpression = Cast(currentSum, resultType) / Cast(currentCount, resultType)
 
-  override def nullable: Boolean = false
+  override def nullable: Boolean = true
   override def dataType: DataType = resultType
   override def children: Seq[Expression] = child :: Nil
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.plans.logical
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate2.AggregateExpression2
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.types._
 import org.apache.spark.util.collection.OpenHashSet
@@ -28,6 +29,7 @@ case class Project(projectList: Seq[NamedExpression], child: LogicalPlan) extend
   override lazy val resolved: Boolean = {
     val hasSpecialExpressions = projectList.exists ( _.collect {
         case agg: AggregateExpression => agg
+        case agg: AggregateExpression2 => agg
         case generator: Generator => generator
         case window: WindowExpression => window
       }.nonEmpty

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate2/Aggregate2Sort.scala

@@ -26,6 +26,8 @@ import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistrib
 import org.apache.spark.sql.execution.{SparkPlan, UnaryNode}
 import org.apache.spark.sql.types.NullType
 
+import scala.collection.mutable.ArrayBuffer
+
 case class Aggregate2Sort(
     preShuffle: Boolean,
     groupingExpressions: Seq[NamedExpression],
@@ -57,40 +59,73 @@ case class Aggregate2Sort(
     child.execute().mapPartitions { iter =>
 
       new Iterator[InternalRow] {
-        private val aggregateFunctions: Array[AggregateFunction2] = {
+        private val aggregateExprsWithBufferOffset = {
           var bufferOffset =
             if (preShuffle) {
               0
             } else {
               groupingExpressions.length
             }
+          val bufferOffsets = new ArrayBuffer[Int]()
           var i = 0
-          val functions = new Array[AggregateFunction2](aggregateExpressions.length)
           while (i < aggregateExpressions.length) {
-            val func = aggregateExpressions(i).aggregateFunction.withBufferOffset(bufferOffset)
-            functions(i) = aggregateExpressions(i).mode match {
-              case Partial | Complete => func
-              case PartialMerge | Final => func
-            }
+            val func = aggregateExpressions(i).aggregateFunction
+            bufferOffsets += bufferOffset
             bufferOffset = aggregateExpressions(i).mode match {
               case Partial | PartialMerge => bufferOffset + func.bufferSchema.length
               case Final | Complete => bufferOffset + 1
             }
             i += 1
           }
+          aggregateExpressions.zip(bufferOffsets)
+        }
 
-          functions.foreach {
-            case ae: AlgebraicAggregate => ae.inputSchema = child.output
-            case _ =>
+        private val algebraicAggregateFunctions: Array[AlgebraicAggregate] = {
+          aggregateExprsWithBufferOffset.collect {
+            case (AggregateExpression2(agg: AlgebraicAggregate, mode, isDistinct), offset) =>
+              agg.inputSchema = child.output
+              agg.withBufferOffset(offset)
+          }.toArray
+        }
+
+        private val nonAlgebraicAggregateFunctions: Array[AggregateFunction2] = {
+          aggregateExprsWithBufferOffset.collect {
+            case (AggregateExpression2(agg: AggregateFunction2, mode, isDistinct), offset)
+              if !agg.isInstanceOf[AlgebraicAggregate] =>
+              val func = agg.withBufferOffset(offset)
+              mode match {
+                case Partial | Complete =>
+                  // Only need to bind reference when the function is not an AlgebraicAggregate
+                  // and the mode is Partial or Complete.
+                  BindReferences.bindReference(func, child.output)
+                case _ => func
+              }
+          }.toArray
+        }
+
+        private val nonAlgebraicAggregateFunctionOrdinals: Array[Int] = {
+          val ordinals = new ArrayBuffer[Int]()
+          var i = 0
+          while (i < aggregateExpressions.length) {
+            aggregateExpressions(i).aggregateFunction match {
+              case agg: AlgebraicAggregate =>
+              case _ => ordinals += i
+            }
+            i += 1
           }
-          functions
+          ordinals.toArray
         }
 
         private val bufferSize: Int = {
-          var i = 0
           var size = 0
-          while (i < aggregateFunctions.length) {
-            size += aggregateFunctions(i).bufferSchema.length
+          var i = 0
+          while (i < algebraicAggregateFunctions.length) {
+            size += algebraicAggregateFunctions(i).bufferSchema.length
+            i += 1
+          }
+          i = 0
+          while (i < nonAlgebraicAggregateFunctions.length) {
+            size += nonAlgebraicAggregateFunctions(i).bufferSchema.length
             i += 1
           }
           if (preShuffle) {
@@ -124,48 +159,49 @@ case class Aggregate2Sort(
         val offsetAttributes = if (preShuffle) Nil else Seq.fill(groupingExpressions.length)(AttributeReference("offset", NullType)())
         val offsetExpressions = if (preShuffle) Nil else Seq.fill(groupingExpressions.length)(NoOp)
 
-        val initialProjection = {
-          val initExpressions = offsetExpressions ++ aggregateFunctions.flatMap {
+        val algebraicInitialProjection = {
+          val initExpressions = offsetExpressions ++ algebraicAggregateFunctions.flatMap {
             case ae: AlgebraicAggregate => ae.initialValues
           }
           // println(initExpressions.mkString(","))
+
           newMutableProjection(initExpressions, Nil)().target(buffer)
         }
 
-        lazy val updateProjection = {
-          val bufferSchema = aggregateFunctions.flatMap {
+        lazy val algebraicUpdateProjection = {
+          val bufferSchema = algebraicAggregateFunctions.flatMap {
             case ae: AlgebraicAggregate => ae.bufferAttributes
           }
-          val updateExpressions = aggregateFunctions.flatMap {
+          val updateExpressions = algebraicAggregateFunctions.flatMap {
             case ae: AlgebraicAggregate => ae.updateExpressions
           }
 
           // println(updateExpressions.mkString(","))
           newMutableProjection(updateExpressions, bufferSchema ++ child.output)().target(buffer)
         }
 
-        lazy val mergeProjection = {
+        lazy val algebraicMergeProjection = {
           val bufferSchemata =
-            offsetAttributes ++ aggregateFunctions.flatMap {
+            offsetAttributes ++ algebraicAggregateFunctions.flatMap {
               case ae: AlgebraicAggregate => ae.bufferAttributes
-            } ++ offsetAttributes ++ aggregateFunctions.flatMap {
+            } ++ offsetAttributes ++ algebraicAggregateFunctions.flatMap {
               case ae: AlgebraicAggregate => ae.rightBufferSchema
             }
-          val mergeExpressions = offsetExpressions ++ aggregateFunctions.flatMap {
+          val mergeExpressions = offsetExpressions ++ algebraicAggregateFunctions.flatMap {
             case ae: AlgebraicAggregate => ae.mergeExpressions
           }
 
           newMutableProjection(mergeExpressions, bufferSchemata)()
         }
 
-        lazy val evalProjection = {
+        lazy val algebraicEvalProjection = {
           val bufferSchemata =
-            offsetAttributes ++ aggregateFunctions.flatMap {
+            offsetAttributes ++ algebraicAggregateFunctions.flatMap {
               case ae: AlgebraicAggregate => ae.bufferAttributes
-            } ++ offsetAttributes ++ aggregateFunctions.flatMap {
+            } ++ offsetAttributes ++ algebraicAggregateFunctions.flatMap {
               case ae: AlgebraicAggregate => ae.rightBufferSchema
             }
-          val evalExpressions = aggregateFunctions.map {
+          val evalExpressions = algebraicAggregateFunctions.map {
             case ae: AlgebraicAggregate => ae.evaluateExpression
           }
 
@@ -190,16 +226,31 @@ case class Aggregate2Sort(
         }
 
         private def initializeBuffer(): Unit = {
-          initialProjection(EmptyRow)
+          algebraicInitialProjection(EmptyRow)
+          var i = 0
+          while (i < nonAlgebraicAggregateFunctions.length) {
+            nonAlgebraicAggregateFunctions(i).initialize(buffer)
+            i += 1
+          }
           // println("initilized: " + buffer)
         }
 
         private def processRow(row: InternalRow): Unit = {
           // The new row is still in the current group.
           if (preShuffle) {
-            updateProjection(joinedRow(buffer, row))
+            algebraicUpdateProjection(joinedRow(buffer, row))
+            var i = 0
+            while (i < nonAlgebraicAggregateFunctions.length) {
+              nonAlgebraicAggregateFunctions(i).update(buffer, row)
+              i += 1
+            }
           } else {
-            mergeProjection.target(buffer)(joinedRow(buffer, row))
+            algebraicMergeProjection.target(buffer)(joinedRow(buffer, row))
+            var i = 0
+            while (i < nonAlgebraicAggregateFunctions.length) {
+              nonAlgebraicAggregateFunctions(i).merge(buffer, row)
+              i += 1
+            }
           }
         }
 
@@ -244,15 +295,15 @@ case class Aggregate2Sort(
                 // If it is preShuffle, we just output the grouping columns and the buffer.
                 joinedRow(currentGroupingKey, buffer).copy()
               } else {
-                /*
+                algebraicEvalProjection.target(aggregateResult)(buffer)
                 var i = 0
-                while (i < aggregateFunctions.length) {
-                  aggregateResult.update(i, aggregateFunctions(i).eval(buffer))
+                while (i < nonAlgebraicAggregateFunctions.length) {
+                  aggregateResult.update(
+                    nonAlgebraicAggregateFunctionOrdinals(i),
+                    nonAlgebraicAggregateFunctions(i).eval(buffer))
                   i += 1
                 }
-                resultProjection(joinedRow(currentGroupingKey, aggregateResult)).copy()
-                */
-                resultProjection(joinedRow(currentGroupingKey, evalProjection.target(aggregateResult)(buffer)))
+                resultProjection(joinedRow(currentGroupingKey, aggregateResult))
 
               }
             initializeBuffer()