do not merge different aggregate implementations and add test

Author: peter-toth

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueries.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.optimizer
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LeafNode, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.{MULTI_SCALAR_SUBQUERY, SCALAR_SUBQUERY}
@@ -134,7 +135,7 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] with PredicateHelper {
           val newOutputMap = createOutputMap(np.projectList, newProjectList)
           Project(distinctExpressions(ep.output ++ newProjectList), mergedChild) -> newOutputMap
         }
-      case (np: Aggregate, ep: Aggregate) =>
+      case (np: Aggregate, ep: Aggregate) if supportedAggregateMerge(np, ep) =>
         tryMergePlans(np.child, ep.child).flatMap { case (mergedChild, outputMap) =>
           val newGroupingExpression = replaceAttributes(np.groupingExpressions, outputMap)
           if (ExpressionSet(newGroupingExpression) == ExpressionSet(ep.groupingExpressions)) {
@@ -147,8 +148,7 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] with PredicateHelper {
             None
           }
         }
-      case _ =>
-        None
+      case _ => None
     }
   }
 
@@ -168,6 +168,29 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] with PredicateHelper {
     ExpressionSet(expressions).toSeq.asInstanceOf[Seq[NamedExpression]]
   }
 
+  // Merging different aggregate implementations could cause performance regression
+  private def supportedAggregateMerge(newPlan: Aggregate, existingPlan: Aggregate) = {
+    val newPlanAggregateExpressions = newPlan.aggregateExpressions.flatMap(_.collect {
+      case a: AggregateExpression => a
+    })
+    val existingPlanAggregateExpressions = existingPlan.aggregateExpressions.flatMap(_.collect {
+      case a: AggregateExpression => a
+    })
+    val newPlanSupportsHashAggregate = Aggregate.supportsHashAggregate(
+      newPlanAggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
+    val existingPlanSupportsHashAggregate = Aggregate.supportsHashAggregate(
+      existingPlanAggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
+    newPlanSupportsHashAggregate && existingPlanSupportsHashAggregate ||
+      !newPlanSupportsHashAggregate && !existingPlanSupportsHashAggregate && {
+        val newPlanSupportsObjectHashAggregate =
+          Aggregate.supportsObjectHashAggregate(newPlanAggregateExpressions)
+        val existingPlanSupportsObjectHashAggregate =
+          Aggregate.supportsObjectHashAggregate(existingPlanAggregateExpressions)
+        newPlanSupportsObjectHashAggregate && existingPlanSupportsObjectHashAggregate ||
+          !newPlanSupportsObjectHashAggregate && !existingPlanSupportsObjectHashAggregate
+      }
+  }
+
   private def removeReferences(
       plan: LogicalPlan,
       mergedSubqueries: ArrayBuffer[LogicalPlan]): LogicalPlan = {

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

@@ -22,7 +22,7 @@ import org.apache.spark.sql.catalyst.analysis.{AnsiTypeCoercion, MultiInstanceRe
 import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable}
 import org.apache.spark.sql.catalyst.catalog.CatalogTable.VIEW_STORING_ANALYZED_PLAN
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, TypedImperativeAggregate}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition}
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
@@ -828,6 +828,24 @@ case class Aggregate(
     copy(child = newChild)
 }
 
+object Aggregate {
+  def supportsAggregationBufferSchema(schema: StructType): Boolean = {
+    schema.forall(f => UnsafeRow.isMutable(f.dataType))
+  }
+
+  def supportsHashAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
+    val aggregationBufferSchema = StructType.fromAttributes(aggregateBufferAttributes)
+    supportsAggregationBufferSchema(aggregationBufferSchema)
+  }
+
+  def supportsObjectHashAggregate(aggregateExpressions: Seq[AggregateExpression]): Boolean = {
+    aggregateExpressions.map(_.aggregateFunction).exists {
+      case _: TypedImperativeAggregate[_] => true
+      case _ => false
+    }
+  }
+}
+
 case class Window(
     windowExpressions: Seq[NamedExpression],
     partitionSpec: Seq[Expression],

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeScalarSubqueriesSuite.scala

@@ -20,18 +20,18 @@ package org.apache.spark.sql.catalyst.optimizer
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions.{GetStructField, MultiScalarSubquery, ScalarSubquery}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 
 class MergeScalarSubqueriesSuite extends PlanTest {
 
   private object Optimize extends RuleExecutor[LogicalPlan] {
-    val batches =
-      Batch("MergeScalarSubqueries", Once, MergeScalarSubqueries) :: Nil
+    val batches = Batch("MergeScalarSubqueries", Once, MergeScalarSubqueries) :: Nil
   }
 
-  val testRelation = LocalRelation('a.int, 'b.int)
+  val testRelation = LocalRelation('a.int, 'b.int, 'c.string)
 
   test("Simple non-correlated scalar subquery merge") {
     val subquery1 = testRelation
@@ -70,4 +70,47 @@ class MergeScalarSubqueriesSuite extends PlanTest {
     // as only `Optimizer` can insert such a node to the plan
     comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer, false)
   }
+
+  test("Do not merge different aggregate implementations") {
+    // supports HashAggregate
+    val subquery1 = testRelation
+      .groupBy('b)(max('a))
+    val subquery2 = testRelation
+      .groupBy('b)(min('a))
+
+    // supports ObjectHashAggregate
+    val subquery3 = testRelation
+      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false))
+    val subquery4 = testRelation
+      .groupBy('b)(CollectSet('a).toAggregateExpression(isDistinct = false))
+
+    // supports SortAggregate
+    val subquery5 = testRelation
+      .groupBy('b)(max('c))
+    val subquery6 = testRelation
+      .groupBy('b)(min('c))
+
+    val originalQuery = testRelation
+      .select(ScalarSubquery(subquery1), ScalarSubquery(subquery2), ScalarSubquery(subquery3),
+        ScalarSubquery(subquery4), ScalarSubquery(subquery5), ScalarSubquery(subquery6))
+
+    val hashAggregates = testRelation
+      .groupBy('b)(max('a), min('a)).analyze
+    val objectHashAggregates = testRelation
+      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false),
+        CollectSet('a).toAggregateExpression(isDistinct = false)).analyze
+    val sortAggregates = testRelation
+      .groupBy('b)(max('c), min('c)).analyze
+    val correctAnswer = testRelation
+      .select(GetStructField(MultiScalarSubquery(hashAggregates), 0).as("scalarsubquery()"),
+        GetStructField(MultiScalarSubquery(hashAggregates), 1).as("scalarsubquery()"),
+        GetStructField(MultiScalarSubquery(objectHashAggregates), 0).as("scalarsubquery()"),
+        GetStructField(MultiScalarSubquery(objectHashAggregates), 1).as("scalarsubquery()"),
+        GetStructField(MultiScalarSubquery(sortAggregates), 0).as("scalarsubquery()"),
+        GetStructField(MultiScalarSubquery(sortAggregates), 1).as("scalarsubquery()"))
+
+    // checkAnalysis is disabled because `Analizer` is not prepared for `MultiScalarSubquery` nodes
+    // as only `Optimizer` can insert such a node to the plan
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer, false)
+  }
 }

sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java

@@ -25,6 +25,7 @@
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.catalyst.plans.logical.Aggregate$;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.unsafe.KVIterator;
@@ -68,12 +69,7 @@ public final class UnsafeFixedWidthAggregationMap {
    *         schema, false otherwise.
    */
   public static boolean supportsAggregationBufferSchema(StructType schema) {
-    for (StructField field: schema.fields()) {
-      if (!UnsafeRow.isMutable(field.dataType())) {
-        return false;
-      }
-    }
-    return true;
+    return Aggregate$.MODULE$.supportsAggregationBufferSchema(schema);
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.aggregate
 
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.plans.logical.Aggregate
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.streaming.{StateStoreRestoreExec, StateStoreSaveExec}
 
@@ -50,7 +51,7 @@ object AggUtils {
       initialInputBufferOffset: Int = 0,
       resultExpressions: Seq[NamedExpression] = Nil,
       child: SparkPlan): SparkPlan = {
-    val useHash = HashAggregateExec.supportsAggregate(
+    val useHash = Aggregate.supportsHashAggregate(
       aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
     if (useHash) {
       HashAggregateExec(
@@ -63,7 +64,7 @@ object AggUtils {
         child = child)
     } else {
       val objectHashEnabled = child.sqlContext.conf.useObjectHashAggregation
-      val useObjectHash = ObjectHashAggregateExec.supportsAggregate(aggregateExpressions)
+      val useObjectHash = Aggregate.supportsObjectHashAggregate(aggregateExpressions)
 
       if (objectHashEnabled && useObjectHash) {
         ObjectHashAggregateExec(

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala

@@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.catalyst.plans.logical.Aggregate
 import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_MILLIS
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.execution._
@@ -55,7 +56,7 @@ case class HashAggregateExec(
   with BlockingOperatorWithCodegen
   with GeneratePredicateHelper {
 
-  require(HashAggregateExec.supportsAggregate(aggregateBufferAttributes))
+  require(Aggregate.supportsHashAggregate(aggregateBufferAttributes))
 
   override lazy val allAttributes: AttributeSeq =
     child.output ++ aggregateBufferAttributes ++ aggregateAttributes ++
@@ -1139,10 +1140,3 @@ case class HashAggregateExec(
   override protected def withNewChildInternal(newChild: SparkPlan): HashAggregateExec =
     copy(child = newChild)
 }
-
-object HashAggregateExec {
-  def supportsAggregate(aggregateBufferAttributes: Seq[Attribute]): Boolean = {
-    val aggregationBufferSchema = StructType.fromAttributes(aggregateBufferAttributes)
-    UnsafeFixedWidthAggregationMap.supportsAggregationBufferSchema(aggregationBufferSchema)
-  }
-}

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/ObjectHashAggregateExec.scala

@@ -142,12 +142,3 @@ case class ObjectHashAggregateExec(
   override protected def withNewChildInternal(newChild: SparkPlan): ObjectHashAggregateExec =
     copy(child = newChild)
 }
-
-object ObjectHashAggregateExec {
-  def supportsAggregate(aggregateExpressions: Seq[AggregateExpression]): Boolean = {
-    aggregateExpressions.map(_.aggregateFunction).exists {
-      case _: TypedImperativeAggregate[_] => true
-      case _ => false
-    }
-  }
-}