sunchao commented on code in PR #5368:
URL: https://github.com/apache/datafusion-comet/pull/5368#discussion_r3816957727
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -107,9 +109,8 @@ private[python] trait CometArrowPythonRunnerBase
private val allocator =
CometArrowAllocator.newChildAllocator(s"stdout writer for
$pythonExec", 0, Long.MaxValue)
private var currentGroup: Iterator[ColumnarBatch] = _
- private var arrowWriter: ArrowStreamWriter = _
+ private var arrowWriter:
CometArrowPythonRunnerBase.DirectArrowStreamWriter = _
Review Comment:
Updated in 64a0114a7. The runner now uses a standard `ArrowStreamWriter` for
the stream schema and end marker, and `DirectArrowStreamWriter` has been
removed.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -167,49 +170,37 @@ 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]
+ }
if (arrowWriter == null) {
- // Build the destination struct root once, sized to the first
batch's child fields.
+ // 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). The field types and child
structure are kept as-is
- // so copyVector still walks the source and destination trees in
lockstep. Keeping the
- // type as-is also means a TimestampType reaches the worker with
Comet's UTC time zone
+ // 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
- val childFields = (0 until cometBatch.numCols()).map { i =>
- val vecField =
-
cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
- renamed(vecField, childNames(i), forceNullable = true)
+ val childFields = sourceVectors.zipWithIndex.map { case (vector, i)
=>
+ renamed(vector.getField, childNames(i), forceNullable = true)
}
startWriter(childFields, dataOut)
}
- var i = 0
- while (i < cometBatch.numCols()) {
- val src = cometBatch
- .column(i)
- .asInstanceOf[CometDecodedVector]
- .getValueVector
- .asInstanceOf[FieldVector]
- val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
- copyVector(src, dst)
- i += 1
- }
- val numRows = cometBatch.numRows()
- structVec.setValueCount(numRows)
- // Mark every row of the struct non-null (all-1 validity). The
validity buffer is freshly
- // allocated and zero-initialised, so without this Python would see an
all-null struct.
- val validityBytes = (numRows + 7) / 8
- Platform.setMemory(
- structVec.getValidityBuffer.memoryAddress(),
- 0xff.toByte,
- validityBytes)
- writeRoot.setRowCount(numRows)
- arrowWriter.writeBatch()
+ CometArrowPythonRunnerBase.writeDirectBatch(
+ arrowWriter,
+ sourceVectors,
+ cometBatch.numRows(),
+ allocator)
Review Comment:
Updated in 64a0114a7. Each batch now gets a fresh `WriteChannel` over
`dataOut` and is written with `MessageSerializer.serialize(writeChannel,
wrappedBatch)`, matching the Spark 4.1 implementation. I kept the batch
assembly and retain/release logic in a small `serializeBatch` helper so the
allocator, nested/null, multi-batch, and failure-cleanup regression tests can
exercise it directly. All five tests pass on both Spark 4.0 and Spark 4.1.
--
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]
