feat: save motions in plan slices

src/ir.rs

@@ -168,6 +168,12 @@ impl Plan {
         self.top.ok_or(QueryPlannerError::InvalidPlan)
     }
 
+    /// Get plan slices.
+    #[must_use]
+    pub fn get_slices(&self) -> Option<Vec<Vec<usize>>> {
+        self.slices.clone()
+    }
+
     /// Construct a plan from the YAML file.
     ///
     /// # Errors
@@ -366,6 +372,11 @@ impl Plan {
             "Node isn't reference type".into(),
         ))
     }
+
+    /// Set slices of the plan.
+    pub fn set_slices(&mut self, slices: Option<Vec<Vec<usize>>>) {
+        self.slices = slices;
+    }
 }
 
 #[cfg(test)]

src/ir/transformation/redistribution.rs

@@ -5,8 +5,8 @@ use crate::ir::operator::{Bool, Relational};
 use crate::ir::{Node, Plan};
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
-use std::collections::{HashMap, HashSet};
-use traversal::DftPost;
+use std::collections::{hash_map::Entry, HashMap, HashSet};
+use traversal::{Bft, DftPost};
 
 /// A motion policy determinate what portion of data to move
 /// between data nodes.
@@ -307,7 +307,33 @@ impl Plan {
             }
         }
 
-        // TODO: gather motions (revert bft)
+        // Gather motions (revert levels in bft)
+        let mut motions: Vec<Vec<usize>> = Vec::new();
+        let top = self.get_top()?;
+        let bft_tree = Bft::new(&top, |node| self.nodes.rel_iter(node));
+        let mut map: HashMap<usize, usize> = HashMap::new();
+        let mut max_level: usize = 0;
+        for (level, node) in bft_tree {
+            if let Node::Relational(Relational::Motion { .. }) = self.get_node(*node)? {
+                let key: usize = match map.entry(level) {
+                    Entry::Occupied(o) => *o.into_mut(),
+                    Entry::Vacant(v) => {
+                        let old_level = max_level;
+                        v.insert(old_level);
+                        max_level += 1;
+                        old_level
+                    }
+                };
+                match motions.get_mut(key) {
+                    Some(list) => list.push(*node),
+                    None => motions.push(vec![*node]),
+                }
+            }
+        }
+        if !motions.is_empty() {
+            self.set_slices(Some(motions.into_iter().rev().collect()));
+        }
+
         Ok(())
     }
 }

tests/artifactory/ir/transformation/redistribution/full_motion_less_for_sub_query.yaml

@@ -201,5 +201,7 @@ relations:
         positions:
           - 0
       name: t1
-slices: ~
+slices:
+  - 
+     - 30
 top: 26

tests/artifactory/ir/transformation/redistribution/full_motion_non_segment_outer_for_sub_query.yaml

@@ -220,5 +220,7 @@ relations:
         positions:
           - 0
       name: t2
-slices: ~
+slices: 
+  -
+    - 32
 top: 28

tests/artifactory/ir/transformation/redistribution/multiple_sub_queries.yaml

@@ -351,5 +351,8 @@ relations:
         positions:
           - 0
       name: t2
-slices: ~
+slices: 
+  -
+    - 50
+    - 54
 top: 46

tests/artifactory/ir/transformation/redistribution/segment_motion_for_sub_query.yaml

@@ -207,5 +207,7 @@ relations:
         positions:
           - 0
       name: t2
-slices: ~
+slices:
+  -
+    - 30
 top: 26