comphead commented on code in PR #6900:
URL: https://github.com/apache/datafusion-comet/pull/6900#discussion_r4277281912
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -267,9 +267,15 @@ object Utils extends CometTypeShim with Logging {
}
val provider = batchProviderOpt.getOrElse(dictionaryProvider)
- val writer = new ArrowStreamWriter(root, provider,
Channels.newChannel(out))
+ // VectorSchemaRootAppender, which coalesces broadcast batches, assumes
offsets that start
+ // at zero, and Arrow Java writes the bytes before the first offset too.
A slice of the root
+ // rebases the offsets to zero and leaves those bytes out.
+ val written =
+ if (fieldVectors.forall(offsetsStartAtZero)) root else root.slice(0,
root.getRowCount)
+ val writer = new ArrowStreamWriter(written, provider,
Channels.newChannel(out))
writer.start()
writer.writeBatch()
+ if (written ne root) written.close()
Review Comment:
Good catch, thanks. Moved the close into a `finally` as you suggested in
63e83b0ac9, and added a `UtilsSuite` test that interrupts the write and checks
that the slice's allocator drains to zero after the batch closes. Without the
`finally` it fails with the same 49,280 bytes you measured.
--
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]
