add TID to logging

native-engine/blaze-serde/proto/blaze.proto

@@ -723,9 +723,9 @@ message ExpandProjection {
 
 // Unique identifier for a materialized partition of data
 message PartitionId {
-  string job_id = 1;
   uint32 stage_id = 2;
   uint32 partition_id = 4;
+  uint64 task_id = 5;
 }
 
 message TaskDefinition {

native-engine/blaze/src/logging.rs

@@ -18,6 +18,7 @@ use log::{Level, LevelFilter, Log, Metadata, Record};
 use once_cell::sync::OnceCell;
 
 thread_local! {
+    pub static THREAD_TID: Cell<usize> = Cell::new(0);
     pub static THREAD_STAGE_ID: Cell<usize> = Cell::new(0);
     pub static THREAD_PARTITION_ID: Cell<usize> = Cell::new(0);
 }
@@ -50,8 +51,9 @@ impl Log for SimpleLogger {
             let elapsed_sec = elapsed.as_secs_f64();
             let stage_id = THREAD_STAGE_ID.get();
             let partition_id = THREAD_PARTITION_ID.get();
+            let tid = THREAD_TID.get();
             eprintln!(
-                "(+{elapsed_sec:.3}s) [{}] (stage: {stage_id}, partition: {partition_id}) - {}",
+                "(+{elapsed_sec:.3}s) [{}] (stage: {stage_id}, partition: {partition_id}, tid: {tid}) - {}",
                 record.level(),
                 record.args()
             );

native-engine/blaze/src/rt.rs

@@ -53,7 +53,7 @@ use tokio::{runtime::Runtime, task::JoinHandle};
 
 use crate::{
     handle_unwinded_scope,
-    logging::{THREAD_PARTITION_ID, THREAD_STAGE_ID},
+    logging::{THREAD_PARTITION_ID, THREAD_STAGE_ID, THREAD_TID},
     metrics::update_spark_metric_node,
 };
 
@@ -80,6 +80,7 @@ impl NativeExecutionRuntime {
         let task_id = &task_definition.task_id.expect("task_id is empty");
         let stage_id = task_id.stage_id as usize;
         let partition_id = task_id.partition_id as usize;
+        let tid = task_id.task_id as usize;
         let plan = &task_definition.plan.expect("plan is empty");
         drop(raw_task_definition);
 
@@ -118,6 +119,7 @@ impl NativeExecutionRuntime {
                 );
                 THREAD_STAGE_ID.set(stage_id);
                 THREAD_PARTITION_ID.set(partition_id);
+                THREAD_TID.set(tid);
             });
         if num_worker_threads > 0 {
             tokio_runtime_builder.worker_threads(num_worker_threads as usize);

spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeCallNativeWrapper.scala

@@ -150,7 +150,7 @@ case class BlazeCallNativeWrapper(
       .newBuilder()
       .setPartitionId(partition.index)
       .setStageId(context.map(_.stageId()).getOrElse(0))
-      .setJobId(partition.index.toString)
+      .setTaskId(context.map(_.taskAttemptId()).getOrElse(0))
       .build()
 
     val taskDefinition = TaskDefinition