1852: NewType pattern for job_id and job_name

.cursor/rules.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

.gemini/rules.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

AGENTS.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

CLAUDE.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

ballista/core/src/execution_plans/distributed_explain_analyze.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::JobId;
 use crate::execution_plans::DistributedQueryExec;
 use crate::extension::SessionConfigExt;
 use crate::serde::protobuf::{
@@ -204,7 +205,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedExplainAnalyzeExec
 
 async fn fetch_job_metrics(
     scheduler_url: &str,
-    job_id: &str,
+    job_id: &JobId,
     session_config: datafusion::prelude::SessionConfig,
 ) -> Result<GetJobMetricsResult> {
     let grpc_interceptor = session_config.ballista_grpc_interceptor();
@@ -248,7 +249,7 @@ async fn fetch_job_metrics(
 
 fn format_metrics_as_record_batch(
     job_metrics: &GetJobMetricsResult,
-    _job_id: &str,
+    _job_id: &JobId,
     schema: SchemaRef,
     _verbose: bool,
 ) -> Result<RecordBatch> {
@@ -309,6 +310,7 @@ mod tests {
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use std::sync::Arc;
 
+    use crate::JobId;
     use crate::serde::protobuf::{
         GetJobMetricsResult, JobStageMetrics, OperatorMetric, OperatorWithMetrics,
         operator_metric,
@@ -357,7 +359,8 @@ mod tests {
         };
 
         let batch =
-            format_metrics_as_record_batch(&response, "job-1", schema, true).unwrap();
+            format_metrics_as_record_batch(&response, &JobId::new("job-1"), schema, true)
+                .unwrap();
 
         let plan_type = batch
             .column(0)

ballista/core/src/execution_plans/distributed_query.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::JobId;
 use crate::client::BallistaClient;
 use crate::config::BallistaConfig;
 use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt};
@@ -84,7 +85,7 @@ pub struct DistributedQueryExec<T: 'static + AsLogicalPlan> {
     /// - job_scheduling_in_ms: Time job waited in scheduler queue (started_at - queued_at)
     metrics: ExecutionPlanMetricsSet,
     /// The scheduler job id after the query has been accepted.
-    job_id: Arc<Mutex<Option<String>>>,
+    job_id: Arc<Mutex<Option<JobId>>>,
 }
 
 impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
@@ -136,7 +137,7 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
     }
 
     /// Returns the scheduler job id after the query has been accepted.
-    pub fn job_id(&self) -> Option<String> {
+    pub fn job_id(&self) -> Option<JobId> {
         self.job_id.lock().clone()
     }
 
@@ -336,7 +337,7 @@ async fn execute_query_pull(
     max_message_size: usize,
     grpc_config: GrpcClientConfig,
     metrics: Arc<ExecutionPlanMetricsSet>,
-    job_id_handle: Arc<Mutex<Option<String>>>,
+    job_id_handle: Arc<Mutex<Option<JobId>>>,
     partition: usize,
     session_config: SessionConfig,
 ) -> Result<impl Stream<Item = Result<RecordBatch>> + Send> {
@@ -394,7 +395,7 @@ async fn execute_query_pull(
         "Session id inconsistent between Client and Server side in DistributedQueryExec."
     );
 
-    let job_id = query_result.job_id;
+    let job_id: JobId = query_result.job_id.into();
     *job_id_handle.lock() = Some(job_id.clone());
     let mut prev_status: Option<job_status::Status> = None;
 
@@ -404,7 +405,7 @@ async fn execute_query_pull(
             flight_proxy,
         } = scheduler
             .get_job_status(GetJobStatusParams {
-                job_id: job_id.clone(),
+                job_id: job_id.clone().into(),
             })
             .await
             .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?
@@ -509,7 +510,7 @@ async fn execute_query_push(
     max_message_size: usize,
     grpc_config: GrpcClientConfig,
     metrics: Arc<ExecutionPlanMetricsSet>,
-    job_id_handle: Arc<Mutex<Option<String>>>,
+    job_id_handle: Arc<Mutex<Option<JobId>>>,
     partition: usize,
     session_config: SessionConfig,
 ) -> Result<impl Stream<Item = Result<RecordBatch>> + Send> {
@@ -568,11 +569,12 @@ async fn execute_query_push(
             status,
             flight_proxy,
         } = item;
-        let job_id = status
+        let job_id: JobId = status
             .as_ref()
             .map(|s| s.job_id.to_owned())
-            .unwrap_or("unknown_job_id".to_string()); // should not happen
-        if !job_id.starts_with("unknown_") {
+            .unwrap_or("unknown_job_id".to_string()) // should not happen
+            .into();
+        if !job_id.as_str().starts_with("unknown_") {
             let mut shared_job_id = job_id_handle.lock();
             if shared_job_id.is_none() {
                 *shared_job_id = Some(job_id.clone());
@@ -761,6 +763,7 @@ async fn fetch_partition(
 
 #[cfg(test)]
 mod test {
+    use crate::JobId;
     use crate::config::BallistaConfig;
     use crate::execution_plans::distributed_query::{
         DistributedQueryExec, get_client_host_port,
@@ -835,14 +838,14 @@ mod test {
             LogicalPlan::default(),
             "session".to_string(),
         ));
-        *exec.job_id.lock() = Some("job-123".to_string());
+        *exec.job_id.lock() = Some("job-123".to_owned().into());
 
         let new_exec = exec.clone().with_new_children(vec![]).unwrap();
         let new_exec = new_exec
             .as_any()
             .downcast_ref::<DistributedQueryExec<LogicalPlanNode>>()
             .unwrap();
 
-        assert_eq!(new_exec.job_id().as_deref(), Some("job-123"));
+        assert_eq!(new_exec.job_id(), Some(JobId::new("job-123")));
     }
 }

ballista/core/src/execution_plans/mod.rs

@@ -41,6 +41,8 @@ pub use shuffle_writer_trait::ShuffleWriter;
 pub use sort_shuffle::SortShuffleWriterExec;
 pub use unresolved_shuffle::UnresolvedShuffleExec;
 
+use crate::JobId;
+
 /// Creates the file path for a shuffle output partition.
 ///
 /// The path structure depends on the shuffle type:
@@ -67,7 +69,7 @@ pub use unresolved_shuffle::UnresolvedShuffleExec;
 /// - `is_sort_shuffle` — selects between sort-shuffle and hash-shuffle path layout
 pub fn create_shuffle_path<P: AsRef<Path>>(
     work_dir: P,
-    job_id: &str,
+    job_id: &JobId,
     stage_id: usize,
     partition_id: usize,
     file_id: Option<u64>,
@@ -76,7 +78,7 @@ pub fn create_shuffle_path<P: AsRef<Path>>(
     let mut path = PathBuf::new();
 
     path.push(work_dir);
-    path.push(job_id);
+    path.push(job_id.as_str());
     path.push(stage_id.to_string());
 
     match (file_id, is_sort_shuffle) {
@@ -106,25 +108,30 @@ mod test {
 
     #[test]
     fn test_regular_shuffle_with_file_id() {
-        let path = create_shuffle_path("/work", "job1", 2, 3, Some(42), false).unwrap();
+        let path =
+            create_shuffle_path("/work", &JobId::new("job1"), 2, 3, Some(42), false)
+                .unwrap();
         assert_eq!(path, PathBuf::from("/work/job1/2/3/data-42.arrow"));
     }
 
     #[test]
     fn test_regular_shuffle_without_file_id() {
-        let path = create_shuffle_path("/work", "job1", 2, 3, None, false).unwrap();
+        let path =
+            create_shuffle_path("/work", &JobId::new("job1"), 2, 3, None, false).unwrap();
         assert_eq!(path, PathBuf::from("/work/job1/2/3/data.arrow"));
     }
 
     #[test]
     fn test_sort_shuffle_with_file_id() {
-        let path = create_shuffle_path("/work", "job1", 2, 3, Some(42), true).unwrap();
+        let path =
+            create_shuffle_path("/work", &JobId::new("job1"), 2, 3, Some(42), true)
+                .unwrap();
         assert_eq!(path, PathBuf::from("/work/job1/2/42/data.arrow"));
     }
 
     #[test]
     fn test_sort_shuffle_without_file_id_returns_error() {
-        let result = create_shuffle_path("/work", "job1", 2, 3, None, true);
+        let result = create_shuffle_path("/work", &JobId::new("job1"), 2, 3, None, true);
         assert!(result.is_err());
     }
 
@@ -135,8 +142,15 @@ mod test {
         for (file_id, is_sort_shuffle) in
             [(Some(1), false), (None, false), (Some(1), true)]
         {
-            let path =
-                create_shuffle_path("/", "job1", 2, 3, file_id, is_sort_shuffle).unwrap();
+            let path = create_shuffle_path(
+                "/",
+                &JobId::new("job1"),
+                2,
+                3,
+                file_id,
+                is_sort_shuffle,
+            )
+            .unwrap();
             assert!(
                 path.parent().is_some(),
                 "path {path:?} (file_id={file_id:?}, is_sort_shuffle={is_sort_shuffle}) has no parent"

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -984,7 +984,7 @@ mod tests {
         ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification,
         PartitionId,
     };
-    use crate::utils;
+    use crate::{JobId, utils};
     use datafusion::arrow::array::{Int32Array, StringArray, UInt32Array};
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::arrow::ipc::writer::StreamWriter;
@@ -1018,7 +1018,7 @@ mod tests {
                 vec![PartitionLocation {
                     map_partition_id: 0,
                     partition_id: PartitionId {
-                        job_id: job_id.to_string(),
+                        job_id: job_id.to_owned().into(),
                         stage_id,
                         partition_id: i,
                     },
@@ -1137,7 +1137,7 @@ mod tests {
             partitions.push(PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: job_id.to_string(),
+                    job_id: job_id.to_owned().into(),
                     stage_id: input_stage_id,
                     partition_id,
                 },
@@ -1188,7 +1188,7 @@ mod tests {
             partitions.push(PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: job_id.to_string(),
+                    job_id: job_id.to_owned().into(),
                     stage_id: input_stage_id,
                     partition_id,
                 },
@@ -1240,7 +1240,7 @@ mod tests {
             partitions.push(PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: job_id.to_string(),
+                    job_id: job_id.to_owned().into(),
                     stage_id: input_stage_id,
                     partition_id,
                 },
@@ -1292,7 +1292,7 @@ mod tests {
             partitions.push(PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: job_id.to_string(),
+                    job_id: job_id.to_owned().into(),
                     stage_id: input_stage_id,
                     partition_id,
                 },
@@ -1350,7 +1350,7 @@ mod tests {
         let task_ctx = session_ctx.task_ctx();
         let work_dir = TempDir::new().unwrap();
         let input = ShuffleWriterExec::try_new(
-            "local_file".to_owned(),
+            "local_file".to_owned().into(),
             1,
             create_test_data_plan().unwrap(),
             work_dir.path().to_str().unwrap().to_owned(),
@@ -1401,7 +1401,8 @@ mod tests {
         let work_dir = tmp_dir.path();
 
         // job name and stage id are hard-coded
-        let file_path = create_shuffle_path(work_dir, "job", 1, 0, None, false).unwrap();
+        let file_path =
+            create_shuffle_path(work_dir, &JobId::new("job"), 1, 0, None, false).unwrap();
 
         std::fs::create_dir_all(file_path.parent().unwrap()).unwrap();
 
@@ -1443,7 +1444,8 @@ mod tests {
         for p in 0..partition_num {
             // job name and stage id are hard-codded
             let file_path =
-                create_shuffle_path(work_dir, "job", 1, p, None, false).unwrap();
+                create_shuffle_path(work_dir, &JobId::new("job"), 1, p, None, false)
+                    .unwrap();
             // this unwrap should not be problem as
             // this function never return root dir
             std::fs::create_dir_all(file_path.parent().unwrap()).unwrap();
@@ -1482,7 +1484,7 @@ mod tests {
             .map(|partition_id| PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: "job".to_string(),
+                    job_id: "job".to_owned().into(),
                     stage_id: 1,
                     partition_id,
                 },
@@ -1718,7 +1720,7 @@ mod tests {
             .map(|partition_id| PartitionLocation {
                 map_partition_id: 0,
                 partition_id: PartitionId {
-                    job_id: "j".to_string(),
+                    job_id: "j".to_owned().into(),
                     stage_id: 7,
                     partition_id,
                 },