diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index 59ce5e553ff..41a755e62fb 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -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::{ @@ -109,6 +110,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiArray<'local>( allow_external_blob_outside_bases: JObject, // Optional blob_pack_file_size_threshold: JObject, // Optional schema_addr: jlong, + session_handle: jlong, // Session handle, 0 means no session ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -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() ) @@ -158,6 +161,7 @@ fn inner_create_with_ffi_array<'local>( allow_external_blob_outside_bases: JObject, // Optional blob_pack_file_size_threshold: JObject, // Optional schema_addr: jlong, + session_handle: jlong, // Session handle, 0 means no session ) -> Result> { let c_array_ptr = arrow_array_addr as *mut FFI_ArrowArray; let c_schema_ptr = arrow_schema_addr as *mut FFI_ArrowSchema; @@ -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, ) } @@ -215,6 +220,7 @@ pub extern "system" fn Java_org_lance_Fragment_createWithFfiStream<'a>( allow_external_blob_outside_bases: JObject, // Optional blob_pack_file_size_threshold: JObject, // Optional schema_addr: jlong, + session_handle: jlong, // Session handle, 0 means no session ) -> JObject<'a> { ok_or_throw_with_return!( env, @@ -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() ) @@ -262,6 +269,7 @@ fn inner_create_with_ffi_stream<'local>( allow_external_blob_outside_bases: JObject, // Optional blob_pack_file_size_threshold: JObject, // Optional schema_addr: jlong, + session_handle: jlong, // Session handle, 0 means no session ) -> Result> { let stream_ptr = arrow_array_stream_addr as *mut FFI_ArrowArrayStream; let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?; @@ -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, ) } @@ -307,6 +316,7 @@ fn create_fragment<'a>( allow_external_blob_outside_bases: JObject, // Optional blob_pack_file_size_threshold: JObject, // Optional schema_addr: jlong, + session_handle: jlong, // Session handle, 0 means no session source: impl StreamingWriteSource, ) -> Result> { let path_str = dataset_uri.extract(env)?; @@ -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 { diff --git a/java/src/main/java/org/lance/Fragment.java b/java/src/main/java/org/lance/Fragment.java index 3b12e158617..0462231f29e 100644 --- a/java/src/main/java/org/lance/Fragment.java +++ b/java/src/main/java/org/lance/Fragment.java @@ -273,10 +273,35 @@ static List create( LanceNamespace namespaceClient, List 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. + * + *

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 create( + String datasetUri, + BufferAllocator allocator, + VectorSchemaRoot root, + WriteParams params, + LanceNamespace namespaceClient, + List 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"); + } try (ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray arrowArray = ArrowArray.allocateNew(allocator)) { Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchema); @@ -301,7 +326,8 @@ static List create( tableId, params.getAllowExternalBlobOutsideBases(), params.getBlobPackFileSizeThreshold(), - lanceSchema.memoryAddress()); + lanceSchema.memoryAddress(), + sessionHandle); } } return createWithFfiArray( @@ -322,7 +348,8 @@ static List create( tableId, params.getAllowExternalBlobOutsideBases(), params.getBlobPackFileSizeThreshold(), - 0L); + 0L, + sessionHandle); } } @@ -345,9 +372,34 @@ static List create( LanceNamespace namespaceClient, List 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. + * + *

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 create( + String datasetUri, + BufferAllocator allocator, + ArrowArrayStream stream, + WriteParams params, + LanceNamespace namespaceClient, + List 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)) { @@ -369,7 +421,8 @@ static List create( tableId, params.getAllowExternalBlobOutsideBases(), params.getBlobPackFileSizeThreshold(), - lanceSchema.memoryAddress()); + lanceSchema.memoryAddress(), + sessionHandle); } } return createWithFfiStream( @@ -389,7 +442,8 @@ static List create( tableId, params.getAllowExternalBlobOutsideBases(), params.getBlobPackFileSizeThreshold(), - 0L); + 0L, + sessionHandle); } /** Create a fragment from the given arrow array and schema. */ @@ -411,7 +465,8 @@ private static native List createWithFfiArray( List tableId, Optional allowExternalBlobOutsideBases, Optional blobPackFileSizeThreshold, - long schemaMemoryAddress); + long schemaMemoryAddress, + long sessionHandle); /** Create a fragment from the given arrow stream. */ private static native List createWithFfiStream( @@ -431,5 +486,6 @@ private static native List createWithFfiStream( List tableId, Optional allowExternalBlobOutsideBases, Optional blobPackFileSizeThreshold, - long schemaMemoryAddress); + long schemaMemoryAddress, + long sessionHandle); } diff --git a/java/src/main/java/org/lance/WriteFragmentBuilder.java b/java/src/main/java/org/lance/WriteFragmentBuilder.java index 2dbef873849..d398d644297 100644 --- a/java/src/main/java/org/lance/WriteFragmentBuilder.java +++ b/java/src/main/java/org/lance/WriteFragmentBuilder.java @@ -51,6 +51,7 @@ public class WriteFragmentBuilder { private WriteParams.Builder writeParamsBuilder; private LanceNamespace namespaceClient; private List tableId; + private Session session; WriteFragmentBuilder() {} @@ -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. + * + *

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. * @@ -302,7 +318,8 @@ public List execute() { finalWriteParams, namespaceClient, tableId, - schema); + schema, + session); } else { return Fragment.create( datasetUri, @@ -311,7 +328,8 @@ public List execute() { finalWriteParams, namespaceClient, tableId, - schema); + schema, + session); } } diff --git a/java/src/test/java/org/lance/FragmentTest.java b/java/src/test/java/org/lance/FragmentTest.java index 29a21b5258a..25ccbb920d7 100644 --- a/java/src/test/java/org/lance/FragmentTest.java +++ b/java/src/test/java/org/lance/FragmentTest.java @@ -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 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 fragments = + Fragment.write() + .datasetUri(datasetPath) + .allocator(allocator) + .data(root) + .schema(dataset.getLanceSchema()) + .mode(WriteParams.WriteMode.APPEND) + .session(session) + .execute(); + assertEquals(1, fragments.size()); + } + } + + // 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(); diff --git a/rust/lance/src/dataset/fragment/write.rs b/rust/lance/src/dataset/fragment/write.rs index 2ab3bd1d92b..4938c340659 100644 --- a/rust/lance/src/dataset/fragment/write.rs +++ b/rust/lance/src/dataset/fragment/write.rs @@ -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