[CARMEL-6077] Use TakeOrderedAndProject if maxRows below the topKSortMaxRowsThreshold (#1015)

* Use TakeOrderedAndProject if maxRows below the topKSortMaxRowsThreshold

* Fix

Author: wangyum

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2151,6 +2151,17 @@ object SQLConf {
       .intConf
       .createWithDefault(ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH)
 
+  val TOP_K_SORT_MAX_ROWS_THRESHOLD = buildConf("spark.sql.execution.topKSortMaxRowsThreshold")
+    .doc("In SQL queries with a SORT and max rows exists like " +
+      "'SELECT DISTINCT x FROM t ORDER BY y', if max rows is under this threshold, do a top-K " +
+      "sort in memory to avoid doing range repartition.")
+    .version("3.4.0")
+    .intConf
+    .checkValue(_ <= ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH,
+      "The top-K sort max rows threshold should be less than or equal to " +
+        s"${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}.")
+    .createWithDefault(655360)
+
   val SORT_MERGE_JOIN_EXEC_BUFFER_SPILL_THRESHOLD =
     buildConf("spark.sql.sortMergeJoinExec.buffer.spill.threshold")
       .internal()

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

@@ -100,8 +100,23 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case Limit(IntegerLiteral(limit), Project(projectList, Sort(order, true, child)))
           if limit < conf.topKSortFallbackThreshold =>
         TakeOrderedAndProjectExec(limit, order, projectList, planLater(child)) :: Nil
+
+      case Sort(order, true, child) if supportTakeOrdered(child) =>
+        val limit = child.maxRows.get.toInt
+        TakeOrderedAndProjectExec(limit, order, child.output, planLater(child)) :: Nil
+      case Project(projectList, Sort(order, true, child)) if supportTakeOrdered(child) =>
+        val limit = child.maxRows.get.toInt
+        TakeOrderedAndProjectExec(limit, order, projectList, planLater(child)) :: Nil
       case _ => Nil
     }
+
+    private def supportTakeOrdered(plan: LogicalPlan): Boolean = {
+      plan.maxRows.exists(_ < conf.getConf(SQLConf.TOP_K_SORT_MAX_ROWS_THRESHOLD)) &&
+        // The plan should not contain global sort, to avoid sorting it again.
+        plan.collectFirst {
+          case s: Sort => s.global
+        }.isEmpty
+    }
   }
 
   /**

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

@@ -53,4 +53,6 @@ case class LogicalQueryStage(
     }
     physicalStats.getOrElse(logicalPlan.stats)
   }
+
+  override def maxRows: Option[Long] = stats.rowCount.map(_.min(Long.MaxValue).toLong)
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Range, Repartition, Sort, Union}
 import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
 import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec}
 import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReusedExchangeExec, ShuffleExchangeExec}
@@ -977,6 +977,36 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
     assert(PartitioningUtil.extractPartitioning[CoalescedPartitioning](pc3) ==
       c :: c :: c :: Nil)
   }
+
+  test("SPARK-39698: Use TakeOrderedAndProject if maxRows below the topKSortFallbackThreshold") {
+    Seq(-1, 10).foreach { threshold =>
+      withSQLConf(SQLConf.TOP_K_SORT_MAX_ROWS_THRESHOLD.key -> threshold.toString,
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+        val df = Seq(1 -> "a", 2 -> "b", 3 -> "c", 4 -> "d", 5 -> "e").toDF("i", "j").orderBy($"i")
+        df.collect()
+        val aeqPlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+        if (threshold < 0) {
+          assert(aeqPlan.currentPhysicalPlan.isInstanceOf[execution.WholeStageCodegenExec])
+        } else {
+          assert(aeqPlan.currentPhysicalPlan.isInstanceOf[TakeOrderedAndProjectExec])
+        }
+      }
+    }
+  }
+
+  test("SPARK-39698: Do not use TakeOrderedAndProject if children contain global sort") {
+    Seq(-1, 10).foreach { threshold =>
+      withSQLConf(SQLConf.TOP_K_SORT_MAX_ROWS_THRESHOLD.key -> threshold.toString,
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
+        val df = Seq(1 -> "a", 2 -> "b", 3 -> "c", 4 -> "d", 5 -> "e").toDF("i", "j")
+          .orderBy($"i").limit(4).orderBy($"i")
+        df.collect()
+        val aeqPlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
+        // The last oder should be removed through RemoveRedundantSorts.
+        assert(aeqPlan.currentPhysicalPlan.isInstanceOf[TakeOrderedAndProjectExec])
+      }
+    }
+  }
 }
 
 // Used for unit-testing EnsureRequirements