address comments

Author: cloud-fan

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

@@ -75,8 +75,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
           // Normally wrapping child with `LocalLimitExec` here is a no-op, because
           // `CollectLimitExec.executeCollect` will call `LogicalLimitExec.executeTake`, which
           // calls `child.executeTake`. If child supports whole stage codegen, adding this
-          // `LocalLimitExec` can break the input consuming loop inside whole stage codegen and
-          // trigger the resource releasing work, after we consume `limit` rows.
+          // `LocalLimitExec` can stop the processing of whole stage codegen and trigger the
+          // resource releasing work, after we consume `limit` rows.
           execution.CollectLimitExec(limit, LocalLimitExec(limit, planLater(child))) :: Nil
         case other => planLater(other) :: Nil
       }

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

@@ -474,10 +474,6 @@ case class CollapseCodegenStages(conf: SQLConf) extends Rule[SparkPlan] {
   }
 
   private def supportCodegen(plan: SparkPlan): Boolean = plan match {
-    // Do not enable whole stage codegen for a single limit.
-    case limit: BaseLimitExec if !limit.child.isInstanceOf[CodegenSupport] ||
-        !limit.child.asInstanceOf[CodegenSupport].supportCodegen =>
-      false
     case plan: CodegenSupport if plan.supportCodegen =>
       val willFallback = plan.expressions.exists(_.find(e => !supportCodegen(e)).isDefined)
       // the generated code will be huge if there are too many columns

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

@@ -54,6 +54,12 @@ trait BaseLimitExec extends UnaryExecNode with CodegenSupport {
   val limit: Int
   override def output: Seq[Attribute] = child.output
 
+  // Do not enable whole stage codegen for a single limit.
+  override def supportCodegen: Boolean = child match {
+    case plan: CodegenSupport => plan.supportCodegen
+    case _ => false
+  }
+
   override def executeTake(n: Int): Array[InternalRow] = child.executeTake(math.min(n, limit))
 
   override def executeCollect(): Array[InternalRow] = child.executeTake(limit)