viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3994139506
##########
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:
Confirmed fixed in `279bbe303`. Returning after each slice gives Spark's
transport loop an opportunity to drain, and retaining the source without
calling upstream `hasNext` or `next` preserves its buffers until the final
slice is serialized.
I ran both new transport regressions successfully. Against the previous
runner, they fail with 16,832,728 and 139,416 pending bytes, respectively. The
interruption and serialization-failure tests also pass. This addresses my
blocker.
##########
.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:
Confirmed. The Scala vector directory pattern now includes
`CometVectorUtils.scala` and preserves coverage of the previously listed files.
Since both push and pull-request triggers use the shared anchor, helper-only
changes are covered.
--
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]
