Merge pull request #41 from lykia-rs/feature/planner

Planner improvements

lykiadb-lang/src/ast/visitor.rs

@@ -1,9 +1,5 @@
 use super::{expr::Expr, stmt::Stmt};
 
-pub trait ExprEvaluator<O, E> {
-    fn eval(&mut self, e: &Expr) -> Result<O, E>;
-}
-
 pub trait Visitor<O, E> {
     fn visit_expr(&self, e: &Expr) -> Result<O, E>;
     fn visit_stmt(&self, s: &Stmt) -> Result<O, E>;

lykiadb-server/src/engine/interpreter.rs

@@ -1,6 +1,6 @@
 use lykiadb_lang::ast::expr::{Expr, Operation, RangeKind};
 use lykiadb_lang::ast::stmt::Stmt;
-use lykiadb_lang::ast::visitor::{ExprEvaluator, VisitorMut};
+use lykiadb_lang::ast::visitor::VisitorMut;
 use lykiadb_lang::parser::program::Program;
 use lykiadb_lang::parser::resolver::Resolver;
 use lykiadb_lang::parser::Parser;
@@ -12,6 +12,7 @@ use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
 use super::error::ExecutionError;
+use super::stdlib::stdlib;
 
 use crate::plan::planner::Planner;
 use crate::util::{alloc_shared, Shared};
@@ -177,9 +178,26 @@ pub struct Interpreter {
     current_program: Option<Arc<Program>>,
 }
 
-impl ExprEvaluator<RV, HaltReason> for Interpreter {
-    fn eval(&mut self, e: &Expr) -> Result<RV, HaltReason> {
-        self.visit_expr(e)
+pub struct ExecutionContext {
+    interpreter: Shared<Interpreter>
+}
+
+impl ExecutionContext {
+    pub fn new(out: Option<Shared<Output>>) -> ExecutionContext {
+        let mut env_man = Environment::new();
+        let native_fns = stdlib(out.clone());
+        let env = env_man.top();
+
+        for (name, value) in native_fns {
+            env_man.declare(env, name.to_string(), value);
+        }
+        ExecutionContext {
+            interpreter: alloc_shared(Interpreter::new(alloc_shared(env_man))),
+        }
+    }
+
+    pub fn eval(&mut self, e: &Expr) -> Result<RV, HaltReason> {
+        self.interpreter.write().unwrap().visit_expr(e)
     }
 }
 
@@ -326,7 +344,7 @@ impl VisitorMut<RV, HaltReason> for Interpreter {
             | Expr::Insert { .. }
             | Expr::Update { .. }
             | Expr::Delete { .. } => {
-                let mut planner = Planner::new();
+                let mut planner = Planner::new(ExecutionContext::new(None));
                 let plan = planner.build(e)?;
                 Ok(RV::Undefined)
             }

lykiadb-server/src/plan/mod.rs

@@ -3,7 +3,7 @@ use std::fmt::Display;
 use lykiadb_lang::{
     ast::{
         expr::Expr,
-        sql::{SqlCollectionIdentifier, SqlJoinType, SqlOrdering},
+        sql::{SqlCollectionIdentifier, SqlCompoundOperator, SqlJoinType, SqlOrdering},
     },
     Identifier,
 };
@@ -27,6 +27,12 @@ pub enum Plan {
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub enum Node {
+    Compound {
+        source: Box<Node>,
+        operator: SqlCompoundOperator,
+        right: Box<Node>,
+    },
+
     Aggregate {
         source: Box<Node>,
         group_by: Vec<Expr>,
@@ -109,6 +115,19 @@ impl Node {
             Node::Scan { source, filter } => {
                 write!(f, "{}- scan [{} as {}]{}", indent_str, source.name, source.alias.as_ref().unwrap_or(&source.name), Self::NEWLINE)
             }
+            Node::Compound { source, operator, right } => {
+                write!(f, "{}- compound [{:?}]{}", indent_str, operator, Self::NEWLINE)?;
+                source._fmt_recursive(f, indent + 1)?;
+                right._fmt_recursive(f, indent + 1)
+            }
+            Node::Limit { source, limit } => {
+                write!(f, "{}- limit {}{}", indent_str, limit, Self::NEWLINE)?;
+                source._fmt_recursive(f, indent + 1)
+            }
+            Node::Offset { source, offset } => {
+                write!(f, "{}- offset {}{}", indent_str, offset, Self::NEWLINE)?;
+                source._fmt_recursive(f, indent + 1)
+            }
             Node::Join {
                 left,
                 join_type,

lykiadb-server/src/plan/planner.rs

@@ -1,21 +1,19 @@
-use crate::engine::interpreter::HaltReason;
+use crate::engine::interpreter::{ExecutionContext, HaltReason};
 use lykiadb_lang::ast::{
     expr::Expr,
-    sql::{SqlFrom, SqlJoinType, SqlSelect},
+    sql::{SqlFrom, SqlJoinType, SqlSelect, SqlSelectCore}
 };
 
 use super::{Node, Plan};
-pub struct Planner;
-
-impl Default for Planner {
-    fn default() -> Self {
-        Self::new()
-    }
+pub struct Planner {
+    context: ExecutionContext
 }
 
 impl Planner {
-    pub fn new() -> Planner {
-        Planner {}
+    pub fn new(ctx: ExecutionContext) -> Planner {
+        Planner {
+            context: ctx
+        }
     }
 
     pub fn build(&mut self, expr: &Expr) -> Result<Plan, HaltReason> {
@@ -32,25 +30,46 @@ impl Planner {
         }
     }
 
-    fn build_select(&mut self, query: &SqlSelect) -> Result<Node, HaltReason> {
+    fn build_select_core(&mut self, core: &SqlSelectCore) -> Result<Node, HaltReason> {
         let mut node: Node = Node::Nothing;
         // FROM/JOIN
-        if let Some(from) = &query.core.from {
+        if let Some(from) = &core.from {
             node = self.build_from(from)?;
         }
         // WHERE
-        if let Some(predicate) = &query.core.r#where {
+        if let Some(predicate) = &core.r#where {
             // TODO: Traverse expression
             node = Node::Filter {
                 source: Box::new(node),
                 predicate: *predicate.clone(),
             }
         }
+        if let Some(compound) = &core.compound {
+            node = Node::Compound { 
+                source: Box::new(node),
+                operator: compound.operator.clone(),
+                right: Box::new(self.build_select_core(&compound.core)?)
+            }
+        }
         // GROUP BY
         // HAVING
         // SELECT
+        Ok(node)
+    }
+
+    fn build_select(&mut self, query: &SqlSelect) -> Result<Node, HaltReason> {
+        let mut node: Node = self.build_select_core(&query.core)?;
+
         // ORDER BY
-        // LIMIT/OFFSET
+
+        if let Some(limit) = &query.limit {
+            if let Some(offset) = &limit.offset {
+                node = Node::Offset { source: Box::new(node), offset: self.context.eval(&offset)?.as_number().expect("Offset is not correct").floor() as usize }
+            }
+
+            node = Node::Limit { source: Box::new(node), limit: self.context.eval(&limit.count)?.as_number().expect("Limit is not correct").floor() as usize }
+        }
+
         Ok(node)
     }
 
@@ -99,10 +118,13 @@ impl Planner {
 pub mod test_helpers {
     use lykiadb_lang::{ast::stmt::Stmt, parser::program::Program};
 
+    use crate::engine::interpreter::ExecutionContext;
+
     use super::Planner;
 
     pub fn expect_plan(query: &str, expected_plan: &str) {
-        let mut planner = Planner::new();
+        let ctx = ExecutionContext::new(None);
+        let mut planner = Planner::new(ctx);
         let program = query.parse::<Program>().unwrap();
         match *program.get_root() {
             Stmt::Program { body, .. } if matches!(body.get(0), Some(Stmt::Expression { .. })) => {

lykiadb-server/tests/planner/compound

@@ -0,0 +1,58 @@
+#[name=simple_union, run=plan]>
+
+SELECT * FROM books
+UNION
+SELECT * FROM books;
+
+---
+
+- compound [Union]
+  - scan [books as books]
+  - scan [books as books]
+
+#[name=simple_intersect, run=plan]>
+
+SELECT * FROM books
+INTERSECT
+SELECT * FROM books;
+
+---
+
+- compound [Intersect]
+  - scan [books as books]
+  - scan [books as books]
+
+
+#[name=simple_except, run=plan]>
+
+SELECT * FROM books
+EXCEPT
+SELECT * FROM books;
+
+---
+
+- compound [Except]
+  - scan [books as books]
+  - scan [books as books]
+
+
+#[name=nested, run=plan]>
+
+SELECT * FROM books where id > 5
+UNION
+SELECT * FROM books
+INTERSECT
+SELECT * FROM books
+EXCEPT
+SELECT * FROM books;
+
+---
+
+- compound [Union]
+  - filter (id Greater Num(5.0)):
+    - scan [books as books]
+  - compound [Intersect]
+    - scan [books as books]
+    - compound [Except]
+      - scan [books as books]
+      - scan [books as books]

lykiadb-server/tests/planner/filter

@@ -0,0 +1,8 @@
+#[name=simple, run=plan]>
+
+SELECT * FROM books b where title like '%hello%';
+
+---
+
+- filter (title Like Str("%hello%")):
+  - scan [books as b]

lykiadb-server/tests/planner/limit_offset

@@ -0,0 +1,16 @@
+#[name=limit, run=plan]>
+
+SELECT * FROM books limit 4 + 4;
+
+---
+- limit 8
+  - scan [books as books]
+
+#[name=limit, run=plan]>
+
+SELECT * FROM books limit 5 + 5 offset 5 + 15;
+
+---
+- limit 10
+  - offset 20
+    - scan [books as books]

lykiadb-server/tests/util.rs

@@ -1,9 +1,10 @@
 use lykiadb_lang::{ast::stmt::Stmt, parser::program::Program};
-use lykiadb_server::{assert_plan, plan::planner::Planner};
+use lykiadb_server::{engine::interpreter::ExecutionContext, plan::planner::Planner};
 use pretty_assertions::assert_eq;
 
 fn expect_plan(query: &str, expected_plan: &str) {
-    let mut planner = Planner::new();
+    let ctx = ExecutionContext::new(None);
+    let mut planner = Planner::new(ctx);
     let program = query.parse::<Program>().unwrap();
     match *program.get_root() {
         Stmt::Program { body, .. } if matches!(body.get(0), Some(Stmt::Expression { .. })) => {