cleanup

Author: adriangb

src/flight_service/do_get.rs

@@ -9,6 +9,7 @@ use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
 use datafusion::execution::{SendableRecordBatchStream, SessionState};
+use datafusion::prelude::SessionContext;
 use futures::TryStreamExt;
 use http::HeaderMap;
 use prost::Message;
@@ -115,10 +116,11 @@ impl ArrowFlightEndpoint {
                     })
                     .await
                     .map_err(|err| datafusion_error_to_tonic_status(&err))?;
+                let ctx = SessionContext::new_with_state(session_state.clone());
 
                 let codec = DistributedCodec::new_combined_with_user(session_state.config());
 
-                let stage = stage_from_proto(stage_proto, &session_state, &self.runtime, &codec)
+                let stage = stage_from_proto(stage_proto, &ctx, &self.runtime, &codec)
                     .map_err(|err| {
                         Status::invalid_argument(format!("Cannot decode stage proto: {err}"))
                     })?;

src/protobuf/stage_proto.rs

@@ -2,7 +2,7 @@ use crate::execution_plans::{ExecutionTask, StageExec};
 use datafusion::{
     common::internal_datafusion_err,
     error::{DataFusionError, Result},
-    execution::{SessionState, context::SessionContext, runtime_env::RuntimeEnv},
+    execution::{runtime_env::RuntimeEnv, context::SessionContext},
     physical_plan::ExecutionPlan,
 };
 use datafusion_proto::{
@@ -100,22 +100,21 @@ pub fn proto_from_stage(
 
 pub fn stage_from_proto(
     msg: StageExecProto,
-    session_state: &SessionState,
+    ctx: &SessionContext,
     runtime: &RuntimeEnv,
     codec: &dyn PhysicalExtensionCodec,
 ) -> Result<StageExec> {
     let plan_node = msg.plan.ok_or(internal_datafusion_err!(
         "ExecutionStageMsg is missing the plan"
     ))?;
 
-    let session_ctx = SessionContext::new_with_state(session_state.clone());
-    let plan = plan_node.try_into_physical_plan(&session_ctx, runtime, codec)?;
+    let plan = plan_node.try_into_physical_plan(&ctx, runtime, codec)?;
 
     let inputs = msg
         .inputs
         .into_iter()
         .map(|s| {
-            stage_from_proto(s, session_state, runtime, codec)
+            stage_from_proto(s, ctx, runtime, codec)
                 .map(|s| Arc::new(s) as Arc<dyn ExecutionPlan>)
         })
         .collect::<Result<Vec<_>>>()?;
@@ -153,9 +152,9 @@ pub fn stage_from_proto(
 mod tests {
     use std::sync::Arc;
 
-    use crate::StageExec;
     use crate::protobuf::stage_proto::StageExecProto;
     use crate::protobuf::{proto_from_stage, stage_from_proto};
+    use crate::StageExec;
     use datafusion::{
         arrow::{
             array::{RecordBatch, StringArray, UInt8Array},
@@ -230,7 +229,7 @@ mod tests {
         // Convert back to ExecutionStage
         let round_trip_stage = stage_from_proto(
             decoded_msg,
-            &ctx.state(),
+            &ctx,
             ctx.runtime_env().as_ref(),
             &DefaultPhysicalExtensionCodec {},
         )?;