diff --git a/src/ir/transformation.rs b/src/ir/transformation.rs index b025a6bddfc9c12a336811200f3872f96a28743b..3935fadc31b6a633b1ad96123fe176abeeed44b1 100644 --- a/src/ir/transformation.rs +++ b/src/ir/transformation.rs @@ -6,6 +6,7 @@ pub mod bool_in; pub mod dnf; pub mod equality_propagation; pub mod redistribution; +pub mod split_columns; use crate::errors::QueryPlannerError; use crate::ir::expression::Expression; diff --git a/src/ir/transformation/split_columns.rs b/src/ir/transformation/split_columns.rs new file mode 100644 index 0000000000000000000000000000000000000000..f8997503f70db4ff9fb20016bc5520ab029c81c2 --- /dev/null +++ b/src/ir/transformation/split_columns.rs @@ -0,0 +1,117 @@ +//! Split the tuples to columns in the expression trees of the plan. +//! +//! The transformation splits the tuples in the boolean operators into +//! the AND-ed chain of the single-column tuples. +//! +//! For example: +//! ```sql +//! select a from t where (a, 2) = (1, b) +//! ``` +//! is transformed to: +//! ```sql +//! select a from t where (a) = (1) and (2) = (b) +//! ``` + +use crate::errors::QueryPlannerError; +use crate::ir::expression::Expression; +use crate::ir::operator::Bool; +use crate::ir::Plan; + +fn call_expr_tree_split_columns( + plan: &mut Plan, + top_id: usize, +) -> Result<usize, QueryPlannerError> { + plan.expr_tree_replace_bool( + top_id, + &call_split_bool, + &[ + Bool::Eq, + Bool::Gt, + Bool::GtEq, + Bool::Lt, + Bool::LtEq, + Bool::NotEq, + ], + ) +} + +fn call_split_bool(plan: &mut Plan, top_id: usize) -> Result<usize, QueryPlannerError> { + plan.split_bool(top_id) +} + +impl Plan { + /// Replace left and right tuples in the boolean operator with the chain + /// of the AND-ed operators constructed from the tuple's columns. + /// + /// # Errors + /// - If the operator is not a boolean operator. + /// - If left and right tuples have different number of columns. + /// - If the plan is invalid for some unknown reason. + fn split_bool(&mut self, expr_id: usize) -> Result<usize, QueryPlannerError> { + let expr = self.get_expression_node(expr_id)?; + let (left_id, right_id, op) = match expr { + Expression::Bool { + left, op, right, .. + } => (*left, *right, op.clone()), + _ => { + return Err(QueryPlannerError::CustomError(format!( + "Node is not a boolean expression: {:?}", + expr + ))); + } + }; + let left_expr = self.get_expression_node(left_id)?; + let right_expr = self.get_expression_node(right_id)?; + if let ( + Expression::Row { + list: left_list, .. + }, + Expression::Row { + list: right_list, .. + }, + ) = (left_expr, right_expr) + { + if left_list.len() != right_list.len() { + return Err(QueryPlannerError::CustomError(format!( + "Left and right rows have different number of columns: {:?}, {:?}", + left_expr, right_expr + ))); + } + let pairs = left_list + .iter() + .zip(right_list.iter()) + .map(|(l, r)| (*l, *r)) + .collect::<Vec<_>>(); + if let Some((first, other)) = pairs.split_first() { + let left_col_id = self.expr_clone(first.0)?; + let right_col_id = self.expr_clone(first.1)?; + let left_row_id = self.nodes.add_row(vec![left_col_id], None); + let right_row_id = self.nodes.add_row(vec![right_col_id], None); + let mut top_id = self.add_cond(left_row_id, op.clone(), right_row_id)?; + + for (left_col_id, right_col_id) in other { + let left_col_id = self.expr_clone(*left_col_id)?; + let right_col_id = self.expr_clone(*right_col_id)?; + let left_row_id = self.nodes.add_row(vec![left_col_id], None); + let right_row_id = self.nodes.add_row(vec![right_col_id], None); + let new_top_id = self.add_cond(left_row_id, op.clone(), right_row_id)?; + top_id = self.concat_and(top_id, new_top_id)?; + } + + return Ok(top_id); + } + } + Ok(expr_id) + } + + /// Split columns in all the boolean operators of the plan. + /// + /// # Errors + /// - If the plan tree is invalid (doesn't contain correct nodes where we expect it to). + pub fn split_columns(&mut self) -> Result<(), QueryPlannerError> { + self.transform_expr_trees(&call_expr_tree_split_columns) + } +} + +#[cfg(test)] +mod tests; diff --git a/src/ir/transformation/split_columns/tests.rs b/src/ir/transformation/split_columns/tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..26787ab27992634b5020f5eaf3b4f5b9d5147262 --- /dev/null +++ b/src/ir/transformation/split_columns/tests.rs @@ -0,0 +1,109 @@ +use pretty_assertions::assert_eq; + +use crate::executor::engine::mock::MetadataMock; +use crate::executor::ir::ExecutionPlan; +use crate::frontend::sql::ast::AbstractSyntaxTree; + +#[test] +fn split_columns1() { + let query = r#"SELECT "a" FROM "t" WHERE ("a", 2) = (1, "b")"#; + + let metadata = &MetadataMock::new(); + let ast = AbstractSyntaxTree::new(query).unwrap(); + let mut plan = ast.to_ir(metadata).unwrap(); + plan.split_columns().unwrap(); + let ex_plan = ExecutionPlan::from(&plan); + + let top_id = plan.get_top().unwrap(); + let sql = ex_plan.subtree_as_sql(top_id).unwrap(); + assert_eq!( + format!( + "{}", + r#"SELECT "t"."a" as "a" FROM "t" WHERE ("t"."a") = (1) and (2) = ("t"."b")"#, + ), + sql + ); +} + +#[test] +fn split_columns2() { + let query = r#"SELECT "a" FROM "t" WHERE "a" = 1"#; + + let metadata = &MetadataMock::new(); + let ast = AbstractSyntaxTree::new(query).unwrap(); + let mut plan = ast.to_ir(metadata).unwrap(); + plan.split_columns().unwrap(); + let ex_plan = ExecutionPlan::from(&plan); + + let top_id = plan.get_top().unwrap(); + let sql = ex_plan.subtree_as_sql(top_id).unwrap(); + assert_eq!( + format!( + "{}", + r#"SELECT "t"."a" as "a" FROM "t" WHERE ("t"."a") = (1)"#, + ), + sql + ); +} + +#[test] +fn split_columns3() { + let query = r#"SELECT "a" FROM "t" WHERE ("a", 2, "b") = (1, "b")"#; + + let metadata = &MetadataMock::new(); + let ast = AbstractSyntaxTree::new(query).unwrap(); + let mut plan = ast.to_ir(metadata).unwrap(); + let plan_err = plan.split_columns().unwrap_err(); + assert_eq!( + format!( + "{} {} {}", + r#"Left and right rows have different number of columns:"#, + r#"Row { list: [8, 9, 10], distribution: None },"#, + r#"Row { list: [12, 13], distribution: None }"#, + ), + format!("{}", plan_err) + ); +} + +#[test] +fn split_columns4() { + let query = r#"SELECT "a" FROM "t" WHERE "a" in (1, 2)"#; + + let metadata = &MetadataMock::new(); + let ast = AbstractSyntaxTree::new(query).unwrap(); + let mut plan = ast.to_ir(metadata).unwrap(); + plan.split_columns().unwrap(); + let ex_plan = ExecutionPlan::from(&plan); + + let top_id = plan.get_top().unwrap(); + let sql = ex_plan.subtree_as_sql(top_id).unwrap(); + assert_eq!( + format!( + "{}", + r#"SELECT "t"."a" as "a" FROM "t" WHERE ("t"."a") in (1, 2)"#, + ), + sql + ); +} + +#[test] +fn split_columns5() { + let query = r#"SELECT "a" FROM "t" WHERE ("a", 2) < (1, "b") and "a" > 2"#; + + let metadata = &MetadataMock::new(); + let ast = AbstractSyntaxTree::new(query).unwrap(); + let mut plan = ast.to_ir(metadata).unwrap(); + plan.split_columns().unwrap(); + let ex_plan = ExecutionPlan::from(&plan); + + let top_id = plan.get_top().unwrap(); + let sql = ex_plan.subtree_as_sql(top_id).unwrap(); + assert_eq!( + format!( + "{} {}", + r#"SELECT "t"."a" as "a" FROM "t" WHERE ("t"."a") < (1) and (2) < ("t"."b")"#, + r#"and ("t"."a") > (2)"#, + ), + sql + ); +}