This is an automated email from the ASF dual-hosted git repository.
liuneng1994 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 85d3ca109e [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream
writer (#12345)
85d3ca109e is described below
commit 85d3ca109ebe38168b2a2e0503490f554df3ed3e
Author: Reema <[email protected]>
AuthorDate: Thu Jul 9 10:02:05 2026 +0000
[GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer (#12345)
What changes are proposed in this pull request?
Fixes #12280.
Fix Spark 4 Arrow Python UDF execution with the Velox backend by keeping
the Arrow stream writer alive across input batches instead of reopening the IPC
stream per batch.
Also adds a regression test for Arrow Python UDF over Parquet scan
How was this patch tested?
Added ArrowEvalPythonExecSuite coverage.
Verified locally on Spark 4.0.2 / Scala 2.13 / linux aarch64. The repro
uses ColumnarArrowPythonRunner, returns max(ship_len) = 7, and no longer fails
with Invalid IPC stream
---
.../api/python/ColumnarArrowEvalPythonExec.scala | 77 +++++++++++++++++++++-
.../python/ArrowEvalPythonExecSuite.scala | 71 ++++++++++++++++++++
2 files changed, 147 insertions(+), 1 deletion(-)
diff --git
a/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
b/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
index 6bb2f50e1f..831259b012 100644
---
a/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
@@ -46,6 +46,7 @@ import java.util.concurrent.atomic.AtomicBoolean
import scala.annotation.nowarn
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
+import scala.util.control.NonFatal
class ColumnarArrowPythonRunner(
funcs: Seq[(ChainedPythonFunctions, Long)],
@@ -59,6 +60,9 @@ class ColumnarArrowPythonRunner(
override val simplifiedTraceback: Boolean =
SQLConf.get.pysparkSimplifiedTraceback
+ // Keep this source compatible with older Spark profiles where
PythonEvalType differs.
+ private val SQL_ARROW_BATCHED_UDF = 101
+
override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
require(
bufferSize >= 4,
@@ -150,6 +154,9 @@ class ColumnarArrowPythonRunner(
PythonRDD.writeUTF(k, dataOut)
PythonRDD.writeUTF(v, dataOut)
}
+ if (SparkVersionUtil.gteSpark41 && evalType == SQL_ARROW_BATCHED_UDF) {
+ PythonRDD.writeUTF(schema.json, dataOut)
+ }
ColumnarArrowPythonRunner.this.writeUdf(dataOut, argMetas)
}
@@ -162,7 +169,75 @@ class ColumnarArrowPythonRunner(
// For Spark 4.0. It overrides the corresponding abstract method in
Writer class.
// We omitted the override keyword for compatibility consideration.
def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
- writeToStreamHelper(dataOut)
+ writeNextInputToStreamHelper(dataOut)
+ }
+
+ private var nextInputRoot: VectorSchemaRoot = _
+ private var nextInputLoader: VectorLoader = _
+ private var nextInputWriter: ArrowStreamWriter = _
+ private var nextInputWriterClosed = false
+
+ context.addTaskCompletionListener[Unit] {
+ _ =>
+ try {
+ closeNextInputWriter()
+ } catch {
+ case NonFatal(_) =>
+ }
+ }
+
+ private def ensureNextInputWriter(dataOut: DataOutputStream): Unit = {
+ if (nextInputWriter == null) {
+ val arrowSchema = SparkSchemaUtil.toArrowSchema(schema, timeZoneId)
+ val allocator = ArrowBufferAllocators.contextInstance()
+ nextInputRoot = VectorSchemaRoot.create(arrowSchema, allocator)
+ nextInputLoader = new VectorLoader(nextInputRoot)
+ nextInputWriter = new ArrowStreamWriter(nextInputRoot, null, dataOut)
+ nextInputWriter.start()
+ }
+ }
+
+ private def closeNextInputWriter(): Unit = {
+ if (!nextInputWriterClosed && nextInputRoot != null) {
+ try {
+ if (nextInputWriter != null) {
+ nextInputWriter.end()
+ }
+ } finally {
+ nextInputRoot.close()
+ nextInputWriterClosed = true
+ }
+ }
+ }
+
+ private def writeNextInputToStreamHelper(dataOut: DataOutputStream):
Boolean = {
+ ensureNextInputWriter(dataOut)
+ if (!inputIterator.hasNext) {
+ closeNextInputWriter()
+ // See https://issues.apache.org/jira/browse/SPARK-44705:
+ // Starting from Spark 4.0, we should return false once the iterator
is drained out,
+ // otherwise Spark won't stop calling this method repeatedly.
+ return false
+ }
+ val nextBatch = inputIterator.next()
+ val cols = (0 until nextBatch.numCols).toList.map(
+ i =>
+ nextBatch
+ .asInstanceOf[ColumnarBatch]
+ .column(i)
+ .asInstanceOf[ArrowWritableColumnVector]
+ .getValueVector)
+ val nextRecordBatch =
+ SparkVectorUtil.toArrowRecordBatch(nextBatch.numRows, cols)
+ try {
+ nextInputLoader.load(nextRecordBatch)
+ nextInputWriter.writeBatch()
+ true
+ } finally {
+ if (nextRecordBatch != null) {
+ nextRecordBatch.close()
+ }
+ }
}
def writeToStreamHelper(dataOut: DataOutputStream): Boolean = {
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala
index 7747bd2193..c6a6ac186e 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala
@@ -21,6 +21,8 @@ import org.apache.gluten.execution.WholeStageTransformerSuite
import org.apache.spark.SparkConf
import org.apache.spark.api.python.ColumnarArrowEvalPythonExec
import org.apache.spark.sql.IntegratedUDFTestUtils
+import org.apache.spark.sql.execution.python.UserDefinedPythonFunction
+import org.apache.spark.sql.functions.max
import org.apache.spark.sql.types.{DataType, LongType, StringType}
import org.apache.spark.util.SparkVersionUtil
@@ -36,6 +38,11 @@ class ArrowEvalPythonExecSuite extends
WholeStageTransformerSuite {
newTestScalarPandasUDF(name = "pyarrowUDF", returnType = Some(StringType))
private val pyarrowTestUDFLong =
newTestScalarPandasUDF(name = "pyarrowUDF", returnType = Some(LongType))
+ // Spark exposes these values through a package-private PythonEvalType
object.
+ private val SQL_BATCHED_UDF = 100
+ private val SQL_ARROW_BATCHED_UDF = 101
+ private lazy val arrowBatchedTestUDFString =
+ newTestArrowBatchedPythonUDF(name = "arrowBatchedUDF", returnType =
Some(StringType))
override def sparkConf: SparkConf = {
super.sparkConf
@@ -109,6 +116,70 @@ class ArrowEvalPythonExecSuite extends
WholeStageTransformerSuite {
checkAnswer(df, expected)
}
+ testWithMinSparkVersion("arrow batched python udf over parquet scan", "4.0")
{
+ withTempPath {
+ f =>
+ Seq(("MAIL", 1), ("RAIL", 2), ("SHIP", 3)).toDF("shipmode",
"id").write.parquet(
+ f.getCanonicalPath)
+ val base = spark.read.parquet(f.getCanonicalPath)
+ val arrowUdfCol =
arrowBatchedTestUDFString(base("shipmode")).as("shipmode_arrow")
+ val df =
base.select(arrowUdfCol).agg(max("shipmode_arrow").as("max_shipmode"))
+ val expected = Seq(Tuple1("SHIP")).toDF("max_shipmode")
+
+ checkAnswer(df, expected)
+ checkSparkPlan[ColumnarArrowEvalPythonExec](df)
+ }
+ }
+
+ private def newTestArrowBatchedPythonUDF(
+ name: String,
+ returnType: Option[DataType] = None): UserDefinedPythonFunction = {
+ val pythonUDF =
+ newTestPythonUDF(name, returnType, Some(SQL_ARROW_BATCHED_UDF))
+ val udf = pythonUDF.getClass
+ .getMethod("udf")
+ .invoke(pythonUDF)
+ .asInstanceOf[UserDefinedPythonFunction]
+ if (SparkVersionUtil.gteSpark41) {
+ udf
+ } else {
+ udf.copy(
+ udf.name,
+ udf.func,
+ udf.dataType,
+ SQL_ARROW_BATCHED_UDF,
+ udf.udfDeterministic)
+ }
+ }
+
+ private def newTestPythonUDF(
+ name: String,
+ returnType: Option[DataType] = None,
+ pythonEvalType: Option[Int] = None): TestPythonUDF = {
+ if (SparkVersionUtil.gteSpark41) {
+ classOf[TestPythonUDF]
+ .getConstructor(
+ classOf[String],
+ classOf[Option[DataType]],
+ Integer.TYPE,
+ java.lang.Boolean.TYPE)
+ .newInstance(
+ name,
+ returnType,
+ pythonEvalType.getOrElse(SQL_BATCHED_UDF).asInstanceOf[Integer],
+ java.lang.Boolean.TRUE)
+ } else if (SparkVersionUtil.gteSpark40) {
+ // After https://github.com/apache/spark/pull/42864 which landed in
Spark 4.0, the return
+ // type of the UDF must be explicitly specified when creating the UDF
instance with column
+ // expressions as parameter.
+ classOf[TestPythonUDF]
+ .getConstructor(classOf[String], classOf[Option[DataType]])
+ .newInstance(name, returnType)
+ } else {
+ TestPythonUDF(name)
+ }
+ }
+
private def newTestScalarPandasUDF(
name: String,
returnType: Option[DataType] = None): TestScalarPandasUDF = {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
