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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
