Scope common execution_plan helpers to src/execution_plans/common.rs

Author: gabotechs

src/common/execution_plan_ops.rs

@@ -1,9 +1,4 @@
-mod execution_plan_ops;
 mod map_last_stream;
-mod partitioning;
-#[allow(unused)]
 pub mod ttl_map;
 
-pub(crate) use execution_plan_ops::*;
 pub(crate) use map_last_stream::map_last_stream;
-pub(crate) use partitioning::{scale_partitioning, scale_partitioning_props};

src/common/mod.rs

@@ -1,7 +1,18 @@
+use datafusion::common::{DataFusionError, plan_err};
 use datafusion::physical_expr::Partitioning;
-use datafusion::physical_plan::PlanProperties;
+use datafusion::physical_plan::{ExecutionPlan, PlanProperties};
+use std::sync::Arc;
 
-pub fn scale_partitioning_props(
+pub(super) fn require_one_child(
+    children: &[Arc<dyn ExecutionPlan>],
+) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+    if children.len() != 1 {
+        return plan_err!("Expected exactly 1 children, got {}", children.len());
+    }
+    Ok(children[0].clone())
+}
+
+pub(super) fn scale_partitioning_props(
     props: &PlanProperties,
     f: impl FnOnce(usize) -> usize,
 ) -> PlanProperties {
@@ -13,7 +24,7 @@ pub fn scale_partitioning_props(
     )
 }
 
-pub fn scale_partitioning(
+pub(super) fn scale_partitioning(
     partitioning: &Partitioning,
     f: impl FnOnce(usize) -> usize,
 ) -> Partitioning {

src/common/partitioning.rs src/execution_plans/common.rssrc/common/partitioning.rs renamed to src/execution_plans/common.rs

@@ -1,6 +1,6 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
-use crate::common::one_child;
 use crate::distributed_physical_optimizer_rule::NetworkBoundaryExt;
+use crate::execution_plans::common::require_one_child;
 use crate::protobuf::DistributedCodec;
 use crate::{ExecutionTask, Stage};
 use datafusion::common::exec_err;
@@ -98,7 +98,7 @@ impl ExecutionPlan for DistributedExec {
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         Ok(Arc::new(DistributedExec {
-            plan: one_child(&children)?,
+            plan: require_one_child(&children)?,
         }))
     }
 

src/execution_plans/distributed.rs

@@ -1,3 +1,4 @@
+mod common;
 mod distributed;
 mod metrics;
 mod network_coalesce;

src/execution_plans/mod.rs

@@ -1,7 +1,7 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
-use crate::common::{one_child, scale_partitioning_props};
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_physical_optimizer_rule::{NetworkBoundary, limit_tasks_err};
+use crate::execution_plans::common::{require_one_child, scale_partitioning_props};
 use crate::flight_service::DoGet;
 use crate::metrics::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
@@ -230,12 +230,12 @@ impl ExecutionPlan for NetworkCoalesceExec {
         match self.as_ref() {
             Self::Pending(v) => {
                 let mut v = v.clone();
-                v.child = one_child(&children)?;
+                v.child = require_one_child(&children)?;
                 Ok(Arc::new(Self::Pending(v)))
             }
             Self::Ready(v) => {
                 let mut v = v.clone();
-                v.input_stage.plan = MaybeEncodedPlan::Decoded(one_child(&children)?);
+                v.input_stage.plan = MaybeEncodedPlan::Decoded(require_one_child(&children)?);
                 Ok(Arc::new(Self::Ready(v)))
             }
         }

src/execution_plans/network_coalesce.rs

@@ -1,7 +1,7 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
-use crate::common::{one_child, scale_partitioning};
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_physical_optimizer_rule::NetworkBoundary;
+use crate::execution_plans::common::{require_one_child, scale_partitioning};
 use crate::flight_service::DoGet;
 use crate::metrics::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
@@ -292,12 +292,12 @@ impl ExecutionPlan for NetworkShuffleExec {
         match self.as_ref() {
             Self::Pending(v) => {
                 let mut v = v.clone();
-                v.repartition_exec = one_child(&children)?;
+                v.repartition_exec = require_one_child(&children)?;
                 Ok(Arc::new(Self::Pending(v)))
             }
             Self::Ready(v) => {
                 let mut v = v.clone();
-                v.input_stage.plan = MaybeEncodedPlan::Decoded(one_child(&children)?);
+                v.input_stage.plan = MaybeEncodedPlan::Decoded(require_one_child(&children)?);
                 Ok(Arc::new(Self::Ready(v)))
             }
         }