peterxcli commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r3873534088
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -254,4 +271,225 @@ private static void evaluateInternal(
}
}
}
+
+ /** Visible to the focused allocator test in this package. */
+ static BufferAllocator taskAllocator(TaskContext taskContext) {
+ return taskState(taskContext).allocator();
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static int taskStateCount() {
+ return TASKS.size();
+ }
+
+ /** Visible to the focused allocator test in this package. */
+ static Runnable beginTaskEvaluation(TaskContext taskContext) {
+ TaskState state = taskState(taskContext);
+ state.beginEvaluation();
+ return state::finishEvaluation;
+ }
+
+ private static TaskState taskState(TaskContext taskContext) {
+ return TASKS.computeIfAbsent(
+ taskContext,
+ context -> {
+ TaskState state = new TaskState(context,
CometTaskContextShim.taskMemoryManager(context));
+ context.addTaskCompletionListener(
+ (TaskCompletionListener) ignored -> state.taskCompleted());
+ return state;
+ });
+ }
+
+ /** Per-task Arrow listener and non-spillable Spark memory consumer. */
+ private static final class TaskState implements AllocationListener {
+ private final TaskContext taskContext;
+ private final long taskAttemptId;
+ private final TaskMemoryManager taskMemoryManager;
+ private final TaskMemoryConsumer consumer;
+ private final ConcurrentHashMap<String, CometUDF> instances = new
ConcurrentHashMap<>();
+
+ private BufferAllocator allocator;
+ // Arrow updates allocator accounting after onPreAllocation returns.
+ private int evaluationsInFlight;
+ private boolean completed;
+ private boolean closed;
+
+ private TaskState(TaskContext taskContext, TaskMemoryManager
taskMemoryManager) {
+ this.taskContext = taskContext;
+ this.taskAttemptId = taskContext.taskAttemptId();
+ this.taskMemoryManager = taskMemoryManager;
+ this.consumer =
+ taskMemoryManager.getTungstenMemoryMode() == MemoryMode.OFF_HEAP
+ ? new TaskMemoryConsumer(taskMemoryManager)
+ : null;
+ }
+
+ private synchronized BufferAllocator allocator() {
+ if (completed) {
+ throw new IllegalStateException(
+ "Cannot allocate JVM UDF memory after task " + taskAttemptId + "
completed");
+ }
+ if (allocator == null) {
+ allocator =
+ ROOT_ALLOCATOR.newChildAllocator(
+ "comet-udf-task-" + taskAttemptId, this, 0L, Long.MAX_VALUE);
+ }
+ return allocator;
+ }
+
+ @Override
+ public void onPreAllocation(long size) {
+ // Spark's executor cleanup also synchronizes on TaskMemoryManager. Keep
that cleanup from
+ // overtaking an admitted allocation, while leaving this TaskState
monitor free for buffer
+ // releases that can satisfy a blocking acquire.
+ synchronized (taskMemoryManager) {
+ synchronized (this) {
+ if (completed) {
+ throw new OutOfMemoryException(
+ "Cannot allocate " + size + " JVM UDF bytes after task
completion");
+ }
+ }
+
+ long acquired = consumer == null ? size : consumer.acquireMemory(size);
Review Comment:
Good catch — fixed in 2fafac953 by transferring accounting ownership across
the FFI boundary at export: the result's buffers move to the root allocator
(zero-copy, FFI addresses unchanged) and the Spark charge is released at that
point, so a retaining native operator's reservation is the only charge. The UDF
is still charged while it holds the memory, preserving the admission control
this PR adds.
Verified per-chunk exact against Arrow 18.3.0's allocator internals:
ownership transfers fire no `AllocationListener` callbacks and the eventual FFI
release only notifies the chunk's owner at destruction time (the root
allocator), so the explicit release cannot double-free; `getAccountedSize()` is
non-zero only on the owning ledger, so pass-through input buffers and
re-exported chunks contribute nothing. One wrinkle: the C schema is exported
from the result's own `Field` (a transferred complex vector rebuilds child
names from runtime data vectors, which Arrow hardcodes to `$data$`), while the
C array — which carries no names — comes from the transferred vector.
The focused test was rewritten to assert the new contract (charge present
while the UDF holds the output, gone at export while buffers are still
readable, unchanged by FFI release), and `CometCodegenSuite` passes end-to-end
(85/85). A side benefit: exported buffers no longer pin `TaskState` past task
completion, so the deferred-release retention and the task-attempt-ID-reuse
concern it covered are gone.
--
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]
