apache / apache/datafusion-comet

Reduce JNI round-trips in the unified memory pools (batching, hysteresis, cheaper call path)

Open
#5,383 1 comment 0 reactions 1 assignee Claimed by @andygrove View on GitHub
area:ffi enhancement performance priority:medium
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 6h
Merged PRs (30d)
190

Description

## What / Why

Every DataFusion `try_grow` / `shrink` against `CometUnifiedMemoryPool` or `CometFairMemoryPool` becomes a JNI round-trip into `CometTaskMemoryManager`, and there is no batching, no slack, and no hysteresis: memory is acquired in exactly the requested amount and released the moment a reservation shrinks, however small the amount.

This breaks out positions 17 and 18 of #5212 into a standalone issue so the work can be sequenced and benchmarked on its own.

## Anatomy of one grow/shrink

`native/core/src/execution/memory_pools/unified_pool.rs:64-78`

```rust
fn acquire_from_spark(&self, additional: usize) -> CometResult {
let handle = self.task_memory_manager_handle.as_obj();
JVMClasses::with_env(|env| unsafe {
jni_call!(env,
comet_task_memory_manager(handle).acquire_memory(additional as i64) -> i64)
})
}
```

Each call pays:

1. `JVMClasses::with_env` (`native/jni-bridge/src/lib.rs:339-368`) — `attach_current_thread_guard` plus `with_local_frame`, i.e. a PushLocalFrame/PopLocalFrame pair, even though `acquire_memory(long) -> long` and `release_memory(long) -> void` create no local references at all.
2. `call_method_unchecked` plus `check_exception` (`native/jni-bridge/src/lib.rs:78-100`).
3. On the JVM side, `internal.acquireExecutionMemory` (`spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java:62`), which takes Spark's executor-wide `synchronized` `MemoryManager` lock. This is the part that scales badly — every concurrent task's grow and shrink contends on it.

Callers hit this per batch: DataFusion's aggregate (`try_resize` per input batch), sort accumulation and repartition, plus Comet's own native shuffle writer, which does `try_grow` per batch and `reservation.free()` per flush (`native/shuffle/src/partitioners/multi_partition.rs:463`, `:528`) — a full release-then-reacquire cycle on every spill.

## Options

### A. Chunked acquire with retained slack

Track two numbers instead of one: `used` (what DataFusion reservations hold) and `granted` (what Spark has handed us), with the invariant `granted >= used`.

- `try_grow(n)`: if `used + n <= granted`, bump `used` and return — no JNI at all. Otherwise ask Spark for `max(need, chunk)`.
- `shrink(n)`: decrement `used`, and call Spark only when `granted - used` exceeds the slack cap, releasing chunk-aligned amounts.
- Release all remaining `granted` in `Drop`. Note that today's `Drop` only warns (`unified_pool.rs:81-91`), so this has to be added regardless.

One knob, something like `spark.comet.exec.memoryPool.acquireChunkSize`, defaulting to a few MB with `0` disabling the behaviour; slack cap equal to the chunk size. Steady-state churn within a chunk becomes free, which is exactly the shuffle writer's access pattern.

Two sharp edges:

- `acquireExecutionMemory` will spill *other* consumers to satisfy an oversized request, so only round small requests up; pass large ones through at their exact size.
- On a short grant, keep it when `grant >= need`, otherwise release it and return `resources_err` as today, so the DataFusion-side spill still triggers. Hoarding memory Spark is actively contending for is worse than the extra round-trip.

Implementation note: the `granted`/`used` bookkeeping and the JNI handle should live in one shared type used by both `unified_pool.rs` and `fair_pool.rs`, so the fair pool gets the same benefit and position 17 of #5212 (mutex held across the JNI round-trip) can be fixed in the same place.

### B. Make the retained slack reclaimable

`NativeMemoryConsumer.spill()` returns 0 unconditionally (`CometTaskMemoryManager.java:110-113`), so Spark can never reclaim anything from Comet. Holding slack makes that worse.

Fix: have `spill(size, trigger)` call back into native — a new `@native def releasePoolSlack(taskAttemptId: Long): Long` resolved through the existing `TASK_SHARED_MEMORY_POOLS` registry (`memory_pools/task_shared.rs:25`) — and return the number of bytes of *free slack* dropped. No data is spilled, so it is cheap and always safe.

This is what makes hysteresis defensible under memory pressure rather than just shifting pain onto other tasks in the executor. It also partly addresses position 5 of #5212, since Comet's consumer would stop being invisible to Spark's spill-victim ordering.

I would not ship A enabled by default without B.

### C. Cheaper per-call path

Add a `with_env` variant that skips `with_local_frame` for calls that create no local references. Two fewer JNI transitions per acquire and release, no behavioural change. Worth doing whether or not A lands.

### D. Pre-reserve the whole per-task budget

Chunk size equal to `memoryLimitPerTask`, then run a native `GreedyMemoryPool` over it and touch Spark only on exhaustion. Nearly eliminates the JNI traffic, but hoards memory from tasks that need little and defeats Spark's dynamic sharing. Reasonable as an opt-in config, not as a default.

### E. Reduce call frequency at the source

Round Comet's own reservations up so growth is stepwise rather than per batch; the native shuffle writer is the main offender. Narrower than A, complementary, and needs no new config.

## Suggested sequencing

1. **Instrumentation.** Count acquire/release calls and bytes, and measure cumulative time spent in the JNI round-trip, then report per task. Without this there is no way to show a win — or to tell whether the win is worth the added complexity. This should be its own PR.
2. **C** — small, unconditional, independent of everything else.
3. **A and B together**, behind a config, with numbers from position 1 above.
4. **E** if the shuffle writer still shows up in the counters afterwards.

D only if a workload turns up that genuinely wants it.

## Related

- #5212 — parent EPIC; this is positions 17 and 18, and B overlaps position 5.
- #4576 — adopting DataFusion's allocator-level accounting would change what triggers these calls, but not what each one costs.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.