beliefer commented on code in PR #36541:
URL: https://github.com/apache/spark/pull/36541#discussion_r873409826
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala:
##########
@@ -215,61 +211,52 @@ case class LocalLimitExec(limit: Int, child: SparkPlan)
extends BaseLimitExec {
}
/**
- * Take the first `limit` elements of the child's single output partition.
- */
-case class GlobalLimitExec(limit: Int, child: SparkPlan) extends BaseLimitExec
{
-
- override def requiredChildDistribution: List[Distribution] = AllTuples :: Nil
-
- override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
- copy(child = newChild)
-}
-
-/**
- * Skip the first `offset` elements then take the first `limit` of the
following elements in
- * the child's single output partition.
+ * Take the first `limit` elements and then drop the first `offset` elements
in the child's single
+ * output partition.
*/
-case class GlobalLimitAndOffsetExec(
- limit: Int = -1,
- offset: Int,
- child: SparkPlan) extends BaseLimitExec {
- assert(offset > 0)
+case class GlobalLimitExec(limit: Int = -1, child: SparkPlan, offset: Int = 0)
+ extends BaseLimitExec {
+ assert(limit >= 0 || (limit == -1 && offset > 0))
override def requiredChildDistribution: List[Distribution] = AllTuples :: Nil
- override def doExecute(): RDD[InternalRow] = if (limit >= 0) {
- child.execute().mapPartitionsInternal(iter => iter.slice(offset, limit +
offset))
- } else {
- child.execute().mapPartitionsInternal(iter => iter.drop(offset))
+ override def doExecute(): RDD[InternalRow] = {
+ if (offset > 0) {
+ if (limit >= 0) {
+ child.execute().mapPartitionsInternal(iter => iter.slice(offset,
limit))
+ } else {
+ child.execute().mapPartitionsInternal(iter => iter.drop(offset))
+ }
+ } else {
+ super.doExecute()
+ }
}
- private lazy val skipTerm = BaseLimitExec.newLimitCountTerm()
-
override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row:
ExprCode): String = {
- ctx.addMutableState(
- CodeGenerator.JAVA_INT, skipTerm, forceInline = true, useFreshName =
false)
- if (limit >= 0) {
- // The counter name is already obtained by the upstream operators via
`limitNotReachedChecks`.
- // Here we have to inline it to not change its name. This is fine as we
won't have many limit
- // operators in one query.
- ctx.addMutableState(
- CodeGenerator.JAVA_INT, countTerm, forceInline = true, useFreshName =
false)
- s"""
- | if ($skipTerm < $offset) {
- | $skipTerm += 1;
- | } else if ($countTerm < $limit) {
- | $countTerm += 1;
- | ${consume(ctx, input)}
- | }
+ if (offset > 0) {
+ val skipTerm = ctx.addMutableState(CodeGenerator.JAVA_INT,
"rowsSkipped", forceInline = true)
+ if (limit > 0) {
+ // In codegen, we skip the first `offset` rows, then take the first
`limit - offset` rows.
+ val finalLimit = limit - offset
+ s"""
+ | if ($skipTerm < $offset) {
+ | $skipTerm += 1;
+ | } else if ($countTerm < $finalLimit) {
+ | $countTerm += 1;
+ | ${consume(ctx, input)}
+ | }
""".stripMargin
- } else {
- s"""
- | if ($skipTerm < $offset) {
- | $skipTerm += 1;
- | } else {
- | ${consume(ctx, input)}
- | }
+ } else {
+ s"""
+ | if ($skipTerm < $offset) {
+ | $skipTerm += 1;
+ | } else {
+ | ${consume(ctx, input)}
+ | }
""".stripMargin
Review Comment:
```suggestion
""".stripMargin
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
