sunchao commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r3868456475
##########
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:
**[P2] Avoid charging exported UDF buffers twice**
The new Spark charge remains until FFI release, but native hash joins and
shuffles reserve the same buffers against Spark again. `JvmScalarUdfExpr`
imports the output without copying, and Arrow's single-batch concatenation also
shares the original buffers. A UDF feeding a single-batch hash-join build can
therefore fail its native reservation even when the physical allocations fit
the pool.
A focused component probe using the exact JVM sources with Spark 4.1.3 and
Arrow 18.3.0 confirmed this: a 32 MiB buffer in a 48 MiB pool receives the full
32 MiB native reservation with the root-allocator control. With this change,
the retained JVM charge leaves only about 16 MiB for that same reservation. The
FFI buffer address is unchanged.
Could we coordinate accounting ownership across the FFI boundary, or exclude
already-charged buffers from native reservations? The native query path was
source-traced; this reproduction was not an end-to-end native query.
--
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]
