Repository navigation
feat(java): opt-in copy of each batch the JNI writer imports (workaround for #46) - #47
Closed
ChristopherWirt wants to merge 2 commits into
Closed
ChristopherWirt wants to merge 2 commits into
ChristopherWirt wants to merge 2 commits into
Conversation
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]>
This was referenced Oct 30, 2026
Signed-off-by: Christopher Wirt <[email protected]>
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 I ported #49 to that branch instead, as
#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. |
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]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_batcheswriter option,"true"or"false"(default off). It's read and removed inNativeWriter_create, so it never reaches the object store.writeBatchdeep-copies the importedArrayData(MutableArrayData) before converting it.VortexWriter.Builder.copyBatches(boolean)andVortexWriter.COPY_BATCHES_OPTIONon 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):
copy_batches=trueThe 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 oncewriteBatchreturns. It then reads the file back.copyBatches(true), every batch is still held (66,776 bytes each). That gives JNI writer keeps imported Arrow batches alive until the file finishes #46 a deterministic repro.writer::tests(Rust):true/false, and is removed from the properties./gradlew :vortex-jni:checkandcargo clippy -p vortex-jni --all-targetspass.🤖 Generated with Claude Code