viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991579134
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -28,23 +28,29 @@ on:
paths: &feature-paths
- "pom.xml"
- "common/pom.xml"
- - "common/src/main/scala/org/apache/comet/CometConf.scala"
+ - "native/shuffle/src/spark_unsafe/row.rs"
- "spark/pom.xml"
+ - "spark/src/main/java/org/apache/comet/vector/**"
+ -
"spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java"
+ - "spark/src/main/scala/org/apache/comet/CometConf.scala"
-
"spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/StreamReader.scala"
Review Comment:
**[P3] Include the new shared decoding helper in the workflow paths**
Now that the runner delegates decoding and cleanup to
`spark/src/main/scala/org/apache/comet/vector/CometVectorUtils.scala`, could we
add that file to this list, or cover the Scala vector directory?
The Java vector glob above does not match it, and the Scala entries
currently list only `NativeUtil.scala` and `StreamReader.scala`. A future
change confined to the shared decoding helper would therefore skip these Python
worker regressions. This PR itself still triggers them through its other
changed files.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -178,44 +182,49 @@ private[python] trait CometArrowPythonRunnerBase
val cometBatch = currentGroup.next()
val startData = dataOut.size()
- val sourceVectors = (0 until cometBatch.numCols()).map { i =>
- cometBatch
- .column(i)
- .asInstanceOf[CometDecodedVector]
- .getValueVector
- .asInstanceOf[FieldVector]
+ val columns = (0 until cometBatch.numCols()).map { i =>
+ cometBatch.column(i).asInstanceOf[CometDecodedVector]
}
- val batchFields = sourceVectors.map(_.getField)
-
- if (arrowWriter == null) {
- // Build the schema-only struct root once from the first batch's
child fields.
- // mapInArrow/mapInPandas exchange the columns under a single
non-nullable struct.
- // Comet's FFI-imported vectors leave the Arrow Field name null, so
restore the real
- // column names from the input schema (the worker reads columns by
name, and shaded
- // Arrow rejects a null field name). Keep the field types and child
structure as-is so
- // the advertised schema matches the source buffers. Keeping the
type as-is also means
- // a TimestampType reaches the worker with Comet's UTC time zone
- // rather than the session zone vanilla Spark would label it with;
this is a documented
- // limitation (see pyarrow-udfs.md), not a value difference, since
the stored instant is
- // identical.
- val childNames = inputStructType.fieldNames
- streamFields = batchFields.zipWithIndex.map { case (field, i) =>
- renamed(field, childNames(i), forceNullable = true)
+ CometArrowPythonRunnerBase.foreachInputBatch(
+ columns,
+ cometBatch.numRows(),
+ arrowMaxRecordsPerBatch,
+ arrowMaxBytesPerBatch,
+ allocator) { (sourceVectors, numRows) =>
Review Comment:
**[P1] Return to Spark between slices so the transport buffer can drain**
`foreachInputBatch` serializes every slice before this
`writeNextInputToStream` call returns. In Spark 4.1, `dataOut` wraps a
`DirectByteBufferOutputStream`;
`ReaderInputStream.writeAdditionalInputToPythonWorker` only writes that buffer
to the worker socket after this call returns ([Spark
source](https://github.com/apache/spark/blob/v4.1.3/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala#L861)).
Consequently, closing each decoded vector does not release the accumulated IPC
bytes.
I reproduced this with 256 rows referencing one 64 KiB dictionary value and
`maxRecordsPerBatch=1`. The source buffers occupied 83,968 bytes and the
decoding allocator peaked at 114,688 bytes, but the transport buffer
accumulated 16,832,512 bytes. With `-XX:MaxDirectMemorySize=8m`, this failed in
`DirectByteBufferOutputStream.grow`; a control that drained/reset the buffer
between slices passed, with at most 65,752 pending IPC bytes.
Could we retain the current source batch and remaining ranges in the writer,
emit one slice per `writeNextInputToStream` call, and return so Spark can drain
the transport buffer? The source must remain alive until all its slices have
been consumed. Calling `flush()` on this stream would not drain it to the
socket.
Please also add a regression covering transport buffering. The existing
allocator-capped test proves that decoding is bounded, but does not detect
accumulation outside the Arrow allocator.
--
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]
