Github user sachouche commented on a diff in the pull request: https://github.com/apache/drill/pull/1237#discussion_r184730050 --- Diff: exec/java-exec/src/main/java/org/apache/drill/exec/record/RawFragmentBatch.java --- @@ -77,4 +83,46 @@ public long getByteCount() { public boolean isAckSent() { return ackSent.get(); } + + /** + * Transfer ownership of this DrillBuf to the target allocator. This is done for better memory + * accounting (that is, the operator should be charged with the body's Drillbuf memory). + * + * <p><b>NOTES -</b> + * <ul> + * <li>This operation is a NOOP when a) the current allocator (associated with the DrillBuf) is not the + * owning allocator or b) the target allocator is already the owner + * <li>When transfer happens, a new RawFragmentBatch instance is allocated; this is done for proper + * DrillBuf reference count accounting + * <li>The RPC handling code caches a reference to this RawFragmentBatch object instance; release() + * calls should be routed to the previous DrillBuf + * </ul> + * + * @param targetAllocator target allocator + * @return a new {@link RawFragmentBatch} object instance on success (where the buffer ownership has + * been switched to the target allocator); otherwise this operation is a NOOP (current instance + * returned) + */ + public RawFragmentBatch transferBodyOwnership(BufferAllocator targetAllocator) { + if (body == null) { + return this; // NOOP + } + + if (!body.getLedger().isOwningLedger() + || body.getLedger().isOwner(targetAllocator)) { + + return this; + } + + int writerIndex = body.writerIndex(); + TransferResult transferResult = body.transferOwnership(targetAllocator); + + // Set the index and increment reference count + transferResult.buffer.writerIndex(writerIndex); + + // Clear the current Drillbuffer since caller will perform release() on the new one + body.release(); + + return new RawFragmentBatch(getHeader(), transferResult.buffer, getSender(), false); --- End diff -- We can take up such an enhancement as as part of another JIRA as any changes within the RPC layer have to be thoroughly tested.
---