Skip to content

perf(session): support fragment session cache #7972

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zhangyue19921010 wants to merge 1 commit into from
Closed

perf(session): support fragment session cache #7972

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions java/lance-jni/src/fragment.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use std::sync::Arc;
use crate::blocking_dataset::extract_namespace_info;
use crate::error::{Error, Result};
use crate::ffi::JNIEnvExt;
use crate::session::session_from_handle;
use crate::traits::{FromJObjectWithEnv, IntoJava, JLance, export_vec, import_vec};
use crate::utils::extract_storage_options;
use crate::{
Expand Down Expand Up @@ -109,6 +110,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiArray<'local>(
allow_external_blob_outside_bases: JObject, // Optional<Boolean>
blob_pack_file_size_threshold: JObject, // Optional<Long>
schema_addr: jlong,
session_handle: jlong, // Session handle, 0 means no session
) -> JObject<'local> {
ok_or_throw_with_return!(
env,
Expand All @@ -132,6 +134,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiArray<'local>(
allow_external_blob_outside_bases,
blob_pack_file_size_threshold,
schema_addr,
session_handle,
),
JObject::default()
)
Expand All @@ -158,6 +161,7 @@ fn inner_create_with_ffi_array<'local>(
allow_external_blob_outside_bases: JObject, // Optional<Boolean>
blob_pack_file_size_threshold: JObject, // Optional<Long>
schema_addr: jlong,
session_handle: jlong, // Session handle, 0 means no session
) -> Result<JObject<'local>> {
let c_array_ptr = arrow_array_addr as *mut FFI_ArrowArray;
let c_schema_ptr = arrow_schema_addr as *mut FFI_ArrowSchema;
Expand Down Expand Up @@ -190,6 +194,7 @@ fn inner_create_with_ffi_array<'local>(
allow_external_blob_outside_bases,
blob_pack_file_size_threshold,
schema_addr,
session_handle,
reader,
)
}
Expand All @@ -215,6 +220,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiStream<'a>(
allow_external_blob_outside_bases: JObject, // Optional<Boolean>
blob_pack_file_size_threshold: JObject, // Optional<Long>
schema_addr: jlong,
session_handle: jlong, // Session handle, 0 means no session
) -> JObject<'a> {
ok_or_throw_with_return!(
env,
Expand All @@ -237,6 +243,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiStream<'a>(
allow_external_blob_outside_bases,
blob_pack_file_size_threshold,
schema_addr,
session_handle,
),
JObject::null()
)
Expand All @@ -262,6 +269,7 @@ fn inner_create_with_ffi_stream<'local>(
allow_external_blob_outside_bases: JObject, // Optional<Boolean>
blob_pack_file_size_threshold: JObject, // Optional<Long>
schema_addr: jlong,
session_handle: jlong, // Session handle, 0 means no session
) -> Result<JObject<'local>> {
let stream_ptr = arrow_array_stream_addr as *mut FFI_ArrowArrayStream;
let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?;
Expand All @@ -284,6 +292,7 @@ fn inner_create_with_ffi_stream<'local>(
allow_external_blob_outside_bases,
blob_pack_file_size_threshold,
schema_addr,
session_handle,
reader,
)
}
Expand All @@ -307,6 +316,7 @@ fn create_fragment<'a>(
allow_external_blob_outside_bases: JObject, // Optional<Boolean>
blob_pack_file_size_threshold: JObject, // Optional<Long>
schema_addr: jlong,
session_handle: jlong, // Session handle, 0 means no session
source: impl StreamingWriteSource,
) -> Result<JObject<'a>> {
let path_str = dataset_uri.extract(env)?;
Expand All @@ -328,6 +338,10 @@ fn create_fragment<'a>(
&blob_pack_file_size_threshold,
)?;

// Shared session (metadata cache + object store registry) so repeated fragment
// writes do not cold-load the manifest or rebuild object stores per call.
write_params.session = session_from_handle(session_handle);

// Set up storage options provider if namespace is provided
let namespace_info = extract_namespace_info(env, &namespace_obj, &table_id_obj)?;
if let Some((namespace, table_id)) = namespace_info {
Expand Down
68 changes: 62 additions & 6 deletions java/src/main/java/org/lance/Fragment.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -273,10 +273,35 @@ static List<FragmentMetadata> create(
LanceNamespace namespaceClient,
List<String> tableId,
LanceSchema schema) {
return create(datasetUri, allocator, root, params, namespaceClient, tableId, schema, null);
}

/**
* Create a fragment from the given arrow array with an explicit target schema and a shared
* session.
*
* <p>When a session is provided, its metadata cache and object store registry are reused across
* fragment writes, so repeated writes against the same dataset avoid cold manifest loads and
* object store re-creation.
*/
static List<FragmentMetadata> create(
String datasetUri,
BufferAllocator allocator,
VectorSchemaRoot root,
WriteParams params,
LanceNamespace namespaceClient,
List<String> tableId,
LanceSchema schema,
Session session) {
Preconditions.checkNotNull(datasetUri);
Preconditions.checkNotNull(allocator);
Preconditions.checkNotNull(root);
Preconditions.checkNotNull(params);
long sessionHandle = 0L;
if (session != null) {
sessionHandle = session.getNativeHandle();
Preconditions.checkArgument(sessionHandle != 0, "Session is closed");
}
Comment on lines +300 to +304

@coderabbitai coderabbitai Bot Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.rs

Repository: 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.java

Repository: 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.java

Repository: 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:

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-L402
  • java/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.

try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator);
ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) {
Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchema);
Expand All @@ -301,7 +326,8 @@ static List<FragmentMetadata> create(
tableId,
params.getAllowExternalBlobOutsideBases(),
params.getBlobPackFileSizeThreshold(),
lanceSchema.memoryAddress());
lanceSchema.memoryAddress(),
sessionHandle);
}
}
return createWithFfiArray(
Expand All @@ -322,7 +348,8 @@ static List<FragmentMetadata> create(
tableId,
params.getAllowExternalBlobOutsideBases(),
params.getBlobPackFileSizeThreshold(),
0L);
0L,
sessionHandle);
}
}

Expand All @@ -345,9 +372,34 @@ static List<FragmentMetadata> create(
LanceNamespace namespaceClient,
List<String> tableId,
LanceSchema schema) {
return create(datasetUri, allocator, stream, params, namespaceClient, tableId, schema, null);
}

/**
* Create a fragment from the given arrow stream with an explicit target schema and a shared
* session.
*
* <p>When a session is provided, its metadata cache and object store registry are reused across
* fragment writes, so repeated writes against the same dataset avoid cold manifest loads and
* object store re-creation.
*/
static List<FragmentMetadata> create(
String datasetUri,
BufferAllocator allocator,
ArrowArrayStream stream,
WriteParams params,
LanceNamespace namespaceClient,
List<String> tableId,
LanceSchema schema,
Session session) {
Preconditions.checkNotNull(datasetUri);
Preconditions.checkNotNull(stream);
Preconditions.checkNotNull(params);
long sessionHandle = 0L;
if (session != null) {
sessionHandle = session.getNativeHandle();
Preconditions.checkArgument(sessionHandle != 0, "Session is closed");
}
if (schema != null) {
Preconditions.checkNotNull(allocator, "allocator is required with schema");
try (ArrowSchema lanceSchema = ArrowSchema.allocateNew(allocator)) {
Expand All @@ -369,7 +421,8 @@ static List<FragmentMetadata> create(
tableId,
params.getAllowExternalBlobOutsideBases(),
params.getBlobPackFileSizeThreshold(),
lanceSchema.memoryAddress());
lanceSchema.memoryAddress(),
sessionHandle);
}
}
return createWithFfiStream(
Expand All @@ -389,7 +442,8 @@ static List<FragmentMetadata> create(
tableId,
params.getAllowExternalBlobOutsideBases(),
params.getBlobPackFileSizeThreshold(),
0L);
0L,
sessionHandle);
}

/** Create a fragment from the given arrow array and schema. */
Expand All @@ -411,7 +465,8 @@ private static native List<FragmentMetadata> createWithFfiArray(
List<String> tableId,
Optional<Boolean> allowExternalBlobOutsideBases,
Optional<Long> blobPackFileSizeThreshold,
long schemaMemoryAddress);
long schemaMemoryAddress,
long sessionHandle);

/** Create a fragment from the given arrow stream. */
private static native List<FragmentMetadata> createWithFfiStream(
Expand All @@ -431,5 +486,6 @@ private static native List<FragmentMetadata> createWithFfiStream(
List<String> tableId,
Optional<Boolean> allowExternalBlobOutsideBases,
Optional<Long> blobPackFileSizeThreshold,
long schemaMemoryAddress);
long schemaMemoryAddress,
long sessionHandle);
}
22 changes: 20 additions & 2 deletions java/src/main/java/org/lance/WriteFragmentBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ public class WriteFragmentBuilder {
private WriteParams.Builder writeParamsBuilder;
private LanceNamespace namespaceClient;
private List<String> tableId;
private Session session;

WriteFragmentBuilder() {}

Expand Down Expand Up @@ -118,6 +119,21 @@ public WriteFragmentBuilder schema(LanceSchema schema) {
return this;
}

/**
* Set a shared session whose metadata cache and object store registry are reused for this write.
*
* <p>The session must be non-null and open, and must remain open for the duration of the write.
* Omit this call to write without a shared session.
*
* @param session the shared session
* @return this builder
*/
public WriteFragmentBuilder session(Session session) {
Preconditions.checkNotNull(session, "session must not be null; omit session() to not use one");
this.session = session;
return this;
}

/**
* Set the write parameters.
*
Expand Down Expand Up @@ -302,7 +318,8 @@ public List<FragmentMetadata> execute() {
finalWriteParams,
namespaceClient,
tableId,
schema);
schema,
session);
} else {
return Fragment.create(
datasetUri,
Expand All @@ -311,7 +328,8 @@ public List<FragmentMetadata> execute() {
finalWriteParams,
namespaceClient,
tableId,
schema);
schema,
session);
}
}

Expand Down
76 changes: 76 additions & 0 deletions java/src/test/java/org/lance/FragmentTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,82 @@ void testWriteFragmentWithSchemaOverride(@TempDir Path tempDir) throws Exception
}
}

@Test
void testWriteFragmentWithSession(@TempDir Path tempDir) throws Exception {
String datasetPath = tempDir.resolve("fragment_with_session").toString();
try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
TestUtils.SimpleTestDataset testDataset =
new TestUtils.SimpleTestDataset(allocator, datasetPath);
try (Dataset dataset = testDataset.createEmptyDataset();
Session session = Session.builder().build()) {
// Write twice with the same session: APPEND-mode schema inference loads the
// manifest through the session's metadata cache instead of a fresh session
// per write.
for (int i = 0; i < 2; i++) {
try (VectorSchemaRoot root = VectorSchemaRoot.create(dataset.getSchema(), allocator)) {
root.allocateNew();
((VarCharVector) root.getVector("name"))
.setSafe(0, ("Person " + i).getBytes(StandardCharsets.UTF_8));
((IntVector) root.getVector("id")).setSafe(0, i);
root.setRowCount(1);

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());
Comment on lines +172 to +205

@coderabbitai coderabbitai Bot Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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

}
}

// A closed session is rejected instead of silently degrading to no session.
Session closedSession = Session.builder().build();
closedSession.close();
try (VectorSchemaRoot root = VectorSchemaRoot.create(testDataset.getSchema(), allocator)) {
root.allocateNew();
root.setRowCount(0);
assertThrows(
IllegalArgumentException.class,
() ->
Fragment.write()
.datasetUri(datasetPath)
.allocator(allocator)
.data(root)
.mode(WriteParams.WriteMode.APPEND)
.session(closedSession)
.execute());
}
}
}

@Test
void commitWithoutVersion(@TempDir Path tempDir) {
String datasetPath = tempDir.resolve("commit_without_version").toString();
Expand Down
5 changes: 5 additions & 0 deletions rust/lance/src/dataset/fragment/write.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -333,6 +333,11 @@ impl<'a> FragmentCreateBuilder<'a> {
if let Some(accessor) = accessor {
builder = builder.with_storage_options_accessor(accessor);
}
// Reuse the caller-provided session so repeated fragment writes against the same
// dataset hit the session's metadata cache instead of cold-loading the manifest.
if let Some(session) = self.write_params.and_then(|p| p.session.clone()) {
builder = builder.with_session(session);
}
match builder.load().await {
Ok(dataset) => {
// Use the schema from the dataset, because it has the correct
Expand Down
Loading