peterxcli commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r3877376452
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -252,6 +302,332 @@ private static void evaluateInternal(
// do not mask the original throwable
}
}
+ if (transferred != null) {
+ try {
+ transferred.close();
+ } catch (RuntimeException ignored) {
+ // do not mask the original throwable
+ }
+ }
+ }
+ }
+
+ /**
+ * Moves the result's buffer accounting from the task allocator to the root
allocator and drops
+ * the corresponding Spark task charge. Neither the ownership transfer nor
the eventual FFI
+ * release is observed by the task allocator's {@link AllocationListener}
(Arrow only notifies the
+ * allocator owning a chunk when the chunk is destroyed), so the charge must
be released here.
+ *
+ * <p>The returned vector shares the original buffers and must be closed by
the caller after
+ * export; the exported FFI array keeps the buffers alive until native
execution releases them.
+ */
+ private static FieldVector transferForExport(
+ TaskState state, BufferAllocator outputAllocator, FieldVector result) {
+ long charged = chargedOutputSize(result, outputAllocator);
+ TransferPair transferPair = result.getTransferPair(result.getField(),
ROOT_ALLOCATOR);
+ transferPair.transfer();
+ state.releaseExportedCharge(charged);
+ return (FieldVector) transferPair.getTo();
+ }
+
+ /**
+ * Bytes of the result's buffers currently accounted against the task
allocator, i.e. the portion
+ * of the Spark task charge that moves to native ownership on export. {@code
getAccountedSize()}
+ * is non-zero only on the ledger that owns a chunk, so pass-through input
buffers (owned by the
+ * root allocator) and chunks whose ownership already moved on a previous
export contribute
+ * nothing.
+ */
+ private static long chargedOutputSize(FieldVector result, BufferAllocator
outputAllocator) {
+ long charged = 0L;
+ Set<ReferenceManager> seen = Collections.newSetFromMap(new
IdentityHashMap<>());
+ for (ArrowBuf buf : result.getBuffers(false)) {
Review Comment:
Good catch — fixed in 9b728163b. `chargedOutputSize` now enumerates each
vector's physical field buffers recursively (`getFieldBuffers()` +
`getChildrenFromFields()`), preserving per-ledger dedup, instead of
`getBuffers(false)`, whose `getBufferSize() == 0` short-circuit omits the
allocated buffers of zero-length children while `transfer()` moves them to the
root allocator all the same. Added
`emptyChildAllocationsReleaseSparkChargeAtExport`: a 1024-row all-empty
`ListVector` with an allocated child data vector previously stranded 32,768
bytes of Spark charge at export; it now moves the full charge and returns to
baseline after the FFI release.
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -252,6 +302,332 @@ private static void evaluateInternal(
// do not mask the original throwable
}
}
+ if (transferred != null) {
+ try {
+ transferred.close();
+ } catch (RuntimeException ignored) {
+ // do not mask the original throwable
+ }
+ }
+ }
+ }
+
+ /**
+ * Moves the result's buffer accounting from the task allocator to the root
allocator and drops
+ * the corresponding Spark task charge. Neither the ownership transfer nor
the eventual FFI
+ * release is observed by the task allocator's {@link AllocationListener}
(Arrow only notifies the
+ * allocator owning a chunk when the chunk is destroyed), so the charge must
be released here.
+ *
+ * <p>The returned vector shares the original buffers and must be closed by
the caller after
+ * export; the exported FFI array keeps the buffers alive until native
execution releases them.
+ */
+ private static FieldVector transferForExport(
+ TaskState state, BufferAllocator outputAllocator, FieldVector result) {
+ long charged = chargedOutputSize(result, outputAllocator);
+ TransferPair transferPair = result.getTransferPair(result.getField(),
ROOT_ALLOCATOR);
+ transferPair.transfer();
+ state.releaseExportedCharge(charged);
Review Comment:
Fixed in 9b728163b by retaining accounting for shared chunks, as you
suggested. `chargedOutputSize` now releases a chunk's charge at export only
when every live reference to its ledger comes from the result tree
(`getRefCount()` equals the distinct result buffers on it). A chunk shared with
retained scratch keeps its Spark charge, so when native releases the FFI result
and Arrow silently returns ownership to the scratch ledger, the eventual
scratch close fires `onRelease` and frees the charge exactly once; if scratch
closes first while native still holds, the retained charge is dropped wholesale
at task completion, the same bound as the pre-export model. Fresh outputs are
unaffected: fixed-width vectors slice data+validity out of one combined chunk
with two retained slices, both visible in `getFieldBuffers()`, so exclusive
ownership is still detected and released at export (the existing export test
guards this). Added `sharedScratchChunkChargeIsReleasedExactlyOnce` reproducing
you
r scenario, including an unrelated allocation surviving the scratch close with
no over-release.
--
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]
