sunchao commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r3876130734
##########
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:
**[P2] Include allocated empty children when releasing export charges**
Could we enumerate each vector's physical field buffers recursively here,
preserving ledger deduplication? `getBuffers(false)` omits allocated buffers
for zero-length children, but `transferPair.transfer()` still moves them to the
root allocator. The omitted bytes never reach `releaseExportedCharge`, and
subsequent FFI release cannot return their Spark reservation.
On this head with Spark 4.1.3, Arrow 18.3.0 and JDK 17, a native
`CometNativeScan -> CometProject` query using `regexp_extract_all(s, '(z)', 1)`
on 8,192 nonmatching input rows left 786,432 bytes reserved after consuming all
output. The matching control left zero. A diagnostic change to recursively
collect physical field buffers also left zero. A separate empty-integer-array
bridge probe exhausted a 1 MiB pool after 31 batches while no Arrow buffers
remained live. Repeated empty/null array batches can therefore fail
otherwise-fitting tasks.
##########
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:
**[P2] Account for ownership returning to retained scratch buffers**
For a custom `CometUDF` that retains a scratch vector and returns an aligned
`splitAndTransfer` slice, this releases the whole chunk's Spark charge even
though a task-allocator ledger still references it. When native releases the
FFI result, Arrow can transfer ownership back to the retained scratch without
an `onPreAllocation` callback. Closing scratch later invokes `onRelease` and
frees the same Spark charge again.
A probe through the current public `CometUdfBridge.evaluate` with Spark
4.1.3 and Arrow 18.3.0 reproduced this: after FFI release, scratch still held
8,192 readable Arrow bytes with zero Spark charge. After allocating an
unrelated 8,192-byte buffer, closing scratch reduced the Spark charge to zero
while that unrelated buffer remained live; closing it then logged an
over-release. Could we retain or coordinate accounting for shared chunks so
returning ownership cannot cause duplicate release?
This is limited to the documented custom-CometUDF scratch-buffer contract.
The in-tree codegen implementation creates fresh outputs; this reproduction
exercised the bridge and Arrow C exports, not a native query using a custom UDF.
--
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]
