Add SessionBuilder fn type to SchedulerServer to allow customized SessionContext creation

Author: Wang

ballista/rust/client/src/prelude.rs

@@ -19,8 +19,8 @@
 
 pub use crate::context::BallistaContext;
 pub use ballista_core::config::BallistaConfig;
-pub use ballista_core::config::BALLISTA_DEFAULT_SHUFFLE_PARTITIONS;
 pub use ballista_core::config::BALLISTA_DEFAULT_BATCH_SIZE;
+pub use ballista_core::config::BALLISTA_DEFAULT_SHUFFLE_PARTITIONS;
 pub use ballista_core::error::{BallistaError, Result};
 
 pub use futures::StreamExt;

ballista/rust/scheduler/src/lib.rs

@@ -105,12 +105,14 @@ use ballista_core::serde::protobuf::execute_query_params::OptionalSessionId;
 use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient;
 use ballista_core::serde::scheduler::to_proto::hash_partitioning_to_proto;
 use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec};
+use datafusion::execution::context::{default_session_builder, SessionState};
 use datafusion::prelude::{SessionConfig, SessionContext};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 use tokio::sync::{mpsc, RwLock};
 use tonic::transport::Channel;
 
 type ExecutorsClient = Arc<RwLock<HashMap<String, ExecutorGrpcClient<Channel>>>>;
+type SessionBuilder = fn(SessionConfig) -> SessionState;
 
 #[derive(Clone)]
 pub struct SchedulerServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
@@ -120,6 +122,8 @@ pub struct SchedulerServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
     scheduler_loop: Option<SchedulerLoop>,
     executors_client: Option<ExecutorsClient>,
     codec: BallistaCodec<T, U>,
+    /// SessionState Builder
+    session_builder: SessionBuilder,
     /// DataFusion session contexts that are registered within the SchedulerServer
     session_context_registry: Arc<SessionContextRegistry>,
 }
@@ -134,13 +138,28 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         config: Arc<dyn ConfigBackendClient>,
         namespace: String,
         codec: BallistaCodec<T, U>,
+    ) -> Self {
+        SchedulerServer::new_with_builder(
+            config,
+            namespace,
+            codec,
+            default_session_builder,
+        )
+    }
+
+    pub fn new_with_builder(
+        config: Arc<dyn ConfigBackendClient>,
+        namespace: String,
+        codec: BallistaCodec<T, U>,
+        session_builder: SessionBuilder,
     ) -> Self {
         SchedulerServer::new_with_policy(
             config,
             namespace,
             TaskSchedulingPolicy::PullStaged,
             None,
             codec,
+            session_builder,
         )
     }
 
@@ -150,6 +169,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         policy: TaskSchedulingPolicy,
         scheduler_loop: Option<SchedulerLoop>,
         codec: BallistaCodec<T, U>,
+        session_builder: SessionBuilder,
     ) -> Self {
         let state = Arc::new(SchedulerState::new(config, namespace, codec.clone()));
 
@@ -168,6 +188,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             scheduler_loop,
             executors_client,
             codec,
+            session_builder,
             session_context_registry: Arc::new(SessionContextRegistry::default()),
         }
     }
@@ -771,7 +792,8 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                     update_datafusion_session_context(session_ctx, &config)
                 }
                 _ => {
-                    let df_session = create_datafusion_session_context(&config);
+                    let df_session =
+                        create_datafusion_session_context(&config, self.session_builder);
                     let session_id = df_session.session_id.clone();
                     self.session_context_registry
                         .register_session(session_id, df_session.clone())
@@ -972,7 +994,8 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                 error!("{}", msg);
                 tonic::Status::internal(msg)
             })?;
-            let df_session = create_datafusion_session_context(&config);
+            let df_session =
+                create_datafusion_session_context(&config, self.session_builder);
             let session_id = df_session.session_id.clone();
             self.session_context_registry
                 .register_session(session_id.clone(), df_session.clone())
@@ -1000,15 +1023,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
 }
 
 /// Create a new DataFusion session context from Ballista Configuration
-pub fn create_datafusion_session_context(config: &BallistaConfig) -> Arc<SessionContext> {
+pub fn create_datafusion_session_context(
+    config: &BallistaConfig,
+    session_builder: SessionBuilder,
+) -> Arc<SessionContext> {
     let config = SessionConfig::new()
         .with_target_partitions(config.default_shuffle_partitions())
         .with_batch_size(config.default_batch_size())
         .with_repartition_joins(config.repartition_joins())
         .with_repartition_aggregations(config.repartition_aggregations())
         .with_repartition_windows(config.repartition_windows())
         .with_parquet_pruning(config.parquet_pruning());
-    Arc::new(SessionContext::with_config(config))
+    let session_state = session_builder(config);
+    Arc::new(SessionContext::with_state(session_state))
 }
 
 /// Update the existing DataFusion session context with Ballista Configuration

ballista/rust/scheduler/src/main.rs

@@ -63,6 +63,7 @@ mod config {
 }
 
 use config::prelude::*;
+use datafusion::execution::context::default_session_builder;
 
 async fn start_server(
     config_backend: Arc<dyn ConfigBackendClient>,
@@ -90,6 +91,7 @@ async fn start_server(
                     policy,
                     Some(SchedulerLoop { tx_job }),
                     BallistaCodec::default(),
+                    default_session_builder,
                 );
                 let task_scheduler =
                     TaskScheduler::new(Arc::new(scheduler_server.clone()));

benchmarks/src/bin/tpch.rs

@@ -30,7 +30,9 @@ use std::{
 };
 
 use ballista::context::BallistaContext;
-use ballista::prelude::{BallistaConfig, BALLISTA_DEFAULT_SHUFFLE_PARTITIONS, BALLISTA_DEFAULT_BATCH_SIZE};
+use ballista::prelude::{
+    BallistaConfig, BALLISTA_DEFAULT_BATCH_SIZE, BALLISTA_DEFAULT_SHUFFLE_PARTITIONS,
+};
 
 use datafusion::datasource::{MemTable, TableProvider};
 use datafusion::error::{DataFusionError, Result};

datafusion/src/execution/context.rs

@@ -1151,8 +1151,16 @@ pub struct SessionState {
     pub runtime: Arc<RuntimeEnv>,
 }
 
+/// Default session builder using the provided configuration
+pub fn default_session_builder(config: SessionConfig) -> SessionState {
+    SessionState::with_config(
+        config,
+        Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()),
+    )
+}
+
 impl SessionState {
-    /// Returns new SessionState using the provided configuration
+    /// Returns new SessionState using the provided configuration and runtime
     pub fn with_config(config: SessionConfig, runtime: Arc<RuntimeEnv>) -> Self {
         let catalog_list = Arc::new(MemoryCatalogList::new()) as Arc<dyn CatalogList>;