Refactor SessionContext, SessionState and SessionConfig to support multi-tenancy configurations - Part 2 (#2029)

* Refactor SessionContext, SessionState add SessionConfig to support multi-tenancy configurations - Part 2

* fix UT

* Resolve review comments

* Fix minor merge issues

Co-authored-by: Andrew Lamb <[email protected]>

Author: mingmwang

ballista/rust/core/src/serde/logical_plan/mod.rs

@@ -228,6 +228,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                 };
 
                 let object_store = ctx
+                    .runtime_env()
                     .object_store(scan.path.as_str())
                     .map_err(|e| {
                         BallistaError::NotImplemented(format!(
@@ -1287,9 +1288,10 @@ mod roundtrip_tests {
         let codec: BallistaCodec<protobuf::LogicalPlanNode, protobuf::PhysicalPlanNode> =
             BallistaCodec::default();
         let custom_object_store = Arc::new(TestObjectStore {});
-        ctx.register_object_store("test", custom_object_store.clone());
+        ctx.runtime_env()
+            .register_object_store("test", custom_object_store.clone());
 
-        let (os, _) = ctx.object_store("test://foo.csv")?;
+        let (os, _) = ctx.runtime_env().object_store("test://foo.csv")?;
 
         println!("Object Store {:?}", os);
 

ballista/rust/core/src/serde/mod.rs

@@ -350,6 +350,7 @@ mod tests {
     use datafusion::datasource::object_store::local::LocalFileSystem;
     use datafusion::error::DataFusionError;
     use datafusion::execution::context::{QueryPlanner, SessionState, TaskContext};
+    use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
     use datafusion::logical_plan::plan::Extension;
     use datafusion::logical_plan::{
         col, DFSchemaRef, Expr, LogicalPlan, LogicalPlanBuilder, UserDefinedLogicalNode,
@@ -699,10 +700,11 @@ mod tests {
     #[tokio::test]
     async fn test_extension_plan() -> crate::error::Result<()> {
         let store = Arc::new(LocalFileSystem {});
-        let config =
-            SessionConfig::new().with_query_planner(Arc::new(TopKQueryPlanner {}));
+        let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap());
+        let session_state = SessionState::with_config(SessionConfig::new(), runtime)
+            .with_query_planner(Arc::new(TopKQueryPlanner {}));
 
-        let ctx = SessionContext::with_config(config);
+        let ctx = SessionContext::with_state(session_state);
 
         let scan = LogicalPlanBuilder::scan_csv(
             store,

ballista/rust/core/src/serde/physical_plan/from_proto.rs

@@ -29,7 +29,7 @@ use chrono::{TimeZone, Utc};
 use datafusion::datasource::object_store::local::LocalFileSystem;
 use datafusion::datasource::object_store::{FileMeta, SizedFile};
 use datafusion::datasource::PartitionedFile;
-use datafusion::execution::context::SessionState;
+use datafusion::execution::context::ExecutionProps;
 
 use datafusion::physical_plan::file_format::FileScanConfig;
 
@@ -153,12 +153,12 @@ impl TryFrom<&protobuf::PhysicalExprNode> for Arc<dyn PhysicalExpr> {
                     .map(|x| x.try_into())
                     .collect::<Result<Vec<_>, _>>()?;
 
-                // TODO Do not create new the SessionState
-                let session_state = SessionState::new();
+                // TODO Do not create new the ExecutionProps
+                let execution_props = ExecutionProps::new();
 
                 let fun_expr = functions::create_physical_fun(
                     &(&scalar_function).into(),
-                    &session_state.execution_props,
+                    &execution_props,
                 )?;
 
                 Arc::new(ScalarFunctionExpr::new(

ballista/rust/core/src/serde/physical_plan/mod.rs

@@ -909,7 +909,7 @@ fn decode_scan_config(
         .collect::<Result<Vec<_>, _>>()?;
 
     let object_store = if let Some(file) = file_groups.get(0).and_then(|h| h.get(0)) {
-        ctx.object_store(file.file_meta.path())?.0
+        ctx.runtime_env().object_store(file.file_meta.path())?.0
     } else {
         Arc::new(LocalFileSystem {})
     };

ballista/rust/core/src/utils.rs

@@ -47,6 +47,7 @@ use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::common::batch_byte_size;
 use datafusion::physical_plan::empty::EmptyExec;
 
+use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
 use datafusion::physical_plan::file_format::{CsvExec, ParquetExec};
 use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::hash_aggregate::HashAggregateExec;
@@ -234,11 +235,16 @@ pub fn create_df_ctx_with_ballista_query_planner<T: 'static + AsLogicalPlan>(
     let scheduler_url = format!("http://{}:{}", scheduler_host, scheduler_port);
     let planner: Arc<BallistaQueryPlanner<T>> =
         Arc::new(BallistaQueryPlanner::new(scheduler_url, config.clone()));
-    let config = SessionConfig::new()
-        .with_query_planner(planner)
+
+    let session_config = SessionConfig::new()
         .with_target_partitions(config.default_shuffle_partitions())
         .with_information_schema(true);
-    SessionContext::with_config(config)
+    let session_state = SessionState::with_config(
+        session_config,
+        Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()),
+    )
+    .with_query_planner(planner);
+    SessionContext::with_state(session_state)
 }
 
 pub struct BallistaQueryPlanner<T: AsLogicalPlan> {

ballista/rust/scheduler/src/scheduler_server/mod.rs

@@ -179,13 +179,23 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     }
 }
 
-/// Create a DataFusion context that is compatible with Ballista
+/// Create a DataFusion session context that is compatible with Ballista Configuration
 pub fn create_datafusion_context(config: &BallistaConfig) -> SessionContext {
     let config =
         SessionConfig::new().with_target_partitions(config.default_shuffle_partitions());
     SessionContext::with_config(config)
 }
 
+/// Update the existing DataFusion session context with Ballista Configuration
+pub fn update_datafusion_context(
+    session_ctx: Arc<SessionContext>,
+    config: &BallistaConfig,
+) -> Arc<SessionContext> {
+    session_ctx.state.write().config.target_partitions =
+        config.default_shuffle_partitions();
+    session_ctx
+}
+
 #[cfg(all(test, feature = "sled"))]
 mod test {
     use std::sync::Arc;

datafusion-proto/src/from_proto.rs

@@ -1168,7 +1168,7 @@ pub fn parse_expr(
         ExprType::ScalarUdfExpr(protobuf::ScalarUdfExprNode { fun_name, args }) => {
             let scalar_fn = ctx
                 .state
-                .lock()
+                .read()
                 .get_function_meta(fun_name.as_str()).ok_or_else(|| Error::General(format!("invalid aggregate function message, function {} is not registered in the ExecutionContext", fun_name)))?;
 
             Ok(Expr::ScalarUDF {
@@ -1182,7 +1182,7 @@ pub fn parse_expr(
         ExprType::AggregateUdfExpr(protobuf::AggregateUdfExprNode { fun_name, args }) => {
             let agg_fn = ctx
                 .state
-                .lock()
+                .read()
                 .get_aggregate_meta(fun_name.as_str()).ok_or_else(|| Error::General(format!("invalid aggregate function message, function {} is not registered in the ExecutionContext", fun_name)))?;
 
             Ok(Expr::AggregateUDF {

datafusion/benches/sort_limit_query_sql.rs

@@ -84,7 +84,7 @@ fn create_context() -> Arc<Mutex<SessionContext>> {
     rt.block_on(async {
         // create local session context
         let mut ctx = SessionContext::new();
-        ctx.state.lock().config.target_partitions = 1;
+        ctx.state.write().config.target_partitions = 1;
 
         let task_ctx = ctx.task_ctx();
         let mem_table = MemTable::load(Arc::new(csv.await), Some(partitions), task_ctx)

datafusion/src/dataframe.rs

@@ -34,13 +34,13 @@ use crate::arrow::datatypes::SchemaRef;
 use crate::arrow::util::pretty;
 use crate::datasource::TableProvider;
 use crate::datasource::TableType;
-use crate::execution::context::{SessionContext, SessionState, TaskContext};
+use crate::execution::context::{SessionState, TaskContext};
 use crate::physical_plan::file_format::{plan_to_csv, plan_to_parquet};
 use crate::physical_plan::{collect, collect_partitioned};
 use crate::physical_plan::{execute_stream, execute_stream_partitioned, ExecutionPlan};
 use crate::scalar::ScalarValue;
 use crate::sql::utils::find_window_exprs;
-use parking_lot::Mutex;
+use parking_lot::RwLock;
 use std::any::Any;
 
 /// DataFrame represents a logical set of rows with the same named columns.
@@ -69,13 +69,13 @@ use std::any::Any;
 /// # }
 /// ```
 pub struct DataFrame {
-    session_state: Arc<Mutex<SessionState>>,
+    session_state: Arc<RwLock<SessionState>>,
     plan: LogicalPlan,
 }
 
 impl DataFrame {
     /// Create a new Table based on an existing logical plan
-    pub fn new(session_state: Arc<Mutex<SessionState>>, plan: &LogicalPlan) -> Self {
+    pub fn new(session_state: Arc<RwLock<SessionState>>, plan: &LogicalPlan) -> Self {
         Self {
             session_state,
             plan: plan.clone(),
@@ -84,10 +84,9 @@ impl DataFrame {
 
     /// Create a physical plan
     pub async fn create_physical_plan(&self) -> Result<Arc<dyn ExecutionPlan>> {
-        let state = self.session_state.lock().clone();
-        let ctx = SessionContext::from(Arc::new(Mutex::new(state)));
-        let plan = ctx.optimize(&self.plan)?;
-        ctx.create_physical_plan(&plan).await
+        let state = self.session_state.read().clone();
+        let optimized_plan = state.optimize(&self.plan)?;
+        state.create_physical_plan(&optimized_plan).await
     }
 
     /// Filter the DataFrame by column. Returns a new DataFrame only containing the
@@ -351,8 +350,8 @@ impl DataFrame {
     /// ```
     pub async fn collect(&self) -> Result<Vec<RecordBatch>> {
         let plan = self.create_physical_plan().await?;
-        let task_ctx = Arc::new(TaskContext::from(&self.session_state.lock().clone()));
-        collect(plan, task_ctx).await
+        let task_ctx = Arc::new(TaskContext::from(&self.session_state.read().clone()));
+        Ok(collect(plan, task_ctx).await?)
     }
 
     /// Print results.
@@ -406,7 +405,7 @@ impl DataFrame {
     /// ```
     pub async fn execute_stream(&self) -> Result<SendableRecordBatchStream> {
         let plan = self.create_physical_plan().await?;
-        let task_ctx = Arc::new(TaskContext::from(&self.session_state.lock().clone()));
+        let task_ctx = Arc::new(TaskContext::from(&self.session_state.read().clone()));
         execute_stream(plan, task_ctx).await
     }
 
@@ -426,8 +425,8 @@ impl DataFrame {
     /// ```
     pub async fn collect_partitioned(&self) -> Result<Vec<Vec<RecordBatch>>> {
         let plan = self.create_physical_plan().await?;
-        let task_ctx = Arc::new(TaskContext::from(&self.session_state.lock().clone()));
-        collect_partitioned(plan, task_ctx).await
+        let task_ctx = Arc::new(TaskContext::from(&self.session_state.read().clone()));
+        Ok(collect_partitioned(plan, task_ctx).await?)
     }
 
     /// Executes this DataFrame and returns one stream per partition.
@@ -447,8 +446,8 @@ impl DataFrame {
         &self,
     ) -> Result<Vec<SendableRecordBatchStream>> {
         let plan = self.create_physical_plan().await?;
-        let task_ctx = Arc::new(TaskContext::from(&self.session_state.lock().clone()));
-        execute_stream_partitioned(plan, task_ctx).await
+        let task_ctx = Arc::new(TaskContext::from(&self.session_state.read().clone()));
+        Ok(execute_stream_partitioned(plan, task_ctx).await?)
     }
 
     /// Returns the schema describing the output of this DataFrame in terms of columns returned,
@@ -511,7 +510,7 @@ impl DataFrame {
     /// # }
     /// ```
     pub fn registry(&self) -> Arc<dyn FunctionRegistry> {
-        let registry = self.session_state.lock().clone();
+        let registry = self.session_state.read().clone();
         Arc::new(registry)
     }
 
@@ -563,9 +562,8 @@ impl DataFrame {
     /// Write a `DataFrame` to a CSV file.
     pub async fn write_csv(&self, path: &str) -> Result<()> {
         let plan = self.create_physical_plan().await?;
-        let state = self.session_state.lock().clone();
-        let ctx = SessionContext::from(Arc::new(Mutex::new(state)));
-        plan_to_csv(&ctx, plan, path).await
+        let state = self.session_state.read().clone();
+        plan_to_csv(&state, plan, path).await
     }
 
     /// Write a `DataFrame` to a Parquet file.
@@ -575,9 +573,8 @@ impl DataFrame {
         writer_properties: Option<WriterProperties>,
     ) -> Result<()> {
         let plan = self.create_physical_plan().await?;
-        let state = self.session_state.lock().clone();
-        let ctx = SessionContext::from(Arc::new(Mutex::new(state)));
-        plan_to_parquet(&ctx, plan, path, writer_properties).await
+        let state = self.session_state.read().clone();
+        plan_to_parquet(&state, plan, path, writer_properties).await
     }
 }
 

datafusion/src/datasource/object_store/mod.rs

@@ -163,7 +163,7 @@ pub trait ObjectStore: Sync + Send + Debug {
 
 static LOCAL_SCHEME: &str = "file";
 
-/// A Registry holds all the object stores at runtime with a scheme for each store.
+/// A Registry holds all the object stores at Runtime with a scheme for each store.
 /// This allows the user to extend DataFusion with different storage systems such as S3 or HDFS
 /// and query data inside these systems.
 pub struct ObjectStoreRegistry {