Skip to content

Register CometArrowAllocator as a Spark MemoryConsumer for JVM-UDF dispatch #4174

Open

Description

@andygrove

Describe the problem

CometUdfBridge.evaluate (common/src/main/java/org/apache/comet/udf/CometUdfBridge.java, on the JVM-scalar-UDF prototype branch) allocates output Arrow vectors via the project-wide CometArrowAllocator. That allocator is a RootAllocator that is not registered with Spark's TaskMemoryManager, so off-heap memory consumed by the UDF dispatch path is invisible to Spark's task memory accounting and back-pressure machinery.

Under workloads with many concurrent JVM-UDF tasks per executor, this can drive native off-heap usage past the operator-level limits Spark would otherwise enforce.

Describe the potential solution

Either:

  1. Register CometArrowAllocator as a MemoryConsumer in Spark's TaskMemoryManager so allocations and frees update the task's accounting.
  2. Allocate UDF output vectors from a child allocator that is itself registered as a per-task consumer, so leakage and accounting stay scoped to the task.

Option (2) is closer to the existing Spark-Arrow integration pattern.

Additional context

Identified during code review of the JVM-scalar-UDF prototype. Filed as a follow-up so the prototype PR can ship without a Spark-integration redesign.

Activity

  1. added
    priority:mediumFunctional bugs, performance regressions, broken features
    area:ffiArrow FFI / JNI boundary
    and removed on May 4, 2026
  2. added this to the 0.17.0 milestone on May 5, 2026
  3. moved this from Todo to In progress in Comet Developmenton May 14, 2026
  4. modified the milestones: 0.17.0, 1.0.0 on Jun 25, 2026
  5. peterxcli commented on Jul 20, 2026

    @peterxcli
    Member

    take

  6. modified the milestones: 1.0.0, 1.1.0 on Aug 17, 2026
  7. modified the milestones: 1.1.0, 1.2.0 on Oct 12, 2026
  8. modified the milestones: 1.2.0, 1.3.0 on Nov 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:expressionsExpression evaluationarea:ffiArrow FFI / JNI boundarypriority:mediumFunctional bugs, performance regressions, broken features

Type

No type

Projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions