This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new d670af7a perf: Cache jstrings during metrics collection (#1029)
d670af7a is described below
commit d670af7aba4aa9bfa99c11ec1d0cd8fee15b4410
Author: Matt Butrovich <[email protected]>
AuthorDate: Sat Oct 26 15:54:26 2024 -0400
perf: Cache jstrings during metrics collection (#1029)
* Attempt at caching Jstrings as GlobalRefs in a HashMap to reduce
reallocations. I need to confirm 1) there's actually a performance benefit to
this, and 2) these GlobalRefs are being released when I want them to be.
* Minor refactor and added more docs.
* Undo import reordering to reduce diff.
* Docs.
* Avoid get() by just cloning the Arc to globalref on insert.
* Store jstring cache in ExecutionContext.
---
native/core/src/execution/jni_api.rs | 8 ++++++--
native/core/src/execution/metrics/utils.rs | 28 +++++++++++++++++++++++++---
2 files changed, 31 insertions(+), 5 deletions(-)
diff --git a/native/core/src/execution/jni_api.rs
b/native/core/src/execution/jni_api.rs
index 8347b794..d7e8ccab 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -87,6 +87,8 @@ struct ExecutionContext {
pub debug_native: bool,
/// Whether to write native plans with metrics to stdout
pub explain_native: bool,
+ /// Map of metrics name -> jstring object to cache jni_NewStringUTF calls.
+ pub metrics_jstrings: HashMap<String, Arc<GlobalRef>>,
}
/// Accept serialized query plan and return the address of the native query
plan.
@@ -178,6 +180,7 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_createPlan(
session_ctx: Arc::new(session),
debug_native,
explain_native,
+ metrics_jstrings: HashMap::new(),
});
Ok(Box::into_raw(exec_context) as i64)
@@ -441,10 +444,11 @@ pub extern "system" fn
Java_org_apache_comet_Native_releasePlan(
}
/// Updates the metrics of the query plan.
-fn update_metrics(env: &mut JNIEnv, exec_context: &ExecutionContext) ->
CometResult<()> {
+fn update_metrics(env: &mut JNIEnv, exec_context: &mut ExecutionContext) ->
CometResult<()> {
let native_query = exec_context.root_op.as_ref().unwrap();
let metrics = exec_context.metrics.as_obj();
- update_comet_metric(env, metrics, native_query)
+ let metrics_jstrings = &mut exec_context.metrics_jstrings;
+ update_comet_metric(env, metrics, native_query, metrics_jstrings)
}
fn convert_datatype_arrays(
diff --git a/native/core/src/execution/metrics/utils.rs
b/native/core/src/execution/metrics/utils.rs
index d759f6a8..9291f32c 100644
--- a/native/core/src/execution/metrics/utils.rs
+++ b/native/core/src/execution/metrics/utils.rs
@@ -15,12 +15,15 @@
// specific language governing permissions and limitations
// under the License.
+use crate::jvm_bridge::jni_new_global_ref;
use crate::{
errors::CometError,
jvm_bridge::{jni_call, jni_new_string},
};
use datafusion::physical_plan::ExecutionPlan;
+use jni::objects::{GlobalRef, JString};
use jni::{objects::JObject, JNIEnv};
+use std::collections::HashMap;
use std::sync::Arc;
/// Updates the metrics of a CometMetricNode. This function is called
recursively to
@@ -30,6 +33,7 @@ pub fn update_comet_metric(
env: &mut JNIEnv,
metric_node: &JObject,
execution_plan: &Arc<dyn ExecutionPlan>,
+ metrics_jstrings: &mut HashMap<String, Arc<GlobalRef>>,
) -> Result<(), CometError> {
update_metrics(
env,
@@ -41,6 +45,7 @@ pub fn update_comet_metric(
.map(|m| m.value())
.map(|m| (m.name(), m.as_usize() as i64))
.collect::<Vec<_>>(),
+ metrics_jstrings,
)?;
unsafe {
@@ -51,7 +56,7 @@ pub fn update_comet_metric(
if child_metric_node.is_null() {
continue;
}
- update_comet_metric(env, &child_metric_node, child_plan)?;
+ update_comet_metric(env, &child_metric_node, child_plan,
metrics_jstrings)?;
}
}
Ok(())
@@ -62,11 +67,28 @@ fn update_metrics(
env: &mut JNIEnv,
metric_node: &JObject,
metric_values: &[(&str, i64)],
+ metrics_jstrings: &mut HashMap<String, Arc<GlobalRef>>,
) -> Result<(), CometError> {
unsafe {
for &(name, value) in metric_values {
- let jname = jni_new_string!(env, &name)?;
- jni_call!(env, comet_metric_node(metric_node).set(&jname, value)
-> ())?;
+ // Perform a lookup in the jstrings cache.
+ if let Some(map_global_ref) = metrics_jstrings.get(name) {
+ // Cache hit. Extract the jstring from the global ref.
+ let jobject = map_global_ref.as_obj();
+ let jstring = JString::from_raw(**jobject);
+ // Update the metrics using the jstring as a key.
+ jni_call!(env, comet_metric_node(metric_node).set(&jstring,
value) -> ())?;
+ } else {
+ // Cache miss. Allocate a new string, promote to global ref,
and insert into cache.
+ let local_jstring = jni_new_string!(env, &name)?;
+ let global_ref = jni_new_global_ref!(env, local_jstring)?;
+ let arc_global_ref = Arc::new(global_ref);
+ metrics_jstrings.insert(name.to_string(),
Arc::clone(&arc_global_ref));
+ let jobject = arc_global_ref.as_obj();
+ let jstring = JString::from_raw(**jobject);
+ // Update the metrics using the jstring as a key.
+ jni_call!(env, comet_metric_node(metric_node).set(&jstring,
value) -> ())?;
+ }
}
}
Ok(())
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
