Refactor ExecutionContext and related conf to support multi-tenancy configurations

Do not use static singleton RuntimeEnv in both DataFusion and Ballista

Add SessionBuilder fn type to SchedulerServer to allow customized SessionContext creation

fix client BallistaContext

Author: Wang

ballista-examples/src/bin/ballista-dataframe.rs

@@ -25,7 +25,7 @@ async fn main() -> Result<()> {
     let config = BallistaConfig::builder()
         .set("ballista.shuffle.partitions", "4")
         .build()?;
-    let ctx = BallistaContext::remote("localhost", 50050, &config);
+    let ctx = BallistaContext::remote("localhost", 50050, &config).await?;
 
     let testdata = datafusion::arrow::util::test_util::parquet_test_data();
 

ballista-examples/src/bin/ballista-sql.rs

@@ -25,7 +25,7 @@ async fn main() -> Result<()> {
     let config = BallistaConfig::builder()
         .set("ballista.shuffle.partitions", "4")
         .build()?;
-    let ctx = BallistaContext::remote("localhost", 50050, &config);
+    let ctx = BallistaContext::remote("localhost", 50050, &config).await?;
 
     let testdata = datafusion::arrow::util::test_util::arrow_test_data();
 

ballista/rust/client/README.md

@@ -106,7 +106,7 @@ async fn main() -> Result<()> {
        .build()?;
 
    // connect to Ballista scheduler
-   let ctx = BallistaContext::remote("localhost", 50050, &config);
+   let ctx = BallistaContext::remote("localhost", 50050, &config).await?;
 
    // register csv file with the execution context
    ctx.register_csv(

ballista/rust/client/src/context.rs

@@ -17,6 +17,7 @@
 
 //! Distributed execution context.
 
+use log::info;
 use parking_lot::Mutex;
 use sqlparser::ast::Statement;
 use std::collections::HashMap;
@@ -25,7 +26,8 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use ballista_core::config::BallistaConfig;
-use ballista_core::serde::protobuf::LogicalPlanNode;
+use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::serde::protobuf::{ExecuteQueryParams, KeyValuePair, LogicalPlanNode};
 use ballista_core::utils::create_df_ctx_with_ballista_query_planner;
 
 use datafusion::catalog::TableReference;
@@ -35,7 +37,7 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::dataframe_impl::DataFrameImpl;
 use datafusion::logical_plan::{CreateExternalTable, LogicalPlan, TableScan};
 use datafusion::prelude::{
-    AvroReadOptions, CsvReadOptions, ExecutionConfig, ExecutionContext,
+    AvroReadOptions, CsvReadOptions, SessionConfig, SessionContext,
 };
 use datafusion::sql::parser::{DFParser, FileType, Statement as DFStatement};
 
@@ -64,26 +66,81 @@ impl BallistaContextState {
         }
     }
 
+    pub fn config(&self) -> &BallistaConfig {
+        &self.config
+    }
+}
+
+pub struct BallistaContext {
+    state: Arc<Mutex<BallistaContextState>>,
+    context: Arc<SessionContext>,
+}
+
+impl BallistaContext {
+    /// Create a context for executing queries against a remote Ballista scheduler instance
+    pub async fn remote(host: &str, port: u16, config: &BallistaConfig) -> Result<Self> {
+        let state = BallistaContextState::new(host.to_owned(), port, config);
+        let scheduler_url =
+            format!("http://{}:{}", &state.scheduler_host, state.scheduler_port);
+        info!(
+            "Connecting to Ballista scheduler at {}",
+            scheduler_url.clone()
+        );
+        let mut scheduler = SchedulerGrpcClient::connect(scheduler_url.clone())
+            .await
+            .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?;
+
+        let remote_session_id = scheduler
+            .execute_query(ExecuteQueryParams {
+                query: None,
+                settings: config
+                    .settings()
+                    .iter()
+                    .map(|(k, v)| KeyValuePair {
+                        key: k.to_owned(),
+                        value: v.to_owned(),
+                    })
+                    .collect::<Vec<_>>(),
+                optional_session_id: None,
+            })
+            .await
+            .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?
+            .into_inner()
+            .session_id;
+
+        info!(
+            "Server side SessionContext created with Session id: {}",
+            remote_session_id
+        );
+
+        let ctx = {
+            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
+                scheduler_url,
+                remote_session_id,
+                state.config(),
+            )
+        };
+
+        Ok(Self {
+            state: Arc::new(Mutex::new(state)),
+            context: Arc::new(ctx),
+        })
+    }
+
     #[cfg(feature = "standalone")]
-    pub async fn new_standalone(
+    pub async fn standalone(
         config: &BallistaConfig,
         concurrent_tasks: usize,
     ) -> ballista_core::error::Result<Self> {
-        use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
         use ballista_core::serde::protobuf::PhysicalPlanNode;
         use ballista_core::serde::BallistaCodec;
 
         log::info!("Running in local mode. Scheduler will be run in-proc");
 
         let addr = ballista_scheduler::standalone::new_standalone_scheduler().await?;
-
-        let scheduler = loop {
-            match SchedulerGrpcClient::connect(format!(
-                "http://localhost:{}",
-                addr.port()
-            ))
-            .await
-            {
+        let scheduler_url = format!("http://localhost:{}", addr.port());
+        let mut scheduler = loop {
+            match SchedulerGrpcClient::connect(scheduler_url.clone()).await {
                 Err(_) => {
                     tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                     log::info!("Attempting to connect to in-proc scheduler...");
@@ -92,6 +149,37 @@ impl BallistaContextState {
             }
         };
 
+        let remote_session_id = scheduler
+            .execute_query(ExecuteQueryParams {
+                query: None,
+                settings: config
+                    .settings()
+                    .iter()
+                    .map(|(k, v)| KeyValuePair {
+                        key: k.to_owned(),
+                        value: v.to_owned(),
+                    })
+                    .collect::<Vec<_>>(),
+                optional_session_id: None,
+            })
+            .await
+            .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?
+            .into_inner()
+            .session_id;
+
+        info!(
+            "Server side SessionContext created with Session id: {}",
+            remote_session_id
+        );
+
+        let ctx = {
+            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
+                scheduler_url,
+                remote_session_id,
+                config,
+            )
+        };
+
         let default_codec: BallistaCodec<LogicalPlanNode, PhysicalPlanNode> =
             BallistaCodec::default();
 
@@ -102,43 +190,12 @@ impl BallistaContextState {
         )
         .await?;
 
-        Ok(Self {
-            config: config.clone(),
-            scheduler_host: "localhost".to_string(),
-            scheduler_port: addr.port(),
-            tables: HashMap::new(),
-        })
-    }
-
-    pub fn config(&self) -> &BallistaConfig {
-        &self.config
-    }
-}
-
-pub struct BallistaContext {
-    state: Arc<Mutex<BallistaContextState>>,
-}
-
-impl BallistaContext {
-    /// Create a context for executing queries against a remote Ballista scheduler instance
-    pub fn remote(host: &str, port: u16, config: &BallistaConfig) -> Self {
-        let state = BallistaContextState::new(host.to_owned(), port, config);
-
-        Self {
-            state: Arc::new(Mutex::new(state)),
-        }
-    }
-
-    #[cfg(feature = "standalone")]
-    pub async fn standalone(
-        config: &BallistaConfig,
-        concurrent_tasks: usize,
-    ) -> ballista_core::error::Result<Self> {
         let state =
-            BallistaContextState::new_standalone(config, concurrent_tasks).await?;
+            BallistaContextState::new("localhost".to_string(), addr.port(), config);
 
         Ok(Self {
             state: Arc::new(Mutex::new(state)),
+            context: Arc::new(ctx),
         })
     }
 
@@ -154,15 +211,10 @@ impl BallistaContext {
         let path = fs::canonicalize(&path)?;
 
         // use local DataFusion context for now but later this might call the scheduler
-        let mut ctx = {
-            let guard = self.state.lock();
-            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
-                &guard.scheduler_host,
-                guard.scheduler_port,
-                guard.config(),
-            )
-        };
-        let df = ctx.read_avro(path.to_str().unwrap(), options).await?;
+        let df = self
+            .context
+            .read_avro(path.to_str().unwrap(), options)
+            .await?;
         Ok(df)
     }
 
@@ -174,15 +226,7 @@ impl BallistaContext {
         let path = fs::canonicalize(&path)?;
 
         // use local DataFusion context for now but later this might call the scheduler
-        let mut ctx = {
-            let guard = self.state.lock();
-            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
-                &guard.scheduler_host,
-                guard.scheduler_port,
-                guard.config(),
-            )
-        };
-        let df = ctx.read_parquet(path.to_str().unwrap()).await?;
+        let df = self.context.read_parquet(path.to_str().unwrap()).await?;
         Ok(df)
     }
 
@@ -198,15 +242,10 @@ impl BallistaContext {
         let path = fs::canonicalize(&path)?;
 
         // use local DataFusion context for now but later this might call the scheduler
-        let mut ctx = {
-            let guard = self.state.lock();
-            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
-                &guard.scheduler_host,
-                guard.scheduler_port,
-                guard.config(),
-            )
-        };
-        let df = ctx.read_csv(path.to_str().unwrap(), options).await?;
+        let df = self
+            .context
+            .read_csv(path.to_str().unwrap(), options)
+            .await?;
         Ok(df)
     }
 
@@ -292,34 +331,30 @@ impl BallistaContext {
     /// This method is `async` because queries of type `CREATE EXTERNAL TABLE`
     /// might require the schema to be inferred.
     pub async fn sql(&self, sql: &str) -> Result<Arc<dyn DataFrame>> {
-        let mut ctx = {
-            let state = self.state.lock();
-            create_df_ctx_with_ballista_query_planner::<LogicalPlanNode>(
-                &state.scheduler_host,
-                state.scheduler_port,
-                state.config(),
-            )
-        };
-
+        let mut ctx = self.context.clone();
         let is_show = self.is_show_statement(sql).await?;
         // the show tables、 show columns sql can not run at scheduler because the tables is store at client
         if is_show {
             let state = self.state.lock();
-            ctx = ExecutionContext::with_config(
-                ExecutionConfig::new().with_information_schema(
+            ctx = Arc::new(SessionContext::with_config(
+                SessionConfig::new().with_information_schema(
                     state.config.default_with_information_schema(),
                 ),
-            );
+            ));
         }
 
         // register tables with DataFusion context
         {
             let state = self.state.lock();
             for (name, prov) in &state.tables {
-                ctx.register_table(
-                    TableReference::Bare { table: name },
-                    Arc::clone(prov),
-                )?;
+                // ctx is shared between queries, check table exists or not before register
+                let table_ref = TableReference::Bare { table: name };
+                if !ctx.table_exist(table_ref)? {
+                    ctx.register_table(
+                        TableReference::Bare { table: name },
+                        Arc::clone(prov),
+                    )?;
+                }
             }
         }
 
@@ -342,16 +377,16 @@ impl BallistaContext {
                             .has_header(*has_header),
                     )
                     .await?;
-                    Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
+                    Ok(Arc::new(DataFrameImpl::new(ctx.state.clone(), &plan)))
                 }
                 FileType::Parquet => {
                     self.register_parquet(name, location).await?;
-                    Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
+                    Ok(Arc::new(DataFrameImpl::new(ctx.state.clone(), &plan)))
                 }
                 FileType::Avro => {
                     self.register_avro(name, location, AvroReadOptions::default())
                         .await?;
-                    Ok(Arc::new(DataFrameImpl::new(ctx.state, &plan)))
+                    Ok(Arc::new(DataFrameImpl::new(ctx.state.clone(), &plan)))
                 }
                 _ => Err(DataFusionError::NotImplemented(format!(
                     "Unsupported file type {:?}.",
@@ -476,17 +511,13 @@ mod tests {
         use datafusion::arrow::datatypes::Schema;
         use datafusion::arrow::util::pretty;
         use datafusion::datasource::file_format::csv::CsvFormat;
-        use datafusion::datasource::file_format::parquet::ParquetFormat;
         use datafusion::datasource::listing::{
             ListingOptions, ListingTable, ListingTableConfig,
         };
 
         use ballista_core::config::{
             BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA,
         };
-        use std::fs::File;
-        use std::io::Write;
-        use tempfile::TempDir;
         let config = BallistaConfigBuilder::default()
             .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true")
             .build()

ballista/rust/client/src/prelude.rs

@@ -19,6 +19,7 @@
 
 pub use crate::context::BallistaContext;
 pub use ballista_core::config::BallistaConfig;
+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};