update

Author: luoluoyuyu

src/sql/logical_planner/optimizers/chaining.rs

@@ -10,99 +10,95 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::mem;
-
-use petgraph::graph::{EdgeIndex, NodeIndex};
 use petgraph::prelude::*;
-use petgraph::visit::NodeRef;
+use petgraph::visit::NodeIndexable;
+use tracing::debug;
 
 use crate::sql::logical_node::logical::{LogicalEdgeType, LogicalGraph, Optimizer};
 
-pub type NodeId = NodeIndex;
-pub type EdgeId = EdgeIndex;
-
 pub struct ChainingOptimizer {}
 
-fn remove_in_place<N, E>(graph: &mut DiGraph<N, E>, node: NodeIndex) {
-    let incoming = graph.edges_directed(node, Incoming).next().unwrap();
-
-    let parent = incoming.source().id();
-    let incoming = incoming.id();
-    graph.remove_edge(incoming);
-
-    let outgoing: Vec<_> = graph
-        .edges_directed(node, Outgoing)
-        .map(|e| (e.id(), e.target().id()))
-        .collect();
-
-    for (edge, target) in outgoing {
-        let weight = graph.remove_edge(edge).unwrap();
-        graph.add_edge(parent, target, weight);
-    }
-
-    graph.remove_node(node);
-}
-
 impl Optimizer for ChainingOptimizer {
     fn optimize_once(&self, plan: &mut LogicalGraph) -> bool {
-        let node_indices: Vec<NodeIndex> = plan.node_indices().collect();
+        let mut match_found = None;
 
-        for &node_idx in &node_indices {
-            let cur = plan.node_weight(node_idx).unwrap();
-
-            // sources can't be chained
-            if cur.operator_chain.is_source() {
+        for node_idx in plan.node_indices() {
+            let mut outgoing = plan.edges_directed(node_idx, Outgoing);
+            let first_out = outgoing.next();
+            if first_out.is_none() || outgoing.next().is_some() {
                 continue;
             }
+            let edge = first_out.unwrap();
 
-            let mut successors = plan.edges_directed(node_idx, Outgoing).collect::<Vec<_>>();
-
-            if successors.len() != 1 {
+            if edge.weight().edge_type != LogicalEdgeType::Forward {
                 continue;
             }
 
-            let edge = successors.remove(0);
-            let edge_type = edge.weight().edge_type;
+            let target_idx = edge.target();
 
-            if edge_type != LogicalEdgeType::Forward {
+            let mut incoming = plan.edges_directed(target_idx, Incoming);
+            let first_in = incoming.next();
+            if first_in.is_none() || incoming.next().is_some() {
                 continue;
             }
 
-            let successor_idx = edge.target();
-
-            let successor_node = plan.node_weight(successor_idx).unwrap();
+            let source_node = plan.node_weight(node_idx).expect("Source node missing");
+            let target_node = plan.node_weight(target_idx).expect("Target node missing");
 
-            // skip if parallelism doesn't match or successor is a sink
-            if cur.parallelism != successor_node.parallelism
-                || successor_node.operator_chain.is_sink()
+            if source_node.operator_chain.is_source()
+                || target_node.operator_chain.is_sink()
+                || source_node.parallelism != target_node.parallelism
             {
                 continue;
             }
 
-            // skip successors with multiple predecessors
-            if plan.edges_directed(successor_idx, Incoming).count() > 1 {
-                continue;
+            match_found = Some((node_idx, target_idx, edge.id()));
+            break;
+        }
+
+        if let Some((source_idx, target_idx, edge_id)) = match_found {
+            let edge_weight = plan.remove_edge(edge_id).expect("Edge should exist");
+
+            let target_outgoing: Vec<_> = plan
+                .edges_directed(target_idx, Outgoing)
+                .map(|e| (e.id(), e.target()))
+                .collect();
+
+            for (e_id, next_target_idx) in target_outgoing {
+                let weight = plan.remove_edge(e_id).expect("Outgoing edge missing");
+                plan.add_edge(source_idx, next_target_idx, weight);
             }
 
-            // construct the new node
-            let mut new_cur = cur.clone();
+            let is_source_last = source_idx.index() == plan.node_bound() - 1;
 
-            new_cur.description = format!("{} -> {}", cur.description, successor_node.description);
+            let target_node = plan
+                .remove_node(target_idx)
+                .expect("Target node should exist");
 
-            new_cur
-                .operator_chain
-                .operators
-                .extend(successor_node.operator_chain.operators.clone());
+            let actual_source_idx = if is_source_last {
+                target_idx
+            } else {
+                source_idx
+            };
 
-            new_cur
-                .operator_chain
-                .edges
-                .push(edge.weight().schema.clone());
+            let source_node = plan
+                .node_weight_mut(actual_source_idx)
+                .expect("Source node missing");
+
+            debug!(
+                "Chaining Optimizer: Fusing '{}' -> '{}'",
+                source_node.description, target_node.description
+            );
 
-            mem::swap(&mut new_cur, plan.node_weight_mut(node_idx).unwrap());
+            source_node.description =
+                format!("{} -> {}", source_node.description, target_node.description);
+
+            source_node
+                .operator_chain
+                .operators
+                .extend(target_node.operator_chain.operators);
+            source_node.operator_chain.edges.push(edge_weight.schema);
 
-            // remove the old successor
-            remove_in_place(plan, successor_idx);
             return true;
         }