[SPARK-38162][SQL] Optimize one row plan in normal and AQE Optimizer

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

@@ -238,6 +238,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
       // PropagateEmptyRelation can change the nullability of an attribute from nullable to
       // non-nullable when an empty relation child of a Union is removed
       UpdateAttributeNullability) :+
+    Batch("Optimize One Row Plan", fixedPoint, OptimizeOneRowPlan) :+
     // The following batch should be executed after batch "Join Reorder" and "LocalRelation".
     Batch("Check Cartesian Products", Once,
       CheckCartesianProducts) :+
@@ -454,6 +455,40 @@ object EliminateAggregateFilter extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * The rule is applied both normal and AQE Optimizer. It optimizes plan using max rows:
+ *   - if the child of sort max rows less than or equal to 1, remove the sort
+ *   - if the child of local sort max rows per partition less than or equal to 1, remove the
+ *     local sort
+ *   - if the child of aggregate max rows less than or equal to 1 and its output is subset of
+ *     its child and it's grouping only(include the rewritten distinct plan), remove the aggregate
+ *   - if the child of aggregate max rows less than or equal to 1, set distinct to false in all
+ *     aggregate expression
+ */
+object OptimizeOneRowPlan extends Rule[LogicalPlan] {
+  private def maxRowNotLargerThanOne(plan: LogicalPlan): Boolean = {
+    plan.maxRows.exists(_ <= 1L)
+  }
+
+  private def maxRowPerPartitionNotLargerThanOne(plan: LogicalPlan): Boolean = {
+    plan.maxRowsPerPartition.exists(_ <= 1L)
+  }
+
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    plan.transformDownWithPruning(_.containsAnyPattern(SORT, AGGREGATE), ruleId) {
+      case Sort(_, _, child) if maxRowNotLargerThanOne(child) => child
+      case Sort(_, false, child) if maxRowPerPartitionNotLargerThanOne(child) => child
+      case agg @ Aggregate(_, _, child) if agg.groupOnly &&
+        agg.outputSet.subsetOf(child.outputSet) && maxRowNotLargerThanOne(child) => child
+      case agg: Aggregate if maxRowNotLargerThanOne(agg.child) =>
+        agg.transformExpressions {
+          case aggExpr: AggregateExpression if aggExpr.isDistinct =>
+            aggExpr.copy(isDistinct = false)
+        }
+    }
+  }
+}
+
 /**
  * An optimizer used in test code.
  *
@@ -1390,15 +1425,14 @@ object CombineFilters extends Rule[LogicalPlan] with PredicateHelper {
  * Removes Sort operations if they don't affect the final output ordering.
  * Note that changes in the final output ordering may affect the file size (SPARK-32318).
  * This rule handles the following cases:
- * 1) if the child maximum number of rows less than or equal to 1
- * 2) if the sort order is empty or the sort order does not have any reference
- * 3) if the Sort operator is a local sort and the child is already sorted
- * 4) if there is another Sort operator separated by 0...n Project, Filter, Repartition or
+ * 1) if the sort order is empty or the sort order does not have any reference
+ * 2) if the Sort operator is a local sort and the child is already sorted
+ * 3) if there is another Sort operator separated by 0...n Project, Filter, Repartition or
  *    RepartitionByExpression, RebalancePartitions (with deterministic expressions) operators
- * 5) if the Sort operator is within Join separated by 0...n Project, Filter, Repartition or
+ * 4) if the Sort operator is within Join separated by 0...n Project, Filter, Repartition or
  *    RepartitionByExpression, RebalancePartitions (with deterministic expressions) operators only
  *    and the Join condition is deterministic
- * 6) if the Sort operator is within GroupBy separated by 0...n Project, Filter, Repartition or
+ * 5) if the Sort operator is within GroupBy separated by 0...n Project, Filter, Repartition or
  *    RepartitionByExpression, RebalancePartitions (with deterministic expressions) operators only
  *    and the aggregate function is order irrelevant
  */
@@ -1407,7 +1441,6 @@ object EliminateSorts extends Rule[LogicalPlan] {
     _.containsPattern(SORT))(applyLocally)
 
   private val applyLocally: PartialFunction[LogicalPlan, LogicalPlan] = {
-    case Sort(_, _, child) if child.maxRows.exists(_ <= 1L) => recursiveRemoveSort(child)
     case s @ Sort(orders, _, child) if orders.isEmpty || orders.exists(_.child.foldable) =>
       val newOrders = orders.filterNot(_.child.foldable)
       if (newOrders.isEmpty) {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala

@@ -121,6 +121,7 @@ object RuleIdCollection {
       "org.apache.spark.sql.catalyst.optimizer.ObjectSerializerPruning" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeCsvJsonExprs" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeIn" ::
+      "org.apache.spark.sql.catalyst.optimizer.OptimizeOneRowPlan" ::
       "org.apache.spark.sql.catalyst.optimizer.Optimizer$OptimizeSubqueries" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeRepartition" ::
       "org.apache.spark.sql.catalyst.optimizer.OptimizeWindowFunctions" ::

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

@@ -422,14 +422,4 @@ class EliminateSortsSuite extends AnalysisTest {
       comparePlans(optimized, correctAnswer)
     }
   }
-
-  test("SPARK-35906: Remove order by if the maximum number of rows less than or equal to 1") {
-    comparePlans(
-      Optimize.execute(testRelation.groupBy()(count(1).as("cnt")).orderBy('cnt.asc)).analyze,
-      testRelation.groupBy()(count(1).as("cnt")).analyze)
-
-    comparePlans(
-      Optimize.execute(testRelation.limit(Literal(1)).orderBy('a.asc).orderBy('a.asc)).analyze,
-      testRelation.limit(Literal(1)).analyze)
-  }
 }

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

@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans._
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+
+class OptimizeOneRowPlanSuite extends PlanTest {
+  object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches =
+      Batch("Replace Operators", Once, ReplaceDistinctWithAggregate) ::
+      Batch("Eliminate Sorts", Once, EliminateSorts) ::
+      Batch("Optimize One Row Plan", FixedPoint(10), OptimizeOneRowPlan) :: Nil
+  }
+
+  private val t1 = LocalRelation.fromExternalRows(Seq($"a".int), data = Seq(Row(1)))
+  private val t2 = LocalRelation.fromExternalRows(Seq($"a".int), data = Seq(Row(1), Row(2)))
+
+  test("SPARK-35906: Remove order by if the maximum number of rows less than or equal to 1") {
+    comparePlans(
+      Optimize.execute(t2.groupBy()(count(1).as("cnt")).orderBy('cnt.asc)).analyze,
+      t2.groupBy()(count(1).as("cnt")).analyze)
+
+    comparePlans(
+      Optimize.execute(t2.limit(Literal(1)).orderBy('a.asc).orderBy('a.asc)).analyze,
+      t2.limit(Literal(1)).analyze)
+  }
+
+  test("Remove sort") {
+    // remove local sort
+    val plan1 = LocalLimit(0, t1).union(LocalLimit(0, t2)).sortBy($"a".desc).analyze
+    val expected = LocalLimit(0, t1).union(LocalLimit(0, t2)).analyze
+    comparePlans(Optimize.execute(plan1), expected)
+
+    // do not remove
+    val plan2 = t2.orderBy($"a".desc).analyze
+    comparePlans(Optimize.execute(plan2), plan2)
+
+    val plan3 = t2.sortBy($"a".desc).analyze
+    comparePlans(Optimize.execute(plan3), plan3)
+  }
+
+  test("Remove group only aggregate") {
+    val plan1 = t1.groupBy($"a")($"a").analyze
+    comparePlans(Optimize.execute(plan1), t1)
+
+    // do not remove
+    val plan2 = t2.groupBy($"a")($"a").analyze
+    comparePlans(Optimize.execute(plan2), plan2)
+
+    val plan3 = t1.groupBy($"a")(sum($"a")).analyze
+    comparePlans(Optimize.execute(plan3), plan3)
+
+    val plan4 = t1.groupBy()(sum($"a")).analyze
+    comparePlans(Optimize.execute(plan4), plan4)
+
+    // do not remove if the output is not subset of its child
+    val plan5 = t1.groupBy($"a" + 1)($"a" + 1).analyze
+    comparePlans(Optimize.execute(plan5), plan5)
+  }
+
+  test("Remove distinct in aggregate expression") {
+    val plan1 = t1.groupBy($"a")(sumDistinct($"a").as("s")).analyze
+    val expected1 = t1.groupBy($"a")(sum($"a").as("s")).analyze
+    comparePlans(Optimize.execute(plan1), expected1)
+
+    val plan2 = t1.groupBy()(sumDistinct($"a").as("s")).analyze
+    val expected2 = t1.groupBy()(sum($"a").as("s")).analyze
+    comparePlans(Optimize.execute(plan2), expected2)
+
+    // do not remove
+    val plan3 = t2.groupBy($"a")(sumDistinct($"a").as("s")).analyze
+    comparePlans(Optimize.execute(plan3), plan3)
+  }
+
+  test("Remove in complex case") {
+    val plan1 = t1.groupBy($"a")($"a").orderBy($"a".asc).analyze
+    val expected1 = t1.analyze
+    comparePlans(Optimize.execute(plan1), expected1)
+
+    val plan2 = t1.groupBy($"a")(sumDistinct($"a").as("s")).orderBy($"s".asc).analyze
+    val expected2 = t1.groupBy($"a")(sum($"a").as("s")).analyze
+    comparePlans(Optimize.execute(plan2), expected2)
+  }
+}

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.execution.adaptive
 
 import org.apache.spark.sql.catalyst.analysis.UpdateAttributeNullability
-import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, EliminateLimits}
+import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, EliminateLimits, OptimizeOneRowPlan}
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, LogicalPlanIntegrity, PlanHelper}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.internal.SQLConf
@@ -40,7 +40,8 @@ class AQEOptimizer(conf: SQLConf) extends RuleExecutor[LogicalPlan] {
       ConvertToLocalRelation,
       UpdateAttributeNullability),
     Batch("Dynamic Join Selection", Once, DynamicJoinSelection),
-    Batch("Eliminate Limits", fixedPoint, EliminateLimits)
+    Batch("Eliminate Limits", fixedPoint, EliminateLimits),
+    Batch("Optimize One Row Plan", fixedPoint, OptimizeOneRowPlan)
   )
 
   final override protected def batches: Seq[Batch] = {

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala

@@ -28,6 +28,7 @@ import org.apache.spark.sql.{Dataset, QueryTest, Row, SparkSession, Strategy}
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
 import org.apache.spark.sql.execution.{CollectLimitExec, CommandResultExec, LocalTableScanExec, PartialReducerPartitionSpec, QueryExecution, ReusedSubqueryExec, ShuffledRowRDD, SortExec, SparkPlan, UnaryExecNode, UnionExec}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
 import org.apache.spark.sql.execution.command.DataWritingCommandExec
 import org.apache.spark.sql.execution.datasources.noop.NoopDataSource
 import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec
@@ -126,6 +127,12 @@ class AdaptiveQueryExecSuite
     }
   }
 
+  private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
+    collect(plan) {
+      case agg: BaseAggregateExec => agg
+    }
+  }
+
   private def findTopLevelLimit(plan: SparkPlan): Seq[CollectLimitExec] = {
     collect(plan) {
       case l: CollectLimitExec => l
@@ -2484,6 +2491,38 @@ class AdaptiveQueryExecSuite
       }
     }
   }
+
+  test("SPARK-38162: Optimize one row plan in AQE Optimizer") {
+    withTempView("v") {
+      spark.sparkContext.parallelize(
+        (1 to 4).map(i => TestData( i, i.toString)), 2)
+        .toDF("c1", "c2").createOrReplaceTempView("v")
+
+      // remove sort
+      val (origin1, adaptive1) = runAdaptiveAndVerifyResult(
+        """
+          |SELECT * FROM v where c1 = 1 order by c1, c2
+          |""".stripMargin)
+      assert(findTopLevelSort(origin1).size == 1)
+      assert(findTopLevelSort(adaptive1).isEmpty)
+
+      // remove group only aggregate
+      val (origin2, adaptive2) = runAdaptiveAndVerifyResult(
+        """
+          |SELECT distinct c1 FROM (SELECT /*+ repartition(c1) */ * FROM v where c1 = 1)
+          |""".stripMargin)
+      assert(findTopLevelAggregate(origin2).size == 2)
+      assert(findTopLevelAggregate(adaptive2).isEmpty)
+
+      // remove distinct in aggregate
+      val (origin3, adaptive3) = runAdaptiveAndVerifyResult(
+        """
+          |SELECT sum(distinct c1) FROM (SELECT /*+ repartition(c1) */ * FROM v where c1 = 1)
+          |""".stripMargin)
+      assert(findTopLevelAggregate(origin3).size == 4)
+      assert(findTopLevelAggregate(adaptive3).size == 2)
+    }
+  }
 }
 
 /**