drop MultiScalarSubquery, use ScalarSubquery(CreateStruct()) instead

Author: peter-toth

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

@@ -22,9 +22,8 @@ import scala.collection.mutable.ArrayBuffer
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan}
-import org.apache.spark.sql.catalyst.trees.LeafLike
 import org.apache.spark.sql.catalyst.trees.TreePattern.{EXISTS_SUBQUERY, LIST_SUBQUERY,
-  MULTI_SCALAR_SUBQUERY, PLAN_EXPRESSION, SCALAR_SUBQUERY, TreePattern}
+  PLAN_EXPRESSION, SCALAR_SUBQUERY, TreePattern}
 import org.apache.spark.sql.types._
 import org.apache.spark.util.collection.BitSet
 
@@ -268,33 +267,6 @@ object ScalarSubquery {
   }
 }
 
-/**
- * A subquery that is capable to return multiple scalar values.
- */
-case class MultiScalarSubquery(
-    plan: LogicalPlan,
-    exprId: ExprId = NamedExpression.newExprId)
-  extends SubqueryExpression(plan, Seq.empty, exprId) with LeafLike[Expression] with Unevaluable {
-  override def dataType: DataType = {
-    assert(plan.schema.nonEmpty, "Multi-column scalar subquery should have columns")
-    plan.schema
-  }
-
-  override def nullable: Boolean = true
-
-  override def withNewPlan(plan: LogicalPlan): MultiScalarSubquery = copy(plan = plan)
-
-  override def toString: String = s"multi-scalar-subquery#${exprId.id}"
-
-  override lazy val canonicalized: Expression = {
-    MultiScalarSubquery(
-      plan.canonicalized,
-      ExprId(0))
-  }
-
-  final override def nodePatternsInternal: Seq[TreePattern] = Seq(MULTI_SCALAR_SUBQUERY)
-}
-
 /**
  * A [[ListQuery]] expression defines the query which we want to search in an IN subquery
  * expression. It should and can only be used in conjunction with an IN expression.

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

@@ -23,7 +23,7 @@ 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}
+import org.apache.spark.sql.catalyst.trees.TreePattern.SCALAR_SUBQUERY
 
 /**
  * This rule tries to merge multiple non-correlated [[ScalarSubquery]]s into a
@@ -72,46 +72,54 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{MULTI_SCALAR_SUBQUERY, S
 object MergeScalarSubqueries extends Rule[LogicalPlan] with PredicateHelper {
   def apply(plan: LogicalPlan): LogicalPlan = {
     if (conf.scalarSubqueryMergeEabled && conf.subqueryReuseEnabled) {
-      val mergedSubqueries = ArrayBuffer.empty[LogicalPlan]
-      removeReferences(mergeAndInsertReferences(plan, mergedSubqueries), mergedSubqueries)
+      val mergedSubqueries = ArrayBuffer.empty[Project]
+      removeReferences(mergeAndInsertReferences(plan, mergedSubqueries))
     } else {
       plan
     }
   }
 
   private def mergeAndInsertReferences(
       plan: LogicalPlan,
-      mergedSubqueries: ArrayBuffer[LogicalPlan]): LogicalPlan = {
-    plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY), ruleId) {
-      case s: ScalarSubquery if s.children.isEmpty =>
-        val (mergedPlan, ordinal) = mergeAndGetReference(s.plan, mergedSubqueries)
-        GetStructField(MultiScalarSubquery(mergedPlan, s.exprId), ordinal)
+      mergedSubqueries: ArrayBuffer[Project]): LogicalPlan = {
+    plan.transformWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY), ruleId) {
+      case o => o.transformExpressionsUpWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY), ruleId) {
+        case s: ScalarSubquery if s.children.isEmpty =>
+          val (mergedPlan, ordinal) = mergeAndGetReference(s.plan, mergedSubqueries)
+          GetStructField(s.copy(plan = mergedPlan), ordinal)
+      }
     }
   }
 
   case class SubqueryReference(
       index: Int,
-      mergedSubqueries: ArrayBuffer[LogicalPlan]) extends LeafNode {
+      mergedSubqueries: ArrayBuffer[Project]) extends LeafNode {
     override def stringArgs: Iterator[Any] = Iterator(index)
 
     override def output: Seq[Attribute] = mergedSubqueries(index).output
   }
 
   private def mergeAndGetReference(
       plan: LogicalPlan,
-      mergedSubqueries: ArrayBuffer[LogicalPlan]): (SubqueryReference, Int) = {
+      mergedSubqueries: ArrayBuffer[Project]): (SubqueryReference, Int) = {
     mergedSubqueries.zipWithIndex.collectFirst {
-      Function.unlift { case (s, i) => tryMergePlans(plan, s).map(_ -> i) }
-    }.map { case ((mergedPlan, outputMap), i) =>
-      mergedSubqueries(i) = mergedPlan
-      SubqueryReference(i, mergedSubqueries) ->
-        mergedPlan.output.indexOf(outputMap(plan.output.head))
+      Function.unlift { case (header, i) => tryMergePlans(plan, header.child).map((header, _, i)) }
+    }.map { case (header, (mergedPlan, outputMap), i) =>
+      if (mergedPlan.output.size > header.child.output.size) {
+        mergedSubqueries(i) = createHeader(mergedPlan)
+      }
+      val ordinal = mergedPlan.output.indexOf(outputMap(plan.output.head))
+      SubqueryReference(i, mergedSubqueries) -> ordinal
     }.getOrElse {
-      mergedSubqueries += plan
+      mergedSubqueries += createHeader(plan)
       SubqueryReference(mergedSubqueries.length - 1, mergedSubqueries) -> 0
     }
   }
 
+  private def createHeader(plan: LogicalPlan) = {
+    Project(Seq(Alias(CreateStruct(plan.output), "mergedValue")()), plan)
+  }
+
   private def tryMergePlans(
       newPlan: LogicalPlan,
       existingPlan: LogicalPlan): Option[(LogicalPlan, AttributeMap[Attribute])] = {
@@ -191,16 +199,14 @@ object MergeScalarSubqueries extends Rule[LogicalPlan] with PredicateHelper {
       }
   }
 
-  private def removeReferences(
-      plan: LogicalPlan,
-      mergedSubqueries: ArrayBuffer[LogicalPlan]): LogicalPlan = {
-    plan.transformAllExpressionsWithPruning(_.containsAnyPattern(MULTI_SCALAR_SUBQUERY), ruleId) {
-      case gsf @ GetStructField(mss @ MultiScalarSubquery(sr: SubqueryReference, _), _, _) =>
-        val dereferencedPlan = removeReferences(mergedSubqueries(sr.index), mergedSubqueries)
-        if (dereferencedPlan.outputSet.size > 1) {
-          gsf.copy(child = mss.copy(plan = dereferencedPlan))
+  private def removeReferences(plan: LogicalPlan): LogicalPlan = {
+    plan.transformAllExpressionsWithPruning(_.containsAnyPattern(SCALAR_SUBQUERY), ruleId) {
+      case gsf @ GetStructField(ss @ ScalarSubquery(sr: SubqueryReference, _, _), _, _) =>
+        val header = sr.mergedSubqueries(sr.index)
+        if (header.child.output.size > 1) {
+          gsf.copy(child = ss.copy(plan = header))
         } else {
-          ScalarSubquery(dereferencedPlan, exprId = mss.exprId)
+          ss.copy(plan = header.child)
         }
     }
   }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala

@@ -48,7 +48,6 @@ object TreePattern extends Enumeration  {
   val LIST_SUBQUERY: Value = Value
   val LITERAL: Value = Value
   val MAP_OBJECTS: Value = Value
-  val MULTI_SCALAR_SUBQUERY: Value = Value
   val NOT: Value = Value
   val NULL_CHECK: Value = Value
   val NULL_LITERAL: Value = Value

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

@@ -19,14 +19,13 @@ 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.{CreateStruct, GetStructField, 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
   }
@@ -35,82 +34,81 @@ class MergeScalarSubqueriesSuite extends PlanTest {
 
   test("Simple non-correlated scalar subquery merge") {
     val subquery1 = testRelation
-      .groupBy('b)(max('a))
+      .groupBy('b)(max('a).as("max_a"))
     val subquery2 = testRelation
-      .groupBy('b)(sum('a))
+      .groupBy('b)(sum('a).as("sum_a"))
     val originalQuery = testRelation
       .select(ScalarSubquery(subquery1), ScalarSubquery(subquery2))
 
     val multiSubquery = testRelation
-      .groupBy('b)(max('a), sum('a)).analyze
+      .groupBy('b)(max('a).as("max_a"), sum('a).as("sum_a"))
+      .select(CreateStruct(Seq('max_a, 'sum_a)).as("mergedValue"))
     val correctAnswer = testRelation
-      .select(GetStructField(MultiScalarSubquery(multiSubquery), 0).as("scalarsubquery()"),
-        GetStructField(MultiScalarSubquery(multiSubquery), 1).as("scalarsubquery()"))
+      .select(GetStructField(ScalarSubquery(multiSubquery), 0).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(multiSubquery), 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)
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
   }
 
   test("Aggregate and group expression merge") {
     val subquery1 = testRelation
-      .groupBy('b)(max('a))
+      .groupBy('b)(max('a).as("max_a"))
     val subquery2 = testRelation
       .groupBy('b)('b)
     val originalQuery = testRelation
       .select(ScalarSubquery(subquery1), ScalarSubquery(subquery2))
 
     val multiSubquery = testRelation
-      .groupBy('b)(max('a), 'b).analyze
+      .groupBy('b)(max('a).as("max_a"), 'b)
+      .select(CreateStruct(Seq('max_a, 'b)).as("mergedValue"))
     val correctAnswer = testRelation
-      .select(GetStructField(MultiScalarSubquery(multiSubquery), 0).as("scalarsubquery()"),
-        GetStructField(MultiScalarSubquery(multiSubquery), 1).as("scalarsubquery()"))
+      .select(GetStructField(ScalarSubquery(multiSubquery), 0).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(multiSubquery), 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)
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
   }
 
   test("Do not merge different aggregate implementations") {
     // supports HashAggregate
     val subquery1 = testRelation
-      .groupBy('b)(max('a))
+      .groupBy('b)(max('a).as("max_a"))
     val subquery2 = testRelation
-      .groupBy('b)(min('a))
+      .groupBy('b)(min('a).as("min_a"))
 
     // supports ObjectHashAggregate
     val subquery3 = testRelation
-      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false))
+      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false).as("collectlist_a"))
     val subquery4 = testRelation
-      .groupBy('b)(CollectSet('a).toAggregateExpression(isDistinct = false))
+      .groupBy('b)(CollectSet('a).toAggregateExpression(isDistinct = false).as("collectset_a"))
 
     // supports SortAggregate
     val subquery5 = testRelation
-      .groupBy('b)(max('c))
+      .groupBy('b)(max('c).as("max_c"))
     val subquery6 = testRelation
-      .groupBy('b)(min('c))
+      .groupBy('b)(min('c).as("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
+      .groupBy('b)(max('a).as("max_a"), min('a).as("min_a"))
+      .select(CreateStruct(Seq('max_a, 'min_a)).as("mergedValue"))
     val objectHashAggregates = testRelation
-      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false),
-        CollectSet('a).toAggregateExpression(isDistinct = false)).analyze
+      .groupBy('b)(CollectList('a).toAggregateExpression(isDistinct = false).as("collectlist_a"),
+        CollectSet('a).toAggregateExpression(isDistinct = false).as("collectset_a"))
+      .select(CreateStruct(Seq('collectlist_a, 'collectset_a)).as("mergedValue"))
     val sortAggregates = testRelation
-      .groupBy('b)(max('c), min('c)).analyze
+      .groupBy('b)(max('c).as("max_c"), min('c).as("min_c"))
+      .select(CreateStruct(Seq('max_c, 'min_c)).as("mergedValue"))
     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)
+      .select(GetStructField(ScalarSubquery(hashAggregates), 0).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(hashAggregates), 1).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(objectHashAggregates), 0).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(objectHashAggregates), 1).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(sortAggregates), 0).as("scalarsubquery()"),
+        GetStructField(ScalarSubquery(sortAggregates), 1).as("scalarsubquery()"))
+
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/PlanTest.scala

@@ -73,8 +73,6 @@ trait PlanTestBase extends PredicateHelper with SQLHelper with SQLConfHelper { s
     plan transformAllExpressions {
       case s: ScalarSubquery =>
         s.copy(plan = normalizeExprIds(s.plan), exprId = ExprId(0))
-      case s: MultiScalarSubquery =>
-        s.copy(plan = normalizeExprIds(s.plan), exprId = ExprId(0))
       case e: Exists =>
         e.copy(exprId = ExprId(0))
       case l: ListQuery =>

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala

@@ -24,8 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.{ListQuery, SubqueryExpression}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{DYNAMIC_PRUNING_SUBQUERY, IN_SUBQUERY,
-  MULTI_SCALAR_SUBQUERY, SCALAR_SUBQUERY}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{DYNAMIC_PRUNING_SUBQUERY, IN_SUBQUERY, SCALAR_SUBQUERY}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
 import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
@@ -115,8 +114,7 @@ case class InsertAdaptiveSparkPlan(
    */
   private def buildSubqueryMap(plan: SparkPlan): Map[Long, BaseSubqueryExec] = {
     val subqueryMap = mutable.HashMap.empty[Long, BaseSubqueryExec]
-    if (!plan.containsAnyPattern(SCALAR_SUBQUERY, MULTI_SCALAR_SUBQUERY, IN_SUBQUERY,
-      DYNAMIC_PRUNING_SUBQUERY)) {
+    if (!plan.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY, DYNAMIC_PRUNING_SUBQUERY)) {
       return subqueryMap.toMap
     }
     plan.foreach(_.expressions.foreach(_.foreach {
@@ -127,13 +125,6 @@ case class InsertAdaptiveSparkPlan(
         val subquery = SubqueryExec.createForScalarSubquery(
           s"subquery#${exprId.id}", executedPlan)
         subqueryMap.put(exprId.id, subquery)
-      case expressions.MultiScalarSubquery(p, exprId)
-        if !subqueryMap.contains(exprId.id) =>
-        val executedPlan = compileSubquery(p)
-        verifyAdaptivePlan(executedPlan, p)
-        val subquery = SubqueryExec.createForScalarSubquery(
-          s"subquery#${exprId.id}", executedPlan)
-        subqueryMap.put(exprId.id, subquery)
       case expressions.InSubquery(_, ListQuery(query, _, exprId, _))
           if !subqueryMap.contains(exprId.id) =>
         val executedPlan = compileSubquery(query)

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala

@@ -20,8 +20,7 @@ package org.apache.spark.sql.execution.adaptive
 import org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.sql.catalyst.expressions.{CreateNamedStruct, DynamicPruningExpression, ListQuery, Literal}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{DYNAMIC_PRUNING_SUBQUERY, IN_SUBQUERY,
-  MULTI_SCALAR_SUBQUERY, SCALAR_SUBQUERY}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{DYNAMIC_PRUNING_SUBQUERY, IN_SUBQUERY, SCALAR_SUBQUERY}
 import org.apache.spark.sql.execution
 import org.apache.spark.sql.execution.{BaseSubqueryExec, InSubqueryExec, SparkPlan}
 
@@ -30,12 +29,9 @@ case class PlanAdaptiveSubqueries(
 
   def apply(plan: SparkPlan): SparkPlan = {
     plan.transformAllExpressionsWithPruning(
-      _.containsAnyPattern(SCALAR_SUBQUERY, MULTI_SCALAR_SUBQUERY, IN_SUBQUERY,
-        DYNAMIC_PRUNING_SUBQUERY)) {
+      _.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY, DYNAMIC_PRUNING_SUBQUERY)) {
       case expressions.ScalarSubquery(_, _, exprId) =>
         execution.ScalarSubquery(subqueryMap(exprId.id), exprId)
-      case expressions.MultiScalarSubquery(_, exprId) =>
-        execution.MultiScalarSubqueryExec(subqueryMap(exprId.id), exprId)
       case expressions.InSubquery(values, ListQuery(_, _, exprId, _)) =>
         val expr = if (values.length == 1) {
           values.head