Skip to content

Commit

Permalink
fix: handle fetch's count in a way that matches roundtrip
Browse files Browse the repository at this point in the history
  • Loading branch information
Blizzara committed Mar 5, 2025
1 parent 123e409 commit eb371c8
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ class ToLogicalPlan(spark: SparkSession) extends DefaultRelVisitor[LogicalPlan]

override def visit(fetch: relation.Fetch): LogicalPlan = {
val child = fetch.getInput.accept(this)
val limit = fetch.getCount.getAsLong.intValue()
val limit = fetch.getCount.orElse(-1).intValue()
val offset = fetch.getOffset.intValue()
val toLiteral = (i: Int) => Literal(i, IntegerType)
if (limit >= 0) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -209,12 +209,15 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging {
}

private def fetch(child: LogicalPlan, offset: Long, limit: Long = -1): relation.Fetch = {
relation.Fetch
val builder = relation.Fetch
.builder()
.input(visit(child))
.offset(offset)
.count(limit)
.build()
if (limit != -1) {
builder.count(limit)
}

builder.build()
}

override def visitGlobalLimit(p: GlobalLimit): relation.Rel = {
Expand Down

0 comments on commit eb371c8

Please sign in to comment.