sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3993592945
##########
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:
Fixed in
[279bbe303](https://github.com/apache/datafusion-comet/commit/279bbe303c06c3833a0a899f542847d37ead5321).
The writer now retains the current source batch and remaining ranges, emits
one slice per `writeNextInputToStream` call, and returns control to Spark's
transport loop. It does not call either upstream iterator's `hasNext`/`next`
until the final slice finishes, since `CometExecIterator.hasNext` can close the
previous batch and reuse its buffers. Temporary slices/decoded vectors remain
scoped to each write; upstream retains source ownership on completion,
interruption, and failure.
Added regressions using the production writer and Spark's actual
`DirectByteBufferOutputStream`, with Spark's threshold-based fill/drain loop.
They cover both large slices and multiple small slices per drain, round-trip
the complete IPC stream, check metrics and source lifetime, and exercise
interruption and serialization failure with pending slices. Against the
original runner, the two transport tests fail with 16,832,728 and 139,416
pending bytes; both pass with this fix. The tests assert pending-byte bounds
directly rather than relying on `MaxDirectMemorySize` enforcement.
All 48 focused JVM tests pass on Spark 4.1.3. All 133 general PyArrow tests
and 6 dictionary-shuffle tests also pass against the rebuilt packaged JAR with
real Python workers. Hosted CI for the new head is pending.
##########
.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:
Fixed in
[279bbe303](https://github.com/apache/datafusion-comet/commit/279bbe303c06c3833a0a899f542847d37ead5321).
Replaced the two individual Scala vector paths with
`spark/src/main/scala/org/apache/comet/vector/**` in the shared path-filter
anchor, so helper-only changes trigger both push and PR runs. Confirmed that
the filter matches `CometVectorUtils.scala` as well as the previously covered
files; Actionlint and YAML formatting pass.
--
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]
