Skip to content

Commit 05ce603

Browse files
committed
fix: handle fetch's count in a way that matches roundtrip
1 parent ab104f5 commit 05ce603

File tree

2 files changed

+7
-4
lines changed

2 files changed

+7
-4
lines changed

spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ class ToLogicalPlan(spark: SparkSession) extends DefaultRelVisitor[LogicalPlan]
215215

216216
override def visit(fetch: relation.Fetch): LogicalPlan = {
217217
val child = fetch.getInput.accept(this)
218-
val limit = fetch.getCount.getAsLong.intValue()
218+
val limit = fetch.getCount.orElse(-1).intValue()
219219
val offset = fetch.getOffset.intValue()
220220
val toLiteral = (i: Int) => Literal(i, IntegerType)
221221
if (limit >= 0) {

spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,12 +206,15 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging {
206206
}
207207

208208
private def fetch(child: LogicalPlan, offset: Long, limit: Long = -1): relation.Fetch = {
209-
relation.Fetch
209+
val builder = relation.Fetch
210210
.builder()
211211
.input(visit(child))
212212
.offset(offset)
213-
.count(limit)
214-
.build()
213+
if (limit != -1) {
214+
builder.count(limit)
215+
}
216+
217+
builder.build()
215218
}
216219

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

0 commit comments

Comments
 (0)