diff --git a/src/ir.rs b/src/ir.rs index 6fab6571044466f4c5dccedbaf4bbbcc3d791169..646b0b9d5dff85a27f0849dc5c64f1e7a104bb72 100644 --- a/src/ir.rs +++ b/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)] diff --git a/src/ir/transformation/redistribution.rs b/src/ir/transformation/redistribution.rs index 5ec36eba53c2e60974c8cc372bc4aae77590e526..bde7ab382e58b911cecf8798d27b47919b7b3b40 100644 --- a/src/ir/transformation/redistribution.rs +++ b/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(()) } } diff --git a/tests/artifactory/ir/transformation/redistribution/full_motion_less_for_sub_query.yaml b/tests/artifactory/ir/transformation/redistribution/full_motion_less_for_sub_query.yaml index 041a3b6570cb4c5fccc3a5f2fbe557fb47e71659..d82f83c30b02ee2903672c9e0274fac45823fc01 100644 --- a/tests/artifactory/ir/transformation/redistribution/full_motion_less_for_sub_query.yaml +++ b/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 diff --git a/tests/artifactory/ir/transformation/redistribution/full_motion_non_segment_outer_for_sub_query.yaml b/tests/artifactory/ir/transformation/redistribution/full_motion_non_segment_outer_for_sub_query.yaml index 55485d6a27cfb461814099484be577c4e428313c..6afe1446c5457b58c9adc39abf9c1aa8c0956254 100644 --- a/tests/artifactory/ir/transformation/redistribution/full_motion_non_segment_outer_for_sub_query.yaml +++ b/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 diff --git a/tests/artifactory/ir/transformation/redistribution/multiple_sub_queries.yaml b/tests/artifactory/ir/transformation/redistribution/multiple_sub_queries.yaml index 7919152d0c3213c011500495ccc84a2be211258b..a55d631166494705b4e03e5dc0376e023130a9ba 100644 --- a/tests/artifactory/ir/transformation/redistribution/multiple_sub_queries.yaml +++ b/tests/artifactory/ir/transformation/redistribution/multiple_sub_queries.yaml @@ -351,5 +351,8 @@ relations: positions: - 0 name: t2 -slices: ~ +slices: + - + - 50 + - 54 top: 46 diff --git a/tests/artifactory/ir/transformation/redistribution/segment_motion_for_sub_query.yaml b/tests/artifactory/ir/transformation/redistribution/segment_motion_for_sub_query.yaml index 64fbaaf4c70bea61e5d154e8d744ce94346928c4..6eae3d01b854dd4d0aa474e7c0d0e3e2eee0ba87 100644 --- a/tests/artifactory/ir/transformation/redistribution/segment_motion_for_sub_query.yaml +++ b/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