resolve review comments

Author: mingmwang

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

@@ -708,7 +708,7 @@ mod tests {
         let store = Arc::new(LocalFileSystem {});
         let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap());
         let session_state =
-            SessionState::with_config(SessionConfig::new(), runtime.clone())
+            SessionState::with_config_rt(SessionConfig::new(), runtime.clone())
                 .with_query_planner(Arc::new(TopKQueryPlanner {}));
 
         let ctx = SessionContext::with_state(session_state);

ballista/rust/core/src/utils.rs

@@ -238,7 +238,7 @@ pub fn create_df_ctx_with_ballista_query_planner<T: 'static + AsLogicalPlan>(
     let session_config = SessionConfig::new()
         .with_target_partitions(config.default_shuffle_partitions())
         .with_information_schema(true);
-    let mut session_state = SessionState::with_config(
+    let mut session_state = SessionState::with_config_rt(
         session_config,
         Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()),
     )

datafusion/src/execution/context.rs

@@ -164,7 +164,7 @@ impl SessionContext {
 
     /// Creates a new session context using the provided configuration and RuntimeEnv.
     pub fn with_config_rt(config: SessionConfig, runtime: Arc<RuntimeEnv>) -> Self {
-        let state = SessionState::with_config(config, runtime);
+        let state = SessionState::with_config_rt(config, runtime);
         Self {
             session_id: state.session_id.clone(),
             session_start_time: chrono::Utc::now(),
@@ -1010,15 +1010,15 @@ pub struct SessionState {
 
 /// Default session builder using the provided configuration
 pub fn default_session_builder(config: SessionConfig) -> SessionState {
-    SessionState::with_config(
+    SessionState::with_config_rt(
         config,
         Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()),
     )
 }
 
 impl SessionState {
     /// Returns new SessionState using the provided configuration and runtime
-    pub fn with_config(config: SessionConfig, runtime: Arc<RuntimeEnv>) -> Self {
+    pub fn with_config_rt(config: SessionConfig, runtime: Arc<RuntimeEnv>) -> Self {
         let session_id = Uuid::new_v4().to_string();
 
         let catalog_list = Arc::new(MemoryCatalogList::new()) as Arc<dyn CatalogList>;
@@ -1287,17 +1287,17 @@ pub enum TaskProperties {
 /// Task Execution Context
 pub struct TaskContext {
     /// Session Id
-    pub session_id: String,
+    session_id: String,
     /// Optional Task Identify
-    pub task_id: Option<String>,
+    task_id: Option<String>,
     /// Task properties
-    pub properties: TaskProperties,
+    properties: TaskProperties,
     /// Scalar functions associated with this task context
-    pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+    scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     /// Aggregate functions associated with this task context
-    pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
     /// Runtime environment associated with this task context
-    pub runtime: Arc<RuntimeEnv>,
+    runtime: Arc<RuntimeEnv>,
 }
 
 impl TaskContext {
@@ -1351,6 +1351,21 @@ impl TaskContext {
             TaskProperties::SessionConfig(session_config) => session_config.clone(),
         }
     }
+
+    /// Return the session_id of this [TaskContext]
+    pub fn session_id(&self) -> String {
+        self.session_id.clone()
+    }
+
+    /// Return the task_id of this [TaskContext]
+    pub fn task_id(&self) -> Option<String> {
+        self.task_id.clone()
+    }
+
+    /// Return the [RuntimeEnv] associated with this [TaskContext]
+    pub fn runtime_env(&self) -> Arc<RuntimeEnv> {
+        self.runtime.clone()
+    }
 }
 
 /// Create a new task context instance from SessionContext
@@ -2961,7 +2976,7 @@ mod tests {
     #[tokio::test]
     async fn custom_query_planner() -> Result<()> {
         let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap());
-        let session_state = SessionState::with_config(SessionConfig::new(), runtime)
+        let session_state = SessionState::with_config_rt(SessionConfig::new(), runtime)
             .with_query_planner(Arc::new(MyQueryPlanner {}));
         let ctx = SessionContext::with_state(session_state);
 

datafusion/src/physical_plan/planner.rs

@@ -1466,7 +1466,7 @@ mod tests {
 
     fn make_session_state() -> SessionState {
         let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap());
-        SessionState::with_config(SessionConfig::new(), runtime)
+        SessionState::with_config_rt(SessionConfig::new(), runtime)
     }
 
     async fn plan(logical_plan: &LogicalPlan) -> Result<Arc<dyn ExecutionPlan>> {

datafusion/src/physical_plan/sorts/sort.rs

@@ -582,9 +582,9 @@ async fn do_sort(
         expr,
         metrics_set,
         Arc::new(context.session_config()),
-        context.runtime.clone(),
+        context.runtime_env(),
     );
-    context.runtime.register_requester(sorter.id());
+    context.runtime_env().register_requester(sorter.id());
     while let Some(batch) = input.next().await {
         let batch = batch?;
         sorter.insert_batch(batch).await?;

datafusion/tests/user_defined_plan.rs

@@ -245,7 +245,7 @@ async fn topk_plan() -> Result<()> {
 fn make_topk_context() -> SessionContext {
     let config = SessionConfig::new().with_target_partitions(48);
     let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap());
-    let state = SessionState::with_config(config, runtime)
+    let state = SessionState::with_config_rt(config, runtime)
         .with_query_planner(Arc::new(TopKQueryPlanner {}))
         .add_optimizer_rule(Arc::new(TopKOptimizerRule {}));
     SessionContext::with_state(state)