andygrove opened a new pull request, #6048:
URL: https://github.com/apache/datafusion-comet/pull/6048
## Which issue does this PR close?
Closes #6047.
## Rationale for this change
A traced run cannot currently account for Arrow memory held on the JVM. The
existing counters cover
native allocations, memory pool reservations and the JVM heap, but Arrow
buffers are off-heap, so
they show up in none of them.
Reporting the root allocator's total on its own would not separate the two
things a reader needs to
tell apart. Comet imports batches from native over the Arrow C Data
Interface, and Arrow charges an
imported buffer to whichever allocator wraps it
(`BaseAllocator.wrapForeignAllocation` calls
`allocateBytes`), so the root's total mixes memory the JVM allocated with
native memory that
`native_allocated` already counts.
## What changes are included in this PR?
- `CometArrowImportAllocator`, a child of `CometArrowAllocator` that owns
buffers imported over the
C Data Interface. It reserves nothing, so every byte still escalates to
the parent and the root
keeps reporting the total. Like the root it is never closed, because
imported buffers are
reference counted and routinely outlive the task that imported them.
- `NativeUtil`'s `ArrowImporter` and the input imports in `CometUdfBridge`
are taken against that
allocator. Struct allocation, exports and Arrow IPC buffers stay on the
root, because those
really are memory the JVM allocated.
- `arrowMemoryMetrics` in the `comet` package object returns the two
counters, and
`CometExecIterator.traceMemoryUsage` emits them alongside `jvm_heap_used`.
- The tracing guide gains both labels and a note on the overlap:
`jvm_arrow_imported` is native
memory that `native_allocated` counts as well, so `jvm_arrow_allocated -
jvm_arrow_imported` is
the Arrow memory the JVM allocated itself.
## How are these changes tested?
New tests in `NativeUtilSuite`:
- a batch imported over FFI is charged to the import allocator, and that
allocator is a child of
the root
- the counters carry both labels, the imported counter is non-zero while an
imported batch is
alive, and the total includes the imported bytes
New `CometUdfBridgeSuite`, registered in both PR build workflows, drives
`CometUdfBridge.evaluate`
directly. The bridge closes its imported inputs before returning, so the
probe UDF reads the import
allocator from inside the call, while its inputs are alive.
For regression coverage, `CometScalaUDFClassLoaderSuite` and
`CometInMemoryCacheSuite` pass (42
tests). The latter already asserts that the root allocator returns to its
prior total, which
exercises the escalation the two counters depend on.
The emission call itself is not unit tested, since that needs a traced
native run. It will be
exercised on a traced TPC-H run.
--
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]
