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

Picked from main
branch(#13938)

PR #13903 (target branch
25.12) needs to use the CSV files in #13938

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>
Co-authored-by: Robert (Bobby) Evans <bobby@apache.org>

Author: res-life

integration_tests/src/main/python/join_test.py

@@ -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)

integration_tests/src/test/resources/t0.csv

@@ -0,0 +1 @@
+c0,c1

integration_tests/src/test/resources/t1.csv

@@ -0,0 +1,2 @@
+c0
+true

sql-plugin/src/main/scala/com/nvidia/spark/rapids/limit.scala

@@ -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._
@@ -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] {
@@ -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()) {
@@ -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:]
@@ -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)
     }
   }
 

tests/src/test/scala/com/nvidia/spark/rapids/LimitRetrySuite.scala

@@ -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.
@@ -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 {
@@ -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,