code refactoring

native-engine/blaze-jni-bridge/src/conf.rs

@@ -30,12 +30,9 @@ macro_rules! define_conf {
 
 define_conf!(IntConf, BATCH_SIZE);
 define_conf!(DoubleConf, MEMORY_FRACTION);
+define_conf!(IntConf, TOKIO_NUM_WORKER_THREADS);
 define_conf!(BooleanConf, SMJ_INEQUALITY_JOIN_ENABLE);
-define_conf!(BooleanConf, BHJ_FALLBACKS_TO_SMJ_ENABLE);
-define_conf!(IntConf, BHJ_FALLBACKS_TO_SMJ_ROWS_THRESHOLD);
-define_conf!(IntConf, BHJ_FALLBACKS_TO_SMJ_MEM_THRESHOLD);
 define_conf!(BooleanConf, CASE_CONVERT_FUNCTIONS_ENABLE);
-define_conf!(IntConf, UDF_WRAPPER_NUM_THREADS);
 define_conf!(BooleanConf, INPUT_BATCH_STATISTICS_ENABLE);
 define_conf!(BooleanConf, IGNORE_CORRUPTED_FILES);
 define_conf!(BooleanConf, PARTIAL_AGG_SKIPPING_ENABLE);

native-engine/blaze/Cargo.toml

@@ -25,7 +25,7 @@ once_cell = "1.20.2"
 panic-message = "0.3.0"
 paste = "1.0.15"
 prost = "0.13.3"
-tokio = "1.41"
+tokio = "=1.41.0"
 
 [target.'cfg(not(windows))'.dependencies]
 jemallocator = { version = "0.5.0", features = ["disable_initial_exec_tls"] }

native-engine/blaze/src/exec.rs

@@ -19,25 +19,21 @@ use blaze_jni_bridge::{
     jni_bridge::JavaClasses,
     *,
 };
-use blaze_serde::protobuf::TaskDefinition;
 use datafusion::{
     common::Result,
     error::DataFusionError,
     execution::{
         disk_manager::DiskManagerConfig,
         runtime_env::{RuntimeConfig, RuntimeEnv},
     },
-    physical_plan::{displayable, ExecutionPlan},
     prelude::{SessionConfig, SessionContext},
 };
-use datafusion_ext_commons::df_execution_err;
 use datafusion_ext_plans::memmgr::MemManager;
 use jni::{
     objects::{JClass, JObject},
     JNIEnv,
 };
 use once_cell::sync::OnceCell;
-use prost::Message;
 
 use crate::{handle_unwinded_scope, logging::init_logging, rt::NativeExecutionRuntime};
 
@@ -81,38 +77,11 @@ pub extern "system" fn Java_org_apache_spark_sql_blaze_JniBridge_callNative(
         })?;
         let native_wrapper = jni_new_global_ref!(native_wrapper)?;
 
-        // decode plan
-        let raw_task_definition = jni_call!(
-            BlazeCallNativeWrapper(native_wrapper.as_obj())
-                .getRawTaskDefinition() -> JObject)?;
-        let task_definition = TaskDefinition::decode(
-            jni_convert_byte_array!(raw_task_definition.as_obj())?.as_slice(),
-        )
-        .or_else(|err| df_execution_err!("cannot decode execution plan: {err:?}"))?;
-
-        let task_id = &task_definition.task_id.expect("task_id is empty");
-        let plan = &task_definition.plan.expect("plan is empty");
-        drop(raw_task_definition);
-
-        // get execution plan
-        let execution_plan: Arc<dyn ExecutionPlan> = plan
-            .try_into()
-            .or_else(|err| df_execution_err!("cannot create execution plan: {err:?}"))?;
-        let execution_plan_displayable = displayable(execution_plan.as_ref())
-            .indent(true)
-            .to_string();
-        log::info!("Creating native execution plan succeeded");
-        log::info!("  task_id={task_id:?}");
-        log::info!("  execution plan:\n{execution_plan_displayable}");
-
-        // execute to stream
+        // create execution runtime
         let runtime = Box::new(NativeExecutionRuntime::start(
             native_wrapper,
-            execution_plan,
-            task_id.partition_id as usize,
             SESSION.get().unwrap().task_ctx(),
         )?);
-        log::info!("Blaze native thread created");
 
         // returns runtime raw pointer
         Ok::<_, DataFusionError>(Box::into_raw(runtime) as usize as i64)

native-engine/blaze/src/logging.rs

@@ -12,11 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::time::Instant;
+use std::{cell::Cell, time::Instant};
 
 use log::{Level, LevelFilter, Log, Metadata, Record};
 use once_cell::sync::OnceCell;
 
+thread_local! {
+    pub static THREAD_STAGE_ID: Cell<usize> = Cell::new(0);
+    pub static THREAD_PARTITION_ID: Cell<usize> = Cell::new(0);
+}
+
 const MAX_LEVEL: Level = Level::Info;
 
 pub fn init_logging() {
@@ -43,8 +48,10 @@ impl Log for SimpleLogger {
         if self.enabled(record.metadata()) {
             let elapsed = Instant::now() - self.start_instant;
             let elapsed_sec = elapsed.as_secs_f64();
+            let stage_id = THREAD_STAGE_ID.get();
+            let partition_id = THREAD_PARTITION_ID.get();
             eprintln!(
-                "(+{elapsed_sec:.3}s) [{}] Blaze - {}",
+                "(+{elapsed_sec:.3}s) [{}] (stage: {stage_id}, partition: {partition_id}) - {}",
                 record.level(),
                 record.args()
             );