perf(session): support fragment session cache#7972
Conversation
Allow passing a shared Session to fragment writes so repeated writes against the same dataset reuse the session's metadata cache (manifest loads during APPEND-mode schema inference) and object store registry, instead of cold-loading the manifest and rebuilding object stores per Fragment.create call.
📝 WalkthroughWalkthroughFragment writing now accepts an optional shared ChangesShared session fragment writes
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant WriteFragmentBuilder
participant Fragment
participant JNI
participant RustFragmentWriter
participant DatasetBuilder
WriteFragmentBuilder->>Fragment: pass optional Session
Fragment->>JNI: pass sessionHandle
JNI->>RustFragmentWriter: forward session handle
RustFragmentWriter->>DatasetBuilder: reuse session for schema loading
DatasetBuilder-->>RustFragmentWriter: return existing dataset schema
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@java/src/main/java/org/lance/Fragment.java`:
- Around line 300-304: Protect the session lifetime across both Fragment write
paths: in Fragment.java at lines 300-304 and 398-402, retain a close-safe guard
or session ownership from handle acquisition until
createWithFfiArray/createWithFfiStream returns, and release it afterward; update
java/lance-jni/src/fragment.rs at lines 341-343 only as needed to consume that
ownership safely. Ensure close() cannot zero or release the native session
before the JNI call completes.
In `@java/src/test/java/org/lance/FragmentTest.java`:
- Around line 172-205: Extend FragmentTest’s shared-session APPEND coverage to
include a write using data(ArrowArrayStream), not only VectorSchemaRoot. Create
an ArrowArrayStream containing the test data, pass it through the same
datasetUri, allocator, schema, mode, and session configuration, execute the
write, and assert the expected fragment result.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: cef3716f-7307-4e3d-8ea9-acbdc594e1cd
📒 Files selected for processing (5)
java/lance-jni/src/fragment.rsjava/src/main/java/org/lance/Fragment.javajava/src/main/java/org/lance/WriteFragmentBuilder.javajava/src/test/java/org/lance/FragmentTest.javarust/lance/src/dataset/fragment/write.rs
| long sessionHandle = 0L; | ||
| if (session != null) { | ||
| sessionHandle = session.getNativeHandle(); | ||
| Preconditions.checkArgument(sessionHandle != 0, "Session is closed"); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline java/src/main/java/org/lance/Session.java --items all
rg -n -C 8 '\b(close|getNativeHandle)\s*\(' java/src/main/java/org/lance/Session.java
rg -n -C 6 'nativeSessionHandle|session_from_handle' \
java/src/main/java/org/lance/Session.java \
java/lance-jni/src/session.rs \
java/lance-jni/src/fragment.rsRepository: lance-format/lance
Length of output: 13032
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Fragment.java relevant sections =="
sed -n '230,320p' java/src/main/java/org/lance/Fragment.java
echo
sed -n '360,420p' java/src/main/java/org/lance/Fragment.java
echo
echo "== Session release/create/native methods =="
sed -n '150,270p' java/src/main/java/org/lance/Session.java
echo
echo "== session_from_handle implementation and Arc field =="
rg -n -C 8 'pub struct LanceSession|LanceSession}|session_from_handle|get_session_handle|releaseNative|napi_async_destroy|create_session' java/lance-jni/src/session.rs java/src/main/java/org/lance/Session.java
echo
echo "== Fragment native signatures and imports =="
sed -n '1,120p' java/src/main/java/org/lance/Fragment.javaRepository: lance-format/lance
Length of output: 21824
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Fragment.java full create implementations around native calls =="
sed -n '320,480p' java/src/main/java/org/lance/Fragment.java
echo
echo "== Fragment.rs native JNI signatures around createWithFfiArray/Stream and session_from_handle =="
rg -n -C 12 'Java_org_lance_Fragment_createWithFfi|Java_org_lance_Fragment_createWithFfiStream|session_from_handle|fn.*create.*ffi' java/lance-jni/src/fragment.rs
echo
echo "== Session.java synchronization/fields broad =="
sed -n '58,90p' java/src/main/java/org/lance/Session.java
echo
echo "== Native session handle fields/locks in Fragment write paths (Java) =="
rg -n -C 6 'synchronized|nativeSessionHandle|getNativeHandle|sessionHandle|createWithFfiArray|createWithFfiStream' java/src/main/java/org/lance/Fragment.java java/src/main/java/org/lance/Session.javaRepository: lance-format/lance
Length of output: 31723
🌐 Web query:
Java JNI global reference ThreadLocalThreadState thread suspension garbage collection
💡 Result:
The interaction between Java Native Interface (JNI) references, thread states, and garbage collection (GC) is governed by the Java Virtual Machine (JVM) mechanics that handle thread transitions between managed and native code [1][2][3]. JNI Global References Global references act as roots for the garbage collector, ensuring that the referenced objects remain reachable and are not prematurely reclaimed [4][5]. Because the GC may compact the heap and move objects, global references (and local references) provide a level of indirection [1][2]. Native code does not hold direct memory addresses to heap objects; instead, it interacts with these references, allowing the JVM to update the underlying object location during GC cycles without breaking the native code's access [1][2]. Thread Suspension and GC Safepoints JVMs typically use a safepoint mechanism to coordinate GC [1][2]. 1. Managed State: When a thread is executing Java bytecode, it is at a safepoint or can reach one, allowing the GC to suspend the thread and safely inspect its stack [1][3]. 2. Native State: When a thread transitions from Java to native code via JNI, it enters a state often referred to as _thread_in_native [1][2]. In this state, the thread is not executing Java bytecode and does not directly access the managed heap, meaning it can continue running concurrently with GC operations [1][2][3]. 3. Safepoint Barriers: The transition between native and managed states acts as a barrier [1][2]. When a native thread attempts to call back into the JVM or use a JNI function, it must check for safepoints [1][2][6]. If a GC is in progress or has requested a suspension, the thread will block until the GC operation completes [1][2][6]. JNI Critical Regions Certain JNI functions, such as GetPrimitiveArrayCritical, allow native code to access raw heap pointers directly [7][2]. These functions enter a "critical" state that often disables GC (or specific GC activities like compaction) for the duration of the region to prevent the memory from moving while the native code holds a direct pointer [7][2]. The JVM ensures that other threads attempting to trigger GC are aware of these active critical regions, often by using mechanisms like the GCLocker [7]. Recent developments in the OpenJDK (e.g., JDK 26 timeframe) have included refinements to how these regions interact with JVM Tool Interface (JVM TI) thread suspension, specifically aiming to defer suspension requests until after a critical region is exited to maintain system stability [8]. In summary, global references keep objects alive, native code execution is generally allowed to continue concurrently with GC, and thread suspension is enforced primarily at JNI boundaries where the thread must synchronize with the JVM's managed state [1][2][3].
Citations:
- 1: https://stackoverflow.com/questions/51321807/jni-code-ie-executed-concurrently-with-garbage-collection
- 2: https://stackoverflow.com/questions/39381339/understanding-safe-access-of-jni-arguments
- 3: https://android.googlesource.com/platform/art/+/f5107e929fdc91891ef2f8a1c3956d73f3ca9b89/runtime/mutator_gc_coord.md
- 4: https://docs.oracle.com/javase/10/docs/specs/jni/design.html
- 5: https://android.googlesource.com/platform/dalvik2/+/refs/heads/main/vm/Jni.cpp
- 6: https://developer.android.com/ndk/guides/jni-tips
- 7: https://shipilev.net/jvm/anatomy-quarks/9-jni-critical-gclocker/
- 8: dholmes-ora/jdk@4341883
Hold a close-safe guard on sessionHandle across the JNI call.
close() can run on another thread, and on most JVMs native code may be preempted after Java snapshots nativeSessionHandle; the native side clones through Arc that will have been released if Java zeroes the field. Keep a close/lifetime guard that is only dropped after the native Fragment write call returns, or move session ownership into that guard across both createWithFfiArray and createWithFfiStream.
📍 Affects 2 files
java/src/main/java/org/lance/Fragment.java#L300-L304(this comment)java/src/main/java/org/lance/Fragment.java#L398-L402java/lance-jni/src/fragment.rs#L341-L343
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@java/src/main/java/org/lance/Fragment.java` around lines 300 - 304, Protect
the session lifetime across both Fragment write paths: in Fragment.java at lines
300-304 and 398-402, retain a close-safe guard or session ownership from handle
acquisition until createWithFfiArray/createWithFfiStream returns, and release it
afterward; update java/lance-jni/src/fragment.rs at lines 341-343 only as needed
to consume that ownership safely. Ensure close() cannot zero or release the
native session before the JNI call completes.
| List<FragmentMetadata> fragments = | ||
| Fragment.write() | ||
| .datasetUri(datasetPath) | ||
| .allocator(allocator) | ||
| .data(root) | ||
| .mode(WriteParams.WriteMode.APPEND) | ||
| .session(session) | ||
| .execute(); | ||
| assertEquals(1, fragments.size()); | ||
| assertEquals(1, fragments.get(0).getPhysicalRows()); | ||
| } | ||
| } | ||
| // Schema inference populated the session's metadata cache. | ||
| assertTrue(session.sizeBytes() > 0); | ||
|
|
||
| // An explicit schema skips inference; the session is still accepted (used for | ||
| // the object store registry). | ||
| try (VectorSchemaRoot root = VectorSchemaRoot.create(dataset.getSchema(), allocator)) { | ||
| root.allocateNew(); | ||
| ((VarCharVector) root.getVector("name")) | ||
| .setSafe(0, "Person 2".getBytes(StandardCharsets.UTF_8)); | ||
| ((IntVector) root.getVector("id")).setSafe(0, 2); | ||
| root.setRowCount(1); | ||
|
|
||
| List<FragmentMetadata> fragments = | ||
| Fragment.write() | ||
| .datasetUri(datasetPath) | ||
| .allocator(allocator) | ||
| .data(root) | ||
| .schema(dataset.getLanceSchema()) | ||
| .mode(WriteParams.WriteMode.APPEND) | ||
| .session(session) | ||
| .execute(); | ||
| assertEquals(1, fragments.size()); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Cover the ArrowArrayStream session route.
This test only exercises VectorSchemaRoot; add a shared-session APPEND write through data(ArrowArrayStream) to cover the separately changed stream JNI path.
As per coding guidelines, “Every bugfix and feature must have corresponding tests.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@java/src/test/java/org/lance/FragmentTest.java` around lines 172 - 205,
Extend FragmentTest’s shared-session APPEND coverage to include a write using
data(ArrowArrayStream), not only VectorSchemaRoot. Create an ArrowArrayStream
containing the test data, pass it through the same datasetUri, allocator,
schema, mode, and session configuration, execute the write, and assert the
expected fragment result.
Source: Coding guidelines
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
closes: #7974