refactor: bind ORM queries through binder plans

Author: KKould

Cargo.toml

@@ -34,7 +34,7 @@ rocksdb = ["dep:rocksdb"]
 unsafe_txdb_checkpoint = ["rocksdb", "dep:librocksdb-sys"]
 lmdb    = ["dep:lmdb", "dep:lmdb-sys"]
 pprof   = ["pprof/criterion", "pprof/flamegraph"]
-python  = ["dep:pyo3"]
+python  = ["parser", "dep:pyo3"]
 shell   = ["parser", "dep:comfy-table", "dep:rustyline"]
 
 [[bench]]

Makefile

@@ -18,7 +18,7 @@ test:
 
 ## Run Python binding API tests implemented with pyo3.
 test-python:
-	PYO3_PYTHON=$(PYO3_PYTHON) $(CARGO) test --features python test_python_
+	PYO3_PYTHON=$(PYO3_PYTHON) $(CARGO) test --features python,decimal test_python_
 
 ## Perform a `cargo check` across the workspace.
 cargo-check:

README.md

@@ -55,7 +55,7 @@ For the full ORM guide, see [`src/orm/README.md`](src/orm/README.md).
 ```rust
 use kite_sql::db::DataBaseBuilder;
 use kite_sql::errors::DatabaseError;
-use kite_sql::orm::{BoundExpressionOps, OrmQueryResultExt};
+use kite_sql::orm::OrmQueryResultExt;
 use kite_sql::Model;
 
 #[derive(Default, Debug, PartialEq, Model)]
@@ -114,8 +114,9 @@ fn main() -> Result<(), DatabaseError> {
             ctx.from::<User>()?
                 .filter(|e| e.column(User::age())?.gte(18))?
                 .project_scalars((User::id(), User::name()))?
-                .asc_by(User::name())?
-                .limit(10)
+                .order_by(User::name())?
+                .limit(10)?
+                .finish()
         })?
         .project_tuple::<(i32, String)>();
 

examples/hello_world.rs

@@ -16,7 +16,6 @@
 mod app {
     use kite_sql::db::{DataBaseBuilder, Database};
     use kite_sql::errors::DatabaseError;
-    use kite_sql::orm::BoundExpressionOps;
     use kite_sql::storage::Storage;
     use kite_sql::Model;
     use std::env;

kite_sql_serde_macros/src/projection.rs

@@ -99,7 +99,9 @@ pub(crate) fn handle(ast: DeriveInput) -> Result<TokenStream, Error> {
                 T: ::kite_sql::storage::Transaction,
                 A: AsRef<[(&'static str, ::kite_sql::types::value::DataValue)]>,
             {
-                Ok(::std::vec![#(#projection_exprs),*])
+                Ok(::std::vec![
+                    #(::kite_sql::orm::IntoOrmScalarExpression::into_orm_scalar(#projection_exprs)),*
+                ])
             }
         }
 

src/binder/expr.rs

@@ -214,10 +214,14 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T
         let (sub_query, column, correlated) =
             self.bind_subquery_plan_with_output(None, arena, build)?;
         let sub_query = ScalarSubqueryOperator::build(sub_query);
-        let (expr, sub_query) = if !self.context.is_step(&QueryBindStep::Where) {
-            self.bind_temp_table(column, sub_query, arena)?
-        } else {
-            (column, sub_query)
+        let (expr, sub_query) = match self.context.step_now() {
+            QueryBindStep::Where => (column, sub_query),
+            QueryBindStep::Project => self.bind_temp_table(column, sub_query, arena)?,
+            _ => {
+                return Err(DatabaseError::UnsupportedStmt(
+                    "scalar subqueries can only appear in `WHERE` or SELECT list".to_string(),
+                ))
+            }
         };
         self.context.sub_query(SubQueryType::SubQuery {
             plan: sub_query,
@@ -238,6 +242,12 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T
             &mut PlanArena<'arena>,
         ) -> Result<LogicalPlan, DatabaseError>,
     {
+        if !self.context.is_step(&QueryBindStep::Where) {
+            return Err(DatabaseError::UnsupportedStmt(
+                "EXISTS subqueries can only appear in `WHERE`".to_string(),
+            ));
+        }
+
         let (sub_query, correlated) = self.bind_subquery_plan(arena, build)?;
         let (_, marker_ref) = self.bind_temp_table_alias(
             ScalarExpression::Constant(DataValue::Boolean(true)),

src/binder/mod.rs

@@ -12,6 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+macro_rules! with_query_bind_step {
+    ($binder:expr, $step:expr, $body:block) => {{
+        let current_step = $binder.context.step_now();
+        $binder.context.step($step);
+        let result = (|| -> Result<_, DatabaseError> { Ok($body) })();
+        $binder.context.step(current_step);
+        result
+    }};
+}
+
+pub(crate) use with_query_bind_step;
+
 pub mod aggregate;
 mod alter_table;
 mod analyze;

src/binder/parser.rs

@@ -17,7 +17,7 @@ use super::select::{
     BindPlanHaving, BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput,
     TableAliasInput,
 };
-use super::{is_valid_identifier, Binder, QueryBindStep, SetOperatorKind};
+use super::{is_valid_identifier, with_query_bind_step, Binder, QueryBindStep, SetOperatorKind};
 #[cfg(feature = "copy")]
 use crate::binder::copy::{ExtSource, FileFormat};
 use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName};
@@ -1376,12 +1376,11 @@ where
         self,
         projection: Vec<SelectItem>,
     ) -> Result<BindPlanSelectList<'s, 'a, 'b, 'arena, T, A>, DatabaseError> {
-        let select_bind_step = self.binder.context.step_now();
-        self.binder.context.step(QueryBindStep::Project);
-        let select_list = self.binder.normalize_select_item(projection, self.arena)?;
-        self.binder.context.step(select_bind_step);
+        let select_list = with_query_bind_step!(self.binder, QueryBindStep::Project, {
+            self.binder.normalize_select_item(projection, self.arena)?
+        });
 
-        Ok(self.select_list(select_list))
+        Ok(self.select_list(select_list?))
     }
 }
 
@@ -1395,8 +1394,9 @@ where
         selection: Option<Expr>,
     ) -> Result<BindPlanFiltered<'s, 'a, 'b, 'arena, T, A>, DatabaseError> {
         let predicate = if let Some(predicate) = selection {
-            self.binder.context.step(QueryBindStep::Where);
-            Some(self.binder.bind_expr(predicate, self.arena)?)
+            Some(with_query_bind_step!(self.binder, QueryBindStep::Where, {
+                self.binder.bind_expr(predicate, self.arena)?
+            })?)
         } else {
             None
         };
@@ -1416,34 +1416,42 @@ where
         having: Option<Expr>,
         orderby: Option<Vec<OrderByExpr>>,
     ) -> Result<BindPlanAggregated<'s, 'a, 'b, 'arena, T, A>, DatabaseError> {
-        let group_by = match group_by {
-            GroupByExpr::Expressions(group_by_exprs, modifiers) => {
-                if !modifiers.is_empty() {
+        let group_by = with_query_bind_step!(self.binder, QueryBindStep::Agg, {
+            match group_by {
+                GroupByExpr::Expressions(group_by_exprs, modifiers) => {
+                    if !modifiers.is_empty() {
+                        return Err(DatabaseError::UnsupportedStmt(
+                            "GROUP BY modifiers are not supported".to_string(),
+                        ));
+                    }
+                    group_by_exprs
+                        .into_iter()
+                        .map(|expr| self.binder.bind_expr(expr, self.arena))
+                        .collect::<Result<Vec<_>, DatabaseError>>()?
+                }
+                GroupByExpr::All(_) => {
                     return Err(DatabaseError::UnsupportedStmt(
-                        "GROUP BY modifiers are not supported".to_string(),
-                    ));
+                        "GROUP BY ALL is not supported".to_string(),
+                    ))
                 }
-                group_by_exprs
-                    .into_iter()
-                    .map(|expr| self.binder.bind_expr(expr, self.arena))
-                    .collect::<Result<Vec<_>, DatabaseError>>()?
             }
-            GroupByExpr::All(_) => {
-                return Err(DatabaseError::UnsupportedStmt(
-                    "GROUP BY ALL is not supported".to_string(),
-                ))
-            }
-        };
+        })?;
         let having = having
-            .map(|having| self.binder.bind_expr(having, self.arena))
+            .map(|having| {
+                with_query_bind_step!(self.binder, QueryBindStep::Having, {
+                    self.binder.bind_expr(having, self.arena)?
+                })
+            })
             .transpose()?;
         self.aggregate(group_by, having, orderby, |binder, arena, orderby| {
             let OrderByExpr { expr, options, .. } = orderby;
-            Ok(SortField::new(
-                binder.bind_expr(expr, arena)?,
-                options.asc.is_none_or(|asc| asc),
-                options.nulls_first.unwrap_or(false),
-            ))
+            with_query_bind_step!(binder, QueryBindStep::Sort, {
+                SortField::new(
+                    binder.bind_expr(expr, arena)?,
+                    options.asc.is_none_or(|asc| asc),
+                    options.nulls_first.unwrap_or(false),
+                )
+            })
         })
     }
 }
@@ -2570,9 +2578,9 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
         predicate: Expr,
         arena: &mut PlanArena,
     ) -> Result<LogicalPlan, DatabaseError> {
-        self.context.step(QueryBindStep::Where);
-
-        let predicate = self.bind_expr(predicate, arena)?;
+        let predicate = with_query_bind_step!(self, QueryBindStep::Where, {
+            self.bind_expr(predicate, arena)?
+        })?;
 
         self.bind_where_expr(children, predicate, arena)
     }

src/binder/select.rs

@@ -363,6 +363,27 @@ where
         })
     }
 
+    pub(crate) fn aggregate_without_group(self) -> Result<Self, DatabaseError> {
+        let sorted = self
+            .filter_expr(None)?
+            .aggregate(
+                Vec::new(),
+                None,
+                None::<Vec<SortField>>,
+                |_binder, _arena, order| Ok(order),
+            )?
+            .having()?
+            .distinct(false)?
+            .order_by()?;
+        Ok(BindPlanSelectList {
+            binder: sorted.binder,
+            arena: sorted.arena,
+            plan: sorted.plan,
+            select_list: sorted.select_list,
+            _marker: std::marker::PhantomData,
+        })
+    }
+
     pub(crate) fn having_expr(mut self, expr: ScalarExpression) -> Result<Self, DatabaseError> {
         self.plan = self.binder.bind_having(self.plan, expr, self.arena)?;
         Ok(self)
@@ -401,6 +422,9 @@ where
     }
 
     pub fn finish(self) -> Result<LogicalPlan, DatabaseError> {
+        if self.select_list.iter().any(ScalarExpression::has_agg_call) {
+            return self.aggregate_without_group()?.finish();
+        }
         self.binder
             .bind_project(self.plan, self.select_list, self.arena)
     }

src/expression/mod.rs

@@ -19,6 +19,7 @@ use crate::expression::function::scala::ScalarFunction;
 use crate::expression::function::table::TableFunction;
 use crate::expression::visitor::{walk_expr, Visitor};
 use crate::expression::visitor_mut::VisitorMut;
+use crate::planner::operator::sort::SortField;
 use crate::planner::{MetaArena, PlanArena};
 use crate::types::evaluator::{
     binary_create, cast_create, unary_create, BinaryEvaluatorRef, CastEvaluatorRef,
@@ -379,6 +380,22 @@ impl Visitor<'_> for HasCountStar {
 }
 
 impl ScalarExpression {
+    pub fn asc(self) -> SortField {
+        SortField::from(self).asc()
+    }
+
+    pub fn desc(self) -> SortField {
+        SortField::from(self).desc()
+    }
+
+    pub fn nulls_first(self) -> SortField {
+        SortField::from(self).nulls_first()
+    }
+
+    pub fn nulls_last(self) -> SortField {
+        SortField::from(self).nulls_last()
+    }
+
     pub fn column_expr(column: ColumnRef, position: usize) -> ScalarExpression {
         ScalarExpression::ColumnRef { column, position }
     }