Merge pull request #37 from KKould/main

KKould · web-flow · commit 6a66644bc63d · 2023-08-01T23:18:38.000+08:00
feat(sort): Implement Sort operator to support Order By clause
diff --git a/src/binder/select.rs b/src/binder/select.rs
@@ -23,6 +23,7 @@ use sqlparser::ast::{
     Expr, Ident, Join, JoinConstraint, JoinOperator, Offset, OrderByExpr, Query, Select,
     SelectItem, SetExpr, TableFactor, TableWithJoins,
 };
+use crate::planner::operator::sort::{SortField, SortOperator};
 
 impl Binder {
     pub(crate) fn bind_query(&mut self, query: &Query) -> Result<LogicalSelectPlan> {
@@ -90,9 +91,9 @@ impl Binder {
         //     plan = self.bind_distinct(plan, select_list.clone())?;
         // }
 
-        // // if let Some(orderby) = having_orderby.1 {
-        // //     plan = self.bind_sort(plan, orderby)?;
-        // // }
+        if let Some(orderby) = having_orderby.1 {
+            plan = self.bind_sort(plan, orderby);
+        }
 
         plan = self.bind_project(plan, select_list);
         Ok(plan)
@@ -281,6 +282,20 @@ impl Binder {
         }
     }
 
+    fn bind_sort(
+        &mut self,
+        children: LogicalSelectPlan,
+        sort_fields: Vec<SortField>,
+    ) -> LogicalSelectPlan {
+        LogicalSelectPlan {
+            operator: Arc::new(Operator::Sort(SortOperator {
+                sort_fields,
+                limit: None,
+            })),
+            children: vec![Arc::new(children)],
+        }
+    }
+
     fn bind_limit(
         &mut self,
         children: LogicalSelectPlan,
diff --git a/src/db.rs b/src/db.rs
@@ -154,14 +154,14 @@ mod test {
 
         tokio_test::block_on(async move {
             let _ = kipsql.run("create table t1 (a int, b int)").await?;
-            let _ = kipsql.run("insert into t1 values (1, 2), (3, 4)").await?;
+            let _ = kipsql.run("insert into t1 values (1, 1), (2, 3), (5, 4)").await?;
 
             println!("full:");
             let vec_batch_full_fields = kipsql.run("select * from t1").await?;
             print_batches(&vec_batch_full_fields)?;
 
             println!("projection_and_filter:");
-            let vec_batch_projection_a = kipsql.run("select a from t1 where a != 3 and a < b ").await?;
+            let vec_batch_projection_a = kipsql.run("select a from t1 where a <= b order by a desc ").await?;
             print_batches(&vec_batch_projection_a)?;
 
             Ok(())
diff --git a/src/execution_v1/physical_plan/mod.rs b/src/execution_v1/physical_plan/mod.rs
@@ -2,6 +2,7 @@ use crate::execution_v1::physical_plan::physical_create_table::PhysicalCreateTab
 use crate::execution_v1::physical_plan::physical_filter::PhysicalFilter;
 use crate::execution_v1::physical_plan::physical_insert::PhysicalInsert;
 use crate::execution_v1::physical_plan::physical_projection::PhysicalProjection;
+use crate::execution_v1::physical_plan::physical_sort::PhysicalSort;
 use crate::execution_v1::physical_plan::physical_table_scan::PhysicalTableScan;
 use crate::execution_v1::physical_plan::physical_values::PhysicalValues;
 
@@ -12,6 +13,7 @@ pub(crate) mod physical_table_scan;
 pub(crate) mod physical_insert;
 pub(crate) mod physical_values;
 pub(crate) mod physical_filter;
+pub(crate) mod physical_sort;
 
 #[derive(Debug)]
 pub enum PhysicalOperator {
@@ -20,5 +22,6 @@ pub enum PhysicalOperator {
     TableScan(PhysicalTableScan),
     Projection(PhysicalProjection),
     Filter(PhysicalFilter),
+    Sort(PhysicalSort),
     Values(PhysicalValues)
 }
diff --git a/src/execution_v1/physical_plan/physical_filter.rs b/src/execution_v1/physical_plan/physical_filter.rs
@@ -3,7 +3,6 @@ use crate::expression::ScalarExpression;
 
 #[derive(Debug)]
 pub struct PhysicalFilter {
-    pub(crate) plan_id: u32,
     pub(crate) predicate: ScalarExpression,
     pub(crate) input: Box<PhysicalOperator>
 }
diff --git a/src/execution_v1/physical_plan/physical_plan_builder.rs b/src/execution_v1/physical_plan/physical_plan_builder.rs
@@ -11,26 +11,20 @@ use anyhow::anyhow;
 use anyhow::Result;
 use crate::execution_v1::physical_plan::physical_filter::PhysicalFilter;
 use crate::execution_v1::physical_plan::physical_insert::PhysicalInsert;
+use crate::execution_v1::physical_plan::physical_sort::PhysicalSort;
 use crate::execution_v1::physical_plan::physical_values::PhysicalValues;
 use crate::planner::logical_insert_plan::LogicalInsertPlan;
 use crate::planner::operator::filter::FilterOperator;
 use crate::planner::operator::insert::InsertOperator;
 use crate::planner::operator::project::ProjectOperator;
+use crate::planner::operator::sort::SortOperator;
 use crate::planner::operator::values::ValuesOperator;
 
-pub struct PhysicalPlanBuilder {
-    plan_id: u32,
-}
+pub struct PhysicalPlanBuilder { }
 
 impl PhysicalPlanBuilder {
     pub fn new() -> Self {
-        PhysicalPlanBuilder { plan_id: 0 }
-    }
-
-    fn next_plan_id(&mut self) -> u32 {
-        let id = self.plan_id;
-        self.plan_id += 1;
-        id
+        PhysicalPlanBuilder { }
     }
 
     pub fn build_plan(&mut self, plan: &LogicalPlan) -> Result<PhysicalOperator> {
@@ -87,6 +81,7 @@ impl PhysicalPlanBuilder {
             Operator::Project(op) => self.build_physical_select_projection(plan, op),
             Operator::Scan(scan) => Ok(self.build_physical_scan(scan.clone())),
             Operator::Filter(op) => self.build_physical_filter(plan, op),
+            Operator::Sort(op) => self.build_physical_sort(plan, op),
             _ => Err(anyhow!(format!(
                 "Unsupported physical plan: {:?}",
                 plan.operator
@@ -98,23 +93,30 @@ impl PhysicalPlanBuilder {
         let input = self.build_select_logical_plan(plan.child(0)?)?;
 
         Ok(PhysicalOperator::Projection(PhysicalProjection {
-            plan_id: self.next_plan_id(),
             exprs: op.columns.clone(),
             input: Box::new(input),
         }))
     }
 
     fn build_physical_scan(&mut self, base: ScanOperator) -> PhysicalOperator {
-        PhysicalOperator::TableScan(PhysicalTableScan { plan_id: self.next_plan_id(), base })
+        PhysicalOperator::TableScan(PhysicalTableScan { base })
     }
 
     fn build_physical_filter(&mut self, plan: &LogicalSelectPlan, base: &FilterOperator) -> Result<PhysicalOperator> {
         let input = self.build_select_logical_plan(plan.child(0)?)?;
 
         Ok(PhysicalOperator::Filter(PhysicalFilter {
-            plan_id: 0,
             predicate: base.predicate.clone(),
             input: Box::new(input),
         }))
     }
+
+    fn build_physical_sort(&mut self, plan: &LogicalSelectPlan, base: &SortOperator) -> Result<PhysicalOperator> {
+        let input = self.build_select_logical_plan(plan.child(0)?)?;
+
+        Ok(PhysicalOperator::Sort(PhysicalSort {
+            op: base.clone(),
+            input: Box::new(input),
+        }))
+    }
 }
diff --git a/src/execution_v1/physical_plan/physical_projection.rs b/src/execution_v1/physical_plan/physical_projection.rs
@@ -3,7 +3,6 @@ use crate::expression::ScalarExpression;
 
 #[derive(Debug)]
 pub struct PhysicalProjection {
-    pub(crate) plan_id: u32,
     pub(crate) exprs: Vec<ScalarExpression>,
     pub(crate) input: Box<PhysicalOperator>
 }
diff --git a/src/execution_v1/physical_plan/physical_sort.rs b/src/execution_v1/physical_plan/physical_sort.rs
@@ -0,0 +1,8 @@
+use crate::execution_v1::physical_plan::PhysicalOperator;
+use crate::planner::operator::sort::SortOperator;
+
+#[derive(Debug)]
+pub struct PhysicalSort {
+    pub(crate) op: SortOperator,
+    pub(crate) input: Box<PhysicalOperator>
+}
diff --git a/src/execution_v1/physical_plan/physical_table_scan.rs b/src/execution_v1/physical_plan/physical_table_scan.rs
@@ -2,6 +2,5 @@ use crate::planner::operator::scan::ScanOperator;
 
 #[derive(Debug)]
 pub struct PhysicalTableScan {
-    pub(crate) plan_id: u32,
     pub(crate) base: ScanOperator
 }
diff --git a/src/execution_v1/volcano_executor/insert.rs b/src/execution_v1/volcano_executor/insert.rs
@@ -23,8 +23,6 @@ impl Insert {
                 let mut arrays = batch.columns().to_vec();
                 let col_len = arrays[0].len();
 
-                arrays.reverse();
-
                 let full_arrays = table.all_columns()
                     .into_iter()
                     .map(|(_, col_catalog)| {
diff --git a/src/execution_v1/volcano_executor/mod.rs b/src/execution_v1/volcano_executor/mod.rs
@@ -4,6 +4,7 @@ mod table_scan;
 mod insert;
 mod values;
 mod filter;
+mod sort;
 
 use crate::execution_v1::physical_plan::physical_projection::PhysicalProjection;
 use crate::execution_v1::physical_plan::PhysicalOperator;
@@ -17,8 +18,10 @@ use futures::stream::BoxStream;
 use futures::TryStreamExt;
 use crate::execution_v1::physical_plan::physical_filter::PhysicalFilter;
 use crate::execution_v1::physical_plan::physical_insert::PhysicalInsert;
+use crate::execution_v1::physical_plan::physical_sort::PhysicalSort;
 use crate::execution_v1::volcano_executor::filter::Filter;
 use crate::execution_v1::volcano_executor::insert::Insert;
+use crate::execution_v1::volcano_executor::sort::Sort;
 use crate::execution_v1::volcano_executor::values::Values;
 
 pub type BoxedExecutor = BoxStream<'static, Result<RecordBatch, ExecutorError>>;
@@ -61,6 +64,11 @@ impl VolcanoExecutor {
 
                 Filter::execute(predicate, input)
             }
+            PhysicalOperator::Sort(PhysicalSort {op, input, ..}) => {
+                let input = self.build(*input);
+
+                Sort::execute(op.sort_fields, op.limit, input)
+            }
         }
     }
 
diff --git a/src/execution_v1/volcano_executor/sort.rs b/src/execution_v1/volcano_executor/sort.rs
@@ -0,0 +1,60 @@
+use arrow::compute::{concat_batches, lexsort_to_indices, SortColumn, SortOptions, take, TakeOptions};
+use arrow::record_batch::RecordBatch;
+use futures_async_stream::try_stream;
+use crate::execution_v1::ExecutorError;
+use crate::execution_v1::volcano_executor::BoxedExecutor;
+use crate::planner::operator::sort::SortField;
+
+pub struct Sort { }
+
+impl Sort {
+    #[try_stream(boxed, ok = RecordBatch, error = ExecutorError)]
+    pub async fn execute(sort_fields: Vec<SortField>, limit: Option<usize>, input: BoxedExecutor) {
+        let mut schema = None;
+        let mut batches = vec![];
+
+        #[for_await]
+        for batch in input {
+            let batch = batch?;
+            if schema.is_none() {
+                schema = Some(batch.schema());
+            }
+            batches.push(batch);
+        }
+
+        let schema = schema.unwrap();
+        let batch = concat_batches(&schema, &batches)?;
+
+        let sort_columns = sort_fields
+            .into_iter()
+            .map(|SortField { expr, desc, nulls_first }| -> Result<SortColumn, ExecutorError> {
+                let values = expr.eval_column(&batch)?;
+                Ok(SortColumn {
+                    values,
+                    options: Some(SortOptions {
+                        descending: desc,
+                        nulls_first,
+                    }),
+                })
+            })
+            .try_collect::<Vec<_>>()?;
+
+        let indices = lexsort_to_indices(&sort_columns, limit)?;
+
+        let columns = batch
+            .columns()
+            .iter()
+            .map(|column| {
+                take(
+                    column,
+                    &indices,
+                    // for sqlrs: disable bound check overhead since indices are already generated from
+                    // the same record batch
+                    Some(TakeOptions { check_bounds: false, })
+                )
+            })
+            .try_collect()?;
+
+        yield RecordBatch::try_new(schema, columns)?;
+    }
+}
diff --git a/src/planner/operator/sort.rs b/src/planner/operator/sort.rs
@@ -2,9 +2,9 @@ use crate::expression::ScalarExpression;
 
 #[derive(Debug, PartialEq, Clone)]
 pub struct SortField {
-    expr: ScalarExpression,
-    desc: bool,
-    nulls_first: bool,
+    pub expr: ScalarExpression,
+    pub desc: bool,
+    pub nulls_first: bool,
 }
 
 impl SortField {

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,6 @@ use crate::expression::ScalarExpression;`
`3`	`3`
`4`	`4`	`#[derive(Debug)]`
`5`	`5`	`pub struct PhysicalFilter {`
`6`		`- pub(crate) plan_id: u32,`
`7`	`6`	`pub(crate) predicate: ScalarExpression,`
`8`	`7`	`pub(crate) input: Box<PhysicalOperator>`
`9`	`8`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,6 @@ use crate::expression::ScalarExpression;`
`3`	`3`
`4`	`4`	`#[derive(Debug)]`
`5`	`5`	`pub struct PhysicalProjection {`
`6`		`- pub(crate) plan_id: u32,`
`7`	`6`	`pub(crate) exprs: Vec<ScalarExpression>,`
`8`	`7`	`pub(crate) input: Box<PhysicalOperator>`
`9`	`8`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,5 @@ use crate::planner::operator::scan::ScanOperator;`
`2`	`2`
`3`	`3`	`#[derive(Debug)]`
`4`	`4`	`pub struct PhysicalTableScan {`
`5`		`- pub(crate) plan_id: u32,`
`6`	`5`	`pub(crate) base: ScanOperator`
`7`	`6`	`}`