apache#73 support limit offset

Author: hn5092

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -823,6 +823,35 @@ case class LocalLimit(limitExpr: Expression, child: LogicalPlan) extends OrderPr
   }
 }
 
+object LimitRange {
+  def apply(startExpr: Expression, endExpr: Expression, child: LogicalPlan): UnaryNode = {
+    LimitRange0(startExpr, endExpr, LocalLimit(endExpr, child))
+  }
+
+  def unapply(p: LimitRange0): Option[(Expression, Expression, LogicalPlan)] = {
+    p match {
+      case LimitRange0(le0, le1, LocalLimit(le2, child)) if le1 == le2 => Some((le0, le1, child))
+      case _ => None
+    }
+  }
+}
+/**
+ * A global (coordinated) limit. This operator can emit at most `limitExpr` number in total.
+ *
+ * See [[Limit]] for more information.
+ */
+case class LimitRange0(startExpr: Expression, endExpr: Expression, child: LogicalPlan)
+  extends OrderPreservingUnaryNode {
+  override def output: Seq[Attribute] = child.output
+  override def maxRows: Option[Long] = {
+    (endExpr, endExpr) match {
+      case (IntegerLiteral(start), IntegerLiteral(end)) => Some(end - start)
+      case _ => None
+    }
+  }
+}
+
+
 /**
  * Aliased subquery.
  *

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -1809,6 +1809,10 @@ class Dataset[T] private[sql](
     Limit(Literal(n), logicalPlan)
   }
 
+  def limitRange(start: Int, end: Int): Dataset[T] = withTypedPlan {
+    LimitRange(Literal(start), Literal(end), logicalPlan)
+  }
+
   /**
    * Returns a new Dataset containing union of rows in this Dataset and another Dataset.
    *

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

@@ -74,8 +74,17 @@ 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 LimitRange(IntegerLiteral(start), IntegerLiteral(end),
+        Project(projectList, Sort(order, true, child)))  if end < conf.topKSortFallbackThreshold =>
+          TakeOrderedRangeAndProjectExec(start, end, order, projectList, planLater(child)) :: Nil
+        case LimitRange(IntegerLiteral(start), IntegerLiteral(end), Sort(order, true, child))
+          if end < conf.topKSortFallbackThreshold =>
+          TakeOrderedRangeAndProjectExec(start, end, order, child.output, planLater(child)) :: Nil
         case Limit(IntegerLiteral(limit), child) =>
           CollectLimitExec(limit, planLater(child)) :: Nil
+        case LimitRange(IntegerLiteral(start),
+        IntegerLiteral(limit), child) =>
+          CollectLimitRangeExec(start, limit, planLater(child)) :: Nil
         case other => planLater(other) :: Nil
       }
       case Limit(IntegerLiteral(limit), Sort(order, true, child))
@@ -84,6 +93,12 @@ 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 LimitRange(IntegerLiteral(start), IntegerLiteral(end),
+      Project(projectList, Sort(order, true, child)))  if end < conf.topKSortFallbackThreshold =>
+        TakeOrderedRangeAndProjectExec(start, end, order, projectList, planLater(child)) :: Nil
+      case LimitRange(IntegerLiteral(start), IntegerLiteral(end), Sort(order, true, child))
+        if end < conf.topKSortFallbackThreshold =>
+        TakeOrderedRangeAndProjectExec(start, end, order, child.output, planLater(child)) :: Nil
       case _ => Nil
     }
   }
@@ -617,6 +632,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.LocalLimitExec(limit, planLater(child)) :: Nil
       case logical.GlobalLimit(IntegerLiteral(limit), child) =>
         execution.GlobalLimitExec(limit, planLater(child)) :: Nil
+      case logical.LimitRange(IntegerLiteral(start),
+      IntegerLiteral(limit), child) =>
+        execution.RangeLimitExec(start, limit, planLater(child)) :: Nil
       case logical.Union(unionChildren) =>
         execution.UnionExec(unionChildren.map(planLater)) :: Nil
       case g @ logical.Generate(generator, _, outer, _, _, child) =>

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

@@ -46,6 +46,26 @@ case class CollectLimitExec(limit: Int, child: SparkPlan) extends UnaryExecNode
   }
 }
 
+/**
+ * Take the first `limit` elements and collect them to a single partition.
+ *
+ * This operator will be used when a logical `Limit` operation is the final operator in an
+ * logical plan, which happens when the user is collecting results back to the driver.
+ */
+case class CollectLimitRangeExec(start: Int, end: Int, child: SparkPlan) extends UnaryExecNode {
+  override def output: Seq[Attribute] = child.output
+  override def outputPartitioning: Partitioning = SinglePartition
+  override def executeCollect(): Array[InternalRow] = child.executeTake(end)
+  private val serializer: Serializer = new UnsafeRowSerializer(child.output.size)
+  protected override def doExecute(): RDD[InternalRow] = {
+    val locallyLimited = child.execute().mapPartitionsInternal(_.take(end))
+    val shuffled = new ShuffledRowRDD(
+      ShuffleExchangeExec.prepareShuffleDependency(
+        locallyLimited, child.output, SinglePartition, serializer))
+    shuffled.mapPartitionsInternal(_.slice(start, end))
+  }
+}
+
 /**
  * Helper trait which defines methods that are shared by both
  * [[LocalLimitExec]] and [[GlobalLimitExec]].
@@ -114,6 +134,43 @@ case class GlobalLimitExec(limit: Int, child: SparkPlan) extends BaseLimitExec {
   override def outputOrdering: Seq[SortOrder] = child.outputOrdering
 }
 
+/**
+ * Take the first `limit` elements of the child's single output partition.
+ */
+case class RangeLimitExec(start: Int, limit: Int, child: SparkPlan) extends BaseLimitExec {
+
+  override def requiredChildDistribution: List[Distribution] = AllTuples :: Nil
+
+  override def outputPartitioning: Partitioning = child.outputPartitioning
+
+  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
+
+  override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {
+    val stopEarly =
+      ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "stopEarly") // init as stopEarly = false
+
+    ctx.addNewFunction("stopEarly", s"""
+      @Override
+      protected boolean stopEarly() {
+        return $stopEarly;
+      }
+    """, inlineToOuterClass = true)
+    val countTerm = ctx.addMutableState(CodeGenerator.JAVA_INT, "count") // init as count = 0
+    s"""
+       | $countTerm += 1;
+       | if ( $countTerm > $start && $countTerm <= $limit) {
+       |   ${consume(ctx, input)}
+       | } if($countTerm > $limit) {
+       |   $stopEarly = true;
+       | }
+     """.stripMargin
+  }
+
+  protected override def doExecute(): RDD[InternalRow] = child.execute().mapPartitions { iter =>
+    iter.slice(start, limit)
+  }
+}
+
 /**
  * Take the first limit elements as defined by the sortOrder, and do projection if needed.
  * This is logically equivalent to having a Limit operator after a [[SortExec]] operator,
@@ -176,3 +233,68 @@ case class TakeOrderedAndProjectExec(
     s"TakeOrderedAndProject(limit=$limit, orderBy=$orderByString, output=$outputString)"
   }
 }
+/**
+ * Take the first limit elements as defined by the sortOrder, and do projection if needed.
+ * This is logically equivalent to having a Limit operator after a [[SortExec]] operator,
+ * or having a [[ProjectExec]] operator between them.
+ * This could have been named TopK, but Spark's top operator does the opposite in ordering
+ * so we name it TakeOrdered to avoid confusion.
+ */
+case class TakeOrderedRangeAndProjectExec(
+  start: Int,
+  end: Int,
+  sortOrder: Seq[SortOrder],
+  projectList: Seq[NamedExpression],
+  child: SparkPlan) extends UnaryExecNode {
+
+  override def output: Seq[Attribute] = {
+    projectList.map(_.toAttribute)
+  }
+
+  override def executeCollect(): Array[InternalRow] = {
+    val ord = new LazilyGeneratedOrdering(sortOrder, child.output)
+    val data = child.execute().map(_.copy()).takeOrdered(end)(ord).drop(start)
+    if (projectList != child.output) {
+      val proj = UnsafeProjection.create(projectList, child.output)
+      data.map(r => proj(r).copy())
+    } else {
+      data
+    }
+  }
+
+  private val serializer: Serializer = new UnsafeRowSerializer(child.output.size)
+
+  protected override def doExecute(): RDD[InternalRow] = {
+    val ord = new LazilyGeneratedOrdering(sortOrder, child.output)
+    val localTopK: RDD[InternalRow] = {
+      child.execute().map(_.copy()).mapPartitions { iter =>
+        org.apache.spark.util.collection.Utils.takeOrdered(iter, end)(ord)
+      }
+    }
+    val shuffled = new ShuffledRowRDD(
+      ShuffleExchangeExec.prepareShuffleDependency(
+        localTopK, child.output, SinglePartition, serializer))
+    shuffled.mapPartitions { iter =>
+      val topK = org.apache.spark.util.collection.Utils.takeOrdered(iter.map(_.copy()), end)(ord)
+        .drop(start)
+      if (projectList != child.output) {
+        val proj = UnsafeProjection.create(projectList, child.output)
+        topK.map(r => proj(r))
+      } else {
+        topK
+      }
+    }
+  }
+
+  override def outputOrdering: Seq[SortOrder] = sortOrder
+
+  override def outputPartitioning: Partitioning = SinglePartition
+
+  override def simpleString: String = {
+    val orderByString = Utils.truncatedString(sortOrder, "[", ",", "]")
+    val outputString = Utils.truncatedString(output, "[", ",", "]")
+
+    s"TakeOrderedRangeAndProject" +
+      s"(start=$start, end=$end, orderBy=$orderByString, output=$outputString)"
+  }
+}