Skip to content
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

chore: Remove some calls to unwrap #598

Merged
merged 4 commits into from
Jun 27, 2024
Merged

chore: Remove some calls to unwrap #598

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ced520f
remove an unwrap from ScanExec
andygrove Jun 27, 2024
c80f9ed
remove another unwrap
andygrove Jun 27, 2024
4bff1dd
remove another unwrap
andygrove Jun 27, 2024
4f50a53
remove more unwraps
andygrove Jun 27, 2024
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
2 changes: 1 addition & 1 deletion core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ strip = "debuginfo"
[lib]
name = "comet"
# "rlib" is for benchmarking with criterion.
crate_type = ["cdylib", "rlib"]
Copy link
Member Author

Choose a reason for hiding this comment

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

drive-by cleanup to fix a warning since we upgraded to Rust 1.79. Unrelated to other changes in this PR.

crate-type = ["cdylib", "rlib"]

[[bench]]
name = "parquet_read"
Expand Down
2 changes: 1 addition & 1 deletion core/src/execution/datafusion/expressions/subquery.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ impl PhysicalExpr for Subquery {
}

fn evaluate(&self, _: &RecordBatch) -> datafusion_common::Result<ColumnarValue> {
let mut env = JVMClasses::get_env();
let mut env = JVMClasses::get_env()?;

unsafe {
let is_null = jni_static_call!(&mut env,
Expand Down
4 changes: 2 additions & 2 deletions core/src/execution/memory_pool.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ impl CometMemoryPool {
}

fn acquire(&self, additional: usize) -> CometResult<i64> {
let mut env = JVMClasses::get_env();
let mut env = JVMClasses::get_env()?;
let handle = self.task_memory_manager_handle.as_obj();
unsafe {
jni_call!(&mut env,
Expand All @@ -68,7 +68,7 @@ impl CometMemoryPool {
}

fn release(&self, size: usize) -> CometResult<()> {
let mut env = JVMClasses::get_env();
let mut env = JVMClasses::get_env()?;
let handle = self.task_memory_manager_handle.as_obj();
unsafe {
jni_call!(&mut env, comet_task_memory_manager(handle).release_memory(size as i64) -> ())
Expand Down
9 changes: 4 additions & 5 deletions core/src/execution/operators/scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ impl ScanExec {
return Ok(InputBatch::EOF);
}

let mut env = JVMClasses::get_env();
let mut env = JVMClasses::get_env()?;

if iter.is_null() {
return Err(CometError::from(ExecutionError::GeneralError(format!(
Expand Down Expand Up @@ -318,13 +318,12 @@ impl ScanStream {
.zip(schema_fields.iter())
.map(|(column, f)| {
if column.data_type() != f.data_type() {
cast_with_options(column, f.data_type(), &cast_options).unwrap()
cast_with_options(column, f.data_type(), &cast_options)
} else {
column.clone()
Ok(column.clone())
}
})
.collect();

.collect::<Result<Vec<_>, _>>()?;
let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
RecordBatch::try_new_with_options(self.schema.clone(), new_columns, &options)
.map_err(|e| arrow_datafusion_err!(e))
Expand Down
2 changes: 1 addition & 1 deletion core/src/jvm_bridge/batch_iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ impl<'a> CometBatchIterator<'a> {

Ok(CometBatchIterator {
class,
method_next: env.get_method_id(Self::JVM_CLASS, "next", "()[J").unwrap(),
method_next: env.get_method_id(Self::JVM_CLASS, "next", "()[J")?,
method_next_ret: ReturnType::Array,
})
}
Expand Down
52 changes: 19 additions & 33 deletions core/src/jvm_bridge/comet_exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,49 +56,35 @@ impl<'a> CometExec<'a> {
let class = env.find_class(Self::JVM_CLASS)?;

Ok(CometExec {
method_get_bool: env
.get_static_method_id(Self::JVM_CLASS, "getBoolean", "(JJ)Z")
.unwrap(),
method_get_bool: env.get_static_method_id(Self::JVM_CLASS, "getBoolean", "(JJ)Z")?,
method_get_bool_ret: ReturnType::Primitive(Primitive::Boolean),
method_get_byte: env
.get_static_method_id(Self::JVM_CLASS, "getByte", "(JJ)B")
.unwrap(),
method_get_byte: env.get_static_method_id(Self::JVM_CLASS, "getByte", "(JJ)B")?,
method_get_byte_ret: ReturnType::Primitive(Primitive::Byte),
method_get_short: env
.get_static_method_id(Self::JVM_CLASS, "getShort", "(JJ)S")
.unwrap(),
method_get_short: env.get_static_method_id(Self::JVM_CLASS, "getShort", "(JJ)S")?,
method_get_short_ret: ReturnType::Primitive(Primitive::Short),
method_get_int: env
.get_static_method_id(Self::JVM_CLASS, "getInt", "(JJ)I")
.unwrap(),
method_get_int: env.get_static_method_id(Self::JVM_CLASS, "getInt", "(JJ)I")?,
method_get_int_ret: ReturnType::Primitive(Primitive::Int),
method_get_long: env
.get_static_method_id(Self::JVM_CLASS, "getLong", "(JJ)J")
.unwrap(),
method_get_long: env.get_static_method_id(Self::JVM_CLASS, "getLong", "(JJ)J")?,
method_get_long_ret: ReturnType::Primitive(Primitive::Long),
method_get_float: env
.get_static_method_id(Self::JVM_CLASS, "getFloat", "(JJ)F")
.unwrap(),
method_get_float: env.get_static_method_id(Self::JVM_CLASS, "getFloat", "(JJ)F")?,
method_get_float_ret: ReturnType::Primitive(Primitive::Float),
method_get_double: env
.get_static_method_id(Self::JVM_CLASS, "getDouble", "(JJ)D")
.unwrap(),
method_get_double: env.get_static_method_id(Self::JVM_CLASS, "getDouble", "(JJ)D")?,
method_get_double_ret: ReturnType::Primitive(Primitive::Double),
method_get_decimal: env
.get_static_method_id(Self::JVM_CLASS, "getDecimal", "(JJ)[B")
.unwrap(),
method_get_decimal: env.get_static_method_id(
Self::JVM_CLASS,
"getDecimal",
"(JJ)[B",
)?,
method_get_decimal_ret: ReturnType::Array,
method_get_string: env
.get_static_method_id(Self::JVM_CLASS, "getString", "(JJ)Ljava/lang/String;")
.unwrap(),
method_get_string: env.get_static_method_id(
Self::JVM_CLASS,
"getString",
"(JJ)Ljava/lang/String;",
)?,
method_get_string_ret: ReturnType::Object,
method_get_binary: env
.get_static_method_id(Self::JVM_CLASS, "getBinary", "(JJ)[B")
.unwrap(),
method_get_binary: env.get_static_method_id(Self::JVM_CLASS, "getBinary", "(JJ)[B")?,
method_get_binary_ret: ReturnType::Array,
method_is_null: env
.get_static_method_id(Self::JVM_CLASS, "isNull", "(JJ)Z")
.unwrap(),
method_is_null: env.get_static_method_id(Self::JVM_CLASS, "isNull", "(JJ)Z")?,
method_is_null_ret: ReturnType::Primitive(Primitive::Boolean),
class,
})
Expand Down
16 changes: 6 additions & 10 deletions core/src/jvm_bridge/comet_metric_node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,17 +38,13 @@ impl<'a> CometMetricNode<'a> {
let class = env.find_class(Self::JVM_CLASS)?;

Ok(CometMetricNode {
method_get_child_node: env
.get_method_id(
Self::JVM_CLASS,
"getChildNode",
format!("(I)L{:};", Self::JVM_CLASS).as_str(),
)
.unwrap(),
method_get_child_node: env.get_method_id(
Self::JVM_CLASS,
"getChildNode",
format!("(I)L{:};", Self::JVM_CLASS).as_str(),
)?,
method_get_child_node_ret: ReturnType::Object,
method_add: env
.get_method_id(Self::JVM_CLASS, "add", "(Ljava/lang/String;J)V")
.unwrap(),
method_add: env.get_method_id(Self::JVM_CLASS, "add", "(Ljava/lang/String;J)V")?,
method_add_ret: ReturnType::Primitive(Primitive::Void),
class,
})
Expand Down
13 changes: 9 additions & 4 deletions core/src/jvm_bridge/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ macro_rules! jni_call {
let ret = $env.call_method_unchecked($obj, method_id, ret_type, args);

// Check if JVM has thrown any exception, and handle it if so.
let result = if let Some(exception) = $crate::jvm_bridge::check_exception($env).unwrap() {
let result = if let Some(exception) = $crate::jvm_bridge::check_exception($env)? {
Err(exception.into())
} else {
$crate::jvm_bridge::jni_map_error!($env, ret)
Expand All @@ -100,7 +100,7 @@ macro_rules! jni_static_call {
let ret = $env.call_static_method_unchecked(clazz, method_id, ret_type, args);

// Check if JVM has thrown any exception, and handle it if so.
let result = if let Some(exception) = $crate::jvm_bridge::check_exception($env).unwrap() {
let result = if let Some(exception) = $crate::jvm_bridge::check_exception($env)? {
Err(exception.into())
} else {
$crate::jvm_bridge::jni_map_error!($env, ret)
Expand Down Expand Up @@ -262,10 +262,15 @@ impl JVMClasses<'_> {
}

/// Gets the JNIEnv for the current thread.
pub fn get_env() -> AttachGuard<'static> {
pub fn get_env() -> CometResult<AttachGuard<'static>> {
unsafe {
let java_vm = JAVA_VM.get_unchecked();
java_vm.attach_current_thread().unwrap()
java_vm.attach_current_thread().map_err(|e| {
CometError::Internal(format!(
"JVMClasses::get_env() failed to attach current thread: {}",
e
))
})
}
}
}
Expand Down
Loading