feat: add storage options for LanceFileWriter (lance-format#3900)

Enhance java `LanceFileWriter`, enable it to access object store(like
s3).

Author: leaves12138

java/core/lance-jni/src/blocking_dataset.rs

@@ -15,7 +15,7 @@
 use crate::error::{Error, Result};
 use crate::ffi::JNIEnvExt;
 use crate::traits::{export_vec, import_vec, FromJObjectWithEnv, FromJString};
-use crate::utils::{extract_storage_options, extract_write_params, get_index_params};
+use crate::utils::{extract_storage_options, extract_write_params, get_index_params, to_rust_map};
 use crate::{traits::IntoJava, RT};
 use arrow::array::RecordBatchReader;
 use arrow::datatypes::Schema;
@@ -439,19 +439,7 @@ pub fn inner_commit_overwrite<'local>(
     let path_str = path.extract(env)?;
     let read_version = env.get_u64_opt(&read_version_obj)?;
     let jmap = JMap::from_env(env, &storage_options_obj)?;
-    let storage_options: HashMap<String, String> = env.with_local_frame(16, |env| {
-        let mut map = HashMap::new();
-        let mut iter = jmap.iter(env)?;
-        while let Some((key, value)) = iter.next(env)? {
-            let key_jstring = JString::from(key);
-            let value_jstring = JString::from(value);
-            let key_string: String = env.get_string(&key_jstring)?.into();
-            let value_string: String = env.get_string(&value_jstring)?.into();
-            map.insert(key_string, value_string);
-        }
-        Ok::<_, Error>(map)
-    })?;
-
+    let storage_options = to_rust_map(env, &jmap)?;
     let dataset = BlockingDataset::commit(&path_str, op, read_version, storage_options)?;
     dataset.into_java(env)
 }
@@ -556,20 +544,7 @@ fn inner_open_native<'local>(
     let version = env.get_int_opt(&version_obj)?;
     let block_size = env.get_int_opt(&block_size_obj)?;
     let jmap = JMap::from_env(env, &storage_options_obj)?;
-    let storage_options: HashMap<String, String> = env.with_local_frame(16, |env| {
-        let mut map = HashMap::new();
-        let mut iter = jmap.iter(env)?;
-
-        while let Some((key, value)) = iter.next(env)? {
-            let key_jstring = JString::from(key);
-            let value_jstring = JString::from(value);
-            let key_string: String = env.get_string(&key_jstring)?.into();
-            let value_string: String = env.get_string(&value_jstring)?.into();
-            map.insert(key_string, value_string);
-        }
-
-        Ok::<_, Error>(map)
-    })?;
+    let storage_options = to_rust_map(env, &jmap)?;
     let dataset = BlockingDataset::open(
         &path_str,
         version,

java/core/lance-jni/src/file_writer.rs

@@ -1,5 +1,6 @@
 use std::sync::{Arc, Mutex};
 
+use crate::utils::to_rust_map;
 use crate::{
     error::{Error, Result},
     traits::IntoJava,
@@ -10,6 +11,7 @@ use arrow::{
     ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema},
 };
 use arrow_schema::DataType;
+use jni::objects::JMap;
 use jni::{
     objects::{JObject, JString},
     sys::jlong,
@@ -20,6 +22,7 @@ use lance_file::{
     v2::writer::{FileWriter, FileWriterOptions},
     version::LanceFileVersion,
 };
+use lance_io::object_store::{ObjectStoreParams, ObjectStoreRegistry};
 
 pub const NATIVE_WRITER: &str = "nativeFileWriterHandle";
 
@@ -61,15 +64,31 @@ pub extern "system" fn Java_com_lancedb_lance_file_LanceFileWriter_openNative<'l
     mut env: JNIEnv<'local>,
     _writer_class: JObject,
     file_uri: JString,
+    storage_options_obj: JObject, // Map<String, String>
 ) -> JObject<'local> {
-    ok_or_throw!(env, inner_open(&mut env, file_uri,))
+    ok_or_throw!(env, inner_open(&mut env, file_uri, storage_options_obj))
 }
 
-fn inner_open<'local>(env: &mut JNIEnv<'local>, file_uri: JString) -> Result<JObject<'local>> {
+fn inner_open<'local>(
+    env: &mut JNIEnv<'local>,
+    file_uri: JString,
+    storage_options_obj: JObject,
+) -> Result<JObject<'local>> {
     let file_uri_str: String = env.get_string(&file_uri)?.into();
+    let jmap = JMap::from_env(env, &storage_options_obj)?;
+    let storage_options = to_rust_map(env, &jmap)?;
 
     let writer = RT.block_on(async move {
-        let (obj_store, path) = ObjectStore::from_uri(&file_uri_str).await?;
+        let object_params = ObjectStoreParams {
+            storage_options: Some(storage_options),
+            ..Default::default()
+        };
+        let (obj_store, path) = ObjectStore::from_uri_and_params(
+            Arc::new(ObjectStoreRegistry::default()),
+            &file_uri_str,
+            &object_params,
+        )
+        .await?;
         let obj_store = Arc::new(obj_store);
         let obj_writer = obj_store.create(&path).await?;
 

java/core/lance-jni/src/utils.rs

@@ -287,3 +287,20 @@ pub fn get_index_params(
         )),
     }
 }
+
+pub fn to_rust_map(env: &mut JNIEnv, jmap: &JMap) -> Result<HashMap<String, String>> {
+    env.with_local_frame(16, |env| {
+        let mut map = HashMap::new();
+        let mut iter = jmap.iter(env)?;
+
+        while let Some((key, value)) = iter.next(env)? {
+            let key_jstring = JString::from(key);
+            let value_jstring = JString::from(value);
+            let key_string: String = env.get_string(&key_jstring)?.into();
+            let value_string: String = env.get_string(&value_jstring)?.into();
+            map.insert(key_string, value_string);
+        }
+
+        Ok::<_, Error>(map)
+    })
+}

java/core/src/main/java/com/lancedb/lance/file/LanceFileWriter.java

@@ -23,6 +23,8 @@
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
 
 public class LanceFileWriter implements AutoCloseable {
 
@@ -34,7 +36,8 @@ public class LanceFileWriter implements AutoCloseable {
   private BufferAllocator allocator;
   private DictionaryProvider dictionaryProvider;
 
-  private static native LanceFileWriter openNative(String fileUri) throws IOException;
+  private static native LanceFileWriter openNative(
+      String fileUri, Map<String, String> storageOptions) throws IOException;
 
   private native void closeNative(long nativeLanceFileReaderHandle) throws IOException;
 
@@ -54,7 +57,25 @@ private LanceFileWriter() {}
   public static LanceFileWriter open(
       String path, BufferAllocator allocator, DictionaryProvider dictionaryProvider)
       throws IOException {
-    LanceFileWriter writer = openNative(path);
+    return open(path, allocator, dictionaryProvider, Collections.emptyMap());
+  }
+
+  /**
+   * Open a LanceFileWriter to write to a given file URI
+   *
+   * @param path the URI of the file to write to
+   * @param allocator the BufferAllocator to use for the writer
+   * @param dictionaryProvider the DictionaryProvider to use for the writer
+   * @param storageOptions additional storage options for the writer
+   * @return a new LanceFileWriter
+   */
+  public static LanceFileWriter open(
+      String path,
+      BufferAllocator allocator,
+      DictionaryProvider dictionaryProvider,
+      Map<String, String> storageOptions)
+      throws IOException {
+    LanceFileWriter writer = openNative(path, storageOptions);
     writer.allocator = allocator;
     writer.dictionaryProvider = dictionaryProvider;
     return writer;