RBO Optimizer implementation (#45)

* feat(optimizer): Complete the addition, deletion and modification of `HepGraph`

TODO: HepGraph::to_plan

* feat(optimizer): implement `HepGraph::to_plan`

* feat(optimizer): implement `HepMatcher` to retrieve the corresponding matching `OptExpr` in the `HepGraph`

* feat(optimizer): complete the general framework of the RBO optimizer: `HepOptimizer` rules are divided into `RuleBatch` for matching and apply

* feat(optimizer): physical operator: `TableScan`, support columns conditional scanning

* refactor(physical_plan): `PhysicalPlanBuilder` -> `PhysicalPlanMapping`, `PhysicalPlanMapping` avoids `clone` and reduces performance overhead

* feat(storage::memory): `InMemoryTransaction` supports `Columns` projection and Limit offset

* fix(insert): fix the table binding of idents when inserting

* feat(rule): `columns_pruning` RuleBatch Added `PUSH_PROJECT_THROUGH_CHILD_RULE` Rule to push down the Project to reduce the IO generated by redundant columns

* feat(rule): `combine_operators` RuleBatch Added `COLLAPSE_PROJECT_RULE` and `COMBINE_FILTERS_RULE` to minimize the calculation amount of Projection and merge Filter calculation

* feat(rule): the new `pushdown_limit` series rules are used to optimize the related structure of the Limit operator

* feat(rule): added Rule `PushPredicateThroughJoin`, which is used to push the Filter condition before the Join, and only works for the JoinType of Inner/Left/Right

* fix(hash_join): add return early

* style(join): code fmt

* feat(graph): use `version` to judge `HepGraph` changes

* docs: rewrite README.md

* style(rule): use `matches!` replace `match`

* style(rule): fix note miss

Author: KKould

Cargo.toml

@@ -41,6 +41,7 @@ async-backtrace = "0.2.6"
 futures = "0.3.25"
 futures-lite = "1.12.0"
 ahash = "0.8.3"
+lazy_static = "1.4.0"
 
 [dev-dependencies]
 tokio-test = "0.4.2"

README.md

@@ -13,10 +13,10 @@ impl Binder {
     pub(crate) fn bind_expr(&mut self, expr: &Expr) -> Result<ScalarExpression> {
         match expr {
             Expr::Identifier(ident) => {
-                self.bind_column_ref_from_identifiers(slice::from_ref(ident))
+                self.bind_column_ref_from_identifiers(slice::from_ref(ident), None)
             }
             Expr::CompoundIdentifier(idents) => {
-                self.bind_column_ref_from_identifiers(idents)
+                self.bind_column_ref_from_identifiers(idents, None)
             }
             Expr::BinaryOp { left, right, op} => {
                 self.bind_binary_op_internal(left, right, op)
@@ -33,6 +33,7 @@ impl Binder {
     pub fn bind_column_ref_from_identifiers(
         &mut self,
         idents: &[Ident],
+        bind_table_name: Option<&String>,
     ) -> Result<ScalarExpression> {
         let idents = idents
             .iter()
@@ -54,7 +55,7 @@ impl Binder {
             }
         };
 
-        if let Some(table) = table_name {
+        if let Some(table) = table_name.or(bind_table_name) {
             let table_catalog = self
                 .context
                 .catalog

src/binder/expr.rs

@@ -33,8 +33,12 @@ impl Binder {
                     .map(|(_, catalog)| catalog.clone())
                     .collect_vec();
             } else {
+                let bind_table_name = Some(table.name.to_string());
                 for ident in idents {
-                    match self.bind_column_ref_from_identifiers(slice::from_ref(ident))? {
+                    match self.bind_column_ref_from_identifiers(
+                        slice::from_ref(ident),
+                        bind_table_name.as_ref()
+                    )? {
                         ScalarExpression::ColumnRef(catalog) => col_catalogs.push(catalog),
                         _ => unreachable!()
                     }

src/binder/insert.rs

@@ -115,3 +115,38 @@ pub enum BindError {
     #[error("catalog error")]
     CatalogError(#[from] CatalogError),
 }
+
+#[cfg(test)]
+pub mod test {
+    use crate::catalog::{ColumnCatalog, ColumnDesc, RootCatalog};
+    use crate::planner::LogicalPlan;
+    use crate::types::LogicalType::Integer;
+    use anyhow::Result;
+    use crate::binder::{Binder, BinderContext};
+
+    fn test_root_catalog() -> Result<RootCatalog> {
+        let mut root = RootCatalog::new();
+
+        let cols_t1 = vec![
+            ColumnCatalog::new("c1".to_string(), false, ColumnDesc::new(Integer, true)),
+            ColumnCatalog::new("c2".to_string(), false, ColumnDesc::new(Integer, false)),
+        ];
+        let _ = root.add_table("t1".to_string(), cols_t1)?;
+
+        let cols_t2 = vec![
+            ColumnCatalog::new("c3".to_string(), false, ColumnDesc::new(Integer, true)),
+            ColumnCatalog::new("c4".to_string(), false, ColumnDesc::new(Integer, false)),
+        ];
+        let _ = root.add_table("t2".to_string(), cols_t2)?;
+        Ok(root)
+    }
+
+    pub fn select_sql_run(sql: &str) -> Result<LogicalPlan> {
+        let root = test_root_catalog()?;
+
+        let binder = Binder::new(BinderContext::new(root));
+        let stmt = crate::parser::parse_sql(sql).unwrap();
+
+        Ok(binder.bind(&stmt[0])?)
+    }
+}

src/binder/mod.rs

@@ -23,6 +23,7 @@ use sqlparser::ast::{
     Expr, Ident, Join, JoinConstraint, JoinOperator, Offset, OrderByExpr, Query, Select,
     SelectItem, SetExpr, TableFactor, TableWithJoins,
 };
+use crate::execution_v1::volcano_executor::join::joins_nullable;
 use crate::expression::BinaryOperator;
 use crate::planner::LogicalPlan;
 use crate::planner::operator::join::JoinCondition;
@@ -154,15 +155,15 @@ impl Binder {
                     )));
                 }
 
-                let table_ref_id = self
+                let table_catalog = self
                     .context
-                    .catalog
-                    .get_table_id_by_name(table)
+                    .catalog.get_table_by_name(table)
                     .ok_or_else(|| anyhow::Error::msg(format!("bind table {}", table)))?;
+                let table_ref_id = table_catalog.id;
 
                 self.context.bind_table.insert(table.into(), (table_ref_id, joint_type));
 
-                (table_ref_id, ScanOperator::new(table_ref_id))
+                (table_ref_id, ScanOperator::new(table_ref_id, &table_catalog))
             }
             _ => unimplemented!(),
         };
@@ -354,13 +355,7 @@ impl Binder {
 
         for (table_id, join_option) in bind_tables.values() {
             if let Some(join_type) = join_option {
-                let (left_force_nullable, right_force_nullable) = match join_type {
-                    JoinType::Inner => (false, false),
-                    JoinType::Left => (false, true),
-                    JoinType::Right => (true, false),
-                    JoinType::Full => (true, true),
-                    JoinType::Cross => (true, true),
-                };
+                let (left_force_nullable, right_force_nullable) = joins_nullable(join_type);
                 table_force_nullable.insert(*table_id, right_force_nullable);
                 left_table_force_nullable = left_force_nullable;
             } else {
@@ -498,36 +493,7 @@ impl Binder {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::binder::BinderContext;
-    use crate::catalog::{ColumnCatalog, ColumnDesc, RootCatalog};
-    use crate::planner::LogicalPlan;
-    use crate::types::LogicalType::Integer;
-
-    fn test_root_catalog() -> Result<RootCatalog> {
-        let mut root = RootCatalog::new();
-
-        let cols_t1 = vec![
-            ColumnCatalog::new("c1".to_string(), false, ColumnDesc::new(Integer, true)),
-            ColumnCatalog::new("c2".to_string(), false, ColumnDesc::new(Integer, false)),
-        ];
-        let _ = root.add_table("t1".to_string(), cols_t1)?;
-
-        let cols_t2 = vec![
-            ColumnCatalog::new("c3".to_string(), false, ColumnDesc::new(Integer, true)),
-            ColumnCatalog::new("c4".to_string(), false, ColumnDesc::new(Integer, false)),
-        ];
-        let _ = root.add_table("t2".to_string(), cols_t2)?;
-        Ok(root)
-    }
-
-    fn select_sql_run(sql: &str) -> Result<LogicalPlan> {
-        let root = test_root_catalog()?;
-
-        let binder = Binder::new(BinderContext::new(root));
-        let stmt = crate::parser::parse_sql(sql).unwrap();
-
-        Ok(binder.bind(&stmt[0])?)
-    }
+    use crate::binder::test::select_sql_run;
 
     #[test]
     fn test_select_bind() -> Result<()> {

src/binder/select.rs

@@ -3,7 +3,7 @@ use sqlparser::ast::{ColumnDef, ColumnOption};
 
 use crate::types::{ColumnId, IdGenerator, LogicalType, TableId};
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone)]
 pub struct ColumnCatalog {
     pub id: ColumnId,
     pub name: String,
@@ -12,6 +12,19 @@ pub struct ColumnCatalog {
     pub desc: ColumnDesc,
 }
 
+// Tips: When there is a Join, the on condition in the Join has a nullable condition,
+// and the nullable in the Projection will change after being affected by the JoinType,
+// so that Eq does not necessarily match, so only use the ID as the matching criterion
+impl PartialEq<Self> for ColumnCatalog {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
+}
+
+impl Eq for ColumnCatalog {
+
+}
+
 impl ColumnCatalog {
     pub(crate) fn new(column_name: String, nullable: bool, column_desc: ColumnDesc) -> ColumnCatalog {
         ColumnCatalog {

src/catalog/column.rs

@@ -6,7 +6,7 @@ pub(crate) use self::root::*;
 pub(crate) use self::table::*;
 
 /// The type of catalog reference.
-pub type CatalogRef = Arc<RootCatalog>;
+pub type RootRef = Arc<RootCatalog>;
 
 pub(crate) static DEFAULT_DATABASE_NAME: &str = "kipsql";
 pub(crate) static DEFAULT_SCHEMA_NAME: &str = "kipsql";

src/catalog/mod.rs

@@ -4,9 +4,13 @@ use arrow::record_batch::RecordBatch;
 use sqlparser::parser::ParserError;
 
 use crate::binder::{BindError, Binder, BinderContext};
-use crate::execution_v1::physical_plan::physical_plan_builder::PhysicalPlanBuilder;
+use crate::execution_v1::physical_plan::physical_plan_mapping::PhysicalPlanMapping;
 use crate::execution_v1::volcano_executor::VolcanoExecutor;
+use crate::optimizer::heuristic::batch::HepBatchStrategy;
+use crate::optimizer::heuristic::optimizer::HepOptimizer;
+use crate::optimizer::rule::RuleImpl;
 use crate::parser::parse_sql;
+use crate::planner::LogicalPlan;
 use crate::storage::memory::InMemoryStorage;
 use crate::storage::{Storage, StorageError, StorageImpl};
 
@@ -44,20 +48,60 @@ impl Database {
         ///   Sort(a)
         ///     Limit(1)
         ///       Project(a,b)
-        let logical_plan = binder.bind(&stmts[0])?;
-        // println!("logic plan: {:#?}", logical_plan);
+        let source_plan = binder.bind(&stmts[0])?;
+        // println!("source_plan plan: {:#?}", source_plan);
 
-        let mut builder = PhysicalPlanBuilder::new();
-        let operator = builder.build_plan(&logical_plan)?;
-        // println!("operator: {:#?}", operator);
+        let best_plan = Self::default_optimizer(source_plan)
+            .find_best();
+        // println!("best_plan plan: {:#?}", best_plan);
+
+        let physical_plan = PhysicalPlanMapping::build_plan(best_plan)?;
+        // println!("physical_plan: {:#?}", physical_plan);
 
         let storage = StorageImpl::InMemoryStorage(self.storage.clone());
         let executor = VolcanoExecutor::new(storage);
 
-        let mut stream = executor.build(operator);
+        let mut stream = executor.build(physical_plan);
 
         Ok(VolcanoExecutor::try_collect(&mut stream).await?)
     }
+
+    fn default_optimizer(source_plan: LogicalPlan) -> HepOptimizer {
+        HepOptimizer::new(source_plan)
+            .batch(
+                "Predicate pushdown".to_string(),
+                HepBatchStrategy::fix_point_topdown(10),
+                vec![
+                    RuleImpl::PushPredicateThroughJoin
+                ]
+            )
+            .batch(
+                "Limit pushdown".to_string(),
+                HepBatchStrategy::fix_point_topdown(10),
+                vec![
+                    RuleImpl::LimitProjectTranspose,
+                    RuleImpl::PushLimitThroughJoin,
+                    RuleImpl::PushLimitIntoTableScan,
+                    RuleImpl::EliminateLimits,
+                ],
+            )
+            .batch(
+                "Column pruning".to_string(),
+                HepBatchStrategy::fix_point_topdown(10),
+                vec![
+                    RuleImpl::PushProjectThroughChild,
+                    RuleImpl::PushProjectIntoScan
+                ]
+            )
+            .batch(
+                "Combine operators".to_string(),
+                HepBatchStrategy::fix_point_topdown(10),
+                vec![
+                    RuleImpl::CollapseProject,
+                    RuleImpl::CombineFilter
+                ]
+            )
+    }
 }
 
 #[derive(thiserror::Error, Debug)]

src/db.rs

@@ -9,7 +9,7 @@ use crate::execution_v1::physical_plan::physical_table_scan::PhysicalTableScan;
 use crate::execution_v1::physical_plan::physical_values::PhysicalValues;
 
 pub(crate) mod physical_create_table;
-pub(crate) mod physical_plan_builder;
+pub(crate) mod physical_plan_mapping;
 pub(crate) mod physical_projection;
 pub(crate) mod physical_table_scan;
 pub(crate) mod physical_insert;