Remove BallistaConfig

systemxlabs · Aug 22, 2024 · effbe19 · effbe19
1 parent 69b61b2
commit effbe19
Show file tree

Hide file tree

Showing 7 changed files with 29 additions and 350 deletions.
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs
@@ -19,7 +19,6 @@
 pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 pub mod client;
-pub mod config;
 pub mod error;
 pub mod event_loop;
 pub mod execution_plans;

diff --git a/ballista/scheduler/src/cluster/kv.rs b/ballista/scheduler/src/cluster/kv.rs
@@ -26,7 +26,6 @@ use crate::state::session_manager::create_datafusion_context;
 use crate::state::task_manager::JobInfoCache;
 use crate::state::{decode_into, decode_protobuf};
 use async_trait::async_trait;
-use ballista_core::config::BallistaConfig;
 use ballista_core::error::{BallistaError, Result};
 use ballista_core::serde::protobuf::job_status::Status;
 use ballista_core::serde::protobuf::{
@@ -67,7 +66,7 @@ pub struct KeyValueState<
     scheduler: String,
     /// In-memory store of queued jobs. Map from Job ID -> queued_at timestamp
     queued_jobs: DashMap<String, u64>,
-    //// `SessionBuilder` for constructing `SessionContext` from stored `BallistaConfig`
+    /// `SessionBuilder` for constructing `SessionContext` from stored `BallistaConfig`
     session_builder: SessionBuilder,
 }
 
@@ -534,19 +533,21 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
 
         let settings: protobuf::SessionSettings = decode_protobuf(&value)?;
 
-        let mut config_builder = BallistaConfig::builder();
-        for kv_pair in &settings.configs {
-            config_builder = config_builder.set(&kv_pair.key, &kv_pair.value);
+        let mut config = HashMap::new();
+        for kv_pair in settings.configs {
+            config.insert(kv_pair.key, kv_pair.value);
         }
-        let config = config_builder.build()?;
 
         Ok(create_datafusion_context(&config, self.session_builder))
     }
 
-    async fn create_session(&self, config: &BallistaConfig) -> Result<Arc<SessionContext>> {
+    async fn create_session(
+        &self,
+        config: &HashMap<String, String>,
+    ) -> Result<Arc<SessionContext>> {
         let mut settings: Vec<KeyValuePair> = vec![];
 
-        for (key, value) in config.settings() {
+        for (key, value) in config {
             settings.push(KeyValuePair {
                 key: key.clone(),
                 value: value.clone(),

diff --git a/ballista/scheduler/src/cluster/mod.rs b/ballista/scheduler/src/cluster/mod.rs
@@ -25,7 +25,6 @@ use datafusion_proto::physical_plan::AsExecutionPlan;
 use futures::Stream;
 use log::{debug, info, warn};
 
-use ballista_core::config::BallistaConfig;
 use ballista_core::error::{BallistaError, Result};
 use ballista_core::serde::protobuf::{
     job_status, AvailableTaskSlots, ExecutorHeartbeat, JobStatus,
@@ -230,7 +229,8 @@ pub trait JobState: Send + Sync {
     async fn get_session(&self, session_id: &str) -> Result<Arc<SessionContext>>;
 
     /// Create a new saved session
-    async fn create_session(&self, config: &BallistaConfig) -> Result<Arc<SessionContext>>;
+    async fn create_session(&self, config: &HashMap<String, String>)
+        -> Result<Arc<SessionContext>>;
 }
 
 pub(crate) async fn bind_task_bias(