File name and folder restructure

src/distributed_ext.rs

@@ -1,8 +1,8 @@
-use crate::channel_manager_ext::set_distributed_channel_resolver;
+use crate::channel_resolver_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
-use crate::user_codec_ext::set_distributed_user_codec;
+use crate::protobuf::set_distributed_user_codec;
 use crate::ChannelResolver;
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;

src/physical_optimizer.rs → src/distributed_physical_optimizer_rule.rs

@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
-use super::stage::ExecutionStage;
-use crate::{plan::PartitionIsolatorExec, ArrowFlightReadExec};
+use super::{ArrowFlightReadExec, PartitionIsolatorExec, StageExec};
 use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::joins::PartitionMode;
@@ -55,7 +54,7 @@ impl PhysicalOptimizerRule for DistributedPhysicalOptimizerRule {
         _config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         // We can only optimize plans that are not already distributed
-        if plan.as_any().is::<ExecutionStage>() {
+        if plan.as_any().is::<StageExec>() {
             return Ok(plan);
         }
 
@@ -106,7 +105,7 @@ impl DistributedPhysicalOptimizerRule {
     pub fn distribute_plan(
         &self,
         plan: Arc<dyn ExecutionPlan>,
-    ) -> Result<ExecutionStage, DataFusionError> {
+    ) -> Result<StageExec, DataFusionError> {
         let query_id = Uuid::new_v4();
         self._distribute_plan_inner(query_id, plan, &mut 1, 0)
     }
@@ -117,7 +116,7 @@ impl DistributedPhysicalOptimizerRule {
         plan: Arc<dyn ExecutionPlan>,
         num: &mut usize,
         depth: usize,
-    ) -> Result<ExecutionStage, DataFusionError> {
+    ) -> Result<StageExec, DataFusionError> {
         let mut inputs = vec![];
 
         let distributed = plan.clone().transform_down(|plan| {
@@ -134,7 +133,7 @@ impl DistributedPhysicalOptimizerRule {
         })?;
 
         let inputs = inputs.into_iter().map(Arc::new).collect();
-        let mut stage = ExecutionStage::new(query_id, *num, distributed.data, inputs);
+        let mut stage = StageExec::new(query_id, *num, distributed.data, inputs);
         *num += 1;
 
         stage = match (self.partitions_per_task, can_be_divided(&plan)?) {
@@ -188,7 +187,7 @@ pub fn can_be_divided(plan: &Arc<dyn ExecutionPlan>) -> Result<bool> {
 #[cfg(test)]
 mod tests {
     use crate::assert_snapshot;
-    use crate::physical_optimizer::DistributedPhysicalOptimizerRule;
+    use crate::distributed_physical_optimizer_rule::DistributedPhysicalOptimizerRule;
     use crate::test_utils::parquet::register_parquet_tables;
     use datafusion::error::DataFusionError;
     use datafusion::execution::SessionStateBuilder;

src/plan/arrow_flight_read.rs → src/execution_plans/arrow_flight_read.rs

@@ -1,9 +1,9 @@
-use crate::channel_manager_ext::get_distributed_channel_resolver;
+use crate::channel_resolver_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::plan::DistributedCodec;
-use crate::stage::{proto_from_stage, ExecutionStage};
+use crate::protobuf::{proto_from_stage, DistributedCodec};
 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

@@ -0,0 +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};