Skip to content

feat(java): opt-in copy of each batch the JNI writer imports (workaround for #46) - #47

Closed
ChristopherWirt wants to merge 2 commits into
developfrom
cwirt/jni-copy-imported-batches
Closed

ChristopherWirt wants to merge 2 commits into
developfrom
cwirt/jni-copy-imported-batches

Conversation

@ChristopherWirt

@ChristopherWirt ChristopherWirt commented Oct 30, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

The JNI writer keeps most of a file's imported Arrow batches alive until the file finishes. A file written in many large batches can hold many times the memory of one: 4.7 GB for one 85-batch file in our case. #46 has the cause and measurements.

This adds an opt-in, temporary workaround: copy each batch at import, so it can be freed as soon as it's converted. It costs a memcpy per batch. The real fix is for the writer to stop holding imported buffers, with no copy, and that's tracked in #46.

Changes

  • vortex.write.copy_batches writer option, "true" or "false" (default off). It's read and removed in NativeWriter_create, so it never reaches the object store.
  • When the option is on, writeBatch deep-copies the imported ArrayData (MutableArrayData) before converting it.
    • Dictionary values and view-array data buffers are still shared. The Java writer produces neither.
  • VortexWriter.Builder.copyBatches(boolean) and VortexWriter.COPY_BATCHES_OPTION on the Java side.
  • createStream (NativeWritable) takes no options and never copies.

Measured on a real file (690k rows, 615 leaf columns, 85 batches of about 56 MB):

Batch bytes held Native memory over baseline
Off (default) 4.7 GB +4.8 GiB
copy_batches=true 257 KB +0.2 GiB

The file is the same size and reads back the same.

Tests

  • CopyBatchesTest (Java) writes 8 batches, each from its own child allocator, and checks that each allocator holds under 1 KB once writeBatch returns. It then reads the file back.
  • writer::tests (Rust):
    • a copied batch shares no buffer with its FFI import
    • the option is parsed, rejects values other than true/false, and is removed from the properties
  • ./gradlew :vortex-jni:check and cargo clippy -p vortex-jni --all-targets pass.

🤖 Generated with Claude Code

The JNI writer keeps most of a file's imported Arrow batches alive until
the file finishes, since all of a batch's buffers share one release and
the layout writers hold slices of earlier batches. A file of many large
batches then holds many times the memory of one.

`vortex.write.copy_batches=true` (`VortexWriter.Builder.copyBatches`)
copies each batch at import so it's freed as soon as it's converted. A
stopgap until the writers stop holding imported buffers: see #46.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Christopher Wirt <[email protected]>
@gitRasheed

Copy link
Copy Markdown
Collaborator

I checked this against the FileFormat connector branch, rasheed/spark-fileformat-v2, and the copy wouldn't help there as written. That branch's VortexOutputWriter opens its writer with createStream on a Hadoop-backed NativeWritable. createStream never copies, and only NativeWriter_create reads vortex.write.copy_batches.

I ported #49 to that branch instead, as rasheed/ff-v2-chris-jni. I measured through the Spark writer with batches of 8,192 rows of a nested-string schema, 2.88 MiB of Arrow memory each. Arrow memory still held before close:

#49 works for both create paths and doesn't copy, so I haven't ported this one. I think this PR can close once #49 merges.

@ChristopherWirt

Copy link
Copy Markdown
Collaborator Author

closing in favor of #49

ChristopherWirt added a commit that referenced this pull request Nov 2, 2026
## Summary

Fixes #46. Provides a zero-copy alternative to the batch-copy workaround
in #47.

A retained validity bitmap or column slice currently keeps an entire
imported Arrow batch alive through its shared C Data Interface release
callback. The new Java-root write path retains each Java allocation
independently, so the writer releases unrelated columns as soon as it no
longer needs them. Spark uses this path by default.

In the nested-string regression workload (16 batches of 8,192 rows),
memory retained before `finish` dropped from about **65 MiB to 8.1
MiB**, with all allocations released at finish. These are synthetic
results; the production 615-column workload has not been rerun.

## Changes

- Add `VortexWriter.writeBatch(VectorSchemaRoot, BufferAllocator)`. It
checks the batch schema, exports the array, and imports it using the
writer's stored schema, avoiding a schema export for every batch.
- Retain Java Arrow reference managers independently and replace the
Rust buffers' owners without copying their data. The last Rust reference
releases its Java allocation, including from native worker threads.
- Preserve Arrow's existing alignment copies for unaligned input. Track
the original exported addresses so missing Java owners still fail import
rather than silently keeping whole batches alive.
- Switch the Spark V2 writer to the new path. The address-based API
remains compatible and keeps the producer's shared release callback; its
callers must use the Java-root overload to get independent allocation
lifetimes. Dictionary-encoded roots still require the address-based API.
- Serialize both write methods with `finish` to prevent access to a
freed native writer, and release the temporary schema used when creating
a writer.

Buffers sharing one underlying Java allocation still share that
allocation's lifetime.

## Testing

- `IndependentBuffersTest`: memory retained before finish and full
release afterward; partial-import and sink-failure cleanup; empty
batches and string views; unaligned integer buffers; schema mismatch
rejection; concurrent write/finish stress for both APIs (16 iterations
each).
- `JNIWriterTest`: both APIs round-trip strings and nested maps. **18
JNI tests pass; one existing test is disabled.**
- Rust ownership tests: unchanged buffer addresses, independent column
release, sliced validity offsets, and release/global-reference cleanup
from a native thread. **3 pass.**
- Spark 2.13 `VortexDataSourceWriteTest`: **11 pass.**
- `cargo clippy -p vortex-jni --all-targets --all-features -- -D
warnings`, pinned Rust formatting, Java formatting checks, and patch
whitespace checks pass.
- Four independent reviews covered reuse, maintainability, efficiency
and correctness; confirmed findings were fixed and tested.

Known separately: the existing Vortex string-view concatenation path
reads metadata from null slots. The view round-trip fixture initializes
those slots; this PR does not fix that issue.

Checked locally against QC `develop`.

I used AI tools to help write and review this change.

Signed-off-by: Christopher Wirt <[email protected]>
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