andygrove commented on code in PR #6900:
URL: https://github.com/apache/datafusion-comet/pull/6900#discussion_r4274400171
##########
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:
If `writer.start()` or `writer.writeBatch()` throws, `written` never gets
closed. Before this change that was harmless, because the batch's vectors are
released when `CometExecIterator` closes the batch. But `splitAndTransfer`
takes its own reference on the validity and data buffers it shares with the
batch, so an unclosed slice keeps them alive after the batch is closed. For a
batch imported from native, that pins the native buffers for the life of the
executor.
It's reachable when a broadcast is cancelled. Broadcast jobs run with
`interruptOnCancel = true`, and the channel from `Channels.newChannel` is
interruptible, so a killed task throws `ClosedByInterruptException` on its next
write. The shuffle reader only checks for a kill when native pulls the next
block, which an aggregate emitting its output slices no longer does. I
reproduced it in `UtilsSuite` by calling `Thread.currentThread().interrupt()`
before serializing a sliced vector. After closing the batch, its allocator
still held 49,280 bytes, against 0 for an unsliced vector.
Could we close the slice in a `finally`, along these lines?
```scala
try {
val writer = new ArrowStreamWriter(written, provider,
Channels.newChannel(out))
writer.start()
writer.writeBatch()
root.clear()
writer.close()
} finally {
if (written ne root) written.close()
}
```
That brought it to 0 for me, and `UtilsSuite` and the new join test still
passed. A test that sets the interrupt flag and checks that a child allocator
drains to zero after the batch closes would keep it that way.
--
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]
