diff --git a/Cargo.toml b/Cargo.toml
index ffb7f0093300a171df212edadf3e0aded8f3172f..a9a29b6b8caa1886a4f26606aa507d459f3e9dd8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,7 +18,6 @@ yaml-rust = "0.4.1"
 
 [dev-dependencies]
 pretty_assertions = "1.0.0"
-itertools = "0.8.2"
 rmp-serde = "0.14"
 
 [lib]
diff --git a/src/errors.rs b/src/errors.rs
index 8009bb27e7f4cc76ed1b98e71471a756a835f740..4f61ff0ada9fddf437148eb4c252218989d8c7c6 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -7,10 +7,12 @@ const DUPLICATE_COLUMN_ERROR: &str = "duplicate column";
 const EMPTY_PLAN_RELATION: &str = "empty plan relations";
 const INCORRECT_BUCKET_ID_ERROR: &str = "incorrect bucket id";
 const INVALID_BOOL_ERROR: &str = "invalid boolean";
+const INVALID_INPUT: &str = "invalid input";
 const INVALID_NAME_ERROR: &str = "invalid name";
 const INVALID_NODE: &str = "invalid node";
 const INVALID_NUMBER_ERROR: &str = "invalid number";
 const INVALID_PLAN_ERROR: &str = "invalid plan";
+const INVALID_REFERENCE: &str = "invalid reference";
 const INVALID_RELATION_ERROR: &str = "invalid relation";
 const INVALID_ROW_ERROR: &str = "invalid row";
 const INVALID_SHARDING_KEY_ERROR: &str = "invalid sharding key";
@@ -30,10 +32,12 @@ pub enum QueryPlannerError {
     EmptyPlanRelations,
     IncorrectBucketIdError,
     InvalidBool,
+    InvalidInput,
     InvalidName,
     InvalidNode,
     InvalidNumber,
     InvalidPlan,
+    InvalidReference,
     InvalidRelation,
     InvalidRow,
     InvalidShardingKey,
@@ -55,10 +59,12 @@ impl fmt::Display for QueryPlannerError {
             QueryPlannerError::EmptyPlanRelations => EMPTY_PLAN_RELATION,
             QueryPlannerError::IncorrectBucketIdError => INCORRECT_BUCKET_ID_ERROR,
             QueryPlannerError::InvalidBool => INVALID_BOOL_ERROR,
+            QueryPlannerError::InvalidInput => INVALID_INPUT,
             QueryPlannerError::InvalidName => INVALID_NAME_ERROR,
             QueryPlannerError::InvalidNode => INVALID_NODE,
             QueryPlannerError::InvalidNumber => INVALID_NUMBER_ERROR,
             QueryPlannerError::InvalidPlan => INVALID_PLAN_ERROR,
+            QueryPlannerError::InvalidReference => INVALID_REFERENCE,
             QueryPlannerError::InvalidRelation => INVALID_RELATION_ERROR,
             QueryPlannerError::InvalidRow => INVALID_ROW_ERROR,
             QueryPlannerError::InvalidShardingKey => INVALID_SHARDING_KEY_ERROR,
diff --git a/src/ir.rs b/src/ir.rs
index fda06da32493a5ce10c4681a47ef2f751f7f51d1..2fb6f345542d1a081691cebcb80a52adddd90b81 100644
--- a/src/ir.rs
+++ b/src/ir.rs
@@ -8,7 +8,7 @@ pub mod relation;
 pub mod value;
 
 use crate::errors::QueryPlannerError;
-use expression::{Branch, Distribution, Expression};
+use expression::Expression;
 use operator::Relational;
 use relation::Table;
 use serde::{Deserialize, Serialize};
@@ -33,156 +33,6 @@ pub enum Node {
     Relational(Relational),
 }
 
-/// Suggested distribution by the child relational node.
-/// A wrapper for `Expression::suggest_distribution()`;
-/// used by relational nodes when they calculate their
-/// distribution.
-///
-/// # Errors
-/// Returns `QueryPlannerError`:
-/// - parent node's output is not a valid tuple
-/// - child node is not relational
-/// - child's output is not a valid tuple
-fn child_dist(
-    output: usize,
-    child: usize,
-    branch: &Branch,
-    plan: &Plan,
-) -> Result<Distribution, QueryPlannerError> {
-    // Get current output tuple column list
-    let aliases: &Vec<usize> =
-        if let Node::Expression(Expression::Row { list, .. }) = plan.get_node(output)? {
-            Ok(list)
-        } else {
-            Err(QueryPlannerError::InvalidRow)
-        }?;
-
-    // Distribution suggested by the child.
-    if let Node::Relational(child_node) = plan.get_node(child)? {
-        if let Node::Expression(child_row) = plan.get_node(child_node.output())? {
-            Ok(child_row.suggest_distribution(branch, aliases, plan)?)
-        } else {
-            Err(QueryPlannerError::InvalidRow)
-        }
-    } else {
-        Err(QueryPlannerError::InvalidPlan)
-    }
-}
-
-/// Set output tuple distribution for the node.
-///
-/// # Errors
-/// Returns `QueryPlannerError`:
-/// - when node position doesn't exist in the plan node arena
-/// - for nodes, that don't produce tuples (all expressions except `Row`)
-/// - for relational nodes with invalid output or children
-pub fn set_distribution(pointer: usize, plan: &mut Plan) -> Result<(), QueryPlannerError> {
-    match plan.get_node(pointer)? {
-        Node::Relational(relational) => {
-            match relational {
-                Relational::ScanRelation {
-                    relation: table_name,
-                    ..
-                } => {
-                    if let Some(relations) = &plan.relations {
-                        if let Some(rel) = relations.get(table_name) {
-                            // Update output tuple distribution to the relation's one.
-                            match rel {
-                                Table::Segment { key, .. } | Table::VirtualSegment { key, .. } => {
-                                    let rel_tuple = relational.output();
-                                    let node = plan
-                                        .nodes
-                                        .get_mut(rel_tuple)
-                                        .ok_or(QueryPlannerError::ValueOutOfRange)?;
-                                    if let Node::Expression(Expression::Row {
-                                        ref mut distribution,
-                                        ..
-                                    }) = node
-                                    {
-                                        *distribution =
-                                            Some(Distribution::Segment { key: key.clone() });
-                                        return Ok(());
-                                    }
-                                    return Err(QueryPlannerError::InvalidRow);
-                                }
-                                Table::Virtual { .. } => {
-                                    let rel_tuple = relational.output();
-                                    let node = plan
-                                        .nodes
-                                        .get_mut(rel_tuple)
-                                        .ok_or(QueryPlannerError::ValueOutOfRange)?;
-                                    if let Node::Expression(Expression::Row {
-                                        ref mut distribution,
-                                        ..
-                                    }) = node
-                                    {
-                                        *distribution = Some(Distribution::Random);
-                                        return Ok(());
-                                    }
-                                    return Err(QueryPlannerError::InvalidRow);
-                                }
-                            }
-                        }
-                        return Err(QueryPlannerError::InvalidRelation);
-                    }
-                    Err(QueryPlannerError::EmptyPlanRelations)
-                }
-                Relational::Projection { child, output, .. }
-                | Relational::Selection { child, output, .. }
-                | Relational::ScanSubQuery { child, output, .. } => {
-                    let dist = child_dist(*output, *child, &Branch::Left, plan)?;
-                    let rel_tuple = relational.output();
-                    let node = plan
-                        .nodes
-                        .get_mut(rel_tuple)
-                        .ok_or(QueryPlannerError::ValueOutOfRange)?;
-
-                    if let Node::Expression(Expression::Row {
-                        ref mut distribution,
-                        ..
-                    }) = node
-                    {
-                        *distribution = Some(dist);
-                        return Ok(());
-                    }
-                    Err(QueryPlannerError::InvalidNode)
-                }
-                Relational::UnionAll {
-                    left,
-                    right,
-                    output,
-                    ..
-                } => {
-                    let left_dist = child_dist(*output, *left, &Branch::Both, plan)?;
-
-                    let right_dist = child_dist(*output, *right, &Branch::Both, plan)?;
-
-                    let rel_tuple = relational.output();
-                    let node = plan
-                        .nodes
-                        .get_mut(rel_tuple)
-                        .ok_or(QueryPlannerError::ValueOutOfRange)?;
-                    if let Node::Expression(Expression::Row {
-                        ref mut distribution,
-                        ..
-                    }) = node
-                    {
-                        *distribution = Some(Distribution::new_union(left_dist, right_dist)?);
-                        return Ok(());
-                    }
-                    Err(QueryPlannerError::InvalidRow)
-                }
-                // TODO: implement it!
-                Relational::InnerJoin { .. } | Relational::Motion { .. } => {
-                    Err(QueryPlannerError::QueryNotImplemented)
-                }
-            }
-        }
-        // TODO: how should we implement it for the `Row`?
-        Node::Expression(_) => Err(QueryPlannerError::QueryNotImplemented),
-    }
-}
-
 /// Plan node "allocator".
 ///
 /// Inserts an element to the array and returns its position,
@@ -294,6 +144,24 @@ impl Plan {
         plan.check()?;
         Ok(plan)
     }
+
+    /// Returns the next node position
+    #[must_use]
+    pub fn next_node_id(&self) -> usize {
+        self.nodes.len()
+    }
+
+    /// Build {logical id: position} map for relational nodes
+    #[must_use]
+    pub fn relational_id_map(&self) -> HashMap<usize, usize> {
+        let mut map: HashMap<usize, usize> = HashMap::new();
+        for (pos, node) in self.nodes.iter().enumerate() {
+            if let Node::Relational(relational) = node {
+                map.insert(relational.logical_id(), pos);
+            }
+        }
+        map
+    }
 }
 
 /// Plan node iterator over its branches.
@@ -327,6 +195,16 @@ impl<'n> Iterator for BranchIterator<'n> {
     type Item = &'n Node;
 
     fn next(&mut self) -> Option<Self::Item> {
+        let get_next_child = |children: &Vec<usize>| -> Option<&Node> {
+            let current_step = *self.step.borrow();
+            let child = children.get(current_step);
+            child.and_then(|pos| {
+                let node = self.plan.nodes.get(*pos);
+                *self.step.borrow_mut() += 1;
+                node
+            })
+        };
+
         match self.node {
             Node::Expression(expr) => match expr {
                 Expression::Alias { child, .. } => {
@@ -360,18 +238,13 @@ impl<'n> Iterator for BranchIterator<'n> {
             },
             Node::Relational(rel) => match rel {
                 Relational::InnerJoin {
-                    left,
-                    right,
+                    children,
                     condition,
                     ..
                 } => {
                     let current_step = *self.step.borrow();
-                    if current_step == 0 {
-                        *self.step.borrow_mut() += 1;
-                        return self.plan.nodes.get(*left);
-                    } else if current_step == 1 {
-                        *self.step.borrow_mut() += 1;
-                        return self.plan.nodes.get(*right);
+                    if current_step == 0 || current_step == 1 {
+                        return get_next_child(children);
                     } else if current_step == 2 {
                         *self.step.borrow_mut() += 1;
                         return self.plan.nodes.get(*condition);
@@ -379,25 +252,14 @@ impl<'n> Iterator for BranchIterator<'n> {
                     None
                 }
                 Relational::ScanRelation { .. } => None,
-                Relational::ScanSubQuery { child, .. }
-                | Relational::Motion { child, .. }
-                | Relational::Selection { child, .. }
-                | Relational::Projection { child, .. } => {
-                    let current_step = *self.step.borrow();
-                    if current_step == 0 {
-                        *self.step.borrow_mut() += 1;
-                        return self.plan.nodes.get(*child);
-                    }
-                    None
-                }
-                Relational::UnionAll { left, right, .. } => {
+                Relational::ScanSubQuery { children, .. }
+                | Relational::Motion { children, .. }
+                | Relational::Selection { children, .. }
+                | Relational::Projection { children, .. } => get_next_child(children),
+                Relational::UnionAll { children, .. } => {
                     let current_step = *self.step.borrow();
-                    if current_step == 0 {
-                        *self.step.borrow_mut() += 1;
-                        return self.plan.nodes.get(*left);
-                    } else if current_step == 1 {
-                        *self.step.borrow_mut() += 1;
-                        return self.plan.nodes.get(*right);
+                    if current_step == 0 || current_step == 1 {
+                        return get_next_child(children);
                     }
                     None
                 }
diff --git a/src/ir/expression.rs b/src/ir/expression.rs
index 1cf579cb71543eeb7dc5c9af850d1a10ab5d8635..2c22ebc4d3b07d6f9c6f6bd8d94c9f36dcba5f80 100644
--- a/src/ir/expression.rs
+++ b/src/ir/expression.rs
@@ -6,11 +6,13 @@
 //! - distribution of the data in the tuple
 
 use super::operator;
+use super::operator::Relational;
+use super::relation::Table;
 use super::value::Value;
 use super::{Node, Plan};
 use crate::errors::QueryPlannerError;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 /// Tuple data chunk distribution policy in the cluster.
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
@@ -81,49 +83,11 @@ impl Distribution {
     }
 }
 
-/// Tree branch.
-///
-/// Reference expressions point to the position in the incoming
-/// tuple. But union and join nodes have multiple incoming tuples
-/// and we need a way to detect which branch the tuple came from.
-/// So, branches act like additional dimension (among with position)
-/// to refer incoming data.
-#[derive(Serialize, Deserialize, PartialEq, Clone, Debug)]
-pub enum Branch {
-    /// Output reference points both to the left and right branches.
-    ///
-    /// Example: union operator.
-    Both,
-    /// Left branch is also the default value for a single-child operator.
-    ///
-    /// Example: join, selection operator.
-    Left,
-    /// Right branch.
-    ///
-    /// Example: join operator.
-    Right,
-}
-
-impl Branch {
-    /// Compare two branches.
-    ///
-    /// When the reference points to both left and right branches,
-    /// it is equivalent to any of them.
-    #[must_use]
-    pub fn equivalent(&self, other: &Branch) -> bool {
-        match self {
-            Branch::Both => true,
-            Branch::Left => !matches!(other, Branch::Right),
-            Branch::Right => !matches!(other, Branch::Left),
-        }
-    }
-}
-
 /// Tuple tree build blocks.
 ///
 /// Tuple describes a single portion of data moved among the cluster.
 /// It consists of the ordered, strictly typed expressions with names
-/// (columns) with additional information about data distribution policy.
+/// (columns) and additional information about data distribution policy.
 ///
 /// Tuple is a tree with a `Row` top (level 0) and a list of the named
 /// `Alias` columns (level 1). This convention is used among the code
@@ -138,7 +102,7 @@ pub enum Expression {
     Alias {
         /// Alias name.
         name: String,
-        /// Child expression node index in the plan node arena (left branch).
+        /// Child expression node index in the plan node arena.
         child: usize,
     },
     /// Binary expression returning boolean result.
@@ -160,15 +124,20 @@ pub enum Expression {
         value: Value,
     },
     /// Reference to the position in the incoming tuple(s).
-    /// Uses child branch and incoming column position as
-    /// a coordinate system.
-    ///
-    // Example: &0 (left)
+    /// Uses a relative pointer as a coordinate system:
+    /// - relational node (containing this reference)
+    /// - target(s) in the relational nodes list of children
+    /// - column position in the child(ren) output tuple
     Reference {
-        /// Branch of the input tuple.
-        branch: Branch,
+        /// Targets in the relational node children list.
+        /// - Leaf nodes (relation scans): None.
+        /// - Union nodes: two elements (left and right).
+        /// - Other: single element.
+        targets: Option<Vec<usize>>,
         /// Expression position in the input tuple.
         position: usize,
+        /// Relational node ID, that contains current reference
+        parent: usize,
     },
     /// Top of the tuple tree with a list of aliases.
     ///
@@ -182,86 +151,246 @@ pub enum Expression {
     },
 }
 
-#[allow(dead_code)]
-impl Expression {
-    /// Suggest possible distribution to the parent tuple.
-    ///
-    /// When a parent tuple deduces its distribution, it builds a reference map for each
-    /// branch and asks the corresponding child for a suggestion about distribution.
-    /// This function is executed on the child's side.
-    ///
-    /// # Errors
-    /// Returns `QueryPlannerError` when aliases are invalid or a node doesn't know its
-    /// distribution yet.
-    pub fn suggest_distribution(
-        &self,
-        my_branch: &Branch,
-        aliases: &[usize],
-        plan: &Plan,
-    ) -> Result<Distribution, QueryPlannerError> {
-        if let Expression::Row {
-            ref distribution, ..
-        } = self
+fn dist_suggested_by_child(
+    child: usize,
+    plan: &Plan,
+    child_pos_map: &HashMap<(usize, usize), usize>,
+) -> Result<Distribution, QueryPlannerError> {
+    if let Node::Relational(relational_op) = plan.get_node(child)? {
+        if let Node::Expression(Expression::Row {
+            distribution: child_dist,
+            ..
+        }) = plan.get_node(relational_op.output())?
         {
-            let dist = match distribution {
-                Some(d) => &*d,
+            match child_dist {
                 None => return Err(QueryPlannerError::UninitializedDistribution),
-            };
-
-            match dist {
-                Distribution::Random => return Ok(Distribution::Random),
-                Distribution::Replicated => return Ok(Distribution::Replicated),
-                Distribution::Segment { ref key } => {
-                    // When expression is a Row, it has a three level structure:
-                    //
-                    // level 0: row itself
-                    // level 1: list of aliases
-                    // level 2: arbitrary expressions (references as well)
-                    // ...
-                    // Now we traverse the level 2 for the reference expressions, as only
-                    // they can contain positions of the distribution key from the input row, then
-                    // build a map <input row list position, self row list position of the reference>.
-                    let mut map: HashMap<usize, usize> = HashMap::new();
-
-                    for (self_pos, alias_id) in aliases.iter().enumerate() {
-                        if let Node::Expression(Expression::Alias { child, .. }) =
-                            plan.get_node(*alias_id)?
-                        {
-                            if let Node::Expression(Expression::Reference {
-                                branch: ref_branch,
-                                position: ref_pos,
-                            }) = plan.get_node(*child)?
-                            {
-                                if my_branch.equivalent(ref_branch)
-                                    && map.insert(*ref_pos, self_pos).is_some()
-                                {
-                                    return Err(QueryPlannerError::InvalidPlan);
-                                }
-                            }
-                        } else {
-                            return Err(QueryPlannerError::InvalidRow);
-                        }
-                    }
-
+                Some(Distribution::Random) => return Ok(Distribution::Random),
+                Some(Distribution::Replicated) => return Ok(Distribution::Replicated),
+                Some(Distribution::Coordinator) => return Ok(Distribution::Coordinator),
+                Some(Distribution::Segment { key }) => {
                     let mut new_key: Vec<usize> = Vec::new();
                     let all_found = key.iter().all(|pos| {
-                        if let Some(new_pos) = map.get(pos) {
-                            new_key.push(*new_pos);
-                            return true;
-                        }
-                        false
+                        child_pos_map.get(&(child, *pos)).map_or(false, |v| {
+                            new_key.push(*v);
+                            true
+                        })
                     });
                     if all_found {
                         return Ok(Distribution::Segment { key: new_key });
                     }
                     return Ok(Distribution::Random);
                 }
-                Distribution::Coordinator => return Ok(Distribution::Coordinator),
             }
         }
-        Err(QueryPlannerError::InvalidRow)
     }
+    Err(QueryPlannerError::InvalidRow)
+}
 
+fn set_scan_tuple_distribution(
+    plan: &mut Plan,
+    table_set: &HashSet<String>,
+    table_pos_map: &HashMap<usize, usize>,
+    row_node: usize,
+) -> Result<(), QueryPlannerError> {
+    if table_set.len() != 1 {
+        return Err(QueryPlannerError::InvalidNode);
+    }
+    if let Some(relations) = &plan.relations {
+        let table_name: &str = table_set
+            .iter()
+            .next()
+            .ok_or(QueryPlannerError::InvalidNode)?;
+        let table: &Table = relations
+            .get(table_name)
+            .ok_or(QueryPlannerError::InvalidRelation)?;
+        match table {
+            Table::Segment { key, .. } | Table::VirtualSegment { key, .. } => {
+                let mut new_key: Vec<usize> = Vec::new();
+                let all_found = key.iter().all(|pos| {
+                    table_pos_map.get(pos).map_or(false, |v| {
+                        new_key.push(*v);
+                        true
+                    })
+                });
+                if all_found {
+                    if let Node::Expression(Expression::Row {
+                        ref mut distribution,
+                        ..
+                    }) = plan
+                        .nodes
+                        .get_mut(row_node)
+                        .ok_or(QueryPlannerError::InvalidRow)?
+                    {
+                        *distribution = Some(Distribution::Segment { key: new_key });
+                    }
+                }
+            }
+            Table::Virtual { .. } => {
+                if let Node::Expression(Expression::Row {
+                    ref mut distribution,
+                    ..
+                }) = plan
+                    .nodes
+                    .get_mut(row_node)
+                    .ok_or(QueryPlannerError::InvalidRow)?
+                {
+                    *distribution = Some(Distribution::Random);
+                }
+            }
+        }
+        Ok(())
+    } else {
+        Err(QueryPlannerError::InvalidPlan)
+    }
+}
+
+fn set_double_children_node_tuple_distribution(
+    plan: &mut Plan,
+    child_set: &HashSet<usize>,
+    child_pos_map: &HashMap<(usize, usize), usize>,
+    parent_node: &Option<usize>,
+    row_node: usize,
+) -> Result<(), QueryPlannerError> {
+    let mut child_set_iter = child_set.iter();
+    let left_child = *child_set_iter
+        .next()
+        .ok_or(QueryPlannerError::InvalidNode)?;
+    let right_child = *child_set_iter
+        .next()
+        .ok_or(QueryPlannerError::InvalidNode)?;
+
+    let is_union_all: bool = matches!(
+        plan.get_node(parent_node.ok_or(QueryPlannerError::InvalidNode)?)?,
+        Node::Relational(Relational::UnionAll { .. })
+    );
+
+    if is_union_all {
+        let left_dist = dist_suggested_by_child(left_child, plan, child_pos_map)?;
+        let right_dist = dist_suggested_by_child(right_child, plan, child_pos_map)?;
+        if let Node::Expression(Expression::Row {
+            ref mut distribution,
+            ..
+        }) = plan
+            .nodes
+            .get_mut(row_node)
+            .ok_or(QueryPlannerError::InvalidRow)?
+        {
+            *distribution = Some(Distribution::new_union(left_dist, right_dist)?);
+        }
+    } else {
+        // TODO: implement join
+        return Err(QueryPlannerError::InvalidNode);
+    }
+    Ok(())
+}
+
+/// Calculate and set tuple distribution.
+///
+/// As the references in the `Row` expression contain only logical ID of the parent relational nodes,
+/// we need at first traverse all the plan nodes and build a "logical id - array position" map with
+/// `relational_id_map()` function and pass its reference to this function.
+///
+/// # Errors
+/// Returns `QueryPlannerError` when current expression is not a `Row` or contains broken references.
+pub fn set_distribution<S: ::std::hash::BuildHasher + Default>(
+    row_node: usize,
+    id_map: &HashMap<usize, usize, S>,
+    plan: &mut Plan,
+) -> Result<(), QueryPlannerError> {
+    let mut child_set: HashSet<usize> = HashSet::new();
+    let mut child_pos_map: HashMap<(usize, usize), usize> = HashMap::new();
+    let mut table_set: HashSet<String> = HashSet::new();
+    let mut table_pos_map: HashMap<usize, usize> = HashMap::default();
+    let mut parent_node: Option<usize> = None;
+
+    if let Node::Expression(Expression::Row { list, .. }) = plan.get_node(row_node)? {
+        // Gather information about children nodes, that are pointed by the row references.
+        for (pos, alias_node) in list.iter().enumerate() {
+            if let Node::Expression(Expression::Alias { child, .. }) = plan.get_node(*alias_node)? {
+                if let Node::Expression(Expression::Reference {
+                    targets,
+                    position,
+                    parent,
+                    ..
+                }) = plan.get_node(*child)?
+                {
+                    // Get the relational node, containing this row
+                    parent_node = Some(*id_map.get(parent).ok_or(QueryPlannerError::InvalidNode)?);
+                    if let Node::Relational(relational_op) =
+                        plan.get_node(parent_node.ok_or(QueryPlannerError::InvalidNode)?)?
+                    {
+                        if let Some(children) = relational_op.children() {
+                            // References in the branch node.
+                            let child_pos_list: &Vec<usize> = targets
+                                .as_ref()
+                                .ok_or(QueryPlannerError::InvalidReference)?;
+                            for target in child_pos_list {
+                                let child_node: usize = *children
+                                    .get(*target)
+                                    .ok_or(QueryPlannerError::ValueOutOfRange)?;
+                                child_set.insert(child_node);
+                                child_pos_map.insert((child_node, *position), pos);
+                            }
+                        } else {
+                            // References in the leaf (relation scan) node.
+                            if targets.is_some() {
+                                return Err(QueryPlannerError::InvalidReference);
+                            }
+                            if let Relational::ScanRelation { relation, .. } = relational_op {
+                                table_set.insert(relation.clone());
+                                table_pos_map.insert(*position, pos);
+                            } else {
+                                return Err(QueryPlannerError::InvalidReference);
+                            }
+                        }
+                    } else {
+                        return Err(QueryPlannerError::InvalidNode);
+                    }
+                }
+            }
+        }
+    }
+
+    match child_set.len() {
+        0 => {
+            // Scan
+            set_scan_tuple_distribution(plan, &table_set, &table_pos_map, row_node)?;
+        }
+        1 => {
+            // Single child
+            let child: usize = *child_set
+                .iter()
+                .next()
+                .ok_or(QueryPlannerError::InvalidNode)?;
+            let suggested_dist = dist_suggested_by_child(child, plan, &child_pos_map)?;
+            if let Node::Expression(Expression::Row {
+                ref mut distribution,
+                ..
+            }) = plan
+                .nodes
+                .get_mut(row_node)
+                .ok_or(QueryPlannerError::InvalidRow)?
+            {
+                *distribution = Some(suggested_dist);
+            }
+        }
+        2 => {
+            // Union, join
+            set_double_children_node_tuple_distribution(
+                plan,
+                &child_set,
+                &child_pos_map,
+                &parent_node,
+                row_node,
+            )?;
+        }
+        _ => return Err(QueryPlannerError::InvalidReference),
+    }
+    Ok(())
+}
+
+#[allow(dead_code)]
+impl Expression {
     /// Get current row distribution.
     ///
     /// # Errors
@@ -295,8 +424,12 @@ impl Expression {
 
     /// Reference expression constructor.
     #[must_use]
-    pub fn new_ref(branch: Branch, position: usize) -> Self {
-        Expression::Reference { branch, position }
+    pub fn new_ref(parent: usize, targets: Option<Vec<usize>>, position: usize) -> Self {
+        Expression::Reference {
+            parent,
+            targets,
+            position,
+        }
     }
 
     // TODO: check that doesn't contain top-level aliases with the same names
diff --git a/src/ir/expression/tests.rs b/src/ir/expression/tests.rs
index 1057cdca828698b5d4821751a73cbe777f7b338d..d066169ebc300e8be4107c0885f8077bb1245001 100644
--- a/src/ir/expression/tests.rs
+++ b/src/ir/expression/tests.rs
@@ -1,81 +1,248 @@
 use super::*;
+use crate::ir::relation::*;
 use crate::ir::*;
 use pretty_assertions::assert_eq;
 use std::fs;
 use std::path::Path;
 
 #[test]
-fn suggest_distribution() {
+fn proj_preserve_dist_key() {
+    let mut plan = Plan::empty();
+
+    let t = Table::new_seg(
+        "t",
+        vec![
+            Column::new("a", Type::Boolean),
+            Column::new("b", Type::Number),
+            Column::new("c", Type::String),
+            Column::new("d", Type::String),
+        ],
+        &["b", "a"],
+    )
+    .unwrap();
+    plan.add_rel(t);
+
+    let scan = Relational::new_scan("t", &mut plan).unwrap();
+    let scan_id = vec_alloc(&mut plan.nodes, Node::Relational(scan));
+
+    let proj = Relational::new_proj(&mut plan, scan_id, &["a", "b"]).unwrap();
+    let proj_id = vec_alloc(&mut plan.nodes, Node::Relational(proj));
+
+    plan.top = Some(proj_id);
+
+    let map = plan.relational_id_map();
+
+    let scan_output: usize = if let Node::Relational(scan) = plan.get_node(scan_id).unwrap() {
+        scan.output()
+    } else {
+        panic!("Invalid plan!");
+    };
+    set_distribution(scan_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_output).unwrap() {
+        assert_eq!(
+            &Distribution::Segment { key: vec![1, 0] },
+            scan_row.distribution().unwrap()
+        );
+    }
+
+    let proj_output: usize = if let Node::Relational(proj) = plan.get_node(proj_id).unwrap() {
+        proj.output()
+    } else {
+        panic!("Invalid plan!");
+    };
+    set_distribution(proj_output, &map, &mut plan).unwrap();
+    if let Node::Expression(proj_row) = plan.get_node(proj_output).unwrap() {
+        assert_eq!(
+            &Distribution::Segment { key: vec![1, 0] },
+            proj_row.distribution().unwrap()
+        );
+    }
+}
+
+#[test]
+fn proj_shuffle_dist_key() {
     // Load a table "t (a, b, c, d)" distributed by ["b", "a"]
-    // with a sec scan and three additional alias-reference pairs
-    // for "a", "b" and "c". We want to see, what suggestions would
-    // sec scan output row make for a new parent row constructed
-    // from different combinations of new "a", "b" and "c".
+    // with projection ["a", "b"].
     let path = Path::new("")
         .join("tests")
         .join("artifactory")
         .join("ir")
         .join("expression")
-        .join("suggest_distribution.yaml");
+        .join("shuffle_dist_key.yaml");
     let s = fs::read_to_string(path).unwrap();
-    let plan = Plan::from_yaml(&s).unwrap();
+    let mut plan = Plan::from_yaml(&s).unwrap();
+
+    let map = plan.relational_id_map();
 
-    let a = 11;
-    let b = 13;
-    let c = 15;
     let scan_output = 8;
+    let proj_output = 14;
 
-    if let Node::Expression(output) = plan.get_node(scan_output).unwrap() {
-        // Same order in distribution key
+    set_distribution(scan_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_output).unwrap() {
         assert_eq!(
-            Distribution::Segment { key: vec![1, 0] },
-            output
-                .suggest_distribution(&Branch::Left, &[a, b, c], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![1, 0] },
+            scan_row.distribution().unwrap()
         );
+    }
 
-        // Shuffle distribution key
+    set_distribution(proj_output, &map, &mut plan).unwrap();
+    if let Node::Expression(proj_row) = plan.get_node(proj_output).unwrap() {
         assert_eq!(
-            Distribution::Segment { key: vec![0, 1] },
-            output
-                .suggest_distribution(&Branch::Left, &[b, a], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![0, 1] },
+            proj_row.distribution().unwrap()
         );
+    }
+}
+
+#[test]
+fn proj_shrink_dist_key_1() {
+    // Load a table "t (a, b, c, d)" distributed by ["b", "a"]
+    // with projection ["c", "a"].
+    let path = Path::new("")
+        .join("tests")
+        .join("artifactory")
+        .join("ir")
+        .join("expression")
+        .join("shrink_dist_key_1.yaml");
+    let s = fs::read_to_string(path).unwrap();
+    let mut plan = Plan::from_yaml(&s).unwrap();
+
+    let map = plan.relational_id_map();
 
-        // Shrink distribution key #1
+    let scan_output = 8;
+    let proj_output = 14;
+
+    set_distribution(scan_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_output).unwrap() {
         assert_eq!(
-            Distribution::Random,
-            output
-                .suggest_distribution(&Branch::Left, &[c, a], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![1, 0] },
+            scan_row.distribution().unwrap()
         );
+    }
 
-        // Shrink distribution key #2
+    set_distribution(proj_output, &map, &mut plan).unwrap();
+    if let Node::Expression(proj_row) = plan.get_node(proj_output).unwrap() {
+        assert_eq!(&Distribution::Random, proj_row.distribution().unwrap());
+    }
+}
+
+#[test]
+fn proj_shrink_dist_key_2() {
+    // Load a table "t (a, b, c, d)" distributed by ["b", "a"]
+    // with projection ["a"].
+    let path = Path::new("")
+        .join("tests")
+        .join("artifactory")
+        .join("ir")
+        .join("expression")
+        .join("shrink_dist_key_2.yaml");
+    let s = fs::read_to_string(path).unwrap();
+    let mut plan = Plan::from_yaml(&s).unwrap();
+
+    let map = plan.relational_id_map();
+
+    let scan_output = 8;
+    let proj_output = 12;
+
+    set_distribution(scan_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_output).unwrap() {
         assert_eq!(
-            Distribution::Random,
-            output
-                .suggest_distribution(&Branch::Left, &[a], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![1, 0] },
+            scan_row.distribution().unwrap()
         );
+    }
+
+    set_distribution(proj_output, &map, &mut plan).unwrap();
+    if let Node::Expression(proj_row) = plan.get_node(proj_output).unwrap() {
+        assert_eq!(&Distribution::Random, proj_row.distribution().unwrap());
+    }
+}
 
-        // Check both branch mode
+#[test]
+fn union_all_fallback_to_random() {
+    // Load table "t1 (a, b)" distributed by ["a"],
+    // table "t2 (a, b)" distributed by ["b"],
+    // union all (t1, t2)
+    let path = Path::new("")
+        .join("tests")
+        .join("artifactory")
+        .join("ir")
+        .join("expression")
+        .join("union_fallback_to_random.yaml");
+    let s = fs::read_to_string(path).unwrap();
+    let mut plan = Plan::from_yaml(&s).unwrap();
+
+    let map = plan.relational_id_map();
+
+    let scan_t1_output = 4;
+    let scan_t2_output = 10;
+    let union_output = 16;
+
+    set_distribution(scan_t1_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_t1_output).unwrap() {
         assert_eq!(
-            Distribution::Segment { key: vec![1, 0] },
-            output
-                .suggest_distribution(&Branch::Both, &[a, b, c], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![0] },
+            scan_row.distribution().unwrap()
         );
+    }
 
-        // Check absent branch int the output
+    set_distribution(scan_t2_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_t2_output).unwrap() {
         assert_eq!(
-            Distribution::Random,
-            output
-                .suggest_distribution(&Branch::Right, &[a, b, c], &plan)
-                .unwrap()
+            &Distribution::Segment { key: vec![1] },
+            scan_row.distribution().unwrap()
         );
+    }
 
-        //TODO: implement checks for Replicated and Single
-    } else {
-        panic!("Wrong output node type!");
+    set_distribution(union_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(union_output).unwrap() {
+        assert_eq!(&Distribution::Random, scan_row.distribution().unwrap());
     }
 }
+
+#[test]
+fn union_preserve_dist() {
+    // Load table "t1 (a, b)" distributed by ["a"],
+    // table "t2 (a, b)" distributed by ["b"],
+    // union all (t1, t2)
+    let path = Path::new("")
+        .join("tests")
+        .join("artifactory")
+        .join("ir")
+        .join("expression")
+        .join("union_preserve_dist.yaml");
+    let s = fs::read_to_string(path).unwrap();
+    let mut plan = Plan::from_yaml(&s).unwrap();
+
+    let map = plan.relational_id_map();
+
+    let scan_t1_output = 4;
+    let scan_t2_output = 10;
+    let union_output = 16;
+
+    set_distribution(scan_t1_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_t1_output).unwrap() {
+        assert_eq!(
+            &Distribution::Segment { key: vec![0] },
+            scan_row.distribution().unwrap()
+        );
+    }
+
+    set_distribution(scan_t2_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(scan_t2_output).unwrap() {
+        assert_eq!(
+            &Distribution::Segment { key: vec![0] },
+            scan_row.distribution().unwrap()
+        );
+    }
+
+    set_distribution(union_output, &map, &mut plan).unwrap();
+    if let Node::Expression(scan_row) = plan.get_node(union_output).unwrap() {
+        assert_eq!(
+            &Distribution::Segment { key: vec![0] },
+            scan_row.distribution().unwrap()
+        );
+    }
+}
+
+//TODO: add other distribution variants to the test cases.
diff --git a/src/ir/operator.rs b/src/ir/operator.rs
index 8b6dd8551573eb1a9984ec247930f14d5bf9fc83..954c483986bddc5af8c794e0fb77517ab5258383 100644
--- a/src/ir/operator.rs
+++ b/src/ir/operator.rs
@@ -1,6 +1,6 @@
 //! Operators for expression transformations.
 
-use super::expression::{Branch, Expression};
+use super::expression::Expression;
 use super::relation::Table;
 use super::{vec_alloc, Node, Plan};
 use crate::errors::QueryPlannerError;
@@ -38,107 +38,201 @@ pub enum Bool {
 pub enum Relational {
     /// Inner Join
     InnerJoin {
+        /// Contains exactly two elements: left and right node indexes
+        /// from the plan node arena.
+        children: Vec<usize>,
         /// Left and right tuple comparison condition.
-        /// In fact - an expression tree top index in plan node arena.
+        /// In fact - an expression tree top index from the plan node arena.
         condition: usize,
-        /// Left branch tuple node index in the plan node arena.
-        left: usize,
-        /// Output tuple node index in the plan node arena.
+        /// Logical node ID
+        id: usize,
+        /// Output tuple node index from the plan node arena.
         output: usize,
-        /// Right branch tuple node index in the plan node arena.
-        right: usize,
     },
     Motion {
-        /// Child tuple node index in the plan node arena (left branch).
-        child: usize,
+        /// Contains exactly a single element: child node index
+        /// from the plan node arena.
+        children: Vec<usize>,
+        /// Logical node ID
+        id: usize,
         /// Output tuple node index in the plan node arena.
         output: usize,
     },
     Projection {
-        /// Child tuple node index in the plan node arena (left branch).
-        child: usize,
+        /// Contains at least a single element: child node index
+        /// from the plan node arena. Every element other that the
+        /// first one should be treated as a `SubQuery` node from
+        /// the output tuple tree.
+        children: Vec<usize>,
+        /// Logical node ID
+        id: usize,
         /// Output tuple node index in the plan node arena.
         output: usize,
     },
     ScanRelation {
         /// Output tuple node index in the plan node arena.
         output: usize,
+        /// Logical node ID
+        id: usize,
         /// Relation name.
         relation: String,
     },
     ScanSubQuery {
         /// SubQuery name
         alias: String,
-        /// Child tuple node index in the plan node arena (left branch).
-        child: usize,
+        /// Contains exactly a single element: child node index
+        /// from the plan node arena.
+        children: Vec<usize>,
+        /// Logical node ID
+        id: usize,
         /// Output tuple node index in the plan node arena.
         output: usize,
     },
     Selection {
-        /// Child tuple node index in the plan node arena (left branch).
-        child: usize,
+        /// Contains at least a single element: child node index
+        /// from the plan node arena. Every element other that the
+        /// first one should be treated as a `SubQuery` node from
+        /// the filter tree.
+        children: Vec<usize>,
         /// Filter expression node index in the plan node arena.
         filter: usize,
+        /// Logical node ID
+        id: usize,
         /// Output tuple node index in the plan node arena.
         output: usize,
     },
     UnionAll {
-        /// Left branch tuple node index in the plan node arena.
-        left: usize,
-        /// Right branch tuple node index in the plan node arena.
-        right: usize,
+        /// Contains exactly two elements: left and right node indexes
+        /// from the plan node arena.
+        children: Vec<usize>,
+        /// Logical node ID
+        id: usize,
         /// Output tuple node index in the plan node arena.
         output: usize,
     },
 }
 
-/// Returns a list of new alias nodes.
-/// Helpful, when construct a new row from the child node
-/// and we have only column names. We can feed this function
-/// with column names and child node pointer to create a new
-/// alias node list.
-///
+/// Create a new tuple from the children nodes output, containing only
+/// a specified list of column names. If the column list is empty then
+/// just copy all the columns to a new tuple.
 /// # Errors
-/// Returns `QueryPlannerError` when child node is not a
-/// relational operator or some of the alias names are
-/// absent in the child node's output.  
-pub fn new_alias_nodes(
+/// Returns `QueryPlannerError`:
+/// - relation node contains invalid `Row` in the output
+/// - targets and children are inconsistent
+/// - column names don't exits
+pub fn new_row_node(
     plan: &mut Plan,
-    node: usize,
+    rel_node_id: usize,
+    children: &[usize],
+    targets: &[usize],
     col_names: &[&str],
-    branch: &Branch,
-) -> Result<Vec<usize>, QueryPlannerError> {
-    if col_names.is_empty() {
+) -> Result<usize, QueryPlannerError> {
+    // We can pass two target children nodes only in a case
+    // of `UnionAll`. Even for a `NaturalJoin` we work with
+    // each child independently. In fact, we need only the
+    // first child in a `UnionAll` operator to get correct
+    // column names for a new tuple (second child aliases
+    // would be shadowed). But each reference should point
+    // to both children to give us additional information
+    // during transformations.
+    if (targets.len() > 2) || targets.is_empty() {
         return Err(QueryPlannerError::InvalidRow);
     }
 
-    if let Node::Relational(relation_node) = plan.get_node(node)? {
-        let map = relation_node.output_alias_position_map(plan)?;
+    if let Some(max) = targets.iter().max() {
+        if *max >= children.len() {
+            return Err(QueryPlannerError::InvalidRow);
+        }
+    }
+
+    let target_child: usize = if let Some(target) = targets.get(0) {
+        *target
+    } else {
+        return Err(QueryPlannerError::InvalidRow);
+    };
+    let child_node: usize = if let Some(child) = children.get(target_child) {
+        *child
+    } else {
+        return Err(QueryPlannerError::InvalidRow);
+    };
+
+    if col_names.is_empty() {
+        let child_row_list: Vec<usize> =
+            if let Node::Relational(relational_op) = plan.get_node(child_node)? {
+                if let Node::Expression(Expression::Row { list, .. }) =
+                    plan.get_node(relational_op.output())?
+                {
+                    list.clone()
+                } else {
+                    return Err(QueryPlannerError::InvalidRow);
+                }
+            } else {
+                return Err(QueryPlannerError::InvalidNode);
+            };
+
         let mut aliases: Vec<usize> = Vec::new();
 
-        let all_found = col_names.iter().all(|col| {
-            map.get(*col).map_or(false, |pos| {
-                // Create new references and aliases. Save them to the plan nodes arena.
-                let r_id = vec_alloc(
-                    &mut plan.nodes,
-                    Node::Expression(Expression::new_ref(branch.clone(), *pos)),
-                );
-                let a_id = vec_alloc(
-                    &mut plan.nodes,
-                    Node::Expression(Expression::new_alias(col, r_id)),
-                );
-                aliases.push(a_id);
-                true
-            })
-        });
-
-        if all_found {
-            return Ok(aliases);
+        for (pos, alias_node) in child_row_list.iter().enumerate() {
+            let name: String = if let Node::Expression(Expression::Alias { ref name, .. }) =
+                plan.get_node(*alias_node)?
+            {
+                String::from(name)
+            } else {
+                return Err(QueryPlannerError::InvalidRow);
+            };
+            let new_targets: Vec<usize> = targets.iter().copied().collect();
+            // Create new references and aliases. Save them to the plan nodes arena.
+            let r_id = vec_alloc(
+                &mut plan.nodes,
+                Node::Expression(Expression::new_ref(rel_node_id, Some(new_targets), pos)),
+            );
+            let a_id = vec_alloc(
+                &mut plan.nodes,
+                Node::Expression(Expression::new_alias(&name, r_id)),
+            );
+            aliases.push(a_id);
         }
 
-        return Err(QueryPlannerError::InvalidRow);
+        let row_node = vec_alloc(
+            &mut plan.nodes,
+            Node::Expression(Expression::new_row(aliases, None)),
+        );
+        return Ok(row_node);
     }
-    Err(QueryPlannerError::InvalidPlan)
+
+    let map = if let Node::Relational(relational_op) = plan.get_node(child_node)? {
+        relational_op.output_alias_position_map(plan)?
+    } else {
+        return Err(QueryPlannerError::InvalidNode);
+    };
+
+    let mut aliases: Vec<usize> = Vec::new();
+
+    let all_found = col_names.iter().all(|col| {
+        map.get(*col).map_or(false, |pos| {
+            let new_targets: Vec<usize> = targets.iter().copied().collect();
+            // Create new references and aliases. Save them to the plan nodes arena.
+            let r_id = vec_alloc(
+                &mut plan.nodes,
+                Node::Expression(Expression::new_ref(rel_node_id, Some(new_targets), *pos)),
+            );
+            let a_id = vec_alloc(
+                &mut plan.nodes,
+                Node::Expression(Expression::new_alias(col, r_id)),
+            );
+            aliases.push(a_id);
+            true
+        })
+    });
+
+    if all_found {
+        let row_node = vec_alloc(
+            &mut plan.nodes,
+            Node::Expression(Expression::new_row(aliases, None)),
+        );
+        return Ok(row_node);
+    }
+    Err(QueryPlannerError::InvalidRow)
 }
 
 #[allow(dead_code)]
@@ -192,29 +286,32 @@ impl Relational {
         }
     }
 
-    /// Return a list of column names from the output tuple.
-    ///
-    /// # Errors
-    /// Returns `QueryPlannerError` if the tuple is invalid.
-    pub fn output_alias_names(&self, nodes: &[Node]) -> Result<Vec<String>, QueryPlannerError> {
-        let mut names: Vec<String> = Vec::new();
+    /// Get logical id of the relational node.
+    #[must_use]
+    pub fn logical_id(&self) -> usize {
+        match self {
+            Relational::InnerJoin { id, .. }
+            | Relational::Motion { id, .. }
+            | Relational::Projection { id, .. }
+            | Relational::ScanRelation { id, .. }
+            | Relational::ScanSubQuery { id, .. }
+            | Relational::Selection { id, .. }
+            | Relational::UnionAll { id, .. } => *id,
+        }
+    }
 
-        if let Some(Node::Expression(Expression::Row { list, .. })) = nodes.get(self.output()) {
-            let valid = list.iter().all(|item| {
-                if let Some(Node::Expression(Expression::Alias { ref name, .. })) = nodes.get(*item)
-                {
-                    names.push(name.clone());
-                    true
-                } else {
-                    false
-                }
-            });
-            if valid {
-                return Ok(names);
-            }
-            return Err(QueryPlannerError::InvalidPlan);
+    // Get a copy of the children nodes.
+    #[must_use]
+    pub fn children(&self) -> Option<Vec<usize>> {
+        match self {
+            Relational::InnerJoin { children, .. }
+            | Relational::Motion { children, .. }
+            | Relational::Projection { children, .. }
+            | Relational::ScanSubQuery { children, .. }
+            | Relational::Selection { children, .. }
+            | Relational::UnionAll { children, .. } => Some(children.clone()),
+            Relational::ScanRelation { .. } => None,
         }
-        Err(QueryPlannerError::ValueOutOfRange)
     }
 
     /// New `ScanRelation` constructor.
@@ -222,7 +319,9 @@ impl Relational {
     /// # Errors
     /// Returns `QueryPlannerError` when relation is invalid.
     pub fn new_scan(table_name: &str, plan: &mut Plan) -> Result<Self, QueryPlannerError> {
+        let scan_id = plan.next_node_id();
         let nodes = &mut plan.nodes;
+
         if let Some(relations) = &plan.relations {
             if let Some(rel) = relations.get(table_name) {
                 match rel {
@@ -231,7 +330,7 @@ impl Relational {
                             .iter()
                             .enumerate()
                             .map(|(pos, col)| {
-                                let r = Expression::new_ref(Branch::Left, pos);
+                                let r = Expression::new_ref(scan_id, None, pos);
                                 let r_id = vec_alloc(nodes, Node::Expression(r));
                                 vec_alloc(
                                     nodes,
@@ -241,6 +340,7 @@ impl Relational {
                             .collect();
 
                         return Ok(Relational::ScanRelation {
+                            id: scan_id,
                             output: vec_alloc(
                                 nodes,
                                 Node::Expression(Expression::new_row(refs, None)),
@@ -256,6 +356,8 @@ impl Relational {
         Err(QueryPlannerError::InvalidRelation)
     }
 
+    // TODO: we need a more flexible projection constructor (constants, etc)
+
     /// New `Projection` constructor.
     ///
     /// # Errors
@@ -266,18 +368,16 @@ impl Relational {
     pub fn new_proj(
         plan: &mut Plan,
         child: usize,
-        output: &[&str],
+        col_names: &[&str],
     ) -> Result<Self, QueryPlannerError> {
-        let aliases = new_alias_nodes(plan, child, output, &Branch::Left)?;
-
-        let new_output = vec_alloc(
-            &mut plan.nodes,
-            Node::Expression(Expression::new_row(aliases, None)),
-        );
+        let id = plan.next_node_id();
+        let children: Vec<usize> = vec![child];
+        let output = new_row_node(plan, id, &children, &[0], col_names)?;
 
         Ok(Relational::Projection {
-            child,
-            output: new_output,
+            children,
+            id,
+            output,
         })
     }
 
@@ -298,23 +398,15 @@ impl Relational {
             return Err(QueryPlannerError::InvalidBool);
         }
 
-        let names: Vec<String> = if let Node::Relational(rel_op) = plan.get_node(child)? {
-            rel_op.output_alias_names(&plan.nodes)?
-        } else {
-            return Err(QueryPlannerError::InvalidRow);
-        };
-        let output: Vec<&str> = names.iter().map(|s| s as &str).collect();
-        let aliases = new_alias_nodes(plan, child, &output, &Branch::Left)?;
-
-        let new_output = vec_alloc(
-            &mut plan.nodes,
-            Node::Expression(Expression::new_row(aliases, None)),
-        );
+        let id = plan.next_node_id();
+        let children: Vec<usize> = vec![child];
+        let output = new_row_node(plan, id, &children, &[0], &[])?;
 
         Ok(Relational::Selection {
-            child,
+            children,
             filter,
-            output: new_output,
+            id,
+            output,
         })
     }
 
@@ -330,37 +422,28 @@ impl Relational {
         left: usize,
         right: usize,
     ) -> Result<Self, QueryPlannerError> {
-        let left_names: Vec<String> = if let Node::Relational(rel_op) = plan.get_node(left)? {
-            rel_op.output_alias_names(&plan.nodes)?
-        } else {
-            return Err(QueryPlannerError::InvalidRow);
-        };
-
-        let right_names: Vec<String> = if let Node::Relational(rel_op) = plan.get_node(right)? {
-            rel_op.output_alias_names(&plan.nodes)?
-        } else {
-            return Err(QueryPlannerError::InvalidRow);
+        let child_row_len = |child: usize, plan: &Plan| -> Result<usize, QueryPlannerError> {
+            if let Node::Relational(relational_op) = plan.get_node(child)? {
+                match plan.get_node(relational_op.output())? {
+                    Node::Expression(Expression::Row { ref list, .. }) => Ok(list.len()),
+                    _ => Err(QueryPlannerError::InvalidRow),
+                }
+            } else {
+                Err(QueryPlannerError::InvalidRow)
+            }
         };
 
-        let equal = (left_names.len() == right_names.len())
-            && left_names.iter().zip(right_names).all(|(l, r)| l.eq(&r));
-
-        if !equal {
+        if child_row_len(left, plan) != child_row_len(right, plan) {
             return Err(QueryPlannerError::NotEqualRows);
         }
 
-        // Generate new output columns.
-        let col_names: Vec<&str> = left_names.iter().map(|s| s as &str).collect();
-        let aliases = new_alias_nodes(plan, left, &col_names, &Branch::Both)?;
-
-        let output = vec_alloc(
-            &mut plan.nodes,
-            Node::Expression(Expression::new_row(aliases, None)),
-        );
+        let id = plan.next_node_id();
+        let children: Vec<usize> = vec![left, right];
+        let output = new_row_node(plan, id, &children, &[0, 1], &[])?;
 
         Ok(Relational::UnionAll {
-            left,
-            right,
+            children,
+            id,
             output,
         })
     }
@@ -377,26 +460,17 @@ impl Relational {
         child: usize,
         alias: &str,
     ) -> Result<Self, QueryPlannerError> {
-        let names: Vec<String> = if let Node::Relational(rel_op) = plan.get_node(child)? {
-            rel_op.output_alias_names(&plan.nodes)?
-        } else {
-            return Err(QueryPlannerError::InvalidRow);
-        };
         if alias.is_empty() {
             return Err(QueryPlannerError::InvalidName);
         }
-
-        let col_names: Vec<&str> = names.iter().map(|s| s as &str).collect();
-        let aliases = new_alias_nodes(plan, child, &col_names, &Branch::Both)?;
-
-        let output = vec_alloc(
-            &mut plan.nodes,
-            Node::Expression(Expression::new_row(aliases, None)),
-        );
+        let id = plan.next_node_id();
+        let children: Vec<usize> = vec![child];
+        let output = new_row_node(plan, id, &children, &[0], &[])?;
 
         Ok(Relational::ScanSubQuery {
             alias: String::from(alias),
-            child,
+            children,
+            id,
             output,
         })
     }
diff --git a/src/ir/operator/tests.rs b/src/ir/operator/tests.rs
index f94ca8d28c718d4e0924dc0d710b95fe38b974d0..71cf3432e1984c1d030fdde742a2730ea824aada 100644
--- a/src/ir/operator/tests.rs
+++ b/src/ir/operator/tests.rs
@@ -4,7 +4,6 @@ use crate::ir::expression::*;
 use crate::ir::relation::*;
 use crate::ir::value::*;
 use crate::ir::*;
-use itertools::Itertools;
 use pretty_assertions::assert_eq;
 use std::fs;
 use std::path::Path;
@@ -26,21 +25,31 @@ fn scan_rel() {
     .unwrap();
     plan.add_rel(t);
 
+    let scan_output = 8;
+    let scan_node = 9;
+
     let scan = Relational::new_scan("t", &mut plan).unwrap();
     assert_eq!(
         Relational::ScanRelation {
-            output: 8,
+            output: scan_output,
+            id: 0,
             relation: String::from("t"),
         },
         scan
     );
-    assert_eq!(9, vec_alloc(&mut plan.nodes, Node::Relational(scan)));
+    assert_eq!(
+        scan_node,
+        vec_alloc(&mut plan.nodes, Node::Relational(scan))
+    );
+    plan.top = Some(scan_node);
+
+    let map = plan.relational_id_map();
 
-    set_distribution(9, &mut plan).unwrap();
-    if let Node::Expression(row) = plan.get_node(8).unwrap() {
+    set_distribution(scan_output, &map, &mut plan).unwrap();
+    if let Node::Expression(row) = plan.get_node(scan_output).unwrap() {
         assert_eq!(
-            *row.distribution().unwrap(),
-            Distribution::Segment { key: vec![1, 0] }
+            row.distribution().unwrap(),
+            &Distribution::Segment { key: vec![1, 0] }
         );
     } else {
         panic!("Wrong output node type!");
@@ -67,7 +76,11 @@ fn scan_rel_serialized() {
     let scan = Relational::new_scan("t", &mut plan).unwrap();
     plan.nodes.push(Node::Relational(scan));
     plan.top = Some(9);
-    set_distribution(plan.top.unwrap(), &mut plan).unwrap();
+
+    let scan_output = 8;
+
+    let map = plan.relational_id_map();
+    set_distribution(scan_output, &map, &mut plan).unwrap();
 
     let path = Path::new("")
         .join("tests")
@@ -98,46 +111,6 @@ fn projection() {
 
     let scan = Relational::new_scan("t", &mut plan).unwrap();
     let scan_id = vec_alloc(&mut plan.nodes, Node::Relational(scan));
-    set_distribution(scan_id, &mut plan).unwrap();
-
-    let proj_seg = Relational::new_proj(&mut plan, scan_id, &["b", "a"]).unwrap();
-    assert_eq!(
-        Relational::Projection {
-            child: scan_id,
-            output: 14
-        },
-        proj_seg
-    );
-    let proj_seg_id = vec_alloc(&mut plan.nodes, Node::Relational(proj_seg));
-    set_distribution(proj_seg_id, &mut plan).unwrap();
-
-    if let Node::Expression(row) = plan.get_node(14).unwrap() {
-        assert_eq!(
-            *row.distribution().unwrap(),
-            Distribution::Segment { key: vec![0, 1] }
-        );
-    }
-
-    let proj_rand = Relational::new_proj(&mut plan, scan_id, &["a", "d"]).unwrap();
-    assert_eq!(
-        Relational::Projection {
-            child: scan_id,
-            output: 20
-        },
-        proj_rand
-    );
-    let proj_rand_id = vec_alloc(&mut plan.nodes, Node::Relational(proj_rand));
-    set_distribution(proj_rand_id, &mut plan).unwrap();
-
-    if let Node::Expression(row) = plan.get_node(20).unwrap() {
-        assert_eq!(*row.distribution().unwrap(), Distribution::Random);
-    }
-
-    // Empty output
-    assert_eq!(
-        QueryPlannerError::InvalidRow,
-        Relational::new_proj(&mut plan, scan_id, &[]).unwrap_err()
-    );
 
     // Invalid alias names in the output
     assert_eq!(
@@ -147,7 +120,7 @@ fn projection() {
 
     // Expression node instead of relational one
     assert_eq!(
-        QueryPlannerError::InvalidPlan,
+        QueryPlannerError::InvalidNode,
         Relational::new_proj(&mut plan, 1, &["a"]).unwrap_err()
     );
 
@@ -190,15 +163,21 @@ fn selection() {
     let scan = Relational::new_scan("t", &mut plan).unwrap();
     let scan_id = vec_alloc(&mut plan.nodes, Node::Relational(scan));
 
-    let new_aliases = new_alias_nodes(&mut plan, scan_id, &["b"], &Branch::Left).unwrap();
-    let a_id = new_aliases.get(0).unwrap();
+    let ref_a_id = vec_alloc(
+        &mut plan.nodes,
+        Node::Expression(Expression::new_ref(scan_id + 1, Some(vec![0]), 0)),
+    );
+    let a_id = vec_alloc(
+        &mut plan.nodes,
+        Node::Expression(Expression::new_alias("a", ref_a_id)),
+    );
     let const_id = vec_alloc(
         &mut plan.nodes,
         Node::Expression(Expression::new_const(Value::number_from_str("10").unwrap())),
     );
     let gt_id = vec_alloc(
         &mut plan.nodes,
-        Node::Expression(Expression::new_bool(*a_id, Bool::Gt, const_id)),
+        Node::Expression(Expression::new_bool(a_id, Bool::Gt, const_id)),
     );
 
     // Correct Selection operator
@@ -212,7 +191,7 @@ fn selection() {
 
     // Non-relational child
     assert_eq!(
-        QueryPlannerError::InvalidRow,
+        QueryPlannerError::InvalidNode,
         Relational::new_select(&mut plan, const_id, gt_id).unwrap_err()
     );
 }
@@ -230,7 +209,7 @@ fn selection_serialize() {
 }
 
 #[test]
-fn union_all() {
+fn union_all_col_amount_mismatch() {
     let mut plan = Plan::empty();
 
     let t1 = Table::new_seg(
@@ -246,103 +225,19 @@ fn union_all() {
 
     let scan_t1 = Relational::new_scan("t1", &mut plan).unwrap();
     let scan_t1_id = vec_alloc(&mut plan.nodes, Node::Relational(scan_t1));
-    set_distribution(scan_t1_id, &mut plan).unwrap();
 
-    // Check fallback to random distribution
-    let t2 = Table::new_seg(
-        "t2",
-        vec![
-            Column::new("a", Type::Boolean),
-            Column::new("b", Type::Number),
-        ],
-        &["b"],
-    )
-    .unwrap();
+    // Check errors for children with different amount of column
+    let t2 = Table::new_seg("t2", vec![Column::new("b", Type::Number)], &["b"]).unwrap();
     plan.add_rel(t2);
 
     let scan_t2 = Relational::new_scan("t2", &mut plan).unwrap();
     let scan_t2_id = vec_alloc(&mut plan.nodes, Node::Relational(scan_t2));
-    set_distribution(scan_t2_id, &mut plan).unwrap();
-
-    let union_all = Relational::new_union_all(&mut plan, scan_t1_id, scan_t2_id).unwrap();
-    let union_all_id = vec_alloc(&mut plan.nodes, Node::Relational(union_all));
-    set_distribution(union_all_id, &mut plan).unwrap();
-
-    if let Node::Relational(union_all) = plan.get_node(union_all_id).unwrap() {
-        if let Node::Expression(row) = plan.get_node(union_all.output()).unwrap() {
-            assert_eq!(Distribution::Random, *row.distribution().unwrap());
-        } else {
-            panic!("Invalid output!");
-        }
-    } else {
-        panic!("Invalid node!");
-    }
-
-    // Check preserving the original distribution
-    let scan_t3 = Relational::new_scan("t1", &mut plan).unwrap();
-    let scan_t3_id = vec_alloc(&mut plan.nodes, Node::Relational(scan_t3));
-    set_distribution(scan_t3_id, &mut plan).unwrap();
-
-    let union_all = Relational::new_union_all(&mut plan, scan_t1_id, scan_t3_id).unwrap();
-    let union_all_id = vec_alloc(&mut plan.nodes, Node::Relational(union_all));
-    set_distribution(union_all_id, &mut plan).unwrap();
-
-    if let Node::Relational(union_all) = plan.get_node(union_all_id).unwrap() {
-        if let Node::Expression(row) = plan.get_node(union_all.output()).unwrap() {
-            assert_eq!(
-                Distribution::Segment { key: vec![0] },
-                *row.distribution().unwrap()
-            );
-        } else {
-            panic!("Invalid output!");
-        }
-    } else {
-        panic!("Invalid node!");
-    }
-
-    // Check errors for children with different column names
-    let t4 = Table::new_seg(
-        "t4",
-        vec![
-            Column::new("c", Type::Boolean),
-            Column::new("b", Type::Number),
-        ],
-        &["b"],
-    )
-    .unwrap();
-    plan.add_rel(t4);
-
-    let scan_t4 = Relational::new_scan("t4", &mut plan).unwrap();
-    let scan_t4_id = vec_alloc(&mut plan.nodes, Node::Relational(scan_t4));
-    assert_eq!(
-        QueryPlannerError::NotEqualRows,
-        Relational::new_union_all(&mut plan, scan_t4_id, scan_t1_id).unwrap_err()
-    );
-
-    // Check errors for children with different amount of column
-    let t5 = Table::new_seg("t5", vec![Column::new("b", Type::Number)], &["b"]).unwrap();
-    plan.add_rel(t5);
-
-    let scan_t5 = Relational::new_scan("t5", &mut plan).unwrap();
-    let scan_t5_id = vec_alloc(&mut plan.nodes, Node::Relational(scan_t5));
     assert_eq!(
         QueryPlannerError::NotEqualRows,
-        Relational::new_union_all(&mut plan, scan_t5_id, scan_t1_id).unwrap_err()
+        Relational::new_union_all(&mut plan, scan_t2_id, scan_t1_id).unwrap_err()
     );
 }
 
-#[test]
-fn union_all_serialize() {
-    let path = Path::new("")
-        .join("tests")
-        .join("artifactory")
-        .join("ir")
-        .join("operator")
-        .join("union_all.yaml");
-    let s = fs::read_to_string(path).unwrap();
-    Plan::from_yaml(&s).unwrap();
-}
-
 #[test]
 fn sub_query() {
     let mut plan = Plan::empty();
@@ -366,7 +261,7 @@ fn sub_query() {
     // Non-relational child node
     let a = 1;
     assert_eq!(
-        QueryPlannerError::InvalidRow,
+        QueryPlannerError::InvalidNode,
         Relational::new_sub_query(&mut plan, a, "sq").unwrap_err()
     );
 
@@ -388,101 +283,3 @@ fn sub_query_serialize() {
     let s = fs::read_to_string(path).unwrap();
     Plan::from_yaml(&s).unwrap();
 }
-
-#[test]
-fn output_alias_position_map() {
-    let path = Path::new("")
-        .join("tests")
-        .join("artifactory")
-        .join("ir")
-        .join("operator")
-        .join("output_aliases.yaml");
-    let s = fs::read_to_string(path).unwrap();
-    let plan = Plan::from_yaml(&s).unwrap();
-
-    let top = plan.nodes.get(plan.top.unwrap()).unwrap();
-    if let Node::Relational(rel) = top {
-        let col_map = rel.output_alias_position_map(&plan).unwrap();
-
-        let expected_keys = vec!["a", "b"];
-        assert_eq!(expected_keys.len(), col_map.len());
-        expected_keys
-            .iter()
-            .zip(col_map.keys().sorted())
-            .for_each(|(e, k)| assert_eq!(e, k));
-
-        let expected_val = vec![0, 1];
-        expected_val
-            .iter()
-            .zip(col_map.values().sorted())
-            .for_each(|(e, v)| assert_eq!(e, v));
-    } else {
-        panic!("Plan top should be a relational operator!");
-    }
-}
-
-#[test]
-fn output_alias_position_map_duplicates() {
-    let path = Path::new("")
-        .join("tests")
-        .join("artifactory")
-        .join("ir")
-        .join("operator")
-        .join("output_aliases_duplicates.yaml");
-    let s = fs::read_to_string(path).unwrap();
-    let plan = Plan::from_yaml(&s).unwrap();
-
-    let top = plan.nodes.get(plan.top.unwrap()).unwrap();
-    if let Node::Relational(rel) = top {
-        assert_eq!(
-            QueryPlannerError::InvalidPlan,
-            rel.output_alias_position_map(&plan).unwrap_err()
-        );
-    } else {
-        panic!("Plan top should be a relational operator!");
-    }
-}
-
-#[test]
-fn output_alias_position_map_unsupported_type() {
-    let path = Path::new("")
-        .join("tests")
-        .join("artifactory")
-        .join("ir")
-        .join("operator")
-        .join("output_aliases_unsupported_type.yaml");
-    let s = fs::read_to_string(path).unwrap();
-    let plan = Plan::from_yaml(&s).unwrap();
-
-    let top = plan.nodes.get(plan.top.unwrap()).unwrap();
-    if let Node::Relational(rel) = top {
-        assert_eq!(
-            QueryPlannerError::InvalidPlan,
-            rel.output_alias_position_map(&plan).unwrap_err()
-        );
-    } else {
-        panic!("Plan top should be a relational operator!");
-    }
-}
-
-#[test]
-fn output_alias_oor() {
-    let path = Path::new("")
-        .join("tests")
-        .join("artifactory")
-        .join("ir")
-        .join("operator")
-        .join("output_aliases_oor.yaml");
-    let s = fs::read_to_string(path).unwrap();
-    let plan = Plan::from_yaml(&s).unwrap();
-
-    let top = plan.nodes.get(plan.top.unwrap()).unwrap();
-    if let Node::Relational(rel) = top {
-        assert_eq!(
-            QueryPlannerError::ValueOutOfRange,
-            rel.output_alias_position_map(&plan).unwrap_err()
-        );
-    } else {
-        panic!("Plan top should be a relational operator!");
-    }
-}
diff --git a/tests/artifactory/ir/expression/suggest_distribution.yaml b/tests/artifactory/ir/expression/shrink_dist_key_1.yaml
similarity index 70%
rename from tests/artifactory/ir/expression/suggest_distribution.yaml
rename to tests/artifactory/ir/expression/shrink_dist_key_1.yaml
index 1e3fc984fd047a4438066cb6a6067abb77f4341c..84b859355b81669b63b8db0ac927f6d137990c60 100644
--- a/tests/artifactory/ir/expression/suggest_distribution.yaml
+++ b/tests/artifactory/ir/expression/shrink_dist_key_1.yaml
@@ -2,32 +2,36 @@
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 0
+        parent: 0
   - Expression:
       Alias:
         name: a
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 1
+        parent: 0
   - Expression:
       Alias:
         name: b
         child: 2
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 2
+        parent: 0
   - Expression:
       Alias:
         name: c
         child: 4
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 3
+        parent: 0
   - Expression:
       Alias:
         name: d
@@ -39,39 +43,44 @@ nodes:
           - 3
           - 5
           - 7
-        distribution:
-          Segment:
-            key:
-              - 1
-              - 0
+        distribution: ~
   - Relational:
       ScanRelation:
         output: 8
+        id: 0
         relation: t
   - Expression:
       Reference:
-        branch: Left
-        position: 0
+        targets:
+          - 0
+        position: 2
+        parent: 10
   - Expression:
       Alias:
-        name: a
+        name: c
         child: 10
   - Expression:
       Reference:
-        branch: Left
-        position: 1
+        targets:
+          - 0
+        position: 0
+        parent: 10
   - Expression:
       Alias:
-        name: b
+        name: a
         child: 12
   - Expression:
-      Reference:
-        branch: Left
-        position: 2
-  - Expression:
-      Alias:
-        name: c
-        child: 14
+      Row:
+        list:
+          - 11
+          - 13
+        distribution: ~
+  - Relational:
+      Projection:
+        children:
+          - 9
+        id: 10
+        output: 14
 relations:
   t:
     Segment:
@@ -89,4 +98,4 @@ relations:
         - 0
       name: t
 slices: ~
-top: 9
+top: 15
diff --git a/tests/artifactory/ir/expression/shrink_dist_key_2.yaml b/tests/artifactory/ir/expression/shrink_dist_key_2.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..902ca79ffe5df28a88443c7a9085c50e7bca1ce1
--- /dev/null
+++ b/tests/artifactory/ir/expression/shrink_dist_key_2.yaml
@@ -0,0 +1,90 @@
+---
+nodes:
+  - Expression:
+      Reference:
+        targets: ~
+        position: 0
+        parent: 0
+  - Expression:
+      Alias:
+        name: a
+        child: 0
+  - Expression:
+      Reference:
+        targets: ~
+        position: 1
+        parent: 0
+  - Expression:
+      Alias:
+        name: b
+        child: 2
+  - Expression:
+      Reference:
+        targets: ~
+        position: 2
+        parent: 0
+  - Expression:
+      Alias:
+        name: c
+        child: 4
+  - Expression:
+      Reference:
+        targets: ~
+        position: 3
+        parent: 0
+  - Expression:
+      Alias:
+        name: d
+        child: 6
+  - Expression:
+      Row:
+        list:
+          - 1
+          - 3
+          - 5
+          - 7
+        distribution: ~
+  - Relational:
+      ScanRelation:
+        output: 8
+        id: 0
+        relation: t
+  - Expression:
+      Reference:
+        targets:
+          - 0
+        position: 0
+        parent: 10
+  - Expression:
+      Alias:
+        name: a
+        child: 10
+  - Expression:
+      Row:
+        list:
+          - 11
+        distribution: ~
+  - Relational:
+      Projection:
+        children:
+          - 9
+        id: 10
+        output: 12
+relations:
+  t:
+    Segment:
+      columns:
+        - name: a
+          type: Boolean
+        - name: b
+          type: Number
+        - name: c
+          type: String
+        - name: d
+          type: String
+      key:
+        - 1
+        - 0
+      name: t
+slices: ~
+top: 13
diff --git a/tests/artifactory/ir/expression/shuffle_dist_key.yaml b/tests/artifactory/ir/expression/shuffle_dist_key.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..8acc09bf0df6b48f3b35ec7237a8ad948fa09710
--- /dev/null
+++ b/tests/artifactory/ir/expression/shuffle_dist_key.yaml
@@ -0,0 +1,101 @@
+---
+nodes:
+  - Expression:
+      Reference:
+        targets: ~
+        position: 0
+        parent: 0
+  - Expression:
+      Alias:
+        name: a
+        child: 0
+  - Expression:
+      Reference:
+        targets: ~
+        position: 1
+        parent: 0
+  - Expression:
+      Alias:
+        name: b
+        child: 2
+  - Expression:
+      Reference:
+        targets: ~
+        position: 2
+        parent: 0
+  - Expression:
+      Alias:
+        name: c
+        child: 4
+  - Expression:
+      Reference:
+        targets: ~
+        position: 3
+        parent: 0
+  - Expression:
+      Alias:
+        name: d
+        child: 6
+  - Expression:
+      Row:
+        list:
+          - 1
+          - 3
+          - 5
+          - 7
+        distribution: ~
+  - Relational:
+      ScanRelation:
+        output: 8
+        id: 0
+        relation: t
+  - Expression:
+      Reference:
+        targets:
+          - 0
+        position: 1
+        parent: 10
+  - Expression:
+      Alias:
+        name: b
+        child: 10
+  - Expression:
+      Reference:
+        targets:
+          - 0
+        position: 0
+        parent: 10
+  - Expression:
+      Alias:
+        name: a
+        child: 12
+  - Expression:
+      Row:
+        list:
+          - 11
+          - 13
+        distribution: ~
+  - Relational:
+      Projection:
+        children:
+          - 9
+        id: 10
+        output: 14
+relations:
+  t:
+    Segment:
+      columns:
+        - name: a
+          type: Boolean
+        - name: b
+          type: Number
+        - name: c
+          type: String
+        - name: d
+          type: String
+      key:
+        - 1
+        - 0
+      name: t
+slices: ~
+top: 15
diff --git a/tests/artifactory/ir/operator/union_all.yaml b/tests/artifactory/ir/expression/union_fallback_to_random.yaml
similarity index 76%
rename from tests/artifactory/ir/operator/union_all.yaml
rename to tests/artifactory/ir/expression/union_fallback_to_random.yaml
index 8d6869f355a5fe2d16dad7212ab306e3ed96c4ba..e2d9e1125afc63852f18c1e802d680a70d6f87b2 100644
--- a/tests/artifactory/ir/operator/union_all.yaml
+++ b/tests/artifactory/ir/expression/union_fallback_to_random.yaml
@@ -1,18 +1,19 @@
-
 ---
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 0
+        parent: 0
   - Expression:
       Alias:
         name: a
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 1
+        parent: 0
   - Expression:
       Alias:
         name: b
@@ -22,26 +23,26 @@ nodes:
         list:
           - 1
           - 3
-        distribution:
-          Segment:
-            key:
-              - 0
+        distribution: ~
   - Relational:
       ScanRelation:
         output: 4
+        id: 0
         relation: t1
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 0
+        parent: 6
   - Expression:
       Alias:
         name: a
         child: 6
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 1
+        parent: 6
   - Expression:
       Alias:
         name: b
@@ -51,26 +52,30 @@ nodes:
         list:
           - 7
           - 9
-        distribution:
-          Segment:
-            key:
-              - 1
+        distribution: ~
   - Relational:
       ScanRelation:
         output: 10
+        id: 6
         relation: t2
   - Expression:
       Reference:
-        branch: Both
+        targets:
+          - 0
+          - 1
         position: 0
+        parent: 12
   - Expression:
       Alias:
         name: a
         child: 12
   - Expression:
       Reference:
-        branch: Both
+        targets:
+          - 0
+          - 1
         position: 1
+        parent: 12
   - Expression:
       Alias:
         name: b
@@ -80,11 +85,13 @@ nodes:
         list:
           - 13
           - 15
-        distribution: Random
+        distribution: ~
   - Relational:
       UnionAll:
-        left: 5
-        right: 11
+        children:
+          - 5
+          - 11
+        id: 12
         output: 16
 relations:
   t1:
diff --git a/tests/artifactory/ir/expression/union_preserve_dist.yaml b/tests/artifactory/ir/expression/union_preserve_dist.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..ca72f859f50d9173cbec2cb84beca3b352f2cd90
--- /dev/null
+++ b/tests/artifactory/ir/expression/union_preserve_dist.yaml
@@ -0,0 +1,118 @@
+---
+nodes:
+  - Expression:
+      Reference:
+        targets: ~
+        position: 0
+        parent: 0
+  - Expression:
+      Alias:
+        name: a
+        child: 0
+  - Expression:
+      Reference:
+        targets: ~
+        position: 1
+        parent: 0
+  - Expression:
+      Alias:
+        name: b
+        child: 2
+  - Expression:
+      Row:
+        list:
+          - 1
+          - 3
+        distribution: ~
+  - Relational:
+      ScanRelation:
+        output: 4
+        id: 0
+        relation: t1
+  - Expression:
+      Reference:
+        targets: ~
+        position: 0
+        parent: 6
+  - Expression:
+      Alias:
+        name: a
+        child: 6
+  - Expression:
+      Reference:
+        targets: ~
+        position: 1
+        parent: 6
+  - Expression:
+      Alias:
+        name: b
+        child: 8
+  - Expression:
+      Row:
+        list:
+          - 7
+          - 9
+        distribution: ~
+  - Relational:
+      ScanRelation:
+        output: 10
+        id: 6
+        relation: t2
+  - Expression:
+      Reference:
+        targets:
+          - 0
+          - 1
+        position: 0
+        parent: 12
+  - Expression:
+      Alias:
+        name: a
+        child: 12
+  - Expression:
+      Reference:
+        targets:
+          - 0
+          - 1
+        position: 1
+        parent: 12
+  - Expression:
+      Alias:
+        name: b
+        child: 14
+  - Expression:
+      Row:
+        list:
+          - 13
+          - 15
+        distribution: ~
+  - Relational:
+      UnionAll:
+        children:
+          - 5
+          - 11
+        id: 12
+        output: 16
+relations:
+  t1:
+    Segment:
+      columns:
+        - name: a
+          type: Boolean
+        - name: b
+          type: Number
+      key:
+        - 0
+      name: t1
+  t2:
+    Segment:
+      columns:
+        - name: a
+          type: Boolean
+        - name: b
+          type: Number
+      key:
+        - 0
+      name: t2
+slices: ~
+top: 17
diff --git a/tests/artifactory/ir/operator/projection.yaml b/tests/artifactory/ir/operator/projection.yaml
index e1e54148f49c8e350d97f79be114c092386ca7d4..3319f0d415e59aa513767b00e0f445cbc4f972a8 100644
--- a/tests/artifactory/ir/operator/projection.yaml
+++ b/tests/artifactory/ir/operator/projection.yaml
@@ -6,7 +6,8 @@
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 0
   - Expression:
       Alias:
@@ -14,7 +15,8 @@ nodes:
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 1
   - Expression:
       Alias:
@@ -32,10 +34,12 @@ nodes:
   - Relational:
       ScanRelation:
         output: 4
+        id: 0
         relation: t
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 5 
         position: 1
   - Expression:
       Alias:
@@ -51,7 +55,9 @@ nodes:
               - 0
   - Relational:
       Projection:
-        child: 5
+        id: 5
+        children:
+          - 5
         output: 8
 relations:
   t:
diff --git a/tests/artifactory/ir/operator/scan_rel.yaml b/tests/artifactory/ir/operator/scan_rel.yaml
index 88232f6e9016874f8b8fb5f7fd1534ffd276a3b7..113978fc1862d8217afedc66605dcff260fac744 100644
--- a/tests/artifactory/ir/operator/scan_rel.yaml
+++ b/tests/artifactory/ir/operator/scan_rel.yaml
@@ -2,32 +2,36 @@
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 0
+        parent: 0
   - Expression:
       Alias:
         name: a
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 1
+        parent: 0
   - Expression:
       Alias:
         name: b
         child: 2
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 2
+        parent: 0
   - Expression:
       Alias:
         name: c
         child: 4
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
         position: 3
+        parent: 0
   - Expression:
       Alias:
         name: d
@@ -43,10 +47,11 @@ nodes:
           Segment:
             key:
               - 1
-              - 0
+              - 0 
   - Relational:
       ScanRelation:
         output: 8
+        id: 0
         relation: t
 relations:
   t:
diff --git a/tests/artifactory/ir/operator/selection.yaml b/tests/artifactory/ir/operator/selection.yaml
index 9d325ee3d647a29e16eb449758f5b68f97ea9974..b20ea8269870dede6afe4f754b56703211b45156 100644
--- a/tests/artifactory/ir/operator/selection.yaml
+++ b/tests/artifactory/ir/operator/selection.yaml
@@ -2,7 +2,8 @@
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 0
   - Expression:
       Alias:
@@ -10,7 +11,8 @@ nodes:
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 1
   - Expression:
       Alias:
@@ -18,7 +20,8 @@ nodes:
         child: 2
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 2
   - Expression:
       Alias:
@@ -26,7 +29,8 @@ nodes:
         child: 4
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 3
   - Expression:
       Alias:
@@ -46,11 +50,13 @@ nodes:
               - 0
   - Relational:
       ScanRelation:
+        id: 0
         output: 8
         relation: t
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 10 
         position: 1
   - Expression:
       Alias:
@@ -67,7 +73,9 @@ nodes:
         right: 12
   - Expression:
       Reference:
-        branch: Left
+        targets:
+          - 0
+        parent: 10 
         position: 0
   - Expression:
       Alias:
@@ -75,7 +83,9 @@ nodes:
         child: 14
   - Expression:
       Reference:
-        branch: Left
+        targets:
+          - 0
+        parent: 10 
         position: 1
   - Expression:
       Alias:
@@ -83,7 +93,9 @@ nodes:
         child: 16
   - Expression:
       Reference:
-        branch: Left
+        targets:
+          - 0
+        parent: 10 
         position: 2
   - Expression:
       Alias:
@@ -91,7 +103,9 @@ nodes:
         child: 18
   - Expression:
       Reference:
-        branch: Left
+        targets:
+          - 0
+        parent: 10 
         position: 3
   - Expression:
       Alias:
@@ -111,7 +125,9 @@ nodes:
               - 0
   - Relational:
       Selection:
-        child: 9
+        children:
+          - 9
+        id: 10
         filter: 13
         output: 22
 relations:
diff --git a/tests/artifactory/ir/operator/sub_query.yaml b/tests/artifactory/ir/operator/sub_query.yaml
index 456ad751ae4c80f9c3c6c04d97bcf4606e2f24d2..7fb0473eae34df809e6952be9c143a53ddcecc6c 100644
--- a/tests/artifactory/ir/operator/sub_query.yaml
+++ b/tests/artifactory/ir/operator/sub_query.yaml
@@ -2,7 +2,8 @@
 nodes:
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 0
   - Expression:
       Alias:
@@ -10,7 +11,8 @@ nodes:
         child: 0
   - Expression:
       Reference:
-        branch: Left
+        targets: ~
+        parent: 0 
         position: 1
   - Expression:
       Alias:
@@ -28,10 +30,12 @@ nodes:
   - Relational:
       ScanRelation:
         output: 4
+        id: 0
         relation: t
   - Expression:
       Reference:
-        branch: Both
+        targets: ~
+        parent: 6
         position: 0
   - Expression:
       Alias:
@@ -39,7 +43,8 @@ nodes:
         child: 6
   - Expression:
       Reference:
-        branch: Both
+        targets: ~
+        parent: 6
         position: 1
   - Expression:
       Alias:
@@ -57,7 +62,9 @@ nodes:
   - Relational:
       ScanSubQuery:
         alias: sq
-        child: 5
+        id: 6
+        children:
+          - 5
         output: 10
 relations:
   t:
diff --git a/tests/artifactory/ir/plan_no_top.yaml b/tests/artifactory/ir/plan_no_top.yaml
index c64b659e6844f771e6caaae164c1100879222764..300337df6e7276f5659177e2ec3f893ea2749bdf 100644
--- a/tests/artifactory/ir/plan_no_top.yaml
+++ b/tests/artifactory/ir/plan_no_top.yaml
@@ -2,20 +2,22 @@
 nodes:
   - Expression:
       Reference:
-        alias: a
-        branch: Left
+        targets: ~
         position: 0
+        parent: 0
+  - Expression:
+      Alias:
+        name: a
+        child: 0
   - Expression:
       Row:
         list:
-          - 0
-        distribution:
-          Segment:
-            key:
-              - 0
+          - 1
+        distribution: ~
   - Relational:
       ScanRelation:
-        output: 1
+        output: 2
+        id: 0
         relation: t
 relations:
   t:
diff --git a/tests/artifactory/ir/plan_oor_top.yaml b/tests/artifactory/ir/plan_oor_top.yaml
index 1400910621a65f871a5da3d39d9249d51596b2d3..be6b86d6cb7d146b3524f63df255935ce1c5d339 100644
--- a/tests/artifactory/ir/plan_oor_top.yaml
+++ b/tests/artifactory/ir/plan_oor_top.yaml
@@ -2,20 +2,22 @@
 nodes:
   - Expression:
       Reference:
-        alias: a
-        branch: Left
+        targets: ~
         position: 0
+        parent: 0
+  - Expression:
+      Alias:
+        name: a
+        child: 0
   - Expression:
       Row:
         list:
-          - 0
-        distribution:
-          Segment:
-            key:
-              - 0
+          - 1
+        distribution: ~
   - Relational:
       ScanRelation:
-        output: 1
+        output: 2
+        id: 0
         relation: t
 relations:
   t: