update

Author: luoluoyuyu

src/sql/logical_planner/optimizers/chaining.rs

@@ -10,93 +10,121 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::mem;
-
-use petgraph::prelude::*;
-use petgraph::visit::NodeRef;
+use petgraph::graph::{EdgeIndex, NodeIndex};
+use petgraph::visit::EdgeRef;
+use petgraph::Direction::{Incoming, Outgoing};
 
 use crate::sql::logical_node::logical::{LogicalEdgeType, LogicalGraph, Optimizer};
 
-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);
+pub type NodeId = NodeIndex;
+pub type EdgeId = EdgeIndex;
 
-    let outgoing: Vec<_> = graph
-        .edges_directed(node, Outgoing)
-        .map(|e| (e.id(), e.target().id()))
-        .collect();
+pub struct ChainingOptimizer;
 
-    for (edge, target) in outgoing {
-        let weight = graph.remove_edge(edge).unwrap();
-        graph.add_edge(parent, target, weight);
-    }
+impl ChainingOptimizer {
+    fn find_fusion_candidate(plan: &LogicalGraph) -> Option<(NodeId, NodeId, EdgeId)> {
+        let node_ids: Vec<NodeId> = plan.node_indices().collect();
 
-    graph.remove_node(node);
-}
+        for upstream_id in node_ids {
+            let upstream_node = plan.node_weight(upstream_id)?;
 
-impl Optimizer for ChainingOptimizer {
-    fn optimize_once(&self, plan: &mut LogicalGraph) -> bool {
-        let node_indices: Vec<NodeIndex> = plan.node_indices().collect();
-
-        for &node_idx in &node_indices {
-            let cur = plan.node_weight(node_idx).unwrap();
-
-            if cur.operator_chain.is_source() {
+            if upstream_node.operator_chain.is_source() {
                 continue;
             }
 
-            let mut successors = plan.edges_directed(node_idx, Outgoing).collect::<Vec<_>>();
+            let outgoing_edges: Vec<_> = plan.edges_directed(upstream_id, Outgoing).collect();
 
-            if successors.len() != 1 {
+            if outgoing_edges.len() != 1 {
                 continue;
             }
 
-            let edge = successors.remove(0);
-            let edge_type = edge.weight().edge_type;
+            let bridging_edge = &outgoing_edges[0];
 
-            if edge_type != LogicalEdgeType::Forward {
+            if bridging_edge.weight().edge_type != LogicalEdgeType::Forward {
                 continue;
             }
 
-            let successor_idx = edge.target();
+            let downstream_id = bridging_edge.target();
+            let downstream_node = plan.node_weight(downstream_id)?;
 
-            let successor_node = plan.node_weight(successor_idx).unwrap();
+            if downstream_node.operator_chain.is_sink() {
+                continue;
+            }
 
-            if cur.parallelism != successor_node.parallelism
-                || successor_node.operator_chain.is_sink()
-            {
+            if upstream_node.parallelism != downstream_node.parallelism {
                 continue;
             }
 
-            if plan.edges_directed(successor_idx, Incoming).count() > 1 {
+            let incoming_edges: Vec<_> = plan.edges_directed(downstream_id, Incoming).collect();
+            if incoming_edges.len() != 1 {
                 continue;
             }
 
-            let mut new_cur = cur.clone();
+            return Some((upstream_id, downstream_id, bridging_edge.id()));
+        }
+
+        None
+    }
 
-            new_cur.description = format!("{} -> {}", cur.description, successor_node.description);
+    fn apply_fusion(
+        plan: &mut LogicalGraph,
+        upstream_id: NodeId,
+        downstream_id: NodeId,
+        bridging_edge_id: EdgeId,
+    ) {
+        let bridging_edge = plan
+            .remove_edge(bridging_edge_id)
+            .expect("Graph Integrity Violation: Bridging edge missing");
+
+        let propagated_schema = bridging_edge.schema.clone();
+
+        let downstream_outgoing: Vec<_> = plan
+            .edges_directed(downstream_id, Outgoing)
+            .map(|e| (e.id(), e.target()))
+            .collect();
+
+        for (edge_id, target_id) in downstream_outgoing {
+            let edge_weight = plan
+                .remove_edge(edge_id)
+                .expect("Graph Integrity Violation: Outgoing edge missing");
+
+            plan.add_edge(upstream_id, target_id, edge_weight);
+        }
 
-            new_cur
-                .operator_chain
-                .operators
-                .extend(successor_node.operator_chain.operators.clone());
+        let downstream_node = plan
+            .remove_node(downstream_id)
+            .expect("Graph Integrity Violation: Downstream node missing");
 
-            new_cur
-                .operator_chain
-                .edges
-                .push(edge.weight().schema.clone());
+        let upstream_node = plan
+            .node_weight_mut(upstream_id)
+            .expect("Graph Integrity Violation: Upstream node missing");
 
-            mem::swap(&mut new_cur, plan.node_weight_mut(node_idx).unwrap());
+        upstream_node.description = format!(
+            "{} -> {}",
+            upstream_node.description, downstream_node.description
+        );
 
-            remove_in_place(plan, successor_idx);
-            return true;
-        }
+        upstream_node
+            .operator_chain
+            .operators
+            .extend(downstream_node.operator_chain.operators);
 
-        false
+        upstream_node
+            .operator_chain
+            .edges
+            .push(propagated_schema);
+    }
+}
+
+impl Optimizer for ChainingOptimizer {
+    fn optimize_once(&self, plan: &mut LogicalGraph) -> bool {
+        if let Some((upstream_id, downstream_id, bridging_edge_id)) =
+            Self::find_fusion_candidate(plan)
+        {
+            Self::apply_fusion(plan, upstream_id, downstream_id, bridging_edge_id);
+            true
+        } else {
+            false
+        }
     }
 }