From 0da4a56a860d0fb7dedb6602774fcdfb667665c8 Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 17:36:59 +0200 Subject: [PATCH 01/11] Newtype macro with implementation for `JobId` and `JobName` exposed from `ballista_core` --- ballista/core/src/ids.rs | 145 +++++++++++++++++++++++++++++++++++++++ ballista/core/src/lib.rs | 5 ++ 2 files changed, 150 insertions(+) create mode 100644 ballista/core/src/ids.rs diff --git a/ballista/core/src/ids.rs b/ballista/core/src/ids.rs new file mode 100644 index 0000000000..e4e26ca603 --- /dev/null +++ b/ballista/core/src/ids.rs @@ -0,0 +1,145 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Strongly-typed string identifiers used throughout Ballista. +//! +//! Many entities in Ballista (jobs, executors, sessions, ...) are identified by +//! strings. Using a bare `String`/`&str` for all of them means the compiler +//! cannot tell, for example, a job id from a job name even though confusing the +//! two is an easy mistake. The newtypes in this module make those +//! values distinct at the type level while remaining as cheap as the `String` +//! they wrap. +//! +//! Each id derives the trait set Ballista relies on: +//! - hashing/equality/ordering so the ids work as map and set keys, +//! - [`std::fmt::Display`] so they format identically to the inner string in logs, +//! - [`From`] conversions and [`AsRef`] for cheap construction and access at +//! the protobuf and DataFusion boundaries, +//! - [`std::borrow::Borrow`] so a `HashMap` can be looked up with a +//! plain `&str` without allocating a temporary id. +//! +//! The newtypes are deliberately incompatible, so confusing one for another is +//! a compile error rather than a silent bug. This is the whole point of the +//! module, and it is enforced by the type system. The example below is +//! expected to *fail* to compile: +//! +//! ```compile_fail +//! use ballista_core::{JobId, JobName}; +//! +//! fn requires_job_id(_: JobId) {} +//! +//! // A `JobName` cannot be passed where a `JobId` is expected. +//! requires_job_id(JobName::new("oops")); +//! ``` + + +/// Defines a transparent newtype wrapping a [`String`], with the full set of +/// conversions and trait impls Ballista needs for an identifier. +macro_rules! string_id { + ($(#[$meta:meta])* $name:ident) => { + $(#[$meta])* + #[derive( + Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, + serde::Serialize, serde::Deserialize, + )] + #[serde(transparent)] + pub struct $name(String); + + impl $name { + #[doc = concat!("Creates a new `", stringify!($name), "` from anything convertible into a `String`.")] + pub fn new(s: impl Into) -> Self { + Self(s.into()) + } + + #[doc = concat!("Returns the underlying string of this `", stringify!($name), "`.")] + pub fn as_str(&self) -> &str { + &self.0 + } + + #[doc = concat!("Consumes the `", stringify!($name), "`, returning the owned inner `String`.")] + pub fn into_inner(self) -> String { + self.0 + } + } + + impl std::fmt::Display for $name { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) + } + } + + impl From for $name { + fn from(s: String) -> Self { + Self(s) + } + } + + impl From<&str> for $name { + fn from(s: &str) -> Self { + Self(s.to_owned()) + } + } + + impl From<$name> for String { + fn from(value: $name) -> Self { + value.0 + } + } + + impl AsRef for $name { + fn as_ref(&self) -> &str { + &self.0 + } + } + + impl std::borrow::Borrow for $name { + fn borrow(&self) -> &str { + &self.0 + } + } + }; +} + +string_id! { + /// Unique identifier for a job, created by the scheduler when a query is accepted. + JobId +} + +string_id! { + /// Human-supplied, non-unique display name for a job. + /// + /// Unlike [`JobId`] this is not guaranteed to be unique and carries no + /// scheduler semantics; it exists purely for display and diagnostics. + JobName +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + + #[test] + fn borrow_enables_str_keyed_lookup() { + // The reason `Borrow` exists: a `JobId`-keyed map can be queried + // with a plain `&str`. This only compiles with the `Borrow` impl, + // and only succeeds if the derived `Hash`/`Eq` agree with the `&str`. + let mut map: HashMap = HashMap::new(); + map.insert(JobId::new("job-123"), 7); + assert_eq!(map.get("job-123"), Some(&7)); + } +} diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs index 7e8fad3252..113f48f358 100644 --- a/ballista/core/src/lib.rs +++ b/ballista/core/src/lib.rs @@ -23,6 +23,9 @@ use std::sync::Arc; use datafusion::{execution::runtime_env::RuntimeEnv, prelude::SessionConfig}; use crate::serde::protobuf::JobStatus; + +pub use crate::ids::{JobId, JobName}; + /// The current version of Ballista, derived from the Cargo package version. pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -47,6 +50,8 @@ pub mod event_loop; pub mod execution_plans; /// Extension traits and utilities for DataFusion integration. pub mod extension; +/// Strongly-typed string identifiers (job ids, job names, ...). +pub mod ids; #[cfg(feature = "build-binary")] /// Object store configuration and utilities for distributed file access. pub mod object_store; From 767d1037a8a5801439fa7c5901340721fbba602f Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 17:37:27 +0200 Subject: [PATCH 02/11] Change all structs from `job_name: String` to `job_name: JobName` --- ballista-cli/src/tui/domain/jobs.rs | 2 +- ballista/scheduler/src/api/handlers.rs | 4 ++-- ballista/scheduler/src/scheduler_server/event.rs | 4 ++-- ballista/scheduler/src/state/aqe/mod.rs | 3 ++- ballista/scheduler/src/state/aqe/planner.rs | 9 +++++---- ballista/scheduler/src/state/execution_graph.rs | 3 ++- ballista/scheduler/src/state/task_manager.rs | 4 ++-- 7 files changed, 16 insertions(+), 13 deletions(-) diff --git a/ballista-cli/src/tui/domain/jobs.rs b/ballista-cli/src/tui/domain/jobs.rs index a7c1c87d0b..dd2d97d944 100644 --- a/ballista-cli/src/tui/domain/jobs.rs +++ b/ballista-cli/src/tui/domain/jobs.rs @@ -24,7 +24,7 @@ use std::collections::BTreeMap; #[derive(Deserialize, Clone, Debug)] pub struct Job { pub job_id: String, - pub job_name: String, + pub job_name: JobName, pub status: String, // Running, Completed, Failed, Canceled pub start_time: i64, pub end_time: i64, diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 6eb15a7d27..9807fded9b 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -22,7 +22,7 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; -use ballista_core::BALLISTA_VERSION; +use ballista_core::{BALLISTA_VERSION, JobName}; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ ExecutorMetric, executor_metric::Metric, task_status, @@ -110,7 +110,7 @@ impl ExecutorMetricResponse { #[derive(Debug, serde::Serialize)] pub struct JobResponse { pub job_id: String, - pub job_name: String, + pub job_name: JobName, pub job_status: String, pub status: String, pub num_stages: usize, diff --git a/ballista/scheduler/src/scheduler_server/event.rs b/ballista/scheduler/src/scheduler_server/event.rs index c6d11fb1bf..0bc82d31db 100644 --- a/ballista/scheduler/src/scheduler_server/event.rs +++ b/ballista/scheduler/src/scheduler_server/event.rs @@ -20,7 +20,7 @@ use std::fmt::{Debug, Formatter}; use datafusion::logical_expr::LogicalPlan; use crate::state::execution_graph::RunningTaskInfo; -use ballista_core::{JobStatusSubscriber, serde::protobuf::TaskStatus}; +use ballista_core::{JobName, JobStatusSubscriber, serde::protobuf::TaskStatus}; use datafusion::prelude::SessionContext; use std::sync::Arc; @@ -32,7 +32,7 @@ pub enum QueryStageSchedulerEvent { /// Unique job identifier. job_id: String, /// Human-readable job name. - job_name: String, + job_name: JobName, /// Session context for the job. session_ctx: Arc, /// Logical plan to execute. diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index f069714e6d..d8910bad6c 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -25,6 +25,7 @@ use crate::state::execution_graph::{ }; use crate::state::execution_stage::RunningStage; use crate::state::task_manager::UpdatedStages; +use ballista_core::JobName; use ballista_core::error::BallistaError; use ballista_core::execution_plans::ShuffleWriter; use ballista_core::serde::protobuf::failed_task::FailedReason; @@ -96,7 +97,7 @@ pub(crate) struct AdaptiveExecutionGraph { /// ID for this job job_id: String, /// Job name, can be empty string - job_name: String, + job_name: JobName, /// Session ID for this job session_id: String, /// Status of this job diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 20ebca3996..1febaefb60 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -24,6 +24,7 @@ use crate::state::aqe::optimizer_rule::{ }; use crate::state::distributed_explain::handle_explain_plan; use crate::state::execution_stage::StageOutput; +use ballista_core::JobName; use ballista_core::execution_plans::ShuffleWriter; use ballista_core::serde::scheduler::PartitionLocation; use datafusion::common; @@ -67,7 +68,7 @@ pub struct AdaptivePlanner { /// caches current runnable stages runnable_stage_cache: HashMap>, /// job name - job_name: String, + job_name: JobName, runnable_stage_output: HashMap, } @@ -95,7 +96,7 @@ impl AdaptivePlanner { pub fn try_new_with_optimizers( session_config: &SessionConfig, plan: Arc, - job_name: String, + job_name: JobName, physical_optimizer_rules: Vec, ) -> common::Result { let session_state = @@ -127,7 +128,7 @@ impl AdaptivePlanner { pub fn try_from_plan( session_config: &SessionConfig, plan: Arc, - job_name: String, + job_name: JobName, ) -> common::Result { let plan_id_generator = Arc::new(AtomicUsize::new(0)); Self::try_new_with_optimizers( @@ -151,7 +152,7 @@ impl AdaptivePlanner { pub async fn try_new( ctx: &SessionContext, logical_plan: &LogicalPlan, - job_name: String, + job_name: JobName, ) -> common::Result { // session state with very limited set of optimizers. // this optimizer set will be executed only once, before diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 97e421756c..e59a32a620 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -22,6 +22,7 @@ use std::iter::FromIterator; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; +use ballista_core::JobName; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanVisitor, accept}; use datafusion::prelude::SessionConfig; @@ -249,7 +250,7 @@ pub struct StaticExecutionGraph { /// ID for this job job_id: String, /// Job name, can be empty string - job_name: String, + job_name: JobName, /// Session ID for this job session_id: String, /// Status of this job diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 73776b90bc..3e0dd6b781 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -25,7 +25,7 @@ use crate::state::execution_graph::{ ExecutionGraphBox, RunningTaskInfo, StaticExecutionGraph, TaskDescription, }; use crate::state::executor_manager::ExecutorManager; -use ballista_core::JobStatusSubscriber; +use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::BallistaError; use ballista_core::error::Result; use ballista_core::extension::{SessionConfigExt, SessionConfigHelperExt}; @@ -826,7 +826,7 @@ pub struct JobOverview { /// Unique identifier for this job. pub job_id: String, /// Human-readable name for this job. - pub job_name: String, + pub job_name: JobName, /// Current status of the job. pub status: JobStatus, /// Timestamp when the job started. From ab1685b49e9b93733d3e0ce0e3e6a842a6269022 Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 18:20:39 +0200 Subject: [PATCH 03/11] Revert in ballista-cli, no dependency on `ballista-core` at present. --- ballista-cli/src/tui/domain/jobs.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ballista-cli/src/tui/domain/jobs.rs b/ballista-cli/src/tui/domain/jobs.rs index dd2d97d944..a7c1c87d0b 100644 --- a/ballista-cli/src/tui/domain/jobs.rs +++ b/ballista-cli/src/tui/domain/jobs.rs @@ -24,7 +24,7 @@ use std::collections::BTreeMap; #[derive(Deserialize, Clone, Debug)] pub struct Job { pub job_id: String, - pub job_name: JobName, + pub job_name: String, pub status: String, // Running, Completed, Failed, Canceled pub start_time: i64, pub end_time: i64, From 5934e94b4932214fed0d3392550f691bc4ea6ed1 Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 18:23:34 +0200 Subject: [PATCH 04/11] (`JobName`): Following the compiler: change all `&str`/`String` to `(&)JobName` in scheduler (untouched: core/executor/cli) --- ballista/scheduler/src/api/handlers.rs | 4 +-- ballista/scheduler/src/cluster/memory.rs | 4 +-- ballista/scheduler/src/cluster/mod.rs | 4 +-- .../scheduler/src/cluster/test_util/mod.rs | 3 ++- .../scheduler/src/scheduler_server/grpc.rs | 6 +++-- .../scheduler/src/scheduler_server/mod.rs | 25 ++++++++++--------- .../scheduler_server/query_stage_scheduler.rs | 7 +++--- ballista/scheduler/src/state/aqe/mod.rs | 12 ++++----- ballista/scheduler/src/state/aqe/planner.rs | 2 +- .../src/state/aqe/test/alter_stages.rs | 10 ++++---- .../src/state/aqe/test/coalesce_rule.rs | 14 +++++------ .../src/state/aqe/test/join_selection.rs | 12 ++++----- .../src/state/aqe/test/plan_to_stages.rs | 20 +++++++-------- .../scheduler/src/state/execution_graph.rs | 14 +++++------ .../src/state/execution_graph_dot.rs | 7 +++--- ballista/scheduler/src/state/mod.rs | 4 +-- ballista/scheduler/src/state/task_manager.rs | 8 +++--- ballista/scheduler/src/test_utils.rs | 20 +++++++-------- 18 files changed, 91 insertions(+), 85 deletions(-) diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 9807fded9b..1c9aa65897 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -347,7 +347,7 @@ pub async fn get_jobs< }; JobResponse { job_id: job.job_id.to_string(), - job_name: job.job_name.to_string(), + job_name: job.job_name.to_owned().into(), job_status, status: plain_status, start_time: job.start_time, @@ -410,7 +410,7 @@ pub async fn get_job< Ok(Json(JobResponse { job_id: job.job_id().to_string(), - job_name: job.job_name().to_string(), + job_name: job.job_name().to_owned().into(), job_status, status: plain_status, start_time: job.start_time(), diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index b4f4e36aa8..d2241b4cc3 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -26,7 +26,7 @@ use ballista_core::serde::protobuf::{ executor_status, }; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata}; -use ballista_core::{ConfigProducer, JobStatusSubscriber}; +use ballista_core::{ConfigProducer, JobName, JobStatusSubscriber}; use dashmap::DashMap; use datafusion::prelude::{SessionConfig, SessionContext}; use tokio::sync::mpsc::error::TrySendError; @@ -477,7 +477,7 @@ impl JobState for InMemoryJobState { Ok(all_jobs) } - fn accept_job(&self, job_id: &str, job_name: &str, queued_at: u64) -> Result<()> { + fn accept_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()> { self.queued_jobs .insert(job_id.to_string(), (job_name.to_string(), queued_at)); diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index b28d7960c7..2ba6317e92 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -28,7 +28,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata, PartitionId}; use ballista_core::utils::{default_config_producer, default_session_builder}; -use ballista_core::{ConfigProducer, JobStatusSubscriber}; +use ballista_core::{ConfigProducer, JobName, JobStatusSubscriber}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{SessionConfig, SessionContext}; use futures::Stream; @@ -282,7 +282,7 @@ pub trait JobState: Send + Sync { /// Accepts a job into the scheduler's queue. /// /// Called when a job is received but before it is planned. - fn accept_job(&self, job_id: &str, job_name: &str, queued_at: u64) -> Result<()>; + fn accept_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()>; /// Returns the number of queued jobs waiting to be scheduled. fn pending_job_number(&self) -> usize; diff --git a/ballista/scheduler/src/cluster/test_util/mod.rs b/ballista/scheduler/src/cluster/test_util/mod.rs index 71693930e7..3d183696d6 100644 --- a/ballista/scheduler/src/cluster/test_util/mod.rs +++ b/ballista/scheduler/src/cluster/test_util/mod.rs @@ -19,6 +19,7 @@ use crate::cluster::{JobState, JobStateEvent}; use crate::scheduler_server::timestamp_millis; use crate::state::execution_graph::ExecutionGraphBox; use crate::test_utils::{await_condition, mock_completed_task, mock_executor}; +use ballista_core::JobName; use ballista_core::error::Result; use ballista_core::serde::protobuf::JobStatus; use ballista_core::serde::protobuf::job_status::Status; @@ -56,7 +57,7 @@ impl JobStateTest { /// Queues a job with the given ID. pub fn queue_job(self, job_id: &str) -> Result { - self.state.accept_job(job_id, "", timestamp_millis())?; + self.state.accept_job(job_id, &JobName::new(""), timestamp_millis())?; Ok(self) } diff --git a/ballista/scheduler/src/scheduler_server/grpc.rs b/ballista/scheduler/src/scheduler_server/grpc.rs index f547b33456..7d25e3bdb1 100644 --- a/ballista/scheduler/src/scheduler_server/grpc.rs +++ b/ballista/scheduler/src/scheduler_server/grpc.rs @@ -391,7 +391,8 @@ impl SchedulerGrpc .iter() .find(|s| s.key == BALLISTA_JOB_NAME) .and_then(|s| s.value.clone()) - .unwrap_or_default(); + .unwrap_or_default() + .into(); info!( "execution query (PUSH) job received - session_id: {session_id}, operation_id: {operation_id}, job_name: {job_name}" @@ -467,7 +468,8 @@ impl SchedulerGrpc .iter() .find(|s| s.key == BALLISTA_JOB_NAME) .and_then(|s| s.value.clone()) - .unwrap_or_default(); + .unwrap_or_default() + .into(); info!( "execution query job received - session_id: {session_id}, operation_id: {operation_id}, job_name: {job_name}" diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 9c9f8c2307..0e43c82b9a 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use ballista_core::JobStatusSubscriber; +use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::Result; use ballista_core::event_loop::{EventLoop, EventSender}; use ballista_core::serde::BallistaCodec; @@ -220,7 +220,7 @@ impl SchedulerServer, plan: &LogicalPlan, subscriber: Option, @@ -231,7 +231,7 @@ impl SchedulerServer u64 { mod test { use std::sync::Arc; - use ballista_core::extension::SessionConfigExt; + use ballista_core::JobName; +use ballista_core::extension::SessionConfigExt; use ballista_core::serde::protobuf::job_status::Status; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; @@ -479,12 +480,12 @@ mod test { scheduler .state .task_manager - .queue_job(job_id, "", timestamp_millis())?; + .queue_job(job_id, &JobName::new(""), timestamp_millis())?; // Submit job scheduler .state - .submit_job(job_id, "", ctx, &plan, 0, None) + .submit_job(job_id, &JobName::new(""), ctx, &plan, 0, None) .await .expect("submitting plan"); @@ -574,7 +575,7 @@ mod test { ) .await?; - let (status, job_id) = test.run("", &plan).await.expect("running plan"); + let (status, job_id) = test.run(&JobName::new(""), &plan).await.expect("running plan"); match status.status { Some(job_status::Status::Successful(SuccessfulJob { @@ -614,7 +615,7 @@ mod test { let (tx, mut rx) = tokio::sync::mpsc::channel(16); let (status, job_id) = test - .run_with_subscriber("", &plan, Some(tx)) + .run_with_subscriber(&JobName::new(""), &plan, Some(tx)) .await .expect("running plan"); @@ -707,7 +708,7 @@ mod test { ) .await?; - let (status, job_id) = test.run("", &plan).await.expect("running plan"); + let (status, job_id) = test.run(&JobName::new(""), &plan).await.expect("running plan"); assert!( matches!( @@ -784,7 +785,7 @@ mod test { .await?; let (tx, mut rx) = tokio::sync::mpsc::channel(16); let (status, job_id) = test - .run_with_subscriber("", &plan, Some(tx)) + .run_with_subscriber(&JobName::new(""), &plan, Some(tx)) .await .expect("running plan"); @@ -848,7 +849,7 @@ mod test { .into_optimized_plan()?; // This should fail when we try and create the physical plan - let (status, job_id) = test.run("", &plan).await?; + let (status, job_id) = test.run(&JobName::new(""), &plan).await?; assert!( matches!( @@ -894,7 +895,7 @@ mod test { .into_optimized_plan()?; let (tx, mut rx) = tokio::sync::mpsc::channel(16); // This should fail when we try and create the physical plan - let (status, job_id) = test.run_with_subscriber("", &plan, Some(tx)).await?; + let (status, job_id) = test.run_with_subscriber(&JobName::new(""), &plan, Some(tx)).await?; assert!( matches!( diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index c2a4e58004..2572df9192 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -129,7 +129,7 @@ impl let timestamp = timestamp_millis(); let job_status = JobStatus { job_id: job_id.clone(), - job_name, + job_name: job_name.into(), status: Some(ballista_core::serde::protobuf::job_status::Status::Failed( FailedJob { error, queued_at, started_at: timestamp, ended_at: timestamp } )) @@ -369,7 +369,8 @@ impl mod tests { use crate::config::SchedulerConfig; use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; - use ballista_core::config::TaskSchedulingPolicy; + use ballista_core::JobName; +use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; @@ -394,7 +395,7 @@ mod tests { ) .await?; - let job_id = test.submit("", &plan).await?; + let job_id = test.submit(&JobName::new(""), &plan).await?; test.tick().await?; diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index d8910bad6c..fb0aa0d279 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -133,7 +133,7 @@ impl AdaptiveExecutionGraph { pub async fn try_new( scheduler_id: &str, job_id: &str, - job_name: &str, + job_name: &JobName, ctx: &SessionContext, logical_plan: &LogicalPlan, queued_at: u64, @@ -174,7 +174,7 @@ impl AdaptiveExecutionGraph { planner, scheduler_id: Some(scheduler_id.to_string()), job_id: job_id.to_string(), - job_name: job_name.to_string(), + job_name: job_name.to_owned(), session_id: session_id.to_string(), status: JobStatus { @@ -508,8 +508,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { self.job_id.as_str() } - fn job_name(&self) -> &str { - self.job_name.as_str() + fn job_name(&self) -> &JobName { + &self.job_name } fn session_id(&self) -> &str { @@ -1177,7 +1177,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { self.status = JobStatus { job_id: self.job_id.clone(), - job_name: self.job_name.clone(), + job_name: self.job_name.clone().into(), status: Some(Status::Failed(FailedJob { error, queued_at: self.queued_at, @@ -1206,7 +1206,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { self.status = JobStatus { job_id: self.job_id.clone(), - job_name: self.job_name.clone(), + job_name: self.job_name.clone().into(), status: Some(job_status::Status::Successful(SuccessfulJob { partition_location, diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 1febaefb60..01429b3d73 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -165,7 +165,7 @@ impl AdaptivePlanner { ); let plan = state.create_physical_plan(logical_plan).await?; - let plan = handle_explain_plan(&job_name, ctx, logical_plan, plan) + let plan = handle_explain_plan(job_name.as_str(), ctx, logical_plan, plan) .await .map_err(|e| DataFusionError::Execution(e.to_string()))?; diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs b/ballista/scheduler/src/state/aqe/test/alter_stages.rs index e38d4f6436..1baaf1b949 100644 --- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs @@ -57,7 +57,7 @@ async fn should_propagate_empty_stage() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); @@ -120,7 +120,7 @@ async fn should_propagate_empty_stage_and_remove() -> datafusion::error::Result< let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -189,7 +189,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), join, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -289,7 +289,7 @@ async fn should_support_cross_join() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), join, - "test_job".to_string(), + "test_job".to_owned().into(), )?; // @@ -395,7 +395,7 @@ async fn should_cancel_the_stage() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), join, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let (stages, cancellable) = planner.actionable_stages()?; diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 2c98003fad..893daa36cd 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -138,7 +138,7 @@ async fn should_attach_coalesce_when_partitions_pack_below_m() let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; // Before any stage finalizes the leaves are unresolved, so the rule @@ -191,7 +191,7 @@ async fn should_skip_coalesce_when_rule_disabled() -> datafusion::error::Result< let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let _ = planner.runnable_stages()?.unwrap(); @@ -228,7 +228,7 @@ async fn should_skip_coalesce_when_partitions_are_full() -> datafusion::error::R let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let _ = planner.runnable_stages()?.unwrap(); @@ -267,7 +267,7 @@ async fn should_attach_coalesce_to_both_sides_of_hash_join() let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); @@ -320,7 +320,7 @@ async fn should_attach_coalesce_to_all_three_legs_of_two_hash_joins() let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); @@ -385,7 +385,7 @@ async fn should_attach_coalesce_to_both_sides_of_sort_merge_join() let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); @@ -432,7 +432,7 @@ async fn shuffle_reader_uses_coalesced_k_when_rule_fires() -> datafusion::error: let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; // Stage 0 is the upstream shuffle writer, partitioning by `c` into M=8. diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index df429b1e3b..4eafdd571b 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -128,7 +128,7 @@ async fn test_hash_join_two_tables_coalesce() -> datafusion::common::Result<()> .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); // TODO: we could probably push datasource to read single partition assert_plan!(plan, @ r" @@ -153,7 +153,7 @@ async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<( .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -214,7 +214,7 @@ async fn test_sort_merge_join_two_tables_repartition() -> datafusion::common::Re .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -280,7 +280,7 @@ async fn test_hash_join_three_tables_collect_left() -> datafusion::common::Resul .unwrap() .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -351,7 +351,7 @@ async fn test_hash_join_three_tables_repartition() -> datafusion::common::Result .unwrap() .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -382,7 +382,7 @@ async fn test_sort_merge_join_three_tables_repartition() -> datafusion::common:: .unwrap() .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_string()).await?; + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs index 46b4fab939..b6808089f1 100644 --- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs @@ -43,7 +43,7 @@ async fn should_add_exchanges() -> datafusion::error::Result<()> { let planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -72,7 +72,7 @@ async fn should_split_plan_into_runnable_stages_internal() -> datafusion::error: let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -124,7 +124,7 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -183,7 +183,7 @@ async fn should_create_initial_plan() -> datafusion::error::Result<()> { let planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; // plan has only two exchanges after initial planning @@ -233,7 +233,7 @@ async fn should_split_stages_resolve_right_branch() -> datafusion::error::Result let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let runnable_stages = planner.identify_runnable_stages()?.unwrap(); @@ -303,7 +303,7 @@ async fn should_split_stages_resolve_left_branch() -> datafusion::error::Result< let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let runnable_stages = planner.identify_runnable_stages()?.unwrap(); @@ -378,7 +378,7 @@ async fn should_split_stages_resolve_both() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let runnable_stages = planner.identify_runnable_stages()?.unwrap(); @@ -438,7 +438,7 @@ async fn should_ignore_inactive_stages() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), exchange_exec, - "test_job".to_string(), + "test_job".to_owned().into(), )?; assert_plan!(planner.current_plan(), @ r" @@ -467,7 +467,7 @@ async fn should_use_sort_shuffle_when_enabled() -> datafusion::error::Result<()> let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); @@ -497,7 +497,7 @@ async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()> { let mut planner = AdaptivePlanner::try_from_plan( ctx.state().config(), plan, - "test_job".to_string(), + "test_job".to_owned().into(), )?; let stages = planner.runnable_stages()?.unwrap(); diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index e59a32a620..022e7ee95c 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -105,7 +105,7 @@ pub trait ExecutionGraph: Debug { fn job_id(&self) -> &str; /// Returns the job name for this execution graph. - fn job_name(&self) -> &str; + fn job_name(&self) -> &JobName; /// Returns the session ID associated with this job. fn session_id(&self) -> &str; @@ -305,7 +305,7 @@ impl StaticExecutionGraph { pub fn new( scheduler_id: &str, job_id: &str, - job_name: &str, + job_name: &JobName, session_id: &str, plan: Arc, queued_at: u64, @@ -324,7 +324,7 @@ impl StaticExecutionGraph { Ok(Self { scheduler_id: Some(scheduler_id.to_string()), job_id: job_id.to_string(), - job_name: job_name.to_string(), + job_name: job_name.to_owned().into(), session_id: session_id.to_string(), status: JobStatus { @@ -639,8 +639,8 @@ impl ExecutionGraph for StaticExecutionGraph { self.job_id.as_str() } - fn job_name(&self) -> &str { - self.job_name.as_str() + fn job_name(&self) -> &JobName { + &self.job_name } fn session_id(&self) -> &str { @@ -1360,7 +1360,7 @@ impl ExecutionGraph for StaticExecutionGraph { self.status = JobStatus { job_id: self.job_id.clone(), - job_name: self.job_name.clone(), + job_name: self.job_name.clone().into(), status: Some(Status::Failed(FailedJob { error, queued_at: self.queued_at, @@ -1389,7 +1389,7 @@ impl ExecutionGraph for StaticExecutionGraph { self.status = JobStatus { job_id: self.job_id.clone(), - job_name: self.job_name.clone(), + job_name: self.job_name.clone().into(), status: Some(job_status::Status::Successful(SuccessfulJob { partition_location, diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index e1a24c4da5..2733805222 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -413,7 +413,8 @@ mod tests { use crate::planner::DefaultDistributedPlanner; use crate::state::execution_graph::StaticExecutionGraph; use crate::state::execution_graph_dot::ExecutionGraphDot; - use ballista_core::error::{BallistaError, Result}; + use ballista_core::JobName; +use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::MemTable; @@ -607,7 +608,7 @@ filter_expr="] StaticExecutionGraph::new( "scheduler_id", "job_id", - "job_name", + &JobName::new("job_name"), "session_id", plan, 0, @@ -644,7 +645,7 @@ filter_expr="] StaticExecutionGraph::new( "scheduler_id", "job_id", - "job_name", + &JobName::new("job_name"), "session_id", plan, 0, diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index b05fdc1caf..eda20d170b 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -22,7 +22,7 @@ use crate::state::execution_graph::TaskDescription; use crate::state::executor_manager::ExecutorManager; use crate::state::session_manager::SessionManager; use crate::state::task_manager::{TaskLauncher, TaskManager}; -use ballista_core::JobStatusSubscriber; +use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::EventSender; use ballista_core::serde::BallistaCodec; @@ -364,7 +364,7 @@ impl SchedulerState, logical_plan: &LogicalPlan, queued_at: u64, diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 3e0dd6b781..0171fcda15 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -254,7 +254,7 @@ impl TaskManager } /// Enqueue a job for scheduling - pub fn queue_job(&self, job_id: &str, job_name: &str, queued_at: u64) -> Result<()> { + pub fn queue_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()> { self.state.accept_job(job_id, job_name, queued_at) } @@ -276,7 +276,7 @@ impl TaskManager pub async fn submit_job( &self, job_id: &str, - job_name: &str, + job_name: &JobName, ctx: Arc, logical_plan: &LogicalPlan, queued_at: u64, @@ -400,7 +400,7 @@ impl TaskManager }; jobs.push(JobOverview { job_id: job_status.job_id.clone(), - job_name: job_status.job_name.clone(), + job_name: job_status.job_name.clone().into(), status: job_status, start_time, end_time, @@ -845,7 +845,7 @@ impl From<&ExecutionGraphBox> for JobOverview { Self { job_id: value.job_id().to_string(), - job_name: value.job_name().to_string(), + job_name: value.job_name().to_owned(), status: value.status().clone(), start_time: value.start_time(), end_time: value.end_time(), diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index b8cc18a8d5..7a962479c0 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use ballista_core::JobStatusSubscriber; +use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; use datafusion::catalog::Session; @@ -499,7 +499,7 @@ impl SchedulerTest { } /// Submits a job and returns its ID. - pub async fn submit(&mut self, job_name: &str, plan: &LogicalPlan) -> Result { + pub async fn submit(&mut self, job_name: &JobName, plan: &LogicalPlan) -> Result { println!("{:?}", self.session_config); let ctx = self .scheduler @@ -626,7 +626,7 @@ impl SchedulerTest { /// Returns job status and job_id pub async fn run( &mut self, - job_name: &str, + job_name: &JobName, plan: &LogicalPlan, ) -> Result<(JobStatus, String)> { self.run_with_subscriber(job_name, plan, None).await @@ -634,7 +634,7 @@ impl SchedulerTest { /// Returns job status and job_id, with provided subscriber pub async fn run_with_subscriber( &mut self, - job_name: &str, + job_name: &JobName, plan: &LogicalPlan, subscriber: Option, ) -> Result<(JobStatus, String)> { @@ -913,7 +913,7 @@ pub async fn test_aggregation_plan_with_job_id( StaticExecutionGraph::new( "localhost:50050", job_id, - "", + &JobName::new(""), "session", plan, 0, @@ -962,7 +962,7 @@ pub async fn test_two_aggregations_plan(partition: usize) -> StaticExecutionGrap StaticExecutionGraph::new( "localhost:50050", "job", - "", + &JobName::new(""), "session", plan, 0, @@ -1003,7 +1003,7 @@ pub async fn test_coalesce_plan(partition: usize) -> StaticExecutionGraph { StaticExecutionGraph::new( "localhost:50050", "job", - "", + &JobName::new(""), "session", plan, 0, @@ -1064,7 +1064,7 @@ pub async fn test_join_plan(partition: usize) -> StaticExecutionGraph { let graph = StaticExecutionGraph::new( "localhost:50050", "job", - "", + &JobName::new(""), "session", plan, 0, @@ -1107,7 +1107,7 @@ pub async fn test_union_all_plan(partition: usize) -> StaticExecutionGraph { let graph = StaticExecutionGraph::new( "localhost:50050", "job", - "", + &JobName::new(""), "session", plan, 0, @@ -1150,7 +1150,7 @@ pub async fn test_union_plan(partition: usize) -> StaticExecutionGraph { let graph = StaticExecutionGraph::new( "localhost:50050", "job", - "", + &JobName::new(""), "session", plan, 0, From 8e82254518b62d93973522e6ea70b1ee477f7f3f Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 20:26:48 +0200 Subject: [PATCH 05/11] fmt --- ballista/core/src/ids.rs | 2 -- ballista/scheduler/src/api/handlers.rs | 2 +- .../scheduler/src/cluster/test_util/mod.rs | 3 ++- .../scheduler/src/scheduler_server/mod.rs | 27 ++++++++++++------- .../scheduler_server/query_stage_scheduler.rs | 2 +- ballista/scheduler/src/state/aqe/mod.rs | 2 +- .../src/state/aqe/test/join_selection.rs | 18 ++++++++----- .../src/state/execution_graph_dot.rs | 2 +- ballista/scheduler/src/state/task_manager.rs | 9 +++++-- ballista/scheduler/src/test_utils.rs | 8 ++++-- 10 files changed, 49 insertions(+), 26 deletions(-) diff --git a/ballista/core/src/ids.rs b/ballista/core/src/ids.rs index e4e26ca603..9a6fe58609 100644 --- a/ballista/core/src/ids.rs +++ b/ballista/core/src/ids.rs @@ -46,7 +46,6 @@ //! requires_job_id(JobName::new("oops")); //! ``` - /// Defines a transparent newtype wrapping a [`String`], with the full set of /// conversions and trait impls Ballista needs for an identifier. macro_rules! string_id { @@ -132,7 +131,6 @@ mod tests { use super::*; use std::collections::HashMap; - #[test] fn borrow_enables_str_keyed_lookup() { // The reason `Borrow` exists: a `JobId`-keyed map can be queried diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 1c9aa65897..b0f943a493 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -22,7 +22,6 @@ use axum::{ extract::{Path, State}, response::{IntoResponse, Response}, }; -use ballista_core::{BALLISTA_VERSION, JobName}; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ ExecutorMetric, executor_metric::Metric, task_status, @@ -31,6 +30,7 @@ use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; use ballista_core::utils::get_current_time; +use ballista_core::{BALLISTA_VERSION, JobName}; use datafusion::DATAFUSION_VERSION; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::displayable; diff --git a/ballista/scheduler/src/cluster/test_util/mod.rs b/ballista/scheduler/src/cluster/test_util/mod.rs index 3d183696d6..1fee57132a 100644 --- a/ballista/scheduler/src/cluster/test_util/mod.rs +++ b/ballista/scheduler/src/cluster/test_util/mod.rs @@ -57,7 +57,8 @@ impl JobStateTest { /// Queues a job with the given ID. pub fn queue_job(self, job_id: &str) -> Result { - self.state.accept_job(job_id, &JobName::new(""), timestamp_millis())?; + self.state + .accept_job(job_id, &JobName::new(""), timestamp_millis())?; Ok(self) } diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 0e43c82b9a..4725f57a60 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -18,11 +18,11 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::Result; use ballista_core::event_loop::{EventLoop, EventSender}; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::TaskStatus; +use ballista_core::{JobName, JobStatusSubscriber}; use datafusion::execution::context::SessionState; use datafusion::logical_expr::LogicalPlan; @@ -412,7 +412,7 @@ mod test { use std::sync::Arc; use ballista_core::JobName; -use ballista_core::extension::SessionConfigExt; + use ballista_core::extension::SessionConfigExt; use ballista_core::serde::protobuf::job_status::Status; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; @@ -477,10 +477,11 @@ use ballista_core::extension::SessionConfigExt; let job_id = "job"; // Enqueue job - scheduler - .state - .task_manager - .queue_job(job_id, &JobName::new(""), timestamp_millis())?; + scheduler.state.task_manager.queue_job( + job_id, + &JobName::new(""), + timestamp_millis(), + )?; // Submit job scheduler @@ -575,7 +576,10 @@ use ballista_core::extension::SessionConfigExt; ) .await?; - let (status, job_id) = test.run(&JobName::new(""), &plan).await.expect("running plan"); + let (status, job_id) = test + .run(&JobName::new(""), &plan) + .await + .expect("running plan"); match status.status { Some(job_status::Status::Successful(SuccessfulJob { @@ -708,7 +712,10 @@ use ballista_core::extension::SessionConfigExt; ) .await?; - let (status, job_id) = test.run(&JobName::new(""), &plan).await.expect("running plan"); + let (status, job_id) = test + .run(&JobName::new(""), &plan) + .await + .expect("running plan"); assert!( matches!( @@ -895,7 +902,9 @@ use ballista_core::extension::SessionConfigExt; .into_optimized_plan()?; let (tx, mut rx) = tokio::sync::mpsc::channel(16); // This should fail when we try and create the physical plan - let (status, job_id) = test.run_with_subscriber(&JobName::new(""), &plan, Some(tx)).await?; + let (status, job_id) = test + .run_with_subscriber(&JobName::new(""), &plan, Some(tx)) + .await?; assert!( matches!( diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 2572df9192..a6642ee2e5 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -370,7 +370,7 @@ mod tests { use crate::config::SchedulerConfig; use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; use ballista_core::JobName; -use ballista_core::config::TaskSchedulingPolicy; + use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index fb0aa0d279..fd32c08c91 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -25,7 +25,6 @@ use crate::state::execution_graph::{ }; use crate::state::execution_stage::RunningStage; use crate::state::task_manager::UpdatedStages; -use ballista_core::JobName; use ballista_core::error::BallistaError; use ballista_core::execution_plans::ShuffleWriter; use ballista_core::serde::protobuf::failed_task::FailedReason; @@ -35,6 +34,7 @@ use ballista_core::serde::protobuf::{ job_status, task_status, }; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; +use ballista_core::JobName; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::ExecutionPlan; diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index 4eafdd571b..ae7c7b52a9 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -128,7 +128,8 @@ async fn test_hash_join_two_tables_coalesce() -> datafusion::common::Result<()> .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); // TODO: we could probably push datasource to read single partition assert_plan!(plan, @ r" @@ -153,7 +154,8 @@ async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<( .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let mut planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -214,7 +216,8 @@ async fn test_sort_merge_join_two_tables_repartition() -> datafusion::common::Re .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let mut planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -280,7 +283,8 @@ async fn test_hash_join_three_tables_collect_left() -> datafusion::common::Resul .unwrap() .into_optimized_plan() .unwrap(); - let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let mut planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -351,7 +355,8 @@ async fn test_hash_join_three_tables_repartition() -> datafusion::common::Result .unwrap() .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" @@ -382,7 +387,8 @@ async fn test_sort_merge_join_three_tables_repartition() -> datafusion::common:: .unwrap() .into_optimized_plan() .unwrap(); - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; + let planner = + AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned().into()).await?; let plan = planner.current_plan(); assert_plan!(plan, @ r" diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index 2733805222..0b39164453 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -414,7 +414,7 @@ mod tests { use crate::state::execution_graph::StaticExecutionGraph; use crate::state::execution_graph_dot::ExecutionGraphDot; use ballista_core::JobName; -use ballista_core::error::{BallistaError, Result}; + use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::MemTable; diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 0171fcda15..4d8184bb3f 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -25,7 +25,6 @@ use crate::state::execution_graph::{ ExecutionGraphBox, RunningTaskInfo, StaticExecutionGraph, TaskDescription, }; use crate::state::executor_manager::ExecutorManager; -use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::BallistaError; use ballista_core::error::Result; use ballista_core::extension::{SessionConfigExt, SessionConfigHelperExt}; @@ -34,6 +33,7 @@ use ballista_core::serde::protobuf::{ JobStatus, MultiTaskDefinition, TaskDefinition, TaskId, TaskStatus, job_status, }; use ballista_core::serde::scheduler::ExecutorMetadata; +use ballista_core::{JobName, JobStatusSubscriber}; use dashmap::DashMap; use datafusion::execution::config::SessionConfig; use datafusion::execution::context::SessionContext; @@ -254,7 +254,12 @@ impl TaskManager } /// Enqueue a job for scheduling - pub fn queue_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()> { + pub fn queue_job( + &self, + job_id: &str, + job_name: &JobName, + queued_at: u64, + ) -> Result<()> { self.state.accept_job(job_id, job_name, queued_at) } diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 7a962479c0..3a953f2b43 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. -use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; +use ballista_core::{JobName, JobStatusSubscriber}; use datafusion::catalog::Session; use std::any::Any; use std::collections::HashMap; @@ -499,7 +499,11 @@ impl SchedulerTest { } /// Submits a job and returns its ID. - pub async fn submit(&mut self, job_name: &JobName, plan: &LogicalPlan) -> Result { + pub async fn submit( + &mut self, + job_name: &JobName, + plan: &LogicalPlan, + ) -> Result { println!("{:?}", self.session_config); let ctx = self .scheduler From 09118b5e53ead6c4db6b42dbb1ae7377932c3ef1 Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 20:32:26 +0200 Subject: [PATCH 06/11] (`JobId`): Change all structs from `job_id: String` to `job_id: JobId` and all function parameters. --- ballista/core/src/execution_plans/shuffle_writer.rs | 5 +++-- .../core/src/execution_plans/sort_shuffle/writer.rs | 5 +++-- ballista/core/src/serde/scheduler/mod.rs | 11 ++++++----- ballista/executor/src/execution_engine.rs | 4 ++-- ballista/executor/src/executor.rs | 2 +- ballista/scheduler/src/api/handlers.rs | 4 ++-- ballista/scheduler/src/cluster/memory.rs | 2 +- ballista/scheduler/src/cluster/mod.rs | 10 +++++----- ballista/scheduler/src/scheduler_server/event.rs | 12 ++++++------ ballista/scheduler/src/scheduler_server/mod.rs | 4 ++-- ballista/scheduler/src/state/aqe/mod.rs | 4 ++-- ballista/scheduler/src/state/execution_graph.rs | 5 +++-- ballista/scheduler/src/state/executor_manager.rs | 6 +++--- ballista/scheduler/src/state/mod.rs | 6 +++--- ballista/scheduler/src/state/task_manager.rs | 6 +++--- 15 files changed, 45 insertions(+), 41 deletions(-) diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index cdb1a16cda..383697c075 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -35,6 +35,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; +use crate::JobId; use crate::execution_plans::create_shuffle_path; use crate::extension::SessionConfigExt; use crate::utils; @@ -77,7 +78,7 @@ pub const DEFAULT_SHUFFLE_CHANNEL_CAPACITY: usize = 8; #[derive(Debug, Clone)] pub struct ShuffleWriterExec { /// Unique ID for the job (query) that this stage is a part of - job_id: String, + job_id: JobId, /// Unique query stage ID within the job stage_id: usize, /// Physical execution plan for this query stage @@ -148,7 +149,7 @@ impl ShuffleWriteMetrics { impl ShuffleWriterExec { /// Create a new shuffle writer pub fn try_new( - job_id: String, + job_id: JobId, stage_id: usize, plan: Arc, work_dir: String, diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 8d88b7562d..bf067329ee 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -35,6 +35,7 @@ use super::config::SortShuffleConfig; use super::index::ShuffleIndex; use super::partitioned_batch_iterator::PartitionedBatchIterator; use super::spill::SpillManager; +use crate::JobId; use crate::execution_plans::create_shuffle_path; use crate::serde::protobuf::ShuffleWritePartition; @@ -75,7 +76,7 @@ type FinalizeResult = (PathBuf, PathBuf, Vec<(usize, u64, u64, u64)>); #[derive(Debug, Clone)] pub struct SortShuffleWriterExec { /// Unique ID for the job (query) that this stage is a part of - job_id: String, + job_id: JobId, /// Unique query stage ID within the job stage_id: usize, /// Physical execution plan for this query stage @@ -131,7 +132,7 @@ impl SortShuffleWriteMetrics { impl SortShuffleWriterExec { /// Create a new sort-based shuffle writer. pub fn try_new( - job_id: String, + job_id: JobId, stage_id: usize, plan: Arc, work_dir: String, diff --git a/ballista/core/src/serde/scheduler/mod.rs b/ballista/core/src/serde/scheduler/mod.rs index eca6700d99..1f72d00368 100644 --- a/ballista/core/src/serde/scheduler/mod.rs +++ b/ballista/core/src/serde/scheduler/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::JobId; use crate::error::BallistaError; use crate::execution_plans::create_shuffle_path; use crate::registry::BallistaFunctionRegistry; @@ -41,7 +42,7 @@ pub enum Action { /// Collect a shuffle partition FetchPartition { /// The job identifier. - job_id: String, + job_id: JobId, /// The stage identifier within the job. stage_id: usize, /// The partition identifier within the stage. @@ -61,7 +62,7 @@ pub enum Action { #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct PartitionId { /// The job identifier. - pub job_id: String, + pub job_id: JobId, /// The stage identifier within the job. pub stage_id: usize, /// The partition identifier within the stage. @@ -387,7 +388,7 @@ impl PartitionStats { #[derive(Debug, Clone)] pub struct ExecutePartition { /// Unique ID representing this query execution - pub job_id: String, + pub job_id: JobId, /// Unique ID representing this query stage within the overall query pub stage_id: usize, /// The partitions to execute. The same plan could be sent to multiple executors and each @@ -404,7 +405,7 @@ pub struct ExecutePartition { impl ExecutePartition { /// Creates a new execute partition task. pub fn new( - job_id: String, + job_id: JobId, stage_id: usize, partition_id: Vec, plan: Arc, @@ -464,7 +465,7 @@ pub struct TaskDefinition { /// Current attempt number for this task. pub task_attempt_num: usize, /// Job identifier this task belongs to. - pub job_id: String, + pub job_id: JobId, /// Stage identifier within the job. pub stage_id: usize, /// Current attempt number for the stage. diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 57ab6d0549..a4fa7aa271 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -48,7 +48,7 @@ pub trait ExecutionEngine: Sync + Send { /// plan partition and writing shuffle output to the specified work directory. fn create_query_stage_exec( &self, - job_id: String, + job_id: JobId, stage_id: usize, partition_id: usize, plan: Arc, @@ -104,7 +104,7 @@ impl DefaultExecutionEngine { impl ExecutionEngine for DefaultExecutionEngine { fn create_query_stage_exec( &self, - job_id: String, + job_id: JobId, stage_id: usize, _partition_id: usize, plan: Arc, diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs index 70e51f284c..d329177440 100644 --- a/ballista/executor/src/executor.rs +++ b/ballista/executor/src/executor.rs @@ -234,7 +234,7 @@ impl Executor { pub async fn cancel_task( &self, task_id: usize, - job_id: String, + job_id: JobId, stage_id: usize, partition_id: usize, ) -> Result { diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index b0f943a493..73e94e1a12 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -30,7 +30,7 @@ use ballista_core::serde::scheduler::{ ExecutorOperatingSystemSpecification, ExecutorSpecification, }; use ballista_core::utils::get_current_time; -use ballista_core::{BALLISTA_VERSION, JobName}; +use ballista_core::{BALLISTA_VERSION, JobId, JobName}; use datafusion::DATAFUSION_VERSION; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::displayable; @@ -109,7 +109,7 @@ impl ExecutorMetricResponse { #[derive(Debug, serde::Serialize)] pub struct JobResponse { - pub job_id: String, + pub job_id: JobId, pub job_name: JobName, pub job_status: String, pub status: String, diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index d2241b4cc3..6f06ee59fe 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -323,7 +323,7 @@ impl ExtendedJobStatus { impl JobState for InMemoryJobState { async fn submit_job( &self, - job_id: String, + job_id: JobId, graph: &ExecutionGraphBox, subscriber: Option, ) -> Result<()> { diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index 2ba6317e92..82eeab4d98 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -28,7 +28,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata, PartitionId}; use ballista_core::utils::{default_config_producer, default_session_builder}; -use ballista_core::{ConfigProducer, JobName, JobStatusSubscriber}; +use ballista_core::{ConfigProducer, JobId, JobName, JobStatusSubscriber}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{SessionConfig, SessionContext}; use futures::Stream; @@ -220,7 +220,7 @@ pub enum JobStateEvent { /// Event when a job status has been updated JobUpdated { /// Job ID of updated job - job_id: String, + job_id: JobId, /// New job status status: JobStatus, }, @@ -230,14 +230,14 @@ pub enum JobStateEvent { /// different scheduler JobAcquired { /// Job ID of the acquired job - job_id: String, + job_id: JobId, /// The scheduler which acquired ownership of the job owner: String, }, /// Event when a scheduler releases ownership of a still active job JobReleased { /// Job ID of the released job - job_id: String, + job_id: JobId, }, /// Event when a new session has been created. SessionAccessed { @@ -292,7 +292,7 @@ pub trait JobState: Send + Sync { /// The submitter is assumed to own the job. async fn submit_job( &self, - job_id: String, + job_id: JobId, graph: &ExecutionGraphBox, subscriber: Option, ) -> Result<()>; diff --git a/ballista/scheduler/src/scheduler_server/event.rs b/ballista/scheduler/src/scheduler_server/event.rs index 0bc82d31db..2fbd7072c1 100644 --- a/ballista/scheduler/src/scheduler_server/event.rs +++ b/ballista/scheduler/src/scheduler_server/event.rs @@ -20,7 +20,7 @@ use std::fmt::{Debug, Formatter}; use datafusion::logical_expr::LogicalPlan; use crate::state::execution_graph::RunningTaskInfo; -use ballista_core::{JobName, JobStatusSubscriber, serde::protobuf::TaskStatus}; +use ballista_core::{JobId, JobName, JobStatusSubscriber, serde::protobuf::TaskStatus}; use datafusion::prelude::SessionContext; use std::sync::Arc; @@ -30,7 +30,7 @@ pub enum QueryStageSchedulerEvent { /// A new job has been queued for execution. JobQueued { /// Unique job identifier. - job_id: String, + job_id: JobId, /// Human-readable job name. job_name: JobName, /// Session context for the job. @@ -45,7 +45,7 @@ pub enum QueryStageSchedulerEvent { /// A job has been submitted for execution. JobSubmitted { /// Unique job identifier. - job_id: String, + job_id: JobId, /// Timestamp when the job was queued. queued_at: u64, /// Timestamp when the job was submitted. @@ -54,7 +54,7 @@ pub enum QueryStageSchedulerEvent { /// A job failed during the planning phase. JobPlanningFailed { /// Unique job identifier. - job_id: String, + job_id: JobId, /// Error message describing the failure. fail_message: String, /// Timestamp when the job was queued. @@ -65,7 +65,7 @@ pub enum QueryStageSchedulerEvent { /// A job has completed successfully. JobFinished { /// Unique job identifier. - job_id: String, + job_id: JobId, /// Timestamp when the job was queued. queued_at: u64, /// Timestamp when the job completed. @@ -74,7 +74,7 @@ pub enum QueryStageSchedulerEvent { /// A job failed during execution. JobRunningFailed { /// Unique job identifier. - job_id: String, + job_id: JobId, /// Error message describing the failure. fail_message: String, /// Timestamp when the job was queued. diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 4725f57a60..4825fd4eed 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -186,7 +186,7 @@ impl SchedulerServer Result<()> { + pub async fn cancel_job(&self, job_id: JobId) -> Result<()> { log::debug!("Received cancellation request for job {job_id}"); self.query_stage_event_loop @@ -199,7 +199,7 @@ impl SchedulerServer Result<()> { log::debug!("Received fail job request for job {job_id}"); diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index fd32c08c91..05d1e38c28 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -34,7 +34,7 @@ use ballista_core::serde::protobuf::{ job_status, task_status, }; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; -use ballista_core::JobName; +use ballista_core::{JobId, JobName}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::ExecutionPlan; @@ -95,7 +95,7 @@ pub(crate) struct AdaptiveExecutionGraph { /// Adaptive Planner to be used with this execution graph planner: AdaptivePlanner, /// ID for this job - job_id: String, + job_id: JobId, /// Job name, can be empty string job_name: JobName, /// Session ID for this job diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 022e7ee95c..0c5817a51f 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -28,6 +28,7 @@ use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanVisitor, accept}; use datafusion::prelude::SessionConfig; use log::{debug, error, info, warn}; +use ballista_core::JobId; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ ShuffleWriter, ShuffleWriterExec, SortShuffleWriterExec, UnresolvedShuffleExec, @@ -248,7 +249,7 @@ pub struct StaticExecutionGraph { #[allow(dead_code)] // not used at the moment, will be used later scheduler_id: Option, /// ID for this job - job_id: String, + job_id: JobId, /// Job name, can be empty string job_name: JobName, /// Session ID for this job @@ -287,7 +288,7 @@ pub struct RunningTaskInfo { /// Unique identifier for this task within the execution graph. pub task_id: usize, /// The job ID this task belongs to. - pub job_id: String, + pub job_id: JobId, /// The stage ID this task belongs to. pub stage_id: usize, /// The partition this task is processing. diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index 4bc3558d8e..d87f0f5bc9 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -177,7 +177,7 @@ impl ExecutorManager { /// Send rpc to Executors to clean up the job data by delayed clean_up_interval seconds pub(crate) fn clean_up_job_data_delayed( &self, - job_id: String, + job_id: JobId, clean_up_interval: u64, ) { if clean_up_interval == 0 { @@ -195,7 +195,7 @@ impl ExecutorManager { } /// Sends RPC requests to executors to clean up job data in a spawned task. - pub fn clean_up_job_data(&self, job_id: String) { + pub fn clean_up_job_data(&self, job_id: JobId) { let executor_manager = self.clone(); tokio::spawn(async move { executor_manager.clean_up_job_data_inner(job_id).await; @@ -204,7 +204,7 @@ impl ExecutorManager { /// 1. Push strategy: Send rpc to Executors to clean up the job data /// 2. Poll strategy: Save cleanup job ids and send them to executors - async fn clean_up_job_data_inner(&self, job_id: String) { + async fn clean_up_job_data_inner(&self, job_id: JobId) { let alive_executors = self.get_alive_executors(); for executor in alive_executors { diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index eda20d170b..76cf5b9a36 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -22,11 +22,11 @@ use crate::state::execution_graph::TaskDescription; use crate::state::executor_manager::ExecutorManager; use crate::state::session_manager::SessionManager; use crate::state::task_manager::{TaskLauncher, TaskManager}; -use ballista_core::{JobName, JobStatusSubscriber}; use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::EventSender; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::TaskStatus; +use ballista_core::{JobName, JobStatusSubscriber}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -391,7 +391,7 @@ impl SchedulerState SchedulerState TaskManager } /// Clean up a failed job in FailedJobs Keyspace by delayed clean_up_interval seconds - pub(crate) fn clean_up_job_delayed(&self, job_id: String, clean_up_interval: u64) { + pub(crate) fn clean_up_job_delayed(&self, job_id: JobId, clean_up_interval: u64) { if clean_up_interval == 0 { info!( "The interval is 0 and the clean up for the failed job state {job_id} will not triggered" @@ -829,7 +829,7 @@ impl TaskManager /// Summary information about a job for display purposes. pub struct JobOverview { /// Unique identifier for this job. - pub job_id: String, + pub job_id: JobId, /// Human-readable name for this job. pub job_name: JobName, /// Current status of the job. From 04f23c1b66d2db03d64d4abf297aedbb038fa51b Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 20:34:54 +0200 Subject: [PATCH 07/11] Replase all `job_id: &str` with `job_id: &JobId` --- .../distributed_explain_analyze.rs | 4 +-- ballista/core/src/execution_plans/mod.rs | 2 +- .../src/execution_plans/sort_shuffle/spill.rs | 2 +- .../execution_plans/sort_shuffle/writer.rs | 2 +- ballista/core/src/serde/scheduler/mod.rs | 2 +- ballista/executor/src/executor_process.rs | 4 +-- ballista/executor/src/metrics/mod.rs | 4 +-- ballista/scheduler/src/cluster/memory.rs | 14 +++++----- ballista/scheduler/src/cluster/mod.rs | 16 +++++------ .../scheduler/src/cluster/test_util/mod.rs | 12 ++++---- ballista/scheduler/src/display.rs | 2 +- ballista/scheduler/src/metrics/mod.rs | 16 +++++------ ballista/scheduler/src/metrics/prometheus.rs | 8 +++--- ballista/scheduler/src/planner.rs | 2 +- ballista/scheduler/src/state/aqe/adapter.rs | 2 +- ballista/scheduler/src/state/aqe/mod.rs | 2 +- .../src/state/distributed_explain.rs | 4 +-- .../scheduler/src/state/execution_graph.rs | 4 +-- ballista/scheduler/src/state/mod.rs | 2 +- ballista/scheduler/src/state/task_manager.rs | 26 ++++++++--------- ballista/scheduler/src/test_utils.rs | 28 +++++++++---------- 21 files changed, 79 insertions(+), 79 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_explain_analyze.rs b/ballista/core/src/execution_plans/distributed_explain_analyze.rs index 0a7ddabb36..08ec5147d2 100644 --- a/ballista/core/src/execution_plans/distributed_explain_analyze.rs +++ b/ballista/core/src/execution_plans/distributed_explain_analyze.rs @@ -204,7 +204,7 @@ impl ExecutionPlan for DistributedExplainAnalyzeExec async fn fetch_job_metrics( scheduler_url: &str, - job_id: &str, + job_id: &JobId, session_config: datafusion::prelude::SessionConfig, ) -> Result { let grpc_interceptor = session_config.ballista_grpc_interceptor(); @@ -248,7 +248,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 { diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 91bfddf450..e2ed649a8a 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -67,7 +67,7 @@ pub use unresolved_shuffle::UnresolvedShuffleExec; /// - `is_sort_shuffle` — selects between sort-shuffle and hash-shuffle path layout pub fn create_shuffle_path>( work_dir: P, - job_id: &str, + job_id: &JobId, stage_id: usize, partition_id: usize, file_id: Option, diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs index 4dbb63146a..f0ff5c05fb 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs @@ -73,7 +73,7 @@ impl SpillManager { /// * `compression` - Compression codec for spill files pub fn new( work_dir: &str, - job_id: &str, + job_id: &JobId, stage_id: usize, input_partition: usize, schema: SchemaRef, diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index bf067329ee..0c9d50cc6c 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -419,7 +419,7 @@ fn spill_all_partitions( #[allow(clippy::too_many_arguments)] fn finalize_output( work_dir: &str, - job_id: &str, + job_id: &JobId, stage_id: usize, input_partition: usize, buffered: &mut BufferedBatches, diff --git a/ballista/core/src/serde/scheduler/mod.rs b/ballista/core/src/serde/scheduler/mod.rs index 1f72d00368..594e4d974a 100644 --- a/ballista/core/src/serde/scheduler/mod.rs +++ b/ballista/core/src/serde/scheduler/mod.rs @@ -71,7 +71,7 @@ pub struct PartitionId { impl PartitionId { /// Creates a new partition ID with the given job, stage, and partition identifiers. - pub fn new(job_id: &str, stage_id: usize, partition_id: usize) -> Self { + pub fn new(job_id: &JobId, stage_id: usize, partition_id: usize) -> Self { Self { job_id: job_id.to_string(), stage_id, diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 131a50549f..70755865c1 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -741,7 +741,7 @@ async fn clean_all_shuffle_data(work_dir: &str) -> ballista_core::error::Result< /// Used by both push-based (gRPC handler) and pull-based (poll loop) cleanup. pub(crate) async fn remove_job_dir( work_dir: &str, - job_id: &str, + job_id: &JobId, ) -> ballista_core::error::Result<()> { let work_path = PathBuf::from(&work_dir); let job_path = work_path.join(job_id); @@ -983,7 +983,7 @@ mod tests { assert!(!is_subdirectory(&job_path, base_dir)); } } - fn prepare_testing_job_directory(base_dir: &Path, job_id: &str) -> PathBuf { + fn prepare_testing_job_directory(base_dir: &Path, job_id: &JobId) -> PathBuf { let mut path = base_dir.to_path_buf(); path.push(job_id); if !path.exists() { diff --git a/ballista/executor/src/metrics/mod.rs b/ballista/executor/src/metrics/mod.rs index b80f75adcc..cf91b3e7d1 100644 --- a/ballista/executor/src/metrics/mod.rs +++ b/ballista/executor/src/metrics/mod.rs @@ -28,7 +28,7 @@ pub trait ExecutorMetricsCollector: Send + Sync { /// Record metrics for stage after it is executed fn record_stage( &self, - job_id: &str, + job_id: &JobId, stage_id: usize, partition: usize, plan: Arc, @@ -43,7 +43,7 @@ pub struct LoggingMetricsCollector {} impl ExecutorMetricsCollector for LoggingMetricsCollector { fn record_stage( &self, - job_id: &str, + job_id: &JobId, stage_id: usize, partition: usize, plan: Arc, diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index 6f06ee59fe..403f6b44c7 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -348,7 +348,7 @@ impl JobState for InMemoryJobState { } } - async fn get_job_status(&self, job_id: &str) -> Result> { + async fn get_job_status(&self, job_id: &JobId) -> Result> { if let Some((job_name, queued_at)) = self.queued_jobs.get(job_id).as_deref() { return Ok(Some(JobStatus { job_id: job_id.to_string(), @@ -372,7 +372,7 @@ impl JobState for InMemoryJobState { async fn get_execution_graph( &self, - job_id: &str, + job_id: &JobId, ) -> Result> { Ok(self .completed_jobs @@ -381,13 +381,13 @@ impl JobState for InMemoryJobState { .and_then(|(_, graph)| graph.as_ref().map(|e| e.cloned()))) } - async fn try_acquire_job(&self, _job_id: &str) -> Result> { + async fn try_acquire_job(&self, _job_id: &JobId) -> Result> { // Always return None. The only state stored here are for completed jobs // which cannot be acquired Ok(None) } - async fn save_job(&self, job_id: &str, graph: &ExecutionGraphBox) -> Result<()> { + async fn save_job(&self, job_id: &JobId, graph: &ExecutionGraphBox) -> Result<()> { let status = graph.status().clone(); // If job is either successful or failed, save to completed jobs if matches!( @@ -451,7 +451,7 @@ impl JobState for InMemoryJobState { Ok(Box::pin(self.job_event_sender.subscribe())) } - async fn remove_job(&self, job_id: &str) -> Result<()> { + async fn remove_job(&self, job_id: &JobId) -> Result<()> { if self.completed_jobs.remove(job_id).is_none() { warn!("Tried to delete non-existent job {job_id} from state"); } @@ -477,7 +477,7 @@ impl JobState for InMemoryJobState { Ok(all_jobs) } - fn accept_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()> { + fn accept_job(&self, job_id: &JobId, job_name: &JobName, queued_at: u64) -> Result<()> { self.queued_jobs .insert(job_id.to_string(), (job_name.to_string(), queued_at)); @@ -488,7 +488,7 @@ impl JobState for InMemoryJobState { self.queued_jobs.len() } - async fn fail_unscheduled_job(&self, job_id: &str, reason: String) -> Result<()> { + async fn fail_unscheduled_job(&self, job_id: &JobId, reason: String) -> Result<()> { if let Some((job_id, (job_name, queued_at))) = self.queued_jobs.remove(job_id) { self.completed_jobs.insert( job_id.clone(), diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index 82eeab4d98..d917ed44bf 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -282,7 +282,7 @@ pub trait JobState: Send + Sync { /// Accepts a job into the scheduler's queue. /// /// Called when a job is received but before it is planned. - fn accept_job(&self, job_id: &str, job_name: &JobName, queued_at: u64) -> Result<()>; + fn accept_job(&self, job_id: &JobId, job_name: &JobName, queued_at: u64) -> Result<()>; /// Returns the number of queued jobs waiting to be scheduled. fn pending_job_number(&self) -> usize; @@ -304,7 +304,7 @@ pub trait JobState: Send + Sync { async fn get_all_jobs(&self) -> Result>; /// Returns the status of the specified job. - async fn get_job_status(&self, job_id: &str) -> Result>; + async fn get_job_status(&self, job_id: &JobId) -> Result>; /// Returns the execution graph for a job. /// @@ -312,27 +312,27 @@ pub trait JobState: Send + Sync { /// by another scheduler after this call returns. async fn get_execution_graph( &self, - job_id: &str, + job_id: &JobId, ) -> Result>; /// Persists the current state of an owned job. /// /// Returns an error if the job is not owned by the caller. - async fn save_job(&self, job_id: &str, graph: &ExecutionGraphBox) -> Result<()>; + async fn save_job(&self, job_id: &JobId, graph: &ExecutionGraphBox) -> Result<()>; /// Marks an unscheduled job as failed. /// /// Called when a job fails during planning before an execution graph is created. - async fn fail_unscheduled_job(&self, job_id: &str, reason: String) -> Result<()>; + async fn fail_unscheduled_job(&self, job_id: &JobId, reason: String) -> Result<()>; /// Deletes a job from the state. - async fn remove_job(&self, job_id: &str) -> Result<()>; + async fn remove_job(&self, job_id: &JobId) -> Result<()>; /// Attempts to acquire ownership of a job. /// /// Returns the execution graph if the job is still running and successfully acquired, /// otherwise returns None. - async fn try_acquire_job(&self, job_id: &str) -> Result>; + async fn try_acquire_job(&self, job_id: &JobId) -> Result>; /// Returns a stream of job state events. async fn job_state_events(&self) -> Result; @@ -720,7 +720,7 @@ mod test { } async fn mock_graph( - job_id: &str, + job_id: &JobId, num_target_partitions: usize, num_pending_task: usize, ) -> Result { diff --git a/ballista/scheduler/src/cluster/test_util/mod.rs b/ballista/scheduler/src/cluster/test_util/mod.rs index 1fee57132a..e98a0fec8b 100644 --- a/ballista/scheduler/src/cluster/test_util/mod.rs +++ b/ballista/scheduler/src/cluster/test_util/mod.rs @@ -56,14 +56,14 @@ impl JobStateTest { } /// Queues a job with the given ID. - pub fn queue_job(self, job_id: &str) -> Result { + pub fn queue_job(self, job_id: &JobId) -> Result { self.state .accept_job(job_id, &JobName::new(""), timestamp_millis())?; Ok(self) } /// Marks a job as failed during planning. - pub async fn fail_planning(self, job_id: &str) -> Result { + pub async fn fail_planning(self, job_id: &JobId) -> Result { self.state .fail_unscheduled_job(job_id, "failed planning".to_string()) .await?; @@ -71,7 +71,7 @@ impl JobStateTest { } /// Asserts the job is in queued status. - pub async fn assert_queued(self, job_id: &str) -> Result { + pub async fn assert_queued(self, job_id: &JobId) -> Result { let status = self.state.get_job_status(job_id).await?; assert!(status.is_some(), "Queued job {} not found", job_id); @@ -97,7 +97,7 @@ impl JobStateTest { } /// Asserts the job is in running status. - pub async fn assert_job_running(self, job_id: &str) -> Result { + pub async fn assert_job_running(self, job_id: &JobId) -> Result { let status = self.state.get_job_status(job_id).await?; assert!(status.is_some(), "Job status not found for {}", job_id); @@ -121,7 +121,7 @@ impl JobStateTest { } /// Asserts the job is in failed status. - pub async fn assert_job_failed(self, job_id: &str) -> Result { + pub async fn assert_job_failed(self, job_id: &JobId) -> Result { let status = self.state.get_job_status(job_id).await?; assert!(status.is_some(), "Job status not found for {}", job_id); @@ -139,7 +139,7 @@ impl JobStateTest { } /// Asserts the job completed successfully. - pub async fn assert_job_successful(self, job_id: &str) -> Result { + pub async fn assert_job_successful(self, job_id: &JobId) -> Result { let status = self.state.get_job_status(job_id).await?; assert!(status.is_some(), "Job status not found for {}", job_id); diff --git a/ballista/scheduler/src/display.rs b/ballista/scheduler/src/display.rs index fb47f42a91..cd929a8a41 100644 --- a/ballista/scheduler/src/display.rs +++ b/ballista/scheduler/src/display.rs @@ -51,7 +51,7 @@ fn merge_stage_metrics( /// Prints the physical plan for a completed stage with its aggregated metrics. pub fn print_stage_metrics( - job_id: &str, + job_id: &JobId, stage_id: usize, plan: &dyn ExecutionPlan, stage_metrics: &[MetricsSet], diff --git a/ballista/scheduler/src/metrics/mod.rs b/ballista/scheduler/src/metrics/mod.rs index 4c8f1149f6..509e319ce2 100644 --- a/ballista/scheduler/src/metrics/mod.rs +++ b/ballista/scheduler/src/metrics/mod.rs @@ -33,21 +33,21 @@ pub trait SchedulerMetricsCollector: Send + Sync { /// on executors. /// When invoked should specify the timestamp in milliseconds when the job was originally /// queued and the timestamp in milliseconds when it was submitted - fn record_submitted(&self, job_id: &str, queued_at: u64, submitted_at: u64); + fn record_submitted(&self, job_id: &JobId, queued_at: u64, submitted_at: u64); /// Record that job with `job_id` has completed successfully. This should only /// be invoked on successful job completion. /// When invoked should specify the timestamp in milliseconds when the job was originally /// queued and the timestamp in milliseconds when it was completed - fn record_completed(&self, job_id: &str, queued_at: u64, completed_at: u64); + fn record_completed(&self, job_id: &JobId, queued_at: u64, completed_at: u64); /// Record that job with `job_id` has failed. /// When invoked should specify the timestamp in milliseconds when the job was originally /// queued and the timestamp in milliseconds when it failed. - fn record_failed(&self, job_id: &str, queued_at: u64, failed_at: u64); + fn record_failed(&self, job_id: &JobId, queued_at: u64, failed_at: u64); /// Record that job with `job_id` was cancelled. - fn record_cancelled(&self, job_id: &str); + fn record_cancelled(&self, job_id: &JobId); /// Set the current number of pending tasks in scheduler. A pending task is a task that is available /// to schedule on an executor but cannot be scheduled because no resources are available. @@ -64,10 +64,10 @@ pub trait SchedulerMetricsCollector: Send + Sync { pub struct NoopMetricsCollector {} impl SchedulerMetricsCollector for NoopMetricsCollector { - fn record_submitted(&self, _job_id: &str, _queued_at: u64, _submitted_at: u64) {} - fn record_completed(&self, _job_id: &str, _queued_at: u64, _completed_att: u64) {} - fn record_failed(&self, _job_id: &str, _queued_at: u64, _failed_at: u64) {} - fn record_cancelled(&self, _job_id: &str) {} + fn record_submitted(&self, _job_id: &JobId, _queued_at: u64, _submitted_at: u64) {} + fn record_completed(&self, _job_id: &JobId, _queued_at: u64, _completed_att: u64) {} + fn record_failed(&self, _job_id: &JobId, _queued_at: u64, _failed_at: u64) {} + fn record_cancelled(&self, _job_id: &JobId) {} fn set_pending_tasks_queue_size(&self, _value: u64) {} fn gather_metrics(&self) -> Result, String)>> { diff --git a/ballista/scheduler/src/metrics/prometheus.rs b/ballista/scheduler/src/metrics/prometheus.rs index 7ffd3e657f..bbddc1fae7 100644 --- a/ballista/scheduler/src/metrics/prometheus.rs +++ b/ballista/scheduler/src/metrics/prometheus.rs @@ -140,23 +140,23 @@ impl PrometheusMetricsCollector { } impl SchedulerMetricsCollector for PrometheusMetricsCollector { - fn record_submitted(&self, _job_id: &str, queued_at: u64, submitted_at: u64) { + fn record_submitted(&self, _job_id: &JobId, queued_at: u64, submitted_at: u64) { self.submitted.inc(); self.planning_time .observe((submitted_at - queued_at) as f64); } - fn record_completed(&self, _job_id: &str, queued_at: u64, completed_at: u64) { + fn record_completed(&self, _job_id: &JobId, queued_at: u64, completed_at: u64) { self.completed.inc(); self.execution_time .observe((completed_at - queued_at) as f64 / 1000_f64) } - fn record_failed(&self, _job_id: &str, _queued_at: u64, _failed_at: u64) { + fn record_failed(&self, _job_id: &JobId, _queued_at: u64, _failed_at: u64) { self.failed.inc() } - fn record_cancelled(&self, _job_id: &str) { + fn record_cancelled(&self, _job_id: &JobId) { self.cancelled.inc(); } diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 04405d9fd2..fbada7ac2a 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -541,7 +541,7 @@ pub fn rollback_resolved_shuffles( } pub(crate) fn create_shuffle_writer_with_config( - job_id: &str, + job_id: &JobId, stage_id: usize, plan: Arc, partitioning: Option, diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index ccc8787cfd..5e851a29cd 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -115,7 +115,7 @@ impl BallistaAdapter { /// ShuffleWriterExec/SortShuffleWriterExec and [ShuffleReaderExec] pub fn adapt_to_ballista( plan: Arc, - job_id: &str, + job_id: &JobId, config: &ConfigOptions, ) -> datafusion::error::Result { if let Some(root) = plan.as_any().downcast_ref::() { diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 05d1e38c28..073a898a19 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -132,7 +132,7 @@ impl AdaptiveExecutionGraph { #[allow(clippy::too_many_arguments)] pub async fn try_new( scheduler_id: &str, - job_id: &str, + job_id: &JobId, job_name: &JobName, ctx: &SessionContext, logical_plan: &LogicalPlan, diff --git a/ballista/scheduler/src/state/distributed_explain.rs b/ballista/scheduler/src/state/distributed_explain.rs index fc31580329..ea08ea17f7 100644 --- a/ballista/scheduler/src/state/distributed_explain.rs +++ b/ballista/scheduler/src/state/distributed_explain.rs @@ -42,7 +42,7 @@ use crate::{ }; pub(crate) async fn generate_distributed_explain_plan( - job_id: &str, + job_id: &JobId, session_ctx: &SessionContext, plan: Arc, ) -> Result { @@ -188,7 +188,7 @@ fn render_stages(stages: HashMap) -> String { } pub(crate) async fn handle_explain_plan( - job_id: &str, + job_id: &JobId, ctx: &SessionContext, logical_plan: &LogicalPlan, plan: Arc, diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 0c5817a51f..7a6fdaf2de 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -305,7 +305,7 @@ impl StaticExecutionGraph { #[allow(clippy::too_many_arguments)] pub fn new( scheduler_id: &str, - job_id: &str, + job_id: &JobId, job_name: &JobName, session_id: &str, plan: Arc, @@ -1733,7 +1733,7 @@ impl TaskDescription { } pub(crate) fn partition_to_location( - job_id: &str, + job_id: &JobId, map_partition_id: usize, stage_id: usize, executor: &ExecutorMetadata, diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index 76cf5b9a36..159c902726 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -363,7 +363,7 @@ impl SchedulerState, logical_plan: &LogicalPlan, diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 3e1680a0e3..3b7b62ee51 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -256,7 +256,7 @@ impl TaskManager /// Enqueue a job for scheduling pub fn queue_job( &self, - job_id: &str, + job_id: &JobId, job_name: &JobName, queued_at: u64, ) -> Result<()> { @@ -280,7 +280,7 @@ impl TaskManager #[allow(clippy::too_many_arguments)] pub async fn submit_job( &self, - job_id: &str, + job_id: &JobId, job_name: &JobName, ctx: Arc, logical_plan: &LogicalPlan, @@ -423,7 +423,7 @@ impl TaskManager /// Get the status of of a job. First look in the active cache. /// If no one found, then in the Active/Completed jobs, and then in Failed jobs - pub async fn get_job_status(&self, job_id: &str) -> Result> { + pub async fn get_job_status(&self, job_id: &JobId) -> Result> { if let Some(graph) = self.get_active_execution_graph(job_id) { let guard = graph.read().await; @@ -437,7 +437,7 @@ impl TaskManager /// If no one found, then in the Active/Completed jobs. pub(crate) async fn get_job_execution_graph( &self, - job_id: &str, + job_id: &JobId, ) -> Result> { if let Some(cached) = self.get_active_execution_graph(job_id) { let guard = cached.read().await; @@ -451,7 +451,7 @@ impl TaskManager } /// Get the session configuration for a job. - pub async fn get_job_config(&self, job_id: &str) -> Result> { + pub async fn get_job_config(&self, job_id: &JobId) -> Result> { let graph = self .get_job_execution_graph(job_id) .await? @@ -510,7 +510,7 @@ impl TaskManager /// Mark a job to success. This will create a key under the CompletedJobs keyspace /// and remove the job from ActiveJobs - pub(crate) async fn succeed_job(&self, job_id: &str) -> Result<()> { + pub(crate) async fn succeed_job(&self, job_id: &JobId) -> Result<()> { debug!("Moving job {job_id} from Active to Success"); if let Some(graph) = self.remove_active_execution_graph(job_id) { @@ -531,7 +531,7 @@ impl TaskManager /// Cancel the job and return a Vec of running tasks need to cancel pub(crate) async fn cancel_job( &self, - job_id: &str, + job_id: &JobId, ) -> Result<(Vec, usize)> { self.abort_job(job_id, "Cancelled".to_owned()).await } @@ -539,7 +539,7 @@ impl TaskManager /// Abort the job and return a Vec of running tasks need to cancel pub(crate) async fn abort_job( &self, - job_id: &str, + job_id: &JobId, failure_reason: String, ) -> Result<(Vec, usize)> { let (tasks_to_cancel, pending_tasks) = if let Some(graph) = @@ -576,7 +576,7 @@ impl TaskManager /// and remove the job from ActiveJobs or QueuedJobs pub async fn fail_unscheduled_job( &self, - job_id: &str, + job_id: &JobId, failure_reason: String, ) -> Result<()> { self.state @@ -585,7 +585,7 @@ impl TaskManager } /// Updates the job state and returns the number of new available tasks. - pub async fn update_job(&self, job_id: &str) -> Result { + pub async fn update_job(&self, job_id: &JobId) -> Result { debug!("Update active job {job_id}"); if let Some(graph) = self.get_active_execution_graph(job_id) { let mut graph = graph.write().await; @@ -632,7 +632,7 @@ impl TaskManager /// Retrieves the number of available tasks for the given job. /// /// The value returned is a point-in-time snapshot and may change immediately. - pub async fn get_available_task_count(&self, job_id: &str) -> Result { + pub async fn get_available_task_count(&self, job_id: &JobId) -> Result { if let Some(graph) = self.get_active_execution_graph(job_id) { let available_tasks = graph.read().await.available_tasks(); Ok(available_tasks) @@ -779,7 +779,7 @@ impl TaskManager /// Get the `ExecutionGraph` for the given job ID from cache pub(crate) fn get_active_execution_graph( &self, - job_id: &str, + job_id: &JobId, ) -> Option>> { self.active_job_cache .get(job_id) @@ -790,7 +790,7 @@ impl TaskManager /// Remove the `ExecutionGraph` for the given job ID from cache pub(crate) fn remove_active_execution_graph( &self, - job_id: &str, + job_id: &JobId, ) -> Option>> { self.active_job_cache .remove(job_id) diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 3a953f2b43..fedad3de0b 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -549,7 +549,7 @@ impl SchedulerTest { } /// Cancels a job by ID. - pub async fn cancel(&self, job_id: &str) -> Result<()> { + pub async fn cancel(&self, job_id: &JobId) -> Result<()> { self.scheduler .query_stage_event_loop .get_sender()? @@ -560,7 +560,7 @@ impl SchedulerTest { /// Waits for job completion with a timeout in milliseconds. pub async fn await_completion_timeout( &self, - job_id: &str, + job_id: &JobId, timeout_ms: u64, ) -> Result { let mut time = 0; @@ -599,7 +599,7 @@ impl SchedulerTest { } /// Waits for job completion indefinitely. - pub async fn await_completion(&self, job_id: &str) -> Result { + pub async fn await_completion(&self, job_id: &JobId) -> Result { let final_status: Result = loop { let status = self .scheduler @@ -728,7 +728,7 @@ pub struct TestMetricsCollector { impl TestMetricsCollector { /// Returns all events for the given job ID. - pub fn job_events(&self, job_id: &str) -> Vec { + pub fn job_events(&self, job_id: &JobId) -> Vec { let guard = self.events.lock(); guard @@ -745,7 +745,7 @@ impl TestMetricsCollector { } impl SchedulerMetricsCollector for TestMetricsCollector { - fn record_submitted(&self, job_id: &str, queued_at: u64, submitted_at: u64) { + fn record_submitted(&self, job_id: &JobId, queued_at: u64, submitted_at: u64) { let mut guard = self.events.lock(); guard.push(MetricEvent::Submitted( job_id.to_owned(), @@ -754,7 +754,7 @@ impl SchedulerMetricsCollector for TestMetricsCollector { )); } - fn record_completed(&self, job_id: &str, queued_at: u64, completed_at: u64) { + fn record_completed(&self, job_id: &JobId, queued_at: u64, completed_at: u64) { let mut guard = self.events.lock(); guard.push(MetricEvent::Completed( job_id.to_owned(), @@ -763,12 +763,12 @@ impl SchedulerMetricsCollector for TestMetricsCollector { )); } - fn record_failed(&self, job_id: &str, queued_at: u64, failed_at: u64) { + fn record_failed(&self, job_id: &JobId, queued_at: u64, failed_at: u64) { let mut guard = self.events.lock(); guard.push(MetricEvent::Failed(job_id.to_owned(), queued_at, failed_at)); } - fn record_cancelled(&self, job_id: &str) { + fn record_cancelled(&self, job_id: &JobId) { let mut guard = self.events.lock(); guard.push(MetricEvent::Cancelled(job_id.to_owned())); } @@ -781,7 +781,7 @@ impl SchedulerMetricsCollector for TestMetricsCollector { } /// Asserts that a submitted event was recorded for the job. -pub fn assert_submitted_event(job_id: &str, collector: &TestMetricsCollector) { +pub fn assert_submitted_event(job_id: &JobId, collector: &TestMetricsCollector) { let found = collector .job_events(job_id) .iter() @@ -791,7 +791,7 @@ pub fn assert_submitted_event(job_id: &str, collector: &TestMetricsCollector) { } /// Asserts that no submitted event was recorded for the job. -pub fn assert_no_submitted_event(job_id: &str, collector: &TestMetricsCollector) { +pub fn assert_no_submitted_event(job_id: &JobId, collector: &TestMetricsCollector) { let found = collector .job_events(job_id) .iter() @@ -801,7 +801,7 @@ pub fn assert_no_submitted_event(job_id: &str, collector: &TestMetricsCollector) } /// Asserts that a completed event was recorded for the job. -pub fn assert_completed_event(job_id: &str, collector: &TestMetricsCollector) { +pub fn assert_completed_event(job_id: &JobId, collector: &TestMetricsCollector) { let found = collector .job_events(job_id) .iter() @@ -811,7 +811,7 @@ pub fn assert_completed_event(job_id: &str, collector: &TestMetricsCollector) { } /// Asserts that a cancelled event was recorded for the job. -pub fn assert_cancelled_event(job_id: &str, collector: &TestMetricsCollector) { +pub fn assert_cancelled_event(job_id: &JobId, collector: &TestMetricsCollector) { let found = collector .job_events(job_id) .iter() @@ -821,7 +821,7 @@ pub fn assert_cancelled_event(job_id: &str, collector: &TestMetricsCollector) { } /// Asserts that a failed event was recorded for the job. -pub fn assert_failed_event(job_id: &str, collector: &TestMetricsCollector) { +pub fn assert_failed_event(job_id: &JobId, collector: &TestMetricsCollector) { let found = collector .job_events(job_id) .iter() @@ -883,7 +883,7 @@ pub async fn test_aggregation_plan(partition: usize) -> StaticExecutionGraph { /// Creates a test execution graph with a simple aggregation plan and custom job ID. pub async fn test_aggregation_plan_with_job_id( partition: usize, - job_id: &str, + job_id: &JobId, ) -> StaticExecutionGraph { let config = SessionConfig::new().with_target_partitions(partition); let ctx = Arc::new(SessionContext::new_with_config(config)); From ba1f845275a42c4086d2b7a3f356e874d41b545a Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Tue, 9 Jun 2026 23:00:50 +0200 Subject: [PATCH 08/11] (`JobId`): Following the compiler: change all `&str`/`String` to `(&)JobId` in core/executor/scheduler. --- .../distributed_explain_analyze.rs | 5 +- .../src/execution_plans/distributed_query.rs | 25 ++++---- ballista/core/src/execution_plans/mod.rs | 28 ++++++-- .../src/execution_plans/shuffle_reader.rs | 24 +++---- .../src/execution_plans/shuffle_writer.rs | 10 +-- .../execution_plans/shuffle_writer_trait.rs | 4 +- .../src/execution_plans/sort_shuffle/spill.rs | 11 ++-- .../execution_plans/sort_shuffle/writer.rs | 14 ++-- ballista/core/src/serde/mod.rs | 4 +- .../core/src/serde/scheduler/from_proto.rs | 10 +-- ballista/core/src/serde/scheduler/mod.rs | 2 +- ballista/core/src/serde/scheduler/to_proto.rs | 4 +- ballista/executor/src/execution_engine.rs | 2 +- ballista/executor/src/execution_loop.rs | 8 +-- ballista/executor/src/executor.rs | 7 +- ballista/executor/src/executor_process.rs | 15 +++-- ballista/executor/src/executor_server.rs | 4 +- ballista/executor/src/lib.rs | 4 +- ballista/executor/src/metrics/mod.rs | 1 + ballista/scheduler/src/api/handlers.rs | 20 +++--- ballista/scheduler/src/cluster/memory.rs | 42 +++++++----- ballista/scheduler/src/cluster/mod.rs | 64 ++++++++++--------- .../scheduler/src/cluster/test_util/mod.rs | 16 ++--- ballista/scheduler/src/display.rs | 1 + ballista/scheduler/src/metrics/mod.rs | 2 +- ballista/scheduler/src/metrics/prometheus.rs | 1 + ballista/scheduler/src/planner.rs | 27 ++++---- .../scheduler/src/scheduler_server/event.rs | 7 +- .../scheduler/src/scheduler_server/grpc.rs | 17 +++-- .../scheduler/src/scheduler_server/mod.rs | 23 ++++--- .../scheduler_server/query_stage_scheduler.rs | 2 +- ballista/scheduler/src/state/aqe/adapter.rs | 1 + ballista/scheduler/src/state/aqe/mod.rs | 10 +-- ballista/scheduler/src/state/aqe/planner.rs | 18 +++--- .../src/state/aqe/test/alter_stages.rs | 4 +- .../src/state/aqe/test/coalesce_rule.rs | 2 +- ballista/scheduler/src/state/aqe/test/mod.rs | 4 +- .../src/state/distributed_explain.rs | 1 + .../scheduler/src/state/execution_graph.rs | 12 ++-- .../src/state/execution_graph_dot.rs | 6 +- .../scheduler/src/state/executor_manager.rs | 14 ++-- ballista/scheduler/src/state/mod.rs | 6 +- ballista/scheduler/src/state/task_manager.rs | 21 +++--- ballista/scheduler/src/test_utils.rs | 42 ++++++------ benchmarks/benches/sort_shuffle.rs | 2 +- benchmarks/src/bin/shuffle_bench.rs | 4 +- 46 files changed, 303 insertions(+), 248 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_explain_analyze.rs b/ballista/core/src/execution_plans/distributed_explain_analyze.rs index 08ec5147d2..841bd58094 100644 --- a/ballista/core/src/execution_plans/distributed_explain_analyze.rs +++ b/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::{ @@ -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) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 2ef4521f67..25fe3592b4 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/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 { /// - 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>>, + job_id: Arc>>, } impl DistributedQueryExec { @@ -136,7 +137,7 @@ impl DistributedQueryExec { } /// Returns the scheduler job id after the query has been accepted. - pub fn job_id(&self) -> Option { + pub fn job_id(&self) -> Option { self.job_id.lock().clone() } @@ -336,7 +337,7 @@ async fn execute_query_pull( max_message_size: usize, grpc_config: GrpcClientConfig, metrics: Arc, - job_id_handle: Arc>>, + job_id_handle: Arc>>, partition: usize, session_config: SessionConfig, ) -> Result> + 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 = 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, - job_id_handle: Arc>>, + job_id_handle: Arc>>, partition: usize, session_config: SessionConfig, ) -> Result> + 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,7 +838,7 @@ 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 @@ -843,6 +846,6 @@ mod test { .downcast_ref::>() .unwrap(); - assert_eq!(new_exec.job_id().as_deref(), Some("job-123")); + assert_eq!(new_exec.job_id(), Some(JobId::new("job-123"))); } } diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index e2ed649a8a..db37427081 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/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: @@ -76,7 +78,7 @@ pub fn create_shuffle_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" diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index 8311c2a6a0..7450c86718 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/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, }, diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 383697c075..684eb176cf 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -178,7 +178,7 @@ impl ShuffleWriterExec { } /// Get the Job ID for this query stage - pub fn job_id(&self) -> &str { + pub fn job_id(&self) -> &JobId { &self.job_id } @@ -487,7 +487,7 @@ impl ExecutionPlan for ShuffleWriterExec { let schema = result_schema(); let schema_captured = schema.clone(); - let job_id = self.job_id.to_string(); + let job_id = self.job_id.clone(); let work_dir = self.work_dir.to_string(); let stage_id = self.stage_id; let fut_stream = self @@ -567,7 +567,7 @@ impl ExecutionPlan for ShuffleWriterExec { } impl ShuffleWriter for ShuffleWriterExec { - fn job_id(&self) -> &str { + fn job_id(&self) -> &JobId { &self.job_id } @@ -730,7 +730,7 @@ mod tests { let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), + "jobOne".to_owned().into(), 1, input_plan, work_dir, @@ -763,7 +763,7 @@ mod tests { let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), + "jobOne".to_owned().into(), 1, input_plan, work_dir, diff --git a/ballista/core/src/execution_plans/shuffle_writer_trait.rs b/ballista/core/src/execution_plans/shuffle_writer_trait.rs index dde69c11af..f73d784cac 100644 --- a/ballista/core/src/execution_plans/shuffle_writer_trait.rs +++ b/ballista/core/src/execution_plans/shuffle_writer_trait.rs @@ -24,13 +24,15 @@ use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use std::fmt::Debug; use std::sync::Arc; +use crate::JobId; + /// Trait for shuffle writer execution plans. /// /// This trait defines the common interface needed by the distributed planner /// and execution graph to work with different shuffle implementations. pub trait ShuffleWriter: ExecutionPlan + Debug + Send + Sync { /// Get the Job ID for this query stage. - fn job_id(&self) -> &str; + fn job_id(&self) -> &JobId; /// Get the Stage ID for this query stage. fn stage_id(&self) -> usize; diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs index f0ff5c05fb..0fa7207f50 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs @@ -21,6 +21,7 @@ //! At finalization, the spill bytes are concatenated verbatim into the //! consolidated output file alongside the in-memory remainder. +use crate::JobId; use crate::error::{BallistaError, Result}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::ipc::reader::StreamReader; @@ -80,7 +81,7 @@ impl SpillManager { compression: CompressionType, ) -> Result { let mut spill_dir = PathBuf::from(work_dir); - spill_dir.push(job_id); + spill_dir.push(job_id.as_str()); spill_dir.push(format!("{stage_id}")); spill_dir.push(format!("{input_partition}")); spill_dir.push("spill"); @@ -273,7 +274,7 @@ mod tests { let mut manager = SpillManager::new( temp_dir.path().to_str().unwrap(), - "job1", + &JobId::new("job1"), 1, 0, schema.clone(), @@ -308,7 +309,7 @@ mod tests { let mut manager = SpillManager::new( temp_dir.path().to_str().unwrap(), - "job1", + &JobId::new("job1"), 1, 0, schema.clone(), @@ -345,7 +346,7 @@ mod tests { let mut manager = SpillManager::new( temp_dir.path().to_str().unwrap(), - "job1", + &JobId::new("job1"), 1, 0, schema.clone(), @@ -391,7 +392,7 @@ mod tests { let mut manager = SpillManager::new( temp_dir.path().to_str().unwrap(), - "job1", + &JobId::new("job1"), 1, 0, schema.clone(), diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 0c9d50cc6c..7af4247ca4 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -169,7 +169,7 @@ impl SortShuffleWriterExec { } /// Get the Job ID for this query stage - pub fn job_id(&self) -> &str { + pub fn job_id(&self) -> &JobId { &self.job_id } @@ -432,7 +432,7 @@ fn finalize_output( let mut partition_stats = Vec::with_capacity(num_partitions); let mut output_dir = PathBuf::from(work_dir); - output_dir.push(job_id); + output_dir.push(job_id.as_str()); output_dir.push(format!("{stage_id}")); output_dir.push(format!("{input_partition}")); std::fs::create_dir_all(&output_dir)?; @@ -589,7 +589,7 @@ impl ExecutionPlan for SortShuffleWriterExec { let schema = result_schema(); let schema_captured = schema.clone(); - let job_id = self.job_id.to_string(); + let job_id = self.job_id.clone(); let work_dir = self.work_dir.to_string(); let stage_id = self.stage_id; let fut_stream = self @@ -669,7 +669,7 @@ impl ExecutionPlan for SortShuffleWriterExec { } impl ShuffleWriter for SortShuffleWriterExec { - fn job_id(&self) -> &str { + fn job_id(&self) -> &JobId { &self.job_id } @@ -840,7 +840,7 @@ mod tests { let config = SortShuffleConfig::default(); let writer = SortShuffleWriterExec::try_new( - "job1".to_string(), + "job1".to_owned().into(), 1, input_plan, work_dir.path().to_str().unwrap().to_string(), @@ -937,7 +937,7 @@ mod tests { let work_dir = TempDir::new()?; let writer = SortShuffleWriterExec::try_new( - "round_trip_job".to_string(), + "round_trip_job".to_owned().into(), 1, input, work_dir.path().to_str().unwrap().to_string(), @@ -1084,7 +1084,7 @@ mod tests { let num_partitions = 8; let writer = SortShuffleWriterExec::try_new( - "empty_partitions_job".to_string(), + "empty_partitions_job".to_owned().into(), 1, input, work_dir.path().to_str().unwrap().to_string(), diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index a231c0a875..46e3430c38 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -394,7 +394,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )?; Ok(Arc::new(ShuffleWriterExec::try_new( - shuffle_writer.job_id.clone(), + shuffle_writer.job_id.clone().into(), shuffle_writer.stage_id as usize, input, "".to_string(), // this is intentional but hacky - the executor will fill this in @@ -430,7 +430,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { ); Ok(Arc::new(SortShuffleWriterExec::try_new( - sort_shuffle_writer.job_id.clone(), + sort_shuffle_writer.job_id.clone().into(), sort_shuffle_writer.stage_id as usize, input, "".to_string(), // executor will fill this in diff --git a/ballista/core/src/serde/scheduler/from_proto.rs b/ballista/core/src/serde/scheduler/from_proto.rs index b99cb274fb..ba8b9b33a0 100644 --- a/ballista/core/src/serde/scheduler/from_proto.rs +++ b/ballista/core/src/serde/scheduler/from_proto.rs @@ -41,8 +41,8 @@ use crate::serde::scheduler::{ PartitionLocation, PartitionStats, TaskDefinition, }; -use crate::RuntimeProducer; use crate::serde::{BallistaCodec, protobuf}; +use crate::{JobId, RuntimeProducer}; use protobuf::{NamedCount, NamedGauge, NamedTime, operator_metric}; impl TryInto for protobuf::Action { @@ -52,7 +52,7 @@ impl TryInto for protobuf::Action { match self.action_type { Some(protobuf::action::ActionType::FetchPartition(fetch)) => { Ok(Action::FetchPartition { - job_id: fetch.job_id, + job_id: fetch.job_id.into(), stage_id: fetch.stage_id as usize, partition_id: fetch.partition_id as usize, file_id: fetch.file_id, @@ -72,7 +72,7 @@ impl TryInto for protobuf::Action { impl Into for protobuf::PartitionId { fn into(self) -> PartitionId { PartitionId::new( - &self.job_id, + &self.job_id.into(), self.stage_id as usize, self.partition_id as usize, ) @@ -371,7 +371,7 @@ pub fn get_task_definition Self { Self { - job_id: job_id.to_string(), + job_id: job_id.to_owned().into(), stage_id, partition_id, } diff --git a/ballista/core/src/serde/scheduler/to_proto.rs b/ballista/core/src/serde/scheduler/to_proto.rs index d09dbad17a..2f804a8ce0 100644 --- a/ballista/core/src/serde/scheduler/to_proto.rs +++ b/ballista/core/src/serde/scheduler/to_proto.rs @@ -46,7 +46,7 @@ impl TryInto for Action { is_sort_shuffle, } => Ok(protobuf::Action { action_type: Some(ActionType::FetchPartition(protobuf::FetchPartition { - job_id, + job_id: job_id.into(), stage_id: stage_id as u32, partition_id: partition_id as u32, host, @@ -64,7 +64,7 @@ impl TryInto for Action { impl Into for PartitionId { fn into(self) -> protobuf::PartitionId { protobuf::PartitionId { - job_id: self.job_id, + job_id: self.job_id.into(), stage_id: self.stage_id as u32, partition_id: self.partition_id as u32, } diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index a4fa7aa271..6ce96fc8be 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -25,7 +25,7 @@ use ballista_core::client_pool::BallistaClientPool; use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec; use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; use ballista_core::serde::protobuf::ShuffleWritePartition; -use ballista_core::utils; +use ballista_core::{JobId, utils}; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::context::TaskContext; diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index c1683d54fd..fdb49a56bb 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -116,7 +116,7 @@ where // Clean up any state related to the listed jobs for cleanup in jobs_to_clean { - let job_id = cleanup.job_id.clone(); + let job_id = cleanup.job_id.clone().into(); let work_dir = executor.work_dir.clone(); // In poll-based cleanup, removing job data is fire-and-forget. @@ -158,7 +158,7 @@ where // let partition_id = PartitionId { - job_id: task.job_id.clone(), + job_id: task.job_id.clone().into(), stage_id: task.stage_id as usize, partition_id: task.partition_id as usize, }; @@ -275,7 +275,7 @@ async fn run_received_task ballista_core::error::Result<()> { let work_path = PathBuf::from(&work_dir); - let job_path = work_path.join(job_id); + let job_path = work_path.join(job_id.as_str()); // Match legacy behavior: If the job path does not exist, return OK if !tokio::fs::try_exists(&job_path).await.unwrap_or(false) { @@ -886,6 +886,7 @@ mod tests { use std::path::{Path, PathBuf}; use super::clean_shuffle_data_loop; + use ballista_core::JobId; use std::fs; use std::fs::File; use std::io::Write; @@ -964,28 +965,28 @@ mod tests { // Normal correct one { - let job_path = prepare_testing_job_directory(base_dir, "job_a"); + let job_path = prepare_testing_job_directory(base_dir, &JobId::new("job_a")); assert!(is_subdirectory(&job_path, base_dir)); } // Empty job id { - let job_path = prepare_testing_job_directory(base_dir, ""); + let job_path = prepare_testing_job_directory(base_dir, &JobId::new("")); assert!(!is_subdirectory(&job_path, base_dir)); - let job_path = prepare_testing_job_directory(base_dir, "."); + let job_path = prepare_testing_job_directory(base_dir, &JobId::new(".")); assert!(!is_subdirectory(&job_path, base_dir)); } // Malicious job id { - let job_path = prepare_testing_job_directory(base_dir, ".."); + let job_path = prepare_testing_job_directory(base_dir, &JobId::new("..")); assert!(!is_subdirectory(&job_path, base_dir)); } } fn prepare_testing_job_directory(base_dir: &Path, job_id: &JobId) -> PathBuf { let mut path = base_dir.to_path_buf(); - path.push(job_id); + path.push(job_id.as_str()); if !path.exists() { fs::create_dir(&path).unwrap(); } diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs index c52998ad59..6ef9fa5e4e 100644 --- a/ballista/executor/src/executor_server.rs +++ b/ballista/executor/src/executor_server.rs @@ -906,7 +906,7 @@ impl ExecutorGrpc .executor .cancel_task( task.task_id as usize, - task.job_id, + task.job_id.into(), task.stage_id as usize, task.partition_id as usize, ) @@ -924,7 +924,7 @@ impl ExecutorGrpc &self, request: Request, ) -> Result, Status> { - let job_id = request.into_inner().job_id; + let job_id = request.into_inner().job_id.into(); remove_job_dir(&self.executor.work_dir, &job_id) .await diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs index 4aeae77106..f0d0f995ec 100644 --- a/ballista/executor/src/lib.rs +++ b/ballista/executor/src/lib.rs @@ -118,7 +118,7 @@ pub fn as_task_status( ); TaskStatus { task_id: task_id as u32, - job_id: partition_id.job_id, + job_id: partition_id.job_id.into(), stage_id: partition_id.stage_id as u32, stage_attempt_num: stage_attempt_num as u32, partition_id: partition_id.partition_id as u32, @@ -138,7 +138,7 @@ pub fn as_task_status( TaskStatus { task_id: task_id as u32, - job_id: partition_id.job_id, + job_id: partition_id.job_id.into(), stage_id: partition_id.stage_id as u32, stage_attempt_num: stage_attempt_num as u32, partition_id: partition_id.partition_id as u32, diff --git a/ballista/executor/src/metrics/mod.rs b/ballista/executor/src/metrics/mod.rs index cf91b3e7d1..97fde7066b 100644 --- a/ballista/executor/src/metrics/mod.rs +++ b/ballista/executor/src/metrics/mod.rs @@ -16,6 +16,7 @@ // under the License. use crate::execution_engine::QueryStageExecutor; +use ballista_core::JobId; use log::debug; use std::{fmt::Display, sync::Arc}; diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 73e94e1a12..1dc07514d6 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -346,8 +346,8 @@ pub async fn get_jobs< ((job.completed_stages as f32 / job.num_stages as f32) * 100_f32) as u8 }; JobResponse { - job_id: job.job_id.to_string(), - job_name: job.job_name.to_owned().into(), + job_id: job.job_id.to_owned(), + job_name: job.job_name.to_owned(), job_status, status: plain_status, start_time: job.start_time, @@ -376,7 +376,7 @@ pub async fn get_job< let graph = data_server .state .task_manager - .get_job_execution_graph(&job_id) + .get_job_execution_graph(&job_id.clone().into()) .await .map_err(|err| { tracing::error!("Error occurred while getting the execution graph for job '{job_id}' reason: {err:?}"); @@ -409,7 +409,7 @@ pub async fn get_job< }; Ok(Json(JobResponse { - job_id: job.job_id().to_string(), + job_id: job.job_id().to_owned(), job_name: job.job_name().to_owned().into(), job_status, status: plain_status, @@ -435,7 +435,7 @@ pub async fn cancel_job< let job_status = data_server .state .task_manager - .get_job_status(&job_id) + .get_job_status(&job_id.clone().into()) .await .map_err(|err| { tracing::error!("Error getting job status: {err:?}"); @@ -457,7 +457,7 @@ pub async fn cancel_job< ); SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) })? - .post_event(QueryStageSchedulerEvent::JobCancel(job_id)) + .post_event(QueryStageSchedulerEvent::JobCancel(job_id.into())) .await .map_err(|_| { SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR) @@ -509,7 +509,7 @@ pub async fn get_query_stages< if let Some(graph) = data_server .state .task_manager - .get_job_execution_graph(&job_id) + .get_job_execution_graph(&job_id.clone().into()) .await .map_err(|e| { tracing::error!("Error occurred while getting the query stages for job '{job_id}' reason: {e:?}"); @@ -833,7 +833,7 @@ pub async fn get_job_dot_graph< if let Some(graph) = data_server .state .task_manager - .get_job_execution_graph(&job_id) + .get_job_execution_graph(&job_id.clone().into()) .await .map_err(|e| { tracing::error!("Error occurred while getting the dot graph for job '{job_id}' reason: {e:?}"); @@ -860,7 +860,7 @@ pub async fn get_query_stage_dot_graph< if let Some(graph) = data_server .state .task_manager - .get_job_execution_graph(&job_id) + .get_job_execution_graph(&job_id.clone().into()) .await .map_err(|_| SchedulerErrorResponse::new(StatusCode::INTERNAL_SERVER_ERROR))? { @@ -936,7 +936,7 @@ pub async fn get_job_config< data_server .state .task_manager - .get_job_config(&job_id) + .get_job_config(&job_id.clone().into()) .await .map(|e| Json(e.to_props())) .map_err(|_| SchedulerErrorResponse::new(StatusCode::NOT_FOUND)) diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index 403f6b44c7..5298c98971 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -26,7 +26,7 @@ use ballista_core::serde::protobuf::{ executor_status, }; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata}; -use ballista_core::{ConfigProducer, JobName, JobStatusSubscriber}; +use ballista_core::{ConfigProducer, JobId, JobName, JobStatusSubscriber}; use dashmap::DashMap; use datafusion::prelude::{SessionConfig, SessionContext}; use tokio::sync::mpsc::error::TrySendError; @@ -67,7 +67,7 @@ impl ClusterState for InMemoryClusterState { async fn bind_schedulable_tasks( &self, distribution: TaskDistributionPolicy, - active_jobs: Arc>, + active_jobs: Arc>, executors: Option>, ) -> Result> { let mut guard = self.task_slots.lock().await; @@ -264,11 +264,11 @@ impl ClusterState for InMemoryClusterState { pub struct InMemoryJobState { scheduler: String, /// Jobs which have either completed successfully or failed - completed_jobs: DashMap)>, + completed_jobs: DashMap)>, /// In-memory store of queued jobs. Map from Job ID -> (Job Name, queued_at timestamp) - queued_jobs: DashMap, + queued_jobs: DashMap, /// In-memory store of running job statuses. Map from Job ID -> JobStatus - running_jobs: DashMap, + running_jobs: DashMap, /// `SessionBuilder` for building DataFusion `SessionContext` from `BallistaConfig` session_builder: SessionBuilder, /// Sender of job events @@ -351,8 +351,8 @@ impl JobState for InMemoryJobState { async fn get_job_status(&self, job_id: &JobId) -> Result> { if let Some((job_name, queued_at)) = self.queued_jobs.get(job_id).as_deref() { return Ok(Some(JobStatus { - job_id: job_id.to_string(), - job_name: job_name.clone(), + job_id: job_id.to_owned().into(), + job_name: job_name.to_owned().into(), status: Some(Status::Queued(QueuedJob { queued_at: *queued_at, })), @@ -381,7 +381,10 @@ impl JobState for InMemoryJobState { .and_then(|(_, graph)| graph.as_ref().map(|e| e.cloned()))) } - async fn try_acquire_job(&self, _job_id: &JobId) -> Result> { + async fn try_acquire_job( + &self, + _job_id: &JobId, + ) -> Result> { // Always return None. The only state stored here are for completed jobs // which cannot be acquired Ok(None) @@ -399,7 +402,7 @@ impl JobState for InMemoryJobState { } self.completed_jobs - .insert(job_id.to_string(), (status.clone(), Some(graph.cloned()))); + .insert(job_id.to_owned(), (status.clone(), Some(graph.cloned()))); } else { // otherwise update running job if let Some(mut job_info) = self.running_jobs.get_mut(job_id) { @@ -417,7 +420,7 @@ impl JobState for InMemoryJobState { // job change event emitted // it is emitting current job status self.job_event_sender.send(&JobStateEvent::JobUpdated { - job_id: job_id.to_string(), + job_id: job_id.to_owned().into(), status, }); @@ -458,7 +461,7 @@ impl JobState for InMemoryJobState { Ok(()) } - async fn get_jobs(&self) -> Result> { + async fn get_jobs(&self) -> Result> { Ok(self .completed_jobs .iter() @@ -466,8 +469,8 @@ impl JobState for InMemoryJobState { .collect()) } - async fn get_all_jobs(&self) -> Result> { - let mut all_jobs: HashSet = self + async fn get_all_jobs(&self) -> Result> { + let mut all_jobs: HashSet = self .queued_jobs .iter() .map(|pair| pair.key().clone()) @@ -477,9 +480,14 @@ impl JobState for InMemoryJobState { Ok(all_jobs) } - fn accept_job(&self, job_id: &JobId, job_name: &JobName, queued_at: u64) -> Result<()> { + fn accept_job( + &self, + job_id: &JobId, + job_name: &JobName, + queued_at: u64, + ) -> Result<()> { self.queued_jobs - .insert(job_id.to_string(), (job_name.to_string(), queued_at)); + .insert(job_id.to_owned(), (job_name.to_owned(), queued_at)); Ok(()) } @@ -494,8 +502,8 @@ impl JobState for InMemoryJobState { job_id.clone(), ( JobStatus { - job_id, - job_name, + job_id: job_id.into(), + job_name: job_name.into(), status: Some(Status::Failed(FailedJob { error: reason, queued_at, diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index d917ed44bf..3771fc1c99 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -169,7 +169,7 @@ pub trait ClusterState: Send + Sync + 'static { async fn bind_schedulable_tasks( &self, distribution: TaskDistributionPolicy, - active_jobs: Arc>, + active_jobs: Arc>, executors: Option>, ) -> Result>; @@ -282,7 +282,12 @@ pub trait JobState: Send + Sync { /// Accepts a job into the scheduler's queue. /// /// Called when a job is received but before it is planned. - fn accept_job(&self, job_id: &JobId, job_name: &JobName, queued_at: u64) -> Result<()>; + fn accept_job( + &self, + job_id: &JobId, + job_name: &JobName, + queued_at: u64, + ) -> Result<()>; /// Returns the number of queued jobs waiting to be scheduled. fn pending_job_number(&self) -> usize; @@ -298,10 +303,10 @@ pub trait JobState: Send + Sync { ) -> Result<()>; /// Returns the set of all active job IDs. - async fn get_jobs(&self) -> Result>; + async fn get_jobs(&self) -> Result>; /// Returns the set of all job IDs including running, queued, and completed jobs. - async fn get_all_jobs(&self) -> Result>; + async fn get_all_jobs(&self) -> Result>; /// Returns the status of the specified job. async fn get_job_status(&self, job_id: &JobId) -> Result>; @@ -353,7 +358,7 @@ pub trait JobState: Send + Sync { pub(crate) async fn bind_task_bias( mut slots: Vec<&mut AvailableTaskSlots>, - running_jobs: Arc>, + running_jobs: Arc>, if_skip: fn(Arc) -> bool, ) -> Vec { let mut schedulable_tasks: Vec = vec![]; @@ -412,7 +417,7 @@ pub(crate) async fn bind_task_bias( *task_info = Some(create_task_info(executor_id.clone(), task_id)); let partition = PartitionId { - job_id: job_id.clone(), + job_id: job_id.clone().into(), stage_id: running_stage.stage_id, partition_id, }; @@ -437,7 +442,7 @@ pub(crate) async fn bind_task_bias( pub(crate) async fn bind_task_round_robin( mut slots: Vec<&mut AvailableTaskSlots>, - running_jobs: Arc>, + running_jobs: Arc>, if_skip: fn(Arc) -> bool, ) -> Vec { let mut schedulable_tasks: Vec = vec![]; @@ -498,7 +503,7 @@ pub(crate) async fn bind_task_round_robin( *task_info = Some(create_task_info(executor_id.clone(), task_id)); let partition = PartitionId { - job_id: job_id.clone(), + job_id: job_id.to_owned().into(), stage_id: running_stage.stage_id, partition_id, }; @@ -551,7 +556,7 @@ pub trait DistributionPolicy: std::fmt::Debug + Send + Sync { async fn bind_tasks( &self, mut slots: Vec<&mut AvailableTaskSlots>, - running_jobs: Arc>, + running_jobs: Arc>, ) -> datafusion::error::Result>; /// Name of [DistributionPolicy] @@ -563,6 +568,7 @@ mod test { use std::collections::HashMap; use std::sync::Arc; + use ballista_core::JobId; use ballista_core::error::Result; use ballista_core::serde::protobuf::AvailableTaskSlots; use ballista_core::serde::scheduler::{ @@ -592,7 +598,7 @@ mod test { let mut expected = Vec::new(); { - let mut expected0 = HashMap::new(); + let mut expected0: HashMap> = HashMap::new(); let mut entry_a = HashMap::new(); entry_a.insert("executor_3".to_string(), 2); @@ -600,21 +606,21 @@ mod test { entry_b.insert("executor_3".to_string(), 5); entry_b.insert("executor_2".to_string(), 2); - expected0.insert("job_a".to_string(), entry_a); - expected0.insert("job_b".to_string(), entry_b); + expected0.insert("job_a".to_owned().into(), entry_a); + expected0.insert("job_b".to_owned().into(), entry_b); expected.push(expected0); } { - let mut expected0 = HashMap::new(); + let mut expected0: HashMap> = HashMap::new(); let mut entry_b = HashMap::new(); entry_b.insert("executor_3".to_string(), 7); let mut entry_a = HashMap::new(); entry_a.insert("executor_2".to_string(), 2); - expected0.insert("job_a".to_string(), entry_a); - expected0.insert("job_b".to_string(), entry_b); + expected0.insert("job_a".to_owned().into(), entry_a); + expected0.insert("job_b".to_owned().into(), entry_b); expected.push(expected0); } @@ -641,9 +647,9 @@ mod test { let result = get_result(bound_tasks); - let mut expected = Vec::new(); + let mut expected: Vec>> = Vec::new(); { - let mut expected0 = HashMap::new(); + let mut expected0: HashMap> = HashMap::new(); let mut entry_a = HashMap::new(); entry_a.insert("executor_3".to_string(), 1); @@ -653,13 +659,13 @@ mod test { entry_b.insert("executor_3".to_string(), 2); entry_b.insert("executor_2".to_string(), 2); - expected0.insert("job_a".to_string(), entry_a); - expected0.insert("job_b".to_string(), entry_b); + expected0.insert("job_a".to_owned().into(), entry_a); + expected0.insert("job_b".to_owned().into(), entry_b); expected.push(expected0); } { - let mut expected0 = HashMap::new(); + let mut expected0: HashMap> = HashMap::new(); let mut entry_b = HashMap::new(); entry_b.insert("executor_3".to_string(), 3); @@ -669,8 +675,8 @@ mod test { entry_a.insert("executor_2".to_string(), 1); entry_a.insert("executor_1".to_string(), 1); - expected0.insert("job_a".to_string(), entry_a); - expected0.insert("job_b".to_string(), entry_b); + expected0.insert("job_a".to_owned().into(), entry_a); + expected0.insert("job_b".to_owned().into(), entry_b); expected.push(expected0); } @@ -683,9 +689,7 @@ mod test { Ok(()) } - fn get_result( - bound_tasks: Vec, - ) -> HashMap> { + fn get_result(bound_tasks: Vec) -> HashMap> { let mut result = HashMap::new(); for bound_task in bound_tasks { @@ -701,18 +705,18 @@ mod test { async fn mock_active_jobs( num_partition: usize, - ) -> Result> { - let graph_a = mock_graph("job_a", num_partition, 2).await?; + ) -> Result> { + let graph_a = mock_graph(&JobId::new("job_a"), num_partition, 2).await?; - let graph_b = mock_graph("job_b", num_partition, 7).await?; + let graph_b = mock_graph(&JobId::new("job_b"), num_partition, 7).await?; let mut active_jobs = HashMap::new(); active_jobs.insert( - graph_a.job_id().to_string(), + graph_a.job_id().to_owned(), JobInfoCache::new(Box::new(graph_a)), ); active_jobs.insert( - graph_b.job_id().to_string(), + graph_b.job_id().to_owned(), JobInfoCache::new(Box::new(graph_b)), ); diff --git a/ballista/scheduler/src/cluster/test_util/mod.rs b/ballista/scheduler/src/cluster/test_util/mod.rs index e98a0fec8b..514c513b97 100644 --- a/ballista/scheduler/src/cluster/test_util/mod.rs +++ b/ballista/scheduler/src/cluster/test_util/mod.rs @@ -19,10 +19,10 @@ use crate::cluster::{JobState, JobStateEvent}; use crate::scheduler_server::timestamp_millis; use crate::state::execution_graph::ExecutionGraphBox; use crate::test_utils::{await_condition, mock_completed_task, mock_executor}; -use ballista_core::JobName; use ballista_core::error::Result; use ballista_core::serde::protobuf::JobStatus; use ballista_core::serde::protobuf::job_status::Status; +use ballista_core::{JobId, JobName}; use futures::StreamExt; use std::sync::Arc; use std::time::Duration; @@ -80,7 +80,7 @@ impl JobStateTest { assert!( matches!(&status, JobStatus { job_id: status_job_id, status: Some(Status::Queued(_)), .. - } if status_job_id.as_str() == job_id), + } if status_job_id.as_str() == job_id.as_str()), "Expected queued status but found {:?}", status ); @@ -91,7 +91,7 @@ impl JobStateTest { /// Submits a job with the given execution graph. pub async fn submit_job(self, graph: &ExecutionGraphBox) -> Result { self.state - .submit_job(graph.job_id().to_string(), graph, None) + .submit_job(graph.job_id().to_owned(), graph, None) .await?; Ok(self) } @@ -106,7 +106,7 @@ impl JobStateTest { assert!( matches!(&status, JobStatus { job_id: status_job_id, status: Some(Status::Running(_)), .. - } if status_job_id.as_str() == job_id), + } if status_job_id.as_str() == job_id.as_str()), "Expected running status but found {:?}", status ); @@ -130,7 +130,7 @@ impl JobStateTest { assert!( matches!(&status, JobStatus { job_id: status_job_id, status: Some(Status::Failed(_)), .. - } if status_job_id.as_str() == job_id), + } if status_job_id.as_str() == job_id.as_str()), "Expected failed status but found {:?}", status ); @@ -147,7 +147,7 @@ impl JobStateTest { assert!( matches!(&status, JobStatus { job_id: status_job_id, status: Some(Status::Successful(_)), .. - } if status_job_id.as_str() == job_id), + } if status_job_id.as_str() == job_id.as_str()), "Expected success status but found {:?}", status ); @@ -178,7 +178,7 @@ pub async fn test_job_lifecycle( ) -> Result<()> { let test = JobStateTest::new(state).await?; - let job_id = graph.job_id().to_string(); + let job_id = graph.job_id().to_owned(); let test = test .queue_job(&job_id)? @@ -207,7 +207,7 @@ pub async fn test_job_planning_failure( ) -> Result<()> { let test = JobStateTest::new(state).await?; - let job_id = graph.job_id().to_string(); + let job_id = graph.job_id().to_owned(); test.queue_job(&job_id)? .fail_planning(&job_id) diff --git a/ballista/scheduler/src/display.rs b/ballista/scheduler/src/display.rs index cd929a8a41..4ea27dce2d 100644 --- a/ballista/scheduler/src/display.rs +++ b/ballista/scheduler/src/display.rs @@ -19,6 +19,7 @@ //! [`datafusion::physical_plan::display`] for examples of how to //! format +use ballista_core::JobId; use ballista_core::utils::collect_plan_metrics; use datafusion::logical_expr::{StringifiedPlan, ToStringifiedPlan}; use datafusion::physical_plan::metrics::MetricsSet; diff --git a/ballista/scheduler/src/metrics/mod.rs b/ballista/scheduler/src/metrics/mod.rs index 509e319ce2..18e06c21cb 100644 --- a/ballista/scheduler/src/metrics/mod.rs +++ b/ballista/scheduler/src/metrics/mod.rs @@ -21,7 +21,7 @@ pub mod prometheus; #[cfg(feature = "prometheus")] use crate::metrics::prometheus::PrometheusMetricsCollector; -use ballista_core::error::Result; +use ballista_core::{JobId, error::Result}; use std::sync::Arc; /// Interface for recording metrics events in the scheduler. An instance of `Arc` diff --git a/ballista/scheduler/src/metrics/prometheus.rs b/ballista/scheduler/src/metrics/prometheus.rs index bbddc1fae7..1d799e8851 100644 --- a/ballista/scheduler/src/metrics/prometheus.rs +++ b/ballista/scheduler/src/metrics/prometheus.rs @@ -16,6 +16,7 @@ // under the License. use crate::metrics::SchedulerMetricsCollector; +use ballista_core::JobId; use ballista_core::error::{BallistaError, Result}; use once_cell::sync::OnceCell; diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index fbada7ac2a..76687cdff4 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -20,6 +20,7 @@ use std::collections::HashMap; use std::sync::Arc; +use ballista_core::JobId; use ballista_core::config::BallistaConfig; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::ShuffleWriter; @@ -60,7 +61,7 @@ pub trait DistributedPlanner { /// partitioning changes. fn plan_query_stages<'a>( &'a mut self, - job_id: &'a str, + job_id: &'a JobId, execution_plan: Arc, config: &ConfigOptions, ) -> Result>>; @@ -101,7 +102,7 @@ impl DistributedPlanner for DefaultDistributedPlanner { /// A shuffle writer is created whenever the partitioning changes. fn plan_query_stages<'a>( &'a mut self, - job_id: &'a str, + job_id: &'a JobId, execution_plan: Arc, config: &ConfigOptions, ) -> Result>> { @@ -125,7 +126,7 @@ impl DefaultDistributedPlanner { /// complete query stage (its parent might also belong to the same stage) fn plan_query_stages_internal<'a>( &'a mut self, - job_id: &'a str, + job_id: &'a JobId, execution_plan: Arc, config: &ConfigOptions, ) -> Result { @@ -644,7 +645,7 @@ mod test { let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); let stages = planner.plan_query_stages( - &job_uuid.to_string(), + &job_uuid.to_string().into(), plan, ctx.state().config().options(), )?; @@ -762,7 +763,8 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); - let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?; + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; for (i, stage) in stages.iter().enumerate() { println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false)); } @@ -905,7 +907,8 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); - let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?; + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; for (i, stage) in stages.iter().enumerate() { println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false)); } @@ -953,7 +956,8 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); - let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?; + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; for stage in &stages { let mut walker: Vec> = @@ -999,7 +1003,8 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); - let stages = planner.plan_query_stages(&job_uuid.to_string(), plan, &options)?; + let stages = + planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; for (i, stage) in stages.iter().enumerate() { println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false)); } @@ -1044,7 +1049,7 @@ order by let make_loc = |partition_id: usize| PartitionLocation { map_partition_id: partition_id, partition_id: PartitionId { - job_id: "job".to_string(), + job_id: "job".to_owned().into(), stage_id: 42, partition_id, }, @@ -1148,7 +1153,7 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); let stages = planner.plan_query_stages( - &job_uuid.to_string(), + &job_uuid.to_string().into(), plan, ctx.state().config().options(), )?; @@ -1259,7 +1264,7 @@ order by let mut planner = DefaultDistributedPlanner::new(); let job_uuid = Uuid::new_v4(); let stages = planner.plan_query_stages( - &job_uuid.to_string(), + &job_uuid.to_string().into(), plan, ctx.state().config().options(), )?; diff --git a/ballista/scheduler/src/scheduler_server/event.rs b/ballista/scheduler/src/scheduler_server/event.rs index 2fbd7072c1..e0852e63fa 100644 --- a/ballista/scheduler/src/scheduler_server/event.rs +++ b/ballista/scheduler/src/scheduler_server/event.rs @@ -83,11 +83,11 @@ pub enum QueryStageSchedulerEvent { failed_at: u64, }, /// A job's execution graph has been updated. - JobUpdated(String), + JobUpdated(JobId), /// Request to cancel a job. - JobCancel(String), + JobCancel(JobId), /// Request to clean up job data. - JobDataClean(String), + JobDataClean(JobId), /// Task status updates received. TaskUpdating(String, Vec), /// Signal to revive task offers. @@ -150,6 +150,7 @@ impl Debug for QueryStageSchedulerEvent { QueryStageSchedulerEvent::JobDataClean(job_id) => { write!(f, "JobDataClean : job_id={job_id}.") } + // TODO: This is not job_id but Executor ID (based on usage). QueryStageSchedulerEvent::TaskUpdating(job_id, status) => { write!(f, "TaskUpdating : job_id={job_id}, status:[{status:?}].") } diff --git a/ballista/scheduler/src/scheduler_server/grpc.rs b/ballista/scheduler/src/scheduler_server/grpc.rs index 7d25e3bdb1..c31cd328db 100644 --- a/ballista/scheduler/src/scheduler_server/grpc.rs +++ b/ballista/scheduler/src/scheduler_server/grpc.rs @@ -162,7 +162,9 @@ impl SchedulerGrpc .executor_manager .drain_pending_cleanup_jobs(&executor_id) .into_iter() - .map(|job_id| CleanJobDataParams { job_id }) + .map(|job_id| CleanJobDataParams { + job_id: job_id.into_inner(), + }) .collect(); Ok(Response::new(PollWorkResult { tasks, @@ -522,7 +524,10 @@ impl SchedulerGrpc Ok(Response::new(ExecuteQueryResult { operation_id, result: Some(execute_query_result::Result::Success( - ExecuteQuerySuccessResult { job_id, session_id }, + ExecuteQuerySuccessResult { + job_id: job_id.into(), + session_id, + }, )), })) } else { @@ -534,7 +539,7 @@ impl SchedulerGrpc &self, request: Request, ) -> Result, Status> { - let job_id = request.into_inner().job_id; + let job_id = request.into_inner().job_id.into(); trace!("Received get_job_status request for job {}", job_id); let flight_proxy = self.flight_proxy_config(); @@ -557,7 +562,7 @@ impl SchedulerGrpc &self, request: Request, ) -> Result, Status> { - let job_id = request.into_inner().job_id; + let job_id = request.into_inner().job_id.into(); trace!("Received get_job_metrics request for job {}", job_id); let graph = self @@ -716,7 +721,7 @@ impl SchedulerGrpc &self, request: Request, ) -> Result, Status> { - let job_id = request.into_inner().job_id; + let job_id = request.into_inner().job_id.into(); info!("Received cancellation request for job {}", job_id); self.cancel_job(job_id).await.map_err(|e| { @@ -732,7 +737,7 @@ impl SchedulerGrpc &self, request: Request, ) -> Result, Status> { - let job_id = request.into_inner().job_id; + let job_id = request.into_inner().job_id.into(); info!("Received clean data request for job {}", job_id); self.query_stage_event_loop diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 4825fd4eed..8f41066d5b 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -22,7 +22,7 @@ use ballista_core::error::Result; use ballista_core::event_loop::{EventLoop, EventSender}; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::TaskStatus; -use ballista_core::{JobName, JobStatusSubscriber}; +use ballista_core::{JobId, JobName, JobStatusSubscriber}; use datafusion::execution::context::SessionState; use datafusion::logical_expr::LogicalPlan; @@ -224,13 +224,13 @@ impl SchedulerServer, plan: &LogicalPlan, subscriber: Option, - ) -> Result { + ) -> Result { log::debug!("Received submit request for job {job_name}"); let job_id = self.state.task_manager.generate_job_id(); self.query_stage_event_loop .get_sender()? .post_event(QueryStageSchedulerEvent::JobQueued { - job_id: job_id.to_owned(), + job_id: job_id.to_owned().into(), job_name: job_name.to_owned().into(), session_ctx: ctx, plan: Box::new(plan.clone()), @@ -411,9 +411,9 @@ pub fn timestamp_millis() -> u64 { mod test { use std::sync::Arc; - use ballista_core::JobName; use ballista_core::extension::SessionConfigExt; use ballista_core::serde::protobuf::job_status::Status; + use ballista_core::{JobId, JobName}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; use datafusion::logical_expr::{LogicalPlan, col}; @@ -423,11 +423,10 @@ mod test { use datafusion_proto::protobuf::LogicalPlanNode; use datafusion_proto::protobuf::PhysicalPlanNode; + use crate::config::SchedulerConfig; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; - use crate::config::SchedulerConfig; - use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::{ ExecutionError, FailedTask, JobStatus, MultiTaskDefinition, @@ -474,11 +473,11 @@ mod test { .create_or_update_session("session_id", &config) .await?; - let job_id = "job"; + let job_id: JobId = "job".to_owned().into(); // Enqueue job scheduler.state.task_manager.queue_job( - job_id, + &job_id, &JobName::new(""), timestamp_millis(), )?; @@ -486,7 +485,7 @@ mod test { // Submit job scheduler .state - .submit_job(job_id, &JobName::new(""), ctx, &plan, 0, None) + .submit_job(&job_id, &JobName::new(""), ctx, &plan, 0, None) .await .expect("submitting plan"); @@ -494,7 +493,7 @@ mod test { while let Some(graph) = scheduler .state .task_manager - .get_active_execution_graph(job_id) + .get_active_execution_graph(&job_id) { let task = { let mut graph = graph.write().await; @@ -519,7 +518,7 @@ mod test { // Complete the task let task_status = TaskStatus { task_id: task.task_id as u32, - job_id: task.partition.job_id.clone(), + job_id: task.partition.job_id.clone().into(), stage_id: task.partition.stage_id as u32, stage_attempt_num: task.stage_attempt_num as u32, partition_id: task.partition.partition_id as u32, @@ -545,7 +544,7 @@ mod test { let final_graph = scheduler .state .task_manager - .get_active_execution_graph(job_id) + .get_active_execution_graph(&job_id) .expect("Fail to find graph in the cache"); let final_graph = final_graph.read().await; diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index a6642ee2e5..580d6a969c 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -128,7 +128,7 @@ impl if let Some(subscriber) = subscriber { let timestamp = timestamp_millis(); let job_status = JobStatus { - job_id: job_id.clone(), + job_id: job_id.clone().into(), job_name: job_name.into(), status: Some(ballista_core::serde::protobuf::job_status::Status::Failed( FailedJob { error, queued_at, started_at: timestamp, ended_at: timestamp } diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 5e851a29cd..f751d81b97 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -18,6 +18,7 @@ use crate::planner::create_shuffle_writer_with_config; use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use crate::state::aqe::planner::AdaptiveStageInfo; +use ballista_core::JobId; use ballista_core::execution_plans::ShuffleReaderExec; use datafusion::common::exec_err; use datafusion::config::ConfigOptions; diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 073a898a19..7ee241d344 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -173,7 +173,7 @@ impl AdaptiveExecutionGraph { Ok(Self { planner, scheduler_id: Some(scheduler_id.to_string()), - job_id: job_id.to_string(), + job_id: job_id.to_owned(), job_name: job_name.to_owned(), session_id: session_id.to_string(), @@ -504,8 +504,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { Box::new(self.clone()) } - fn job_id(&self) -> &str { - self.job_id.as_str() + fn job_id(&self) -> &JobId { + &self.job_id } fn job_name(&self) -> &JobName { @@ -1176,7 +1176,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { self.end_time = timestamp_millis(); self.status = JobStatus { - job_id: self.job_id.clone(), + job_id: self.job_id.clone().into(), job_name: self.job_name.clone().into(), status: Some(Status::Failed(FailedJob { error, @@ -1205,7 +1205,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { self.end_time = timestamp_millis(); self.status = JobStatus { - job_id: self.job_id.clone(), + job_id: self.job_id.clone().into(), job_name: self.job_name.clone().into(), status: Some(job_status::Status::Successful(SuccessfulJob { partition_location, diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 01429b3d73..fcaaa26582 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -24,9 +24,9 @@ use crate::state::aqe::optimizer_rule::{ }; use crate::state::distributed_explain::handle_explain_plan; use crate::state::execution_stage::StageOutput; -use ballista_core::JobName; use ballista_core::execution_plans::ShuffleWriter; use ballista_core::serde::scheduler::PartitionLocation; +use ballista_core::{JobId, JobName}; use datafusion::common; use datafusion::common::{HashMap, exec_err}; use datafusion::error::DataFusionError; @@ -165,7 +165,11 @@ impl AdaptivePlanner { ); let plan = state.create_physical_plan(logical_plan).await?; - let plan = handle_explain_plan(job_name.as_str(), ctx, logical_plan, plan) + + // Note: the signature requires a JobId, but we are passing a JobName. The below is a + // dirty fix but this seems like a bug or a design flaw. + let job_id: JobId = job_name.clone().into_inner().into(); + let plan = handle_explain_plan(&job_id, ctx, logical_plan, plan) .await .map_err(|e| DataFusionError::Execution(e.to_string()))?; @@ -360,12 +364,10 @@ impl AdaptivePlanner { // that would arise if the rule walked the entire residual // plan in `default_optimizers()`. let plan = CoalescePartitionsRule.optimize(plan, config)?; - BallistaAdapter::adapt_to_ballista( - plan, - self.job_name.as_str(), - config, - ) - .map(|w| (w.plan.stage_id(), w)) + // adapt_to_ballista takes an job_id, we are passing a job_name. Need to transform to fix compiler. + let job_id = self.job_name.clone().into_inner().into(); + BallistaAdapter::adapt_to_ballista(plan, &job_id, config) + .map(|w| (w.plan.stage_id(), w)) }) .collect::, Vec<_>)>>()?; diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs b/ballista/scheduler/src/state/aqe/test/alter_stages.rs index 1baaf1b949..83ba01a005 100644 --- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs @@ -530,7 +530,7 @@ fn small_statistics_exchange() -> Vec> { // next few properties are generic values map_partition_id: 0, partition_id: PartitionId { - job_id: "".to_string(), + job_id: "".to_owned().into(), stage_id: 0, partition_id: 0, }, @@ -562,7 +562,7 @@ fn big_statistics_exchange() -> Vec> { // next few properties are generic values map_partition_id: 0, partition_id: PartitionId { - job_id: "".to_string(), + job_id: "".to_owned().into(), stage_id: 0, partition_id: 0, }, diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 893daa36cd..8d67616002 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -97,7 +97,7 @@ fn partitions_with_byte_sizes( vec![PartitionLocation { map_partition_id: 0, partition_id: PartitionId { - job_id: "".to_string(), + job_id: "".to_owned().into(), stage_id: 0, partition_id: idx, }, diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index 939347359f..ed68d2912d 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -43,7 +43,7 @@ pub(crate) fn mock_partitions_with_statistics() -> Vec> { // next few properties are generic values map_partition_id: 0, partition_id: PartitionId { - job_id: "".to_string(), + job_id: "".to_owned().into(), stage_id: 0, partition_id: 0, }, @@ -68,7 +68,7 @@ pub(crate) fn mock_partitions_with_statistics_no_data() -> Vec; /// publish its outputs to the `ExecutionGraph`s `output_locations` representing the final query results. pub trait ExecutionGraph: Debug { /// Returns the job ID for this execution graph. - fn job_id(&self) -> &str; + fn job_id(&self) -> &JobId; /// Returns the job name for this execution graph. fn job_name(&self) -> &JobName; @@ -324,7 +324,7 @@ impl StaticExecutionGraph { Ok(Self { scheduler_id: Some(scheduler_id.to_string()), - job_id: job_id.to_string(), + job_id: job_id.to_owned().into(), job_name: job_name.to_owned().into(), session_id: session_id.to_string(), @@ -636,8 +636,8 @@ impl ExecutionGraph for StaticExecutionGraph { Box::new(self.clone()) } - fn job_id(&self) -> &str { - self.job_id.as_str() + fn job_id(&self) -> &JobId { + &self.job_id } fn job_name(&self) -> &JobName { @@ -1360,7 +1360,7 @@ impl ExecutionGraph for StaticExecutionGraph { self.end_time = timestamp_millis(); self.status = JobStatus { - job_id: self.job_id.clone(), + job_id: self.job_id.clone().into(), job_name: self.job_name.clone().into(), status: Some(Status::Failed(FailedJob { error, @@ -1389,7 +1389,7 @@ impl ExecutionGraph for StaticExecutionGraph { self.end_time = timestamp_millis(); self.status = JobStatus { - job_id: self.job_id.clone(), + job_id: self.job_id.clone().into(), job_name: self.job_name.clone().into(), status: Some(job_status::Status::Successful(SuccessfulJob { partition_location, diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index 0b39164453..81754cfb1e 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -413,9 +413,9 @@ mod tests { use crate::planner::DefaultDistributedPlanner; use crate::state::execution_graph::StaticExecutionGraph; use crate::state::execution_graph_dot::ExecutionGraphDot; - use ballista_core::JobName; use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; + use ballista_core::{JobId, JobName}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::MemTable; use datafusion::prelude::{SessionConfig, SessionContext}; @@ -607,7 +607,7 @@ filter_expr="] let mut planner = DefaultDistributedPlanner::new(); StaticExecutionGraph::new( "scheduler_id", - "job_id", + &JobId::new("job_id"), &JobName::new("job_name"), "session_id", plan, @@ -644,7 +644,7 @@ filter_expr="] let mut planner = DefaultDistributedPlanner::new(); StaticExecutionGraph::new( "scheduler_id", - "job_id", + &JobId::new("job_id"), &JobName::new("job_name"), "session_id", plan, diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index d87f0f5bc9..1871514a7b 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -17,6 +17,7 @@ use std::time::Duration; +use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::error::Result; use ballista_core::serde::protobuf; @@ -66,7 +67,7 @@ pub struct ExecutorManager { /// Cached gRPC clients for communicating with executors. clients: ExecutorClients, /// Jobs pending cleanup on each executor. - pending_cleanup_jobs: Arc>>, + pending_cleanup_jobs: Arc>>, /// Configuration for gRPC client connections. grpc_client_config: GrpcClientConfig, } @@ -106,7 +107,7 @@ impl ExecutorManager { /// Returns a list of bound tasks that can be launched on executors. pub async fn bind_schedulable_tasks( &self, - running_jobs: Arc>, + running_jobs: Arc>, ) -> Result> { if running_jobs.is_empty() { debug!("There's no active jobs for binding tasks"); @@ -142,7 +143,7 @@ impl ExecutorManager { let infos = tasks_to_cancel.entry(task_info.executor_id).or_default(); infos.push(protobuf::RunningTaskInfo { task_id: task_info.task_id as u32, - job_id: task_info.job_id, + job_id: task_info.job_id.into(), stage_id: task_info.stage_id as u32, partition_id: task_info.partition_id as u32, }); @@ -208,7 +209,7 @@ impl ExecutorManager { let alive_executors = self.get_alive_executors(); for executor in alive_executors { - let job_id_clone = job_id.to_owned(); + let job_id_clone = job_id.to_owned().into_inner(); if self.config.is_push_staged_scheduling() { if let Ok(mut client) = @@ -386,10 +387,7 @@ impl ExecutorManager { Ok(()) } - pub(crate) fn drain_pending_cleanup_jobs( - &self, - executor_id: &str, - ) -> HashSet { + pub(crate) fn drain_pending_cleanup_jobs(&self, executor_id: &str) -> HashSet { self.pending_cleanup_jobs .remove(executor_id) .map(|(_, jobs)| jobs) diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index 159c902726..dad3cb3d7f 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -26,7 +26,7 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::EventSender; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::TaskStatus; -use ballista_core::{JobName, JobStatusSubscriber}; +use ballista_core::{JobId, JobName, JobStatusSubscriber}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -264,7 +264,7 @@ impl SchedulerState>, + HashMap<(JobId, usize), Vec>, > = HashMap::new(); for (executor_id, task) in bound_tasks.into_iter() { let stage_key = (task.partition.job_id.clone(), task.partition.stage_id); @@ -276,7 +276,7 @@ impl SchedulerState, > = HashMap::new(); executor_stage_tasks.insert(stage_key, vec![task]); diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 3b7b62ee51..c2ac1b6143 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -52,7 +52,7 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::RwLock; -type ActiveJobCache = Arc>; +type ActiveJobCache = Arc>; /// Trait for launching tasks on executors. /// @@ -333,7 +333,7 @@ impl TaskManager info!("Submitting execution graph for job_id: {job_id}:\n\n{string_plan}"); self.state - .submit_job(job_id.to_string(), &graph, subscriber) + .submit_job(job_id.to_owned(), &graph, subscriber) .await?; graph.revive(); self.active_job_cache @@ -343,7 +343,7 @@ impl TaskManager } /// Returns a snapshot of currently running jobs from the cache. - pub fn get_running_job_cache(&self) -> Arc> { + pub fn get_running_job_cache(&self) -> Arc> { let ret = self .active_job_cache .iter() @@ -404,7 +404,7 @@ impl TaskManager _ => (0, 0), }; jobs.push(JobOverview { - job_id: job_status.job_id.clone(), + job_id: job_status.job_id.clone().into(), job_name: job_status.job_name.clone().into(), status: job_status, start_time, @@ -483,7 +483,7 @@ impl TaskManager // let graph = self.get_active_execution_graph(&job_id).await; let job_events = if let Some(cached) = - self.get_active_execution_graph(&job_id) + self.get_active_execution_graph(&job_id.clone().into()) { let mut graph = cached.write().await; graph.update_task_status( @@ -663,7 +663,7 @@ impl TaskManager let task_definition = TaskDefinition { task_id: task.task_id as u32, task_attempt_num: task.task_attempt as u32, - job_id, + job_id: job_id.into(), stage_id: stage_id as u32, stage_attempt_num: task.stage_attempt_num as u32, partition_id: task.partition.partition_id as u32, @@ -754,7 +754,7 @@ impl TaskManager .collect(); multi_tasks.push(MultiTaskDefinition { task_ids, - job_id, + job_id: job_id.into(), stage_id: stage_id as u32, stage_attempt_num: stage_attempt_num as u32, plan, @@ -798,13 +798,14 @@ impl TaskManager } /// Generates a new random 7-character alphanumeric job ID. - pub fn generate_job_id(&self) -> String { + pub fn generate_job_id(&self) -> JobId { let mut rng = rng(); std::iter::repeat(()) .map(|()| rng.sample(Alphanumeric)) .map(char::from) .take(7) - .collect() + .collect::() + .into() } /// Clean up a failed job in FailedJobs Keyspace by delayed clean_up_interval seconds @@ -849,7 +850,7 @@ impl From<&ExecutionGraphBox> for JobOverview { let completed_stages = value.completed_stages(); Self { - job_id: value.job_id().to_string(), + job_id: value.job_id().to_owned(), job_name: value.job_name().to_owned(), status: value.status().clone(), start_time: value.start_time(), diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index fedad3de0b..f62d3f1431 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -17,7 +17,7 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; -use ballista_core::{JobName, JobStatusSubscriber}; +use ballista_core::{JobId, JobName, JobStatusSubscriber}; use datafusion::catalog::Session; use std::any::Any; use std::collections::HashMap; @@ -503,7 +503,7 @@ impl SchedulerTest { &mut self, job_name: &JobName, plan: &LogicalPlan, - ) -> Result { + ) -> Result { println!("{:?}", self.session_config); let ctx = self .scheduler @@ -632,7 +632,7 @@ impl SchedulerTest { &mut self, job_name: &JobName, plan: &LogicalPlan, - ) -> Result<(JobStatus, String)> { + ) -> Result<(JobStatus, JobId)> { self.run_with_subscriber(job_name, plan, None).await } /// Returns job status and job_id, with provided subscriber @@ -641,7 +641,7 @@ impl SchedulerTest { job_name: &JobName, plan: &LogicalPlan, subscriber: Option, - ) -> Result<(JobStatus, String)> { + ) -> Result<(JobStatus, JobId)> { let ctx = self .scheduler .state @@ -698,23 +698,23 @@ impl SchedulerTest { #[derive(Clone)] pub enum MetricEvent { /// Job submitted event (job_id, queued_at, submitted_at). - Submitted(String, u64, u64), + Submitted(JobId, u64, u64), /// Job completed event (job_id, queued_at, completed_at). - Completed(String, u64, u64), + Completed(JobId, u64, u64), /// Job cancelled event (job_id). - Cancelled(String), + Cancelled(JobId), /// Job failed event (job_id, queued_at, failed_at). - Failed(String, u64, u64), + Failed(JobId, u64, u64), } impl MetricEvent { /// Returns the job ID associated with this event. - pub fn job_id(&self) -> &str { + pub fn job_id(&self) -> &JobId { match self { - MetricEvent::Submitted(job, _, _) => job.as_str(), - MetricEvent::Completed(job, _, _) => job.as_str(), - MetricEvent::Cancelled(job) => job.as_str(), - MetricEvent::Failed(job, _, _) => job.as_str(), + MetricEvent::Submitted(job, _, _) => job, + MetricEvent::Completed(job, _, _) => job, + MetricEvent::Cancelled(job) => job, + MetricEvent::Failed(job, _, _) => job, } } } @@ -877,7 +877,7 @@ pub fn revive_graph_and_complete_next_stage_with_executor( /// Creates a test execution graph with a simple aggregation plan. pub async fn test_aggregation_plan(partition: usize) -> StaticExecutionGraph { - test_aggregation_plan_with_job_id(partition, "job").await + test_aggregation_plan_with_job_id(partition, &JobId::new("job")).await } /// Creates a test execution graph with a simple aggregation plan and custom job ID. @@ -965,7 +965,7 @@ pub async fn test_two_aggregations_plan(partition: usize) -> StaticExecutionGrap StaticExecutionGraph::new( "localhost:50050", - "job", + &JobId::new("job"), &JobName::new(""), "session", plan, @@ -1006,7 +1006,7 @@ pub async fn test_coalesce_plan(partition: usize) -> StaticExecutionGraph { StaticExecutionGraph::new( "localhost:50050", - "job", + &JobId::new("job"), &JobName::new(""), "session", plan, @@ -1067,7 +1067,7 @@ pub async fn test_join_plan(partition: usize) -> StaticExecutionGraph { let mut planner = DefaultDistributedPlanner::new(); let graph = StaticExecutionGraph::new( "localhost:50050", - "job", + &JobId::new("job"), &JobName::new(""), "session", plan, @@ -1110,7 +1110,7 @@ pub async fn test_union_all_plan(partition: usize) -> StaticExecutionGraph { let mut planner = DefaultDistributedPlanner::new(); let graph = StaticExecutionGraph::new( "localhost:50050", - "job", + &JobId::new("job"), &JobName::new(""), "session", plan, @@ -1153,7 +1153,7 @@ pub async fn test_union_plan(partition: usize) -> StaticExecutionGraph { let mut planner = DefaultDistributedPlanner::new(); let graph = StaticExecutionGraph::new( "localhost:50050", - "job", + &JobId::new("job"), &JobName::new(""), "session", plan, @@ -1201,7 +1201,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat // Complete the task protobuf::TaskStatus { task_id: task.task_id as u32, - job_id: task.partition.job_id.clone(), + job_id: task.partition.job_id.clone().into(), stage_id: task.partition.stage_id as u32, stage_attempt_num: task.stage_attempt_num as u32, partition_id: task.partition.partition_id as u32, @@ -1236,7 +1236,7 @@ pub fn mock_failed_task(task: TaskDescription, failed_task: FailedTask) -> TaskS // Fail the task protobuf::TaskStatus { task_id: task.task_id as u32, - job_id: task.partition.job_id.clone(), + job_id: task.partition.job_id.clone().into(), stage_id: task.partition.stage_id as u32, stage_attempt_num: task.stage_attempt_num as u32, partition_id: task.partition.partition_id as u32, diff --git a/benchmarks/benches/sort_shuffle.rs b/benchmarks/benches/sort_shuffle.rs index ddfcf29a0e..90fe054104 100644 --- a/benchmarks/benches/sort_shuffle.rs +++ b/benchmarks/benches/sort_shuffle.rs @@ -239,7 +239,7 @@ fn run_sort_shuffle( let config = SortShuffleConfig::new(true, CompressionType::LZ4_FRAME, 8192); let writer = SortShuffleWriterExec::try_new( - "bench_job".to_string(), + "bench_job".to_owned().into(), 1, input, work_dir.to_string(), diff --git a/benchmarks/src/bin/shuffle_bench.rs b/benchmarks/src/bin/shuffle_bench.rs index d4f248b461..934f80a236 100644 --- a/benchmarks/src/bin/shuffle_bench.rs +++ b/benchmarks/src/bin/shuffle_bench.rs @@ -274,7 +274,7 @@ async fn execute_shuffle_write( let metrics: MetricsSet = match writer_kind { WriterKind::Hash => { let exec = ShuffleWriterExec::try_new( - format!("bench_job_{task_id}"), + format!("bench_job_{task_id}").into(), 1, input, work_dir_str, @@ -289,7 +289,7 @@ async fn execute_shuffle_write( let cfg = SortShuffleConfig::new(true, CompressionType::LZ4_FRAME, args.batch_size); let exec = SortShuffleWriterExec::try_new( - format!("bench_job_{task_id}"), + format!("bench_job_{task_id}").into(), 1, input, work_dir_str, From 67936763fc039ce18914d10a81df5a1550cc8a77 Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Wed, 10 Jun 2026 08:51:20 +0200 Subject: [PATCH 09/11] Fixup in feature-gated code --- ballista/core/src/execution_plans/shuffle_writer.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 684eb176cf..fe4d3fe2cf 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -622,7 +622,7 @@ mod tests { let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); let work_dir = TempDir::new()?; let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), + JobId::new("jobOne"), 1, input_plan, work_dir.path().to_str().unwrap().to_owned(), @@ -680,7 +680,7 @@ mod tests { let input_plan = create_input_plan()?; let work_dir = TempDir::new()?; let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), + JobId::new("jobOne"), 1, input_plan, work_dir.path().to_str().unwrap().to_owned(), From 310020290f9e2f67884e9c1675e12f0bddfb1fbe Mon Sep 17 00:00:00 2001 From: JarroVGIT Date: Wed, 10 Jun 2026 08:56:25 +0200 Subject: [PATCH 10/11] Clippy --- ballista/core/src/serde/scheduler/mod.rs | 2 +- ballista/scheduler/src/api/handlers.rs | 2 +- ballista/scheduler/src/cluster/memory.rs | 2 +- ballista/scheduler/src/cluster/mod.rs | 4 ++-- ballista/scheduler/src/scheduler_server/mod.rs | 4 ++-- ballista/scheduler/src/state/execution_graph.rs | 4 ++-- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/ballista/core/src/serde/scheduler/mod.rs b/ballista/core/src/serde/scheduler/mod.rs index 7d490f53f6..cb357e0037 100644 --- a/ballista/core/src/serde/scheduler/mod.rs +++ b/ballista/core/src/serde/scheduler/mod.rs @@ -73,7 +73,7 @@ impl PartitionId { /// Creates a new partition ID with the given job, stage, and partition identifiers. pub fn new(job_id: &JobId, stage_id: usize, partition_id: usize) -> Self { Self { - job_id: job_id.to_owned().into(), + job_id: job_id.to_owned(), stage_id, partition_id, } diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 1dc07514d6..8f77cdb2f2 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -410,7 +410,7 @@ pub async fn get_job< Ok(Json(JobResponse { job_id: job.job_id().to_owned(), - job_name: job.job_name().to_owned().into(), + job_name: job.job_name().to_owned(), job_status, status: plain_status, start_time: job.start_time(), diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index 5298c98971..3cbdab53a7 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -420,7 +420,7 @@ impl JobState for InMemoryJobState { // job change event emitted // it is emitting current job status self.job_event_sender.send(&JobStateEvent::JobUpdated { - job_id: job_id.to_owned().into(), + job_id: job_id.to_owned(), status, }); diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs index 3771fc1c99..61e939a767 100644 --- a/ballista/scheduler/src/cluster/mod.rs +++ b/ballista/scheduler/src/cluster/mod.rs @@ -417,7 +417,7 @@ pub(crate) async fn bind_task_bias( *task_info = Some(create_task_info(executor_id.clone(), task_id)); let partition = PartitionId { - job_id: job_id.clone().into(), + job_id: job_id.clone(), stage_id: running_stage.stage_id, partition_id, }; @@ -503,7 +503,7 @@ pub(crate) async fn bind_task_round_robin( *task_info = Some(create_task_info(executor_id.clone(), task_id)); let partition = PartitionId { - job_id: job_id.to_owned().into(), + job_id: job_id.to_owned(), stage_id: running_stage.stage_id, partition_id, }; diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 8f41066d5b..7a77b7ad5d 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -230,8 +230,8 @@ impl SchedulerServer Date: Wed, 10 Jun 2026 13:37:54 +0300 Subject: [PATCH 11/11] Add AI config files --- .cursor/rules.md | 5 +++++ .gemini/rules.md | 5 +++++ AGENTS.md | 5 +++++ CLAUDE.md | 5 +++++ 4 files changed, 20 insertions(+) create mode 100644 .cursor/rules.md create mode 100644 .gemini/rules.md create mode 100644 AGENTS.md create mode 100644 CLAUDE.md diff --git a/.cursor/rules.md b/.cursor/rules.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/.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! + diff --git a/.gemini/rules.md b/.gemini/rules.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/.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! + diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/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! + diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/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! +