Merge branch 'main' into lei/scale_ray_sink

Author: eddyxu

java/core/lance-jni/Cargo.toml

@@ -23,3 +23,5 @@ tokio.workspace = true
 jni = "0.21.1"
 snafu.workspace = true
 lazy_static.workspace = true
+serde = { version = "^1" }
+serde_json = { version = "1" }

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

@@ -12,14 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::{traits::IntoJava, Result, RT};
+use crate::{traits::IntoJava, Error, Result, RT};
 use arrow::array::RecordBatchReader;
 use arrow::ffi::FFI_ArrowSchema;
 use arrow_schema::Schema;
 use jni::sys::jlong;
 use jni::{objects::JObject, JNIEnv};
+use lance::dataset::fragment::FileFragment;
+use lance::dataset::transaction::Operation;
 use lance::dataset::{Dataset, WriteParams};
-
+use snafu::{location, Location};
 pub const NATIVE_DATASET: &str = "nativeDatasetHandle";
 
 #[derive(Clone)]
@@ -42,6 +44,11 @@ impl BlockingDataset {
         Ok(Self { inner })
     }
 
+    pub fn commit(uri: &str, operation: Operation, read_version: Option<u64>) -> Result<Self> {
+        let inner = RT.block_on(Dataset::commit(uri, operation, read_version, None, None))?;
+        Ok(Self { inner })
+    }
+
     pub fn count_rows(&self, filter: Option<String>) -> Result<usize> {
         Ok(RT.block_on(self.inner.count_rows(filter))?)
     }
@@ -100,7 +107,7 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_releaseNativeDataset(
 }
 
 #[no_mangle]
-pub extern "system" fn Java_com_lancedb_lance_Dataset_getFragmentsIds<'a>(
+pub extern "system" fn Java_com_lancedb_lance_Dataset_getJsonFragments<'a>(
     mut env: JNIEnv<'a>,
     jdataset: JObject,
 ) -> JObject<'a> {
@@ -111,13 +118,7 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_getFragmentsIds<'a>(
         dataset.inner.get_fragments()
     };
 
-    let array_list = env
-        .new_int_array(fragments.len() as i32)
-        .expect("Failed to create int array");
-    let fragment_ids = fragments.iter().map(|f| f.id() as i32).collect::<Vec<_>>();
-    env.set_int_array_region(&array_list, 0, &fragment_ids)
-        .expect("Failed to set int array region");
-    array_list.into()
+    ok_or_throw!(env, create_json_fragment_list(&mut env, fragments))
 }
 
 #[no_mangle]
@@ -133,20 +134,34 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_importFfiSchema(
         Schema::from(dataset.inner.schema())
     };
     let out_c_schema = arrow_schema_addr as *mut FFI_ArrowSchema;
-    let c_schema = match FFI_ArrowSchema::try_from(&schema) {
-        Ok(schema) => schema,
-        Err(err) => {
-            env.throw_new(
-                "java/lang/RuntimeException",
-                format!("Failed to convert Arrow schema: {}", err),
-            )
-            .expect("Error throwing exception");
-            return;
-        }
-    };
+    let c_schema = ok_or_throw_without_return!(env, FFI_ArrowSchema::try_from(&schema));
 
     unsafe {
         std::ptr::copy(std::ptr::addr_of!(c_schema), out_c_schema, 1);
         std::mem::forget(c_schema);
     };
 }
+
+fn create_json_fragment_list<'a>(
+    env: &mut JNIEnv<'a>,
+    fragments: Vec<FileFragment>,
+) -> Result<JObject<'a>> {
+    let array_list_class = env.find_class("java/util/ArrayList")?;
+
+    let array_list = env.new_object(array_list_class, "()V", &[])?;
+
+    for fragment in fragments {
+        let json_string = serde_json::to_string(fragment.metadata()).map_err(|e| Error::JSON {
+            message: e.to_string(),
+            location: location!(),
+        })?;
+        let jstring = env.new_string(json_string)?;
+        env.call_method(
+            &array_list,
+            "add",
+            "(Ljava/lang/Object;)Z",
+            &[(&jstring).into()],
+        )?;
+    }
+    Ok(array_list)
+}

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

@@ -14,8 +14,10 @@
 
 use std::str::Utf8Error;
 
+use arrow_schema::ArrowError;
 use jni::errors::Error as JniError;
-use snafu::{Location, Snafu};
+use serde_json::Error as JsonError;
+use snafu::{location, Location, Snafu};
 
 /// Java Exception types
 pub enum JavaException {
@@ -47,6 +49,8 @@ pub enum Error {
     Arrow { message: String, location: Location },
     #[snafu(display("Index error: {}, location", message))]
     Index { message: String, location: Location },
+    #[snafu(display("JSON error: {}, location: {}", message, location))]
+    JSON { message: String, location: Location },
     #[snafu(display("Dataset not found error: {}, location {}", path, location))]
     DatasetNotFound { path: String, location: Location },
     #[snafu(display("Dataset already exists error: {}, location {}", uri, location))]
@@ -66,6 +70,7 @@ impl Error {
             Self::Arrow { .. }
             | Self::DatasetNotFound { .. }
             | Self::DatasetAlreadyExists { .. }
+            | Self::JSON { .. }
             | Self::Other { .. }
             | Self::Jni { .. } => self.throw_as(env, JavaException::RuntimeException),
         }
@@ -96,6 +101,24 @@ impl From<Utf8Error> for Error {
     }
 }
 
+impl From<ArrowError> for Error {
+    fn from(source: ArrowError) -> Self {
+        Self::Arrow {
+            message: source.to_string(),
+            location: location!(),
+        }
+    }
+}
+
+impl From<JsonError> for Error {
+    fn from(source: JsonError) -> Self {
+        Self::JSON {
+            message: source.to_string(),
+            location: location!(),
+        }
+    }
+}
+
 impl From<lance::Error> for Error {
     fn from(source: lance::Error) -> Self {
         match source {

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

@@ -19,6 +19,9 @@ use crate::Result;
 
 /// Extend JNIEnv with helper functions.
 pub trait JNIEnvExt {
+    /// Get integers from Java List<Integer> object.
+    fn get_integers(&mut self, obj: &JObject) -> Result<Vec<i32>>;
+
     /// Get strings from Java List<String> object.
     fn get_strings(&mut self, obj: &JObject) -> Result<Vec<String>>;
 
@@ -34,12 +37,27 @@ pub trait JNIEnvExt {
     /// Get Option<i64> from Java Optional<Long>.
     fn get_long_opt(&mut self, obj: &JObject) -> Result<Option<i64>>;
 
+    /// Get Option<u64> from Java Optional<Long>.
+    fn get_u64_opt(&mut self, obj: &JObject) -> Result<Option<u64>>;
+
     fn get_optional<T, F>(&mut self, obj: &JObject, f: F) -> Result<Option<T>>
     where
         F: FnOnce(&mut JNIEnv, &JObject) -> Result<T>;
 }
 
 impl JNIEnvExt for JNIEnv<'_> {
+    fn get_integers(&mut self, obj: &JObject) -> Result<Vec<i32>> {
+        let list = self.get_list(obj)?;
+        let mut iter = list.iter(self)?;
+        let mut results = Vec::with_capacity(list.size(self)? as usize);
+        while let Some(elem) = iter.next(self)? {
+            let int_obj = self.call_method(elem, "intValue", "()I", &[])?;
+            let int_value = int_obj.i()?;
+            results.push(int_value);
+        }
+        Ok(results)
+    }
+
     fn get_strings(&mut self, obj: &JObject) -> Result<Vec<String>> {
         let list = self.get_list(obj)?;
         let mut iter = list.iter(self)?;
@@ -90,6 +108,16 @@ impl JNIEnvExt for JNIEnv<'_> {
         })
     }
 
+    fn get_u64_opt(&mut self, obj: &JObject) -> Result<Option<u64>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_long_obj = java_obj_gen.l()?;
+            let long_obj = env.call_method(java_long_obj, "longValue", "()J", &[])?;
+            let long_value = long_obj.j()?;
+            Ok(long_value as u64)
+        })
+    }
+
     fn get_optional<T, F>(&mut self, obj: &JObject, f: F) -> Result<Option<T>>
     where
         F: FnOnce(&mut JNIEnv, &JObject) -> Result<T>,