Move stage along with the other execution plan implementations

Author: gabotechs

src/execution_plans/arrow_flight_read.rs

@@ -1,9 +1,9 @@
 use crate::channel_manager_ext::get_distributed_channel_resolver;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error};
+use crate::execution_plans::StageExec;
 use crate::flight_service::{DoGet, StageKey};
 use crate::protobuf::{proto_from_stage, DistributedCodec};
-use crate::stage::ExecutionStage;
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
@@ -161,7 +161,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
         let stage = context
             .session_config()
-            .get_extension::<ExecutionStage>()
+            .get_extension::<StageExec>()
             .ok_or(internal_datafusion_err!(
                 "ArrowFlightReadExec requires an ExecutionStage in the session config"
             ))?;
@@ -218,7 +218,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
             );
 
             async move {
-                let url = task.url()?.ok_or(internal_datafusion_err!(
+                let url = task.url.ok_or(internal_datafusion_err!(
                     "ArrowFlightReadExec: task is unassigned, cannot proceed"
                 ))?;
 

src/execution_plans/mod.rs

@@ -1,5 +1,7 @@
 mod arrow_flight_read;
 mod partition_isolator;
+mod stage;
 
 pub use arrow_flight_read::ArrowFlightReadExec;
 pub use partition_isolator::{PartitionGroup, PartitionIsolatorExec};
+pub use stage::{display_stage_graphviz, ExecutionTask, StageExec};