andygrove commented on code in PR #5998:
URL: https://github.com/apache/datafusion-comet/pull/5998#discussion_r4047683933
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala:
##########
@@ -87,21 +87,23 @@ private[codegen] object CometBatchKernelCodegenOutput
extends CometTypeShim {
* Closes the vector on any failure so a partially-initialized tree doesn't
leak buffers.
*/
def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int):
FieldVector = {
+ // JVM-owned codegen output, accounted to the task that is running the
kernel.
+ val allocator = CometTaskArrowAllocator.forCurrentTask()
Review Comment:
Confirmed, and you are right that the previous revision only moved the
exporter. `CometUdfBridge` switched to the root for its import and its
`Data.exportVector`, but the output buffers were allocated before either of
those, and `CometBatchKernelCodegenOutput.allocateOutput` still selected the
task allocator. So the JVM charge stood while native reserved the same buffers.
It allocates from the listener-less root now. The only production caller is
`CometScalaUDFCodegen.evaluate`, which is the only `CometUDF` implementation,
and its result is consumed solely by `CometUdfBridge`, which exports it over
the C Data Interface and then closes its own reference. Every buffer this
method hands out is therefore allocated to be handed to native, which puts it
under the same rule as `NativeUtil`, the imported inputs and
`CometNativeArrowSource.stream`: native's pool is the authority for bytes
native holds, so charging them here reserved the same memory twice.
That should close the 2 MiB JVM reservation you measured alongside the
1,064,960-byte native one, since nothing on the JVM side reserves those buffers
any more. The contributor guide listed codegen output under the task allocator
and is corrected.
One side effect worth recording: this also unstacks the overlap with #5027.
With codegen output no longer charged here, the per-task consumer in that PR
can own the UDF output path outright rather than adding a second charge on top
of this one.
Latest main is merged in, which clears the conflict that was holding CI to
the label check.
--
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]
