Cover Spark 4.2 Arrow Python UDF columnar input [fast-ut] [reduced-it][databricks]

integration_tests/src/main/python/datasourcev2_read_test.py

@@ -13,10 +13,16 @@
 # limitations under the License.
 
 import pytest
+import pyspark.sql.functions as f
+from pyspark.sql.types import IntegerType
 
-from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_and_cpu_row_counts_equal
+from asserts import (
+    assert_cpu_and_gpu_are_equal_collect_with_capture,
+    assert_gpu_and_cpu_are_equal_collect,
+    assert_gpu_and_cpu_row_counts_equal)
 from data_gen import non_utc_allow, copy_and_update
 from marks import *
+from spark_session import is_spark_420_or_later
 
 columnarClass = 'com.nvidia.spark.rapids.tests.datasourcev2.parquet.ArrowColumnarDataSourceV2'
 
@@ -65,3 +71,65 @@ def test_read_arrow_off():
     assert_gpu_and_cpu_are_equal_collect(
         readTable("int,bool,byte,short,long,string,float,double,date,timestamp", columnarClass),
             conf=conf)
+
+
+arrow_udf_conf = copy_and_update(aqe_disabled, {
+    'spark.sql.execution.arrow.pyspark.enabled': 'true',
+})
+
+
+def _arrow_int_df(spark):
+    return spark.read.option("arrowTypes", "int").format(columnarClass).load()
+
+
+@allow_non_gpu('BatchScanExec')
+def test_arrow_source_pandas_udf():
+    pytest.importorskip('pandas')
+    pytest.importorskip('pyarrow')
+
+    def add_one(a):
+        return a + 1
+
+    my_udf = f.pandas_udf(add_one, returnType=IntegerType())
+
+    def do_it(spark):
+        return _arrow_int_df(spark).select(
+            f.col('col1'), my_udf(f.col('col1')).alias('u')).orderBy('col1')
+
+    # Spark 4.2 SPARK-56350 can skip ColumnarToRow for Arrow-backed CPU input.
+    # That CPU path is what required the test Arrow source to keep batches alive.
+    # The GPU path does not use Spark's Arrow pass-through; it still ingests
+    # Arrow via HostColumnarToGpu and evaluates the UDF with GpuArrowEvalPythonExec.
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        do_it,
+        exist_classes='HostColumnarToGpu,GpuArrowEvalPythonExec',
+        conf=arrow_udf_conf)
+
+
+@allow_non_gpu('BatchScanExec', 'PythonUDF')
+@pytest.mark.skipif(not is_spark_420_or_later(),
+                    reason='Arrow-optimized regular Python UDFs use ArrowEvalPythonExec from Spark 4.2')
+def test_arrow_source_regular_udf():
+    pytest.importorskip('pandas')
+    pytest.importorskip('pyarrow')
+
+    def add_one(a):
+        return a + 1
+
+    my_udf = f.udf(add_one, returnType=IntegerType())
+
+    def do_it(spark):
+        return _arrow_int_df(spark).select(
+            f.col('col1'), my_udf(f.col('col1')).alias('u')).orderBy('col1')
+
+    # SPARK-58241 is a Spark CPU bug on 4.2.0: evalType=101 hangs when Arrow
+    # columnar input is enabled. It is not a GPU bug. Disable that CPU path
+    # so this test compares GpuArrowEvalPythonExec against a stable row-based
+    # CPU ArrowEvalPythonExec baseline.
+    conf = copy_and_update(arrow_udf_conf, {
+        'spark.sql.execution.arrow.pythonUDF.columnarInput.enabled': 'false',
+    })
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        do_it,
+        exist_classes='HostColumnarToGpu,GpuArrowEvalPythonExec',
+        conf=conf)

integration_tests/src/main/python/udf_test.py

@@ -18,7 +18,8 @@
 from conftest import is_at_least_precommit_run, is_databricks_runtime
 from spark_session import (is_before_spark_331, is_before_spark_350,
                            is_spark_400_or_later,
-                           is_spark_411_or_later, is_spark_420_or_later)
+                           is_spark_411_or_later, is_spark_420_or_later,
+                           with_cpu_session)
 
 from pyspark.sql.pandas.utils import require_minimum_pyarrow_version, require_minimum_pandas_version
 
@@ -587,3 +588,31 @@ def test_func(spark):
         return df.groupby("id").applyInPandas(sum_udf, schema="v long")
 
     assert_gpu_and_cpu_are_equal_collect(test_func, conf=arrow_udf_conf_unsafe)
+
+
+# SPARK-56350 lets Spark 4.2 CPU ArrowEvalPythonExec consume columnar batches.
+# If the plugin leaves that CPU exec above a GPU scan, a host transition is
+# required so GpuColumnVector is not handed to Spark's CPU Arrow UDF path.
+@allow_non_gpu('ArrowEvalPythonExec', 'PythonUDF')
+def test_pandas_udf_cpu_arrow_eval_after_gpu_scan(spark_tmp_path):
+    data_path = spark_tmp_path + '/PARQUET_DATA'
+
+    def add_one(a):
+        return a + 1
+
+    my_udf = f.pandas_udf(add_one, returnType=IntegerType())
+    # Spark 4.x converts pandas UDF output with Arrow's safe checker by default.
+    # int_gen includes Integer.MAX_VALUE, and pandas a+1 becomes float64 2^31
+    # which cannot be cast back to int32. Keep values that stay in int32 after +1.
+    plus_one_int_gen = IntegerGen(min_val=-1000, max_val=1000, special_cases=[0, 1, -1])
+    with_cpu_session(
+        lambda spark: unary_op_df(spark, plus_one_int_gen, length=200).write.parquet(data_path))
+
+    conf = copy_and_update(arrow_udf_conf, {
+        'spark.rapids.sql.exec.ArrowEvalPythonExec': 'false',
+    })
+    assert_gpu_fallback_collect(
+        lambda spark: spark.read.parquet(data_path).select(
+            f.col('a'), my_udf(f.col('a')).alias('u')),
+        'ArrowEvalPythonExec',
+        conf=conf)

integration_tests/src/main/scala/com/nvidia/spark/rapids/tests/datasourcev2/parquet/TestingV2Source.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,13 +19,14 @@ import java.util
 
 import scala.collection.JavaConverters._
 
-import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
 import org.apache.arrow.vector._
 import org.apache.arrow.vector.complex.MapVector
 import org.apache.arrow.vector.types.{DateUnit, FloatingPointPrecision, TimeUnit}
 import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
 import org.apache.arrow.vector.util.Text;
 
+import org.apache.spark.TaskContext
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.catalog._
 import org.apache.spark.sql.connector.catalog.TableCapability.BATCH_READ
@@ -147,10 +148,33 @@ object ColumnarReaderFactory extends PartitionReaderFactory {
   override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = {
     val ArrowInputPartition(dataTypes, numRows, startNum) = partition
     new PartitionReader[ColumnarBatch] {
+      // Spark 4.2 DataSourceRDD closes this reader as soon as next() returns
+      // false, while CPU ArrowEvalPythonExec (SPARK-56350) may still hold
+      // pass-through ArrowColumnVector refs from get(). Do not close batches
+      // here: Spark/GPU consumers close those ColumnarBatches. Close the
+      // allocator only once nothing is allocated. See SPARK-56350 /
+      // cudf-spark#15663.
+      private val rootAllocator = new RootAllocator(Long.MaxValue)
+      private val allocator: BufferAllocator =
+        rootAllocator.newChildAllocator(s"arrow-test-reader-$startNum", 0, Long.MaxValue)
+      private val allBatches = new util.ArrayList[ColumnarBatch]()
       private var batch: ColumnarBatch = _
+      private var allocatorsClosed = false
+
+      Option(TaskContext.get()).foreach { ctx =>
+        ctx.addTaskCompletionListener[Unit](_ => tryCloseAllocators())
+      }
 
       private var current = 0
 
+      private def tryCloseAllocators(): Unit = {
+        if (!allocatorsClosed && allocator.getAllocatedMemory == 0) {
+          allocator.close()
+          rootAllocator.close()
+          allocatorsClosed = true
+        }
+      }
+
       override def next(): Boolean = {
         val batchSize = if (current < numRows) {
           if (current + BATCH_SIZE > numRows) {
@@ -167,22 +191,25 @@ object ColumnarReaderFactory extends PartitionReaderFactory {
         } else {
           var dtypeNum = 0
           val vecs = dataTypes.map { dtype =>
-            val vector = setupArrowVector(s"v$current$dtypeNum", dtype)
+            val vector = setupArrowVector(allocator, s"v$current$dtypeNum", dtype)
             val startVal = current + startNum * (dtypeNum + 2)
-            fillArrowVec(dtype, vector, startVal, numRows)
+            fillArrowVec(dtype, vector, startVal, batchSize)
             dtypeNum += 1
             new ArrowColumnVector(vector)
           }
           batch = new ColumnarBatch(vecs.toArray)
           batch.setNumRows(batchSize)
+          allBatches.add(batch)
           current += batchSize
           true
         }
       }
 
       override def get(): ColumnarBatch = batch
 
-      override def close(): Unit = batch.close()
+      override def close(): Unit = {
+        tryCloseAllocators()
+      }
     }
   }
 
@@ -323,12 +350,11 @@ object ColumnarReaderFactory extends PartitionReaderFactory {
       throw new UnsupportedOperationException(s"Unsupported data type: ${dt.catalogString}")
   }
 
-  private def setupArrowVector(name: String, dataType: DataType): ValueVector = {
-    val rootAllocator = new RootAllocator(Long.MaxValue)
-    val allocator = rootAllocator.newChildAllocator(s"$name", 0, Long.MaxValue)
-    val vector = toArrowField(s"field$name", dataType, nullable = true, "Utc")
-      .createVector(allocator)
-    vector
+  private def setupArrowVector(
+      allocator: BufferAllocator,
+      name: String,
+      dataType: DataType): ValueVector = {
+    toArrowField(s"field$name", dataType, nullable = true, "Utc").createVector(allocator)
   }
 }