Skip to content

feat: account JVM UDF Arrow allocations in Spark task memory - #5027

Open
peterxcli wants to merge 12 commits into
apache:mainfrom
peterxcli:feat/arrow-allocator-as-spark=memory-consumer-for-jvmdispatch
Open

feat: account JVM UDF Arrow allocations in Spark task memory#5027
peterxcli wants to merge 12 commits into
apache:mainfrom
peterxcli:feat/arrow-allocator-as-spark=memory-consumer-for-jvmdispatch

Conversation

@peterxcli

@peterxcli peterxcli commented Jul 24, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #4174.

Rationale for this change

CometScalaUDFCodegen, currently the only CometUDF implementation, creates output Arrow vectors from the process-wide CometArrowAllocator. These JVM allocations bypass Spark’s task memory accounting.

FFI input vectors only wrap native-owned buffers that are already accounted by Comet, so they must not be charged again. Output buffers, however, are created on the JVM and may remain alive through the FFI release callback after the Java vector closes.

What changes are included in this PR?

  • Add a per-task Arrow child allocator backed by a non-spillable Spark MemoryConsumer.
  • Pass the allocator through the generic CometUDF interface and codegen output path.
  • Keep imported native buffers on the root allocator to avoid double-accounting.
  • Release Spark accounting with the task while retaining allocator state until the final FFI buffer release.

How are these changes tested?

Added CometUdfBridgeTest covering allocation accounting, FFI retention, task completion, and final cleanup. The focused test passes with Spark 4.1 and Spark 3.5/Scala 2.12.

Also ran few scala script with spark shell to verify:

Check main branch
Small output task peak 0 262,144 bytes
32 MiB output task peak 0 50,397,184 bytes
Peak increase 0 50,135,040 bytes
256 MiB batch under 256 MiB pool Succeeds, bypassing Spark Rejected by TaskMemoryManager
Same workload with batch size 32 Succeeds, 268,435,456 bytes processed

@peterxcli

Copy link
Copy Markdown
Member Author

@sunchao would you like to take a look at this? TIA!

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two P2 findings in task-allocator concurrency, detailed inline. Validated with targeted JVM reproductions against this exact head and Arrow 18.3.0. The native producer/cleanup paths were source-traced, not run end-to-end.

Comment thread spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java Outdated
Comment thread spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java Outdated
@peterxcli
peterxcli requested a review from sunchao August 23, 2026 12:53
@sunchao

sunchao commented Aug 23, 2026

Copy link
Copy Markdown
Member

@peterxcli can you check the CI failures?

@peterxcli

Copy link
Copy Markdown
Member Author

@peterxcli can you check the CI failures?

@sunchao I opened a fix for it

#5439

@sunchao

sunchao commented Aug 23, 2026

Copy link
Copy Markdown
Member

Ah thanks! Approved the PR and will merge soon

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Register CometArrowAllocator as a Spark MemoryConsumer for JVM-UDF dispatch

2 participants