You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@gluten.apache.org by "ReemaAlzaid (via GitHub)" <gi...@apache.org> on 2026/06/23 19:47:27 UTC

[PR] [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer [gluten]

ReemaAlzaid opened a new pull request, #12345:
URL: https://github.com/apache/gluten/pull/12345

   <!--
   Thank you for submitting a pull request! Here are some tips:
   
   1. For first-time contributors, please read our contributing guide:
      https://github.com/apache/gluten/blob/main/CONTRIBUTING.md
   2. If necessary, create a GitHub issue for discussion beforehand to avoid duplicate work.
   3. If the PR is specific to a single backend, include [VL] or [CH] in the PR title to indicate the
      Velox or ClickHouse backend, respectively.
   4. If the PR is not ready for review, please mark it as a draft.
   -->
   
   ## 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
   
   <!--
   Provide a clear and concise description of the changes introduced in this PR.
   Ensure the PR description aligns with the code changes, especially after updates.
   If applicable, include "Fixes #<GitHub_Issue_ID>" to automatically close the corresponding issue
   when the PR is merged.
   -->
   
   ## 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`
   
   <!--
   Describe how the changes were tested, if applicable.
   Include new tests to validate the functionality, if necessary.
   For UI-related changes, attach screenshots to demonstrate the updates.
   -->


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org


Re: [PR] [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer [gluten]

Posted by "liuneng1994 (via GitHub)" <gi...@apache.org>.
liuneng1994 commented on PR #12345:
URL: https://github.com/apache/gluten/pull/12345#issuecomment-4912154955

   ArrowEvalPythonExecSuite:
   03:27:19.744 WARN org.apache.spark.sql.execution.GlutenFallbackReporter: Validation failed for plan: LocalTableScan[QueryId=5798], due to: Gluten does not touch it or does not support it
   
   - arrow batched python udf over parquet scan *** FAILED ***
     java.lang.NoSuchMethodException: org.apache.spark.sql.IntegratedUDFTestUtils$TestPythonUDF.<init>(java.lang.String,scala.Option)
     at java.base/java.lang.Class.getConstructor0(Class.java:3587)
     at java.base/java.lang.Class.getConstructor(Class.java:2273)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.newTestPythonUDF(ArrowEvalPythonExecSuite.scala:157)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.newTestArrowBatchedPythonUDF(ArrowEvalPythonExecSuite.scala:136)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.arrowBatchedTestUDFString$lzycompute(ArrowEvalPythonExecSuite.scala:44)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.arrowBatchedTestUDFString(ArrowEvalPythonExecSuite.scala:43)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.$anonfun$new$8(ArrowEvalPythonExecSuite.scala:124)
     at org.apache.gluten.execution.python.ArrowEvalPythonExecSuite.$anonfun$new$8$adapted(ArrowEvalPythonExecSuite.scala:120)
     at org.apache.spark.sql.catalyst.plans.SQLHelper.withTempPath(SQLHelper.scala:41)
     at org.apache.spark.sql.catalyst.plans.SQLHelper.withTempPath$(SQLHelper.scala:38)
     ... 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org


Re: [PR] [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer [gluten]

Posted by "ReemaAlzaid (via GitHub)" <gi...@apache.org>.
ReemaAlzaid commented on PR #12345:
URL: https://github.com/apache/gluten/pull/12345#issuecomment-4919612542

   Hey @liuneng1994 I have fixed the errors 
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org


Re: [PR] [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer [gluten]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #12345:
URL: https://github.com/apache/gluten/pull/12345#discussion_r3485524538


##########
backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala:
##########
@@ -162,7 +163,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()

Review Comment:
   `writeNextInputToStreamHelper` currently calls `ensureNextInputWriter(dataOut)` before checking `inputIterator.hasNext`. This means an Arrow IPC stream (schema/start) can be opened even for empty partitions, which differs from the previous behavior (and adds unnecessary bytes/overhead). Consider checking `hasNext` first, and only creating/starting the `ArrowStreamWriter` when there is at least one batch to write.



##########
backends-velox/src/test/scala/org/apache/gluten/execution/python/ArrowEvalPythonExecSuite.scala:
##########
@@ -36,6 +38,9 @@ class ArrowEvalPythonExecSuite extends WholeStageTransformerSuite {
     newTestScalarPandasUDF(name = "pyarrowUDF", returnType = Some(StringType))
   private val pyarrowTestUDFLong =
     newTestScalarPandasUDF(name = "pyarrowUDF", returnType = Some(LongType))
+  private val SQL_ARROW_BATCHED_UDF = 101

Review Comment:
   The eval type is hard-coded as `101`, which is a Spark internal constant (`PythonEvalType.SQL_ARROW_BATCHED_UDF`). Using the named constant makes the test more self-documenting and avoids coupling to a magic number.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org


Re: [PR] [GLUTEN-12280][VL] Fix Spark 4 Arrow Python UDF stream writer [gluten]

Posted by "liuneng1994 (via GitHub)" <gi...@apache.org>.
liuneng1994 merged PR #12345:
URL: https://github.com/apache/gluten/pull/12345


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org