Merge pull request apache#36 from marmbrus/partialAgg

Implement partial aggregation.

Author: rxin

src/main/scala/catalyst/execution/PlanningStrategies.scala

@@ -119,6 +119,60 @@ trait PlanningStrategies {
       expr.references subsetOf plan.outputSet
   }
 
+  object PartialAggregation extends Strategy {
+    def apply(plan: LogicalPlan): Seq[SharkPlan] = plan match {
+      case logical.Aggregate(groupingExpressions, aggregateExpressions, child) =>
+        // Collect all aggregate expressions.
+        val allAggregates =
+          aggregateExpressions.flatMap(_ collect { case a: AggregateExpression => a})
+        // Collect all aggregate expressions that can be computed partially.
+        val partialAggregates =
+          aggregateExpressions.flatMap(_ collect { case p: PartialAggregate => p})
+
+        // Only do partial aggregation if supported by all aggregate expressions.
+        if (allAggregates.size == partialAggregates.size) {
+          // Create a map of expressions to their partial evaluations for all aggregate expressions.
+          val partialEvaluations: Map[Long, SplitEvaluation] =
+            partialAggregates.map(a => (a.id, a.asPartial)).toMap
+
+          // We need to pass all grouping expressions though so the grouping can happen a second
+          // time. However some of them might be unnamed so we alias them allowing them to be
+          // referenced in the second aggregation.
+          val namedGroupingExpressions: Map[Expression, NamedExpression] = groupingExpressions.map {
+            case n: NamedExpression => (n, n)
+            case other => (other, Alias(other, "PartialGroup")())
+          }.toMap
+
+          // Replace aggregations with a new expression that computes the result from the already
+          // computed partial evaluations and grouping values.
+          val rewrittenAggregateExpressions = aggregateExpressions.map(_.transformUp {
+            case e: Expression if partialEvaluations.contains(e.id) =>
+              partialEvaluations(e.id).finalEvaluation
+            case e: Expression if namedGroupingExpressions.contains(e) =>
+              namedGroupingExpressions(e).toAttribute
+          }).asInstanceOf[Seq[NamedExpression]]
+
+          val partialComputation =
+            (namedGroupingExpressions.values ++
+             partialEvaluations.values.flatMap(_.partialEvaluations)).toSeq
+
+          // Construct two phased aggregation.
+          execution.Aggregate(
+            partial = false,
+            namedGroupingExpressions.values.map(_.toAttribute).toSeq,
+            rewrittenAggregateExpressions,
+            execution.Aggregate(
+              partial = true,
+              groupingExpressions,
+              partialComputation,
+              planLater(child))(sc))(sc) :: Nil
+        } else {
+          Nil
+        }
+      case _ => Nil
+    }
+  }
+
   object BroadcastNestedLoopJoin extends Strategy {
     def apply(plan: LogicalPlan): Seq[SharkPlan] = plan match {
       case logical.Join(left, right, joinType, condition) =>
@@ -143,7 +197,8 @@ trait PlanningStrategies {
   object BasicOperators extends Strategy {
     def apply(plan: LogicalPlan): Seq[SharkPlan] = plan match {
       case logical.Distinct(child) =>
-        execution.Aggregate(child.output, child.output, planLater(child))(sc) :: Nil
+        execution.Aggregate(
+          partial = false, child.output, child.output, planLater(child))(sc) :: Nil
       case logical.Sort(sortExprs, child) =>
         // This sort is a global sort. Its requiredDistribution will be an OrderedDistribution.
         execution.Sort(sortExprs, global = true, planLater(child)):: Nil
@@ -156,7 +211,7 @@ trait PlanningStrategies {
       case logical.Filter(condition, child) =>
         execution.Filter(condition, planLater(child)) :: Nil
       case logical.Aggregate(group, agg, child) =>
-        execution.Aggregate(group, agg, planLater(child))(sc) :: Nil
+        execution.Aggregate(partial = false, group, agg, planLater(child))(sc) :: Nil
       case logical.Sample(fraction, withReplacement, seed, child) =>
         execution.Sample(fraction, withReplacement, seed, planLater(child)) :: Nil
       case logical.LocalRelation(output, data) =>

src/main/scala/catalyst/execution/SharkInstance.scala

@@ -73,6 +73,7 @@ abstract class SharkInstance extends Logging {
   object TrivialPlanner extends QueryPlanner[SharkPlan] with PlanningStrategies {
     val sc = self.sc
     val strategies =
+      PartialAggregation ::
       SparkEquiInnerJoin ::
       PartitionPrunings ::
       HiveTableScans ::

src/main/scala/catalyst/execution/aggregates.scala

@@ -5,22 +5,37 @@ import org.apache.hadoop.hive.ql.udf.generic.{GenericUDAFEvaluator, AbstractGene
 
 import catalyst.errors._
 import catalyst.expressions._
-import catalyst.plans.physical.{ClusteredDistribution, AllTuples}
+import catalyst.plans.physical.{UnspecifiedDistribution, ClusteredDistribution, AllTuples}
 
 /* Implicits */
 import org.apache.spark.rdd.SharkPairRDDFunctions._
 
+/**
+ * Groups input data by `groupingExpressions` and computes the `aggregateExpressions` for each
+ * group.
+ *
+ * @param partial if true then aggregation is done partially on local data without shuffling to
+ *                ensure all values where `groupingExpressions` are equal are present.
+ * @param groupingExpressions expressions that are evaluated to determine grouping.
+ * @param aggregateExpressions expressions that are computed for each group.
+ * @param child the input data source.
+ */
 case class Aggregate(
+    partial: Boolean,
     groupingExpressions: Seq[Expression],
     aggregateExpressions: Seq[NamedExpression],
     child: SharkPlan)(@transient sc: SharkContext)
   extends UnaryNode {
 
   override def requiredChildDistribution =
-    if (groupingExpressions == Nil) {
-      AllTuples :: Nil
+    if (partial) {
+      UnspecifiedDistribution :: Nil
     } else {
-      ClusteredDistribution(groupingExpressions) :: Nil
+      if (groupingExpressions == Nil) {
+        AllTuples :: Nil
+      } else {
+        ClusteredDistribution(groupingExpressions) :: Nil
+      }
     }
 
   override def otherCopyArgs = sc :: Nil

src/main/scala/catalyst/expressions/aggregates.scala

@@ -8,9 +8,34 @@ abstract class AggregateExpression extends Expression {
 
 }
 
+/**
+ * Represents an aggregation that has been rewritten to be performed in two steps.
+ *
+ * @param finalEvaluation an aggregate expression that evaluates to same final result as the
+ *                        original aggregation.
+ * @param partialEvaluations A sequence of [[NamedExpression]]s that can be computed on partial
+ *                           data sets and are required to compute the `finalEvaluation`.
+ */
+case class SplitEvaluation(
+    finalEvaluation: Expression,
+    partialEvaluations: Seq[NamedExpression])
+
+/**
+ * An [[AggregateExpression]] that can be partially computed without seeing all relevent tuples.
+ * These partial evaluations can then be combined to compute the actual answer.
+ */
+abstract class PartialAggregate extends AggregateExpression {
+  self: Product =>
+
+  /**
+   * Returns a [[SplitEvaluation]] that computes this aggregation using partial aggregation.
+   */
+  def asPartial: SplitEvaluation
+}
+
 /**
  * A specific implementation of an aggregate function. Used to wrap a generic
- * [[AggregateExpression]] with an algorithm that will be used to compute the result.
+ * [[AggregateExpression]] with an algorithm that will be used to compute one specific result.
  */
 abstract class AggregateFunction
   extends AggregateExpression with Serializable with trees.LeafNode[Expression] {
@@ -26,11 +51,16 @@ abstract class AggregateFunction
   def result: Any
 }
 
-case class Count(child: Expression) extends AggregateExpression with trees.UnaryNode[Expression] {
+case class Count(child: Expression) extends PartialAggregate with trees.UnaryNode[Expression] {
   def references = child.references
   def nullable = false
   def dataType = IntegerType
   override def toString = s"COUNT($child)"
+
+  def asPartial: SplitEvaluation = {
+    val partialCount = Alias(Count(child), "PartialCount")()
+    SplitEvaluation(Sum(partialCount.toAttribute), partialCount :: Nil)
+  }
 }
 
 case class CountDistinct(expressions: Seq[Expression]) extends AggregateExpression {
@@ -41,23 +71,48 @@ case class CountDistinct(expressions: Seq[Expression]) extends AggregateExpressi
   override def toString = s"COUNT(DISTINCT ${expressions.mkString(",")}})"
 }
 
-case class Average(child: Expression) extends AggregateExpression with trees.UnaryNode[Expression] {
+case class Average(child: Expression) extends PartialAggregate with trees.UnaryNode[Expression] {
   def references = child.references
   def nullable = false
   def dataType = DoubleType
   override def toString = s"AVG($child)"
+
+  override def asPartial: SplitEvaluation = {
+    val partialSum = Alias(Sum(child), "PartialSum")()
+    val partialCount = Alias(Count(child), "PartialCount")()
+    val castedSum = Cast(Sum(partialSum.toAttribute), dataType)
+    val castedCount = Cast(Sum(partialCount.toAttribute), dataType)
+
+    SplitEvaluation(
+      Divide(castedSum, castedCount),
+      partialCount :: partialSum :: Nil)
+  }
 }
 
-case class Sum(child: Expression) extends AggregateExpression with trees.UnaryNode[Expression] {
+case class Sum(child: Expression) extends PartialAggregate with trees.UnaryNode[Expression] {
   def references = child.references
   def nullable = false
   def dataType = child.dataType
   override def toString = s"SUM($child)"
+
+  override def asPartial: SplitEvaluation = {
+    val partialSum = Alias(Sum(child), "PartialSum")()
+    SplitEvaluation(
+      Sum(partialSum.toAttribute),
+      partialSum :: Nil)
+  }
 }
 
-case class First(child: Expression) extends AggregateExpression with trees.UnaryNode[Expression] {
+case class First(child: Expression) extends PartialAggregate with trees.UnaryNode[Expression] {
   def references = child.references
   def nullable = child.nullable
   def dataType = child.dataType
   override def toString = s"FIRST($child)"
+
+  override def asPartial: SplitEvaluation = {
+    val partialFirst = Alias(First(child), "PartialFirst")()
+    SplitEvaluation(
+      First(partialFirst.toAttribute),
+      partialFirst :: Nil)
+  }
 }

src/test/scala/catalyst/execution/PlannerSuite.scala

@@ -0,0 +1,42 @@
+package catalyst
+package execution
+
+import org.scalatest.FunSuite
+
+import catalyst.expressions._
+import catalyst.plans.logical
+import catalyst.dsl._
+
+class PlannerSuite extends FunSuite {
+  import TestData._
+
+  test("unions are collapsed") {
+    val query = testData.unionAll(testData).unionAll(testData)
+    val planned = TestShark.TrivialPlanner.BasicOperators(query).head
+    val logicalUnions = query collect { case u: logical.Union => u}
+    val physicalUnions = planned collect { case u: execution.Union => u}
+
+    assert(logicalUnions.size === 2)
+    assert(physicalUnions.size === 1)
+  }
+
+  test("count is partially aggregated") {
+    val query = testData.groupBy('value)(Count('key)).analyze
+    val planned = TestShark.TrivialPlanner.PartialAggregation(query).head
+    val aggregations = planned.collect { case a: Aggregate => a }
+
+    assert(aggregations.size === 2)
+  }
+
+  test("count distinct is not partially aggregated") {
+    val query = testData.groupBy('value)(CountDistinct('key :: Nil)).analyze
+    val planned = TestShark.TrivialPlanner.PartialAggregation(query)
+    assert(planned.isEmpty)
+  }
+
+  test("mixed aggregates are not partially aggregated") {
+    val query = testData.groupBy('value)(Count('value), CountDistinct('key :: Nil)).analyze
+    val planned = TestShark.TrivialPlanner.PartialAggregation(query)
+    assert(planned.isEmpty)
+  }
+}