use optimized plan

Author: turboFei

externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelper.scala

@@ -25,8 +25,9 @@ import org.apache.spark.network.util.{ByteUnit, JavaUtils}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.plans.logical.GlobalLimit
 import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
-import org.apache.spark.sql.execution.{CollectLimitExec, HiveResult, LocalTableScanExec, SparkPlan, SQLExecution, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.{CollectLimitExec, HiveResult, LocalTableScanExec, QueryExecution, SparkPlan, SQLExecution}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
 import org.apache.spark.sql.execution.arrow.KyuubiArrowConverters
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
@@ -294,18 +295,17 @@ object SparkDatasetHelper extends Logging {
     SQLMetrics.postDriverMetricUpdates(sc, executionId, metrics.values.toSeq)
   }
 
-  private[kyuubi] def planLimit(plan: SparkPlan): Option[Int] = plan match {
-    case tp: TakeOrderedAndProjectExec => Option(tp.limit)
-    case c: CollectLimitExec => Option(c.limit)
-    case ap: AdaptiveSparkPlanExec => planLimit(ap.inputPlan)
-    case _ => None
-  }
+  private[kyuubi] def optimizedPlanLimit(queryExecution: QueryExecution): Option[Long] =
+    queryExecution.optimizedPlan match {
+      case globalLimit: GlobalLimit => globalLimit.maxRows
+      case _ => None
+    }
 
   def shouldSaveResultToFs(resultMaxRows: Int, minSize: Long, result: DataFrame): Boolean = {
     if (isCommandExec(result.queryExecution.executedPlan.nodeName)) {
       return false
     }
-    val finalLimit = planLimit(result.queryExecution.sparkPlan) match {
+    val finalLimit = optimizedPlanLimit(result.queryExecution) match {
       case Some(limit) if resultMaxRows > 0 => math.min(limit, resultMaxRows)
       case Some(limit) => limit
       case None => resultMaxRows

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/spark/sql/kyuubi/SparkDatasetHelperSuite.scala

@@ -33,13 +33,13 @@ class SparkDatasetHelperSuite extends WithSparkSQLEngine {
         " SELECT * FROM VALUES(1),(2),(3),(4) AS t(id)")
 
       val topKStatement = s"SELECT * FROM(SELECT * FROM tv ORDER BY id LIMIT ${topKThreshold - 1})"
-      assert(SparkDatasetHelper.planLimit(
-        spark.sql(topKStatement).queryExecution.sparkPlan) === Option(topKThreshold - 1))
+      assert(SparkDatasetHelper.optimizedPlanLimit(
+        spark.sql(topKStatement).queryExecution) === Option(topKThreshold - 1))
 
       val collectLimitStatement =
         s"SELECT * FROM (SELECT * FROM tv ORDER BY id LIMIT $topKThreshold)"
-      assert(SparkDatasetHelper.planLimit(
-        spark.sql(collectLimitStatement).queryExecution.sparkPlan) === Option(topKThreshold))
+      assert(SparkDatasetHelper.optimizedPlanLimit(
+        spark.sql(collectLimitStatement).queryExecution) === Option(topKThreshold))
     }
   }
 }