Skip to content

perf(jni): bound completed scan outputs with background prefetch - #78

Merged
ChristopherWirt merged 4 commits into
developfrom
cwirt/spark-peak-memory-investigation
Nov 4, 2026
Merged

ChristopherWirt merged 4 commits into
developfrom
cwirt/spark-peak-memory-investigation

Conversation

@ChristopherWirt

@ChristopherWirt ChristopherWirt commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

Problem

  • A JNI scan hands batches to Java through block_on_stream_thread_safe.
  • That bridge queues one completed batch per CPU core, but each JNI stream has only one consumer.
  • On wide or nested Spark scans those queued Arrow batches add up, so peak memory grows with host core count.

Fix

  • Add CurrentThreadRuntime::block_on_stream_prefetch to vortex-io. It is a single-consumer bridge that:
    • runs the stream in a background task, so native work still overlaps Java consumption;
    • keeps at most one completed item ahead of the consumer;
    • cancels the producer on drop and re-raises its panic at end of stream.
  • Use it for the default JNI scan path. Scans with an explicit splitConcurrency keep the direct iterator.

The new API is additive. JNI is its only caller, and other bindings are unchanged.

Results

Local Spark measurements, baseline QC develop 29348ee39. Values are medians of macOS peak physical footprint; n is runs per variant.

Workload n Peak GiB, baseline → candidate Change
Flat-layout read 3 1.46 → 1.05 −28%
List-layout read 2 3.25 → 2.22 −32%
Ordered Vortex rewrite 2 1.86 → 1.43 −23%
  • Batch sizes, compression and decode concurrency are unchanged, and the rewrite output was byte-identical in size and row count.
  • Timing was too noisy on a shared host to claim a speedup or rule out a regression.
  • Only one Spark task with four native workers was measured.
  • The benchmark harness is not part of this PR.

Notes for reviewers

  • Once a batch is queued, the producer stops pulling the next one until Java takes it. Conversions already running keep going, but with a single native worker there is little overlap.
  • This PR does not change how a native panic during a scan is handled. It is still re-raised inside the Arrow C stream callback, which aborts the JVM. A follow-up could surface it as an ArrowError instead.

Tests

  • vortex-io: one-item-ahead bound and driver panic propagation.
  • vortex-jni: the default scan keeps one completed output ahead.
  • Not run yet: no tests, clippy or fmt have been run locally on this branch.

🤖 Generated with Claude Code

ChristopherWirt and others added 4 commits November 3, 2026 15:49
- Add `CurrentThreadRuntime::block_on_stream_prefetch`, a single-consumer
  bridge that keeps one completed item ahead and skips executor entry when
  an item is already queued. JNI's default scan path now calls it.
- Move the producer panic test next to the bridge.
- Simplify the memory harness: pass only the arguments that vary, read the
  buffered peak from the writer summary instead of reflection, and dedupe
  javac/JVM flag/variant handling.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Christopher Wirt <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Christopher Wirt <[email protected]>
@ChristopherWirt
ChristopherWirt marked this pull request as ready for review November 3, 2026 21:42
@ChristopherWirt

Copy link
Copy Markdown
Collaborator Author

@gitRasheed what do you think of this one? feel free to test the crap out of it w/benchmarks and everything.

@ChristopherWirt
ChristopherWirt merged commit 0519dda into develop Nov 4, 2026
2 checks passed
@ChristopherWirt

Copy link
Copy Markdown
Collaborator Author

Ran a broader testing pass on 032fe1ff1 against its parent 29348ee39: 14 Spark scenarios and 38,400 JNI scans, covering projections, nullable lists, concurrent tasks, rewrites and different native worker counts.

The memory saving holds up. On the 12-core Mac, with four native workers and a fixed, pre-touched 512 MiB Java heap, these are median peak physical footprints from three fresh JVMs per variant:

Workload Baseline → PR (GiB) Change
Flat-layout read, one task 1.71 → 1.31 −23%
List-layout read, one task 3.07 → 2.18 −29%
Flat-layout read, four tasks 4.42 → 2.53 −43%
List-layout read, four tasks 10.18 → 5.73 −44%

RSS fell too. Narrow projections had little memory to save, so these percentages won't apply everywhere.

The 18 runtime/JNI unit tests passed, full-row digests matched across 3.67 million input rows, and the ordered rewrite outputs were byte-identical.

Timing is still inconclusive. Some fast JNI scans initially looked 20–30% slower, but extra repetitions were inconsistent and even unchanged control paths swung badly while other builds were running. I'd want the same checks on an idle Linux executor before calling this regression-free.

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.

2 participants