Skip to content

Fix a special case in limit where it could return an empty batch with the wrong number of columns (#13938) #13971

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
res-life merged 1 commit into from
Dec 8, 2025
Merged

Fix a special case in limit where it could return an empty batch with the wrong number of columns (#13938) #13971

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions integration_tests/src/main/python/join_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -476,6 +476,15 @@ def do_join(spark):
conf = {kudo_enabled_conf_key: kudo_enabled}
assert_gpu_and_cpu_are_equal_collect(do_join, conf = conf)

# because this infers the schema for CSV we need to allow some ops to be on the CPU
@allow_non_gpu("CollectLimitExec", "FileSourceScanExec", "DeserializeToObjectExec")
def test_empty_cross_side_with_limit(std_input_path):
def do_join(spark):
t0 = spark.read.csv(std_input_path + '/t0.csv', header=True, inferSchema=True)
t1 = spark.read.csv(std_input_path + '/t1.csv', header=True, inferSchema=True)
return t0.crossJoin(t1).limit(21)
assert_gpu_and_cpu_are_equal_collect(do_join)

# local sort because of https://github.com/NVIDIA/spark-rapids/issues/84
# After 3.1.0 is the min spark version we can drop this
@ignore_order(local=True)
Expand Down
1 change: 1 addition & 0 deletions integration_tests/src/test/resources/t0.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
c0,c1
2 changes: 2 additions & 0 deletions integration_tests/src/test/resources/t1.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
c0
true
13 changes: 8 additions & 5 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/limit.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@

package com.nvidia.spark.rapids

import scala.collection.mutable.ArrayBuffer

import ai.rapids.cudf.Table
import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
import com.nvidia.spark.rapids.GpuMetric._
Expand All @@ -35,12 +33,14 @@ import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, Distribution, Pa
import org.apache.spark.sql.catalyst.util.truncatedString
import org.apache.spark.sql.execution.{CollectLimitExec, LimitExec, SparkPlan, TakeOrderedAndProjectExec}
import org.apache.spark.sql.execution.exchange.ENSURE_REQUIREMENTS
import org.apache.spark.sql.types.{DataType, StructField, StructType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}

class GpuBaseLimitIterator(
input: Iterator[ColumnarBatch],
limit: Int,
offset: Int,
dataTypes : Array[DataType],
opTime: GpuMetric,
numOutputBatches: GpuMetric,
numOutputRows: GpuMetric) extends Iterator[ColumnarBatch] {
Expand All @@ -55,7 +55,6 @@ class GpuBaseLimitIterator(
}

var batch = input.next()
val numCols = batch.numCols()

// In each partition, we need to skip `offset` rows
while (batch != null && remainingOffset >= batch.numRows()) {
Expand All @@ -71,7 +70,10 @@ class GpuBaseLimitIterator(
// If the last batch is null, then we have offset >= numRows in this partition.
// In such case, we should return an empty batch
if (batch == null || batch.numRows() == 0) {
return new ColumnarBatch(new ArrayBuffer[GpuColumnVector](numCols).toArray, 0)
val fields = dataTypes.zipWithIndex.map {
case (dt, idx) => StructField(s"_col$idx", dt, nullable = true)
}
return GpuColumnVector.emptyBatch(StructType(fields))
}

// Here 0 <= remainingOffset < batch.numRow(), we need to get batch[remainingOffset:]
Expand Down Expand Up @@ -156,7 +158,8 @@ trait GpuBaseLimitExec extends LimitExec with GpuExec with ShimUnaryExecNode {
val numOutputRows = gpuLongMetric(NUM_OUTPUT_ROWS)
val numOutputBatches = gpuLongMetric(NUM_OUTPUT_BATCHES)
rdd.mapPartitions { iter =>
new GpuBaseLimitIterator(iter, limit, offset, opTime, numOutputBatches, numOutputRows)
new GpuBaseLimitIterator(iter, limit, offset, output.map(_.dataType).toArray,
opTime, numOutputBatches, numOutputRows)
}
}

Expand Down
7 changes: 4 additions & 3 deletions tests/src/test/scala/com/nvidia/spark/rapids/LimitRetrySuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
* Copyright (c) 2023-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -21,7 +21,7 @@ import com.nvidia.spark.rapids.Arm.withResource
import com.nvidia.spark.rapids.jni.RmmSpark

import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, ExprId, SortOrder}
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.{DataType, IntegerType}
import org.apache.spark.sql.vectorized.ColumnarBatch

class LimitRetrySuite extends RmmSparkRetrySuiteBase {
Expand Down Expand Up @@ -79,7 +79,8 @@ class LimitRetrySuite extends RmmSparkRetrySuiteBase {
val limitIter = new GpuBaseLimitIterator(
// 3 batches as input, and each has 8 rows
(0 until totalRows).grouped(8).map(buildBatch(_)).toList.toIterator,
limit, offset, NoopMetric, NoopMetric, NoopMetric)
limit, offset, Array[DataType](IntegerType),
NoopMetric, NoopMetric, NoopMetric)
var leftRows = if (limit > totalRows) totalRows - offset else limit - offset
var curValue = offset
RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1,
Expand Down