Move out of single-node/distributed into pending/ready

Author: gabotechs

src/physical_optimizer.rs

@@ -109,7 +109,7 @@ impl DistributedPhysicalOptimizerRule {
                 };
 
                 return Ok(Transformed::yes(Arc::new(
-                    ArrowFlightReadExec::new_single_node(
+                    ArrowFlightReadExec::new_pending(
                         Arc::clone(&maybe_isolated_plan),
                         maybe_isolated_plan.output_partitioning().clone(),
                     ),

src/plan/arrow_flight_read.rs

@@ -8,7 +8,7 @@ use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::Ticket;
 use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::common::{internal_datafusion_err, internal_err, plan_err};
+use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
@@ -22,40 +22,39 @@ use std::fmt::Formatter;
 use std::sync::Arc;
 use url::Url;
 
-/// This node can operate in two modes:
-/// 1. Single-node: it acts as a pure passthrough the child plan, as if it was not there.
-/// 2. Distributed: runs within a distributed stage and queries the next input stage over
-///     the network using Arrow Flight.
+/// This node has two variants.
+/// 1. Pending: it acts as a placeholder for the distributed optimization step to mark it as ready.
+/// 2. Ready: runs within a distributed stage and queries the next input stage over the network
+///     using Arrow Flight.
 #[derive(Debug, Clone)]
 pub enum ArrowFlightReadExec {
-    SingleNode(ArrowFlightReadSingleNodeExec),
-    Distributed(ArrowFlightReadDistributedExec),
+    Pending(ArrowFlightReadPendingExec),
+    Ready(ArrowFlightReadReadyExec),
 }
 
-/// Single-node version of the [ArrowFlightReadExec] node. Users can choose to place
-/// this node wherever they want in their plan wherever they want, and the distributed
-/// optimization step will replace it with the appropriate [ArrowFlightReadDistributedExec]
-/// node.
+/// Placeholder version of the [ArrowFlightReadExec] node. It acts as a marker for the
+/// distributed optimization step, which will replace it with the appropriate
+/// [ArrowFlightReadReadyExec] node.
 #[derive(Debug, Clone)]
-pub struct ArrowFlightReadSingleNodeExec {
+pub struct ArrowFlightReadPendingExec {
     properties: PlanProperties,
     child: Arc<dyn ExecutionPlan>,
 }
 
-/// Distributed version of the [ArrowFlightReadExec] node. This node can be created in
+/// Ready version of the [ArrowFlightReadExec] node. This node can be created in
 /// just two ways:
-/// - by the distributed optimization step based on an original [ArrowFlightReadSingleNodeExec]
+/// - by the distributed optimization step based on an original [ArrowFlightReadPendingExec]
 /// - deserialized from a protobuf plan sent over the network.
 #[derive(Debug, Clone)]
-pub struct ArrowFlightReadDistributedExec {
+pub struct ArrowFlightReadReadyExec {
     /// the properties we advertise for this execution plan
     properties: PlanProperties,
     pub(crate) stage_num: usize,
 }
 
 impl ArrowFlightReadExec {
-    pub fn new_single_node(child: Arc<dyn ExecutionPlan>, partitioning: Partitioning) -> Self {
-        Self::SingleNode(ArrowFlightReadSingleNodeExec {
+    pub fn new_pending(child: Arc<dyn ExecutionPlan>, partitioning: Partitioning) -> Self {
+        Self::Pending(ArrowFlightReadPendingExec {
             properties: PlanProperties::new(
                 EquivalenceProperties::new(child.schema()),
                 partitioning,
@@ -66,7 +65,7 @@ impl ArrowFlightReadExec {
         })
     }
 
-    pub(crate) fn new_distributed(
+    pub(crate) fn new_ready(
         partitioning: Partitioning,
         schema: SchemaRef,
         stage_num: usize,
@@ -77,15 +76,15 @@ impl ArrowFlightReadExec {
             EmissionType::Incremental,
             Boundedness::Bounded,
         );
-        Self::Distributed(ArrowFlightReadDistributedExec {
+        Self::Ready(ArrowFlightReadReadyExec {
             properties,
             stage_num,
         })
     }
 
     pub(crate) fn to_distributed(&self, stage_num: usize) -> Result<Self, DataFusionError> {
         match self {
-            ArrowFlightReadExec::SingleNode(p) => Ok(Self::new_distributed(
+            ArrowFlightReadExec::Pending(p) => Ok(Self::new_ready(
                 p.properties.partitioning.clone(),
                 p.child.schema(),
                 stage_num,
@@ -98,8 +97,8 @@ impl ArrowFlightReadExec {
 impl DisplayAs for ArrowFlightReadExec {
     fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
         match self {
-            ArrowFlightReadExec::SingleNode(_) => write!(f, "ArrowFlightReadExec"),
-            ArrowFlightReadExec::Distributed(v) => {
+            ArrowFlightReadExec::Pending(_) => write!(f, "ArrowFlightReadExec"),
+            ArrowFlightReadExec::Ready(v) => {
                 write!(f, "ArrowFlightReadExec: Stage {:<3}", v.stage_num)
             }
         }
@@ -117,15 +116,15 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
     fn properties(&self) -> &PlanProperties {
         match self {
-            ArrowFlightReadExec::SingleNode(v) => &v.properties,
-            ArrowFlightReadExec::Distributed(v) => &v.properties,
+            ArrowFlightReadExec::Pending(v) => &v.properties,
+            ArrowFlightReadExec::Ready(v) => &v.properties,
         }
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         match self {
-            ArrowFlightReadExec::SingleNode(v) => vec![&v.child],
-            ArrowFlightReadExec::Distributed(_) => vec![],
+            ArrowFlightReadExec::Pending(v) => vec![&v.child],
+            ArrowFlightReadExec::Ready(_) => vec![],
         }
     }
 
@@ -147,9 +146,8 @@ impl ExecutionPlan for ArrowFlightReadExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream, DataFusionError> {
-        let this = match self {
-            ArrowFlightReadExec::SingleNode(this) => return this.child.execute(partition, context),
-            ArrowFlightReadExec::Distributed(this) => this,
+        let ArrowFlightReadExec::Ready(this) = self else {
+            return exec_err!("ArrowFlightReadExec is not ready, was the distributed optimization step performed?");
         };
 
         // get the channel manager and current stage from our context

src/plan/codec.rs

@@ -52,7 +52,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
                 )?
                 .ok_or(proto_error("ArrowFlightReadExec is missing partitioning"))?;
 
-                Ok(Arc::new(ArrowFlightReadExec::new_distributed(
+                Ok(Arc::new(ArrowFlightReadExec::new_ready(
                     partioning,
                     Arc::new(schema),
                     stage_num as usize,
@@ -84,9 +84,9 @@ impl PhysicalExtensionCodec for DistributedCodec {
         buf: &mut Vec<u8>,
     ) -> datafusion::common::Result<()> {
         if let Some(node) = node.as_any().downcast_ref::<ArrowFlightReadExec>() {
-            let ArrowFlightReadExec::Distributed(ready_node) = node else {
+            let ArrowFlightReadExec::Ready(ready_node) = node else {
                 return Err(proto_error(
-                    "deserialized an ArrowFlightReadExec that is not distributed",
+                    "deserialized an ArrowFlightReadExec that is not ready",
                 ));
             };
             ArrowFlightReadExecProto {

tests/common/plan.rs

@@ -36,7 +36,7 @@ pub fn distribute_aggregate(
                 let child = node.children()[0].clone();
 
                 let node = node.with_new_children(vec![Arc::new(
-                    ArrowFlightReadExec::new_single_node(child, Partitioning::Hash(expr, 1)),
+                    ArrowFlightReadExec::new_pending(child, Partitioning::Hash(expr, 1)),
                 )])?;
                 Ok(Transformed::yes(node))
             }
@@ -54,7 +54,7 @@ pub fn distribute_aggregate(
                 let child = node.children()[0].clone();
 
                 let node = node.with_new_children(vec![Arc::new(
-                    ArrowFlightReadExec::new_single_node(child, Partitioning::RoundRobinBatch(8)),
+                    ArrowFlightReadExec::new_pending(child, Partitioning::RoundRobinBatch(8)),
                 )])?;
                 Ok(Transformed::yes(node))
             }

tests/custom_extension_codec.rs

@@ -126,7 +126,7 @@ mod tests {
         )?);
 
         if distributed {
-            plan = Arc::new(ArrowFlightReadExec::new_single_node(
+            plan = Arc::new(ArrowFlightReadExec::new_pending(
                 plan.clone(),
                 Partitioning::Hash(vec![col("numbers", &plan.schema())?], 1),
             ));
@@ -141,7 +141,7 @@ mod tests {
         ));
 
         if distributed {
-            plan = Arc::new(ArrowFlightReadExec::new_single_node(
+            plan = Arc::new(ArrowFlightReadExec::new_pending(
                 plan.clone(),
                 Partitioning::RoundRobinBatch(10),
             ));

tests/error_propagation.rs

@@ -51,7 +51,7 @@ mod tests {
         let mut plan: Arc<dyn ExecutionPlan> = Arc::new(ErrorExec::new("something failed"));
 
         for size in [1, 2, 3] {
-            plan = Arc::new(ArrowFlightReadExec::new_single_node(
+            plan = Arc::new(ArrowFlightReadExec::new_pending(
                 plan,
                 Partitioning::RoundRobinBatch(size),
             ));

tests/highly_distributed_query.rs

@@ -31,7 +31,7 @@ mod tests {
 
         let mut physical_distributed = physical.clone();
         for size in [1, 10, 5] {
-            physical_distributed = Arc::new(ArrowFlightReadExec::new_single_node(
+            physical_distributed = Arc::new(ArrowFlightReadExec::new_pending(
                 physical_distributed,
                 Partitioning::RoundRobinBatch(size),
             ));