First draft of codegen

Author: marmbrus

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/BoundAttribute.scala

@@ -31,7 +31,7 @@ import org.apache.spark.sql.types._
 case class BoundReference(ordinal: Int, dataType: DataType, nullable: Boolean)
   extends NamedExpression with trees.LeafNode[Expression] {
 
-  override def toString: String = s"input[$ordinal]"
+  override def toString: String = s"input[$ordinal, $dataType]"
 
   override def eval(input: InternalRow): Any = input(ordinal)
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala

@@ -72,7 +72,7 @@ abstract class Expression extends TreeNode[Expression] {
     val primitive = ctx.freshName("primitive")
     val ve = GeneratedExpressionCode("", isNull, primitive)
     ve.code = genCode(ctx, ve)
-    ve
+    ve.copy(s"/* $this */\n" + ve.code)
   }
 
   /**

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

@@ -34,6 +34,13 @@ private[sql] case object Final extends AggregateMode
 
 private[sql] case object Complete extends AggregateMode
 
+case object NoOp extends Expression {
+  override def nullable: Boolean = true
+  override def eval(input: expressions.InternalRow): Any = ???
+  override def dataType: DataType = NullType
+  override def children: Seq[Expression] = Nil
+}
+
 /**
  * A container of a Aggregate Function, Aggregate Mode, and a field (`isDistinct`) indicating
  * if DISTINCT keyword is specified for this function.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.expressions.codegen
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate2.NoOp
 
 // MutableProjection is not accessible in Java
 abstract class BaseMutableProjection extends MutableProjection
@@ -36,15 +37,18 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
 
   protected def create(expressions: Seq[Expression]): (() => MutableProjection) = {
     val ctx = newCodeGenContext()
-    val projectionCode = expressions.zipWithIndex.map { case (e, i) =>
-      val evaluationCode = e.gen(ctx)
-      evaluationCode.code +
-        s"""
-          if(${evaluationCode.isNull})
-            mutableRow.setNullAt($i);
-          else
-            ${ctx.setColumn("mutableRow", e.dataType, i, evaluationCode.primitive)};
-        """
+    val projectionCode = expressions.zipWithIndex.map {
+      case (NoOp, _) => ""
+      case (e, i) =>
+        val evaluationCode = e.gen(ctx)
+        evaluationCode.code +
+          s"""
+            /** output[$i] = $e */
+            if(${evaluationCode.isNull})
+              mutableRow.setNullAt($i);
+            else
+              ${ctx.setColumn("mutableRow", e.dataType, i, evaluationCode.primitive)};
+          """
     }.mkString("\n")
     val code = s"""
       public Object generate($exprType[] expr) {
@@ -80,7 +84,7 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
       }
     """
 
-    logDebug(s"code for ${expressions.mkString(",")}:\n$code")
+    logWarning(s"code for ${expressions.mkString(",")}:\n$code")
 
     val c = compile(code)
     () => {

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

@@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.errors._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate2._
 import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, AllTuples, UnspecifiedDistribution, Distribution}
+import org.apache.spark.sql.types.NullType
 
 case class Aggregate2Sort(
     preShuffle: Boolean,
@@ -66,7 +67,7 @@ case class Aggregate2Sort(
           while (i < aggregateExpressions.length) {
             val func = aggregateExpressions(i).aggregateFunction.withBufferOffset(bufferOffset)
             functions(i) = aggregateExpressions(i).mode match {
-              case Partial | Complete => BindReferences.bindReference(func, child.output)
+              case Partial | Complete => func
               case PartialMerge | Final => func
             }
             bufferOffset = aggregateExpressions(i).mode match {
@@ -118,6 +119,43 @@ case class Aggregate2Sort(
           new InterpretedMutableProjection(
             resultExpressions, groupingExpressions.map(_.toAttribute) ++ aggregateAttributes)
 
+        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 {
+            case ae: AlgebraicAggregate => ae.initialValues
+          }
+          println(initExpressions.mkString(","))
+          newMutableProjection(initExpressions, Nil)().target(buffer)
+        }
+
+        lazy val updateProjection = {
+          val bufferSchema = aggregateFunctions.flatMap {
+            case ae: AlgebraicAggregate => ae.bufferSchema
+          }
+          val updateExpressions = aggregateFunctions.flatMap {
+            case ae: AlgebraicAggregate => ae.updateExpressions
+          }
+
+          println(updateExpressions.mkString(","))
+          newMutableProjection(updateExpressions, bufferSchema ++ child.output)().target(buffer)
+        }
+
+        val mergeProjection = {
+          val bufferSchemata =
+            offsetAttributes ++ aggregateFunctions.flatMap {
+              case ae: AlgebraicAggregate => ae.bufferSchema
+            } ++ offsetAttributes ++ aggregateFunctions.flatMap {
+              case ae: AlgebraicAggregate => ae.rightBufferSchema
+            }
+            val mergeExpressions = offsetExpressions ++ aggregateFunctions.flatMap {
+              case ae: AlgebraicAggregate => ae.mergeExpressions
+            }
+
+          newMutableProjection(mergeExpressions, bufferSchemata)()
+        }
+
         // Initialize this iterator.
         initialize()
 
@@ -136,28 +174,16 @@ case class Aggregate2Sort(
         }
 
         private def initializeBuffer(): Unit = {
-          var i = 0
-          while (i < aggregateFunctions.length) {
-            aggregateFunctions(i).initialize(buffer)
-            i += 1
-          }
+          initialProjection(EmptyRow)
+          println("initilized: " + buffer)
         }
 
         private def processRow(row: InternalRow): Unit = {
           // The new row is still in the current group.
           if (preShuffle) {
-            var i = 0
-            while (i < aggregateFunctions.length) {
-              aggregateFunctions(i).update(buffer, row)
-              i += 1
-            }
+            updateProjection(joinedRow(buffer, row))
           } else {
-            var i = 0
-            println("post shuffle: " + buffer + " " + row)
-            while (i < aggregateFunctions.length) {
-              aggregateFunctions(i).merge(buffer, row)
-              i += 1
-            }
+            mergeProjection.target(buffer)(joinedRow(buffer, row))
           }
         }