peterxcli commented on code in PR #5368:
URL: https://github.com/apache/datafusion-comet/pull/5368#discussion_r3815605210
##########
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:
```suggestion
private var arrowWriter: ArrowStreamWriter = _
```
##########
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:
then use `MessageSerializer.serialize`
```suggestion
// This root only borrows the Comet vectors; do not close it.
val sourceRoot =
new VectorSchemaRoot(sourceVectors.map(_.getField).asJava,
sourceVectors.asJava, numRows)
val sourceBatch = new VectorUnloader(sourceRoot).getRecordBatch
try {
val validityBytes = (numRows.toLong + 7L) / 8L
val structValidity = allocator.buffer(validityBytes)
try {
if (validityBytes > 0) {
structValidity.setOne(0L, validityBytes)
}
structValidity.writerIndex(validityBytes)
val nodes = new
ArrayList[ArrowFieldNode](sourceBatch.getNodes.size() + 1)
nodes.add(new ArrowFieldNode(numRows, 0))
nodes.addAll(sourceBatch.getNodes)
val buffers = new
ArrayList[ArrowBuf](sourceBatch.getBuffers.size() + 1)
buffers.add(structValidity)
buffers.addAll(sourceBatch.getBuffers)
val wrappedBatch = new ArrowRecordBatch(
numRows,
nodes,
buffers,
sourceBatch.getBodyCompression,
sourceBatch.getVariadicBufferCounts,
true)
try {
MessageSerializer.serialize(
new WriteChannel(Channels.newChannel(dataOut)),
wrappedBatch)
} finally {
wrappedBatch.close()
}
} finally {
structValidity.close()
}
} finally {
sourceBatch.close()
}
pythonMetrics("pythonDataSent") += dataOut.size() - startData
true
}
```
ref:
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala#L134-L166
--
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]
