feat(limit): support physical limit operator

Author: loloxwg

src/binder/select.rs

@@ -394,6 +394,18 @@ mod tests {
             plan_6
         );
 
+        let plan_7 = select_sql_run("select * from t1 limit 1")?;
+        println!(
+            "limit:\n {:#?}",
+            plan_7
+        );
+
+        let plan_8 = select_sql_run("select * from t1 offset 2")?;
+        println!(
+            "offset:\n {:#?}",
+            plan_8
+        );
+
         Ok(())
     }
 }

src/db.rs

@@ -164,6 +164,10 @@ mod test {
             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)?;
 
+            println!("limit:");
+            let vec_batch_limit=kipsql.run("select * from t1  limit 2 offset 1").await?;
+            print_batches(&vec_batch_limit)?;
+
             Ok(())
         })
     }

src/execution/physical/physical_plan_builder.rs

@@ -59,7 +59,7 @@ impl PhysicalPlanBuilder {
                 Ok(PhysicalOperator::Limit(PhysicalLimit {
                     plan_id: self.next_plan_id(),
                     input: Arc::new(input),
-                    limit: limit.count,
+                    limit: limit.limit,
                     offset: limit.offset,
                 }))
             }

src/execution_v1/physical_plan/mod.rs

@@ -1,6 +1,7 @@
 use crate::execution_v1::physical_plan::physical_create_table::PhysicalCreateTable;
 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_limit::PhysicalLimit;
 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;
@@ -14,6 +15,7 @@ pub(crate) mod physical_insert;
 pub(crate) mod physical_values;
 pub(crate) mod physical_filter;
 pub(crate) mod physical_sort;
+pub(crate) mod physical_limit;
 
 #[derive(Debug)]
 pub enum PhysicalOperator {
@@ -23,5 +25,6 @@ pub enum PhysicalOperator {
     Projection(PhysicalProjection),
     Filter(PhysicalFilter),
     Sort(PhysicalSort),
-    Values(PhysicalValues)
+    Values(PhysicalValues),
+    Limit(PhysicalLimit),
 }

src/execution_v1/physical_plan/physical_limit.rs

@@ -0,0 +1,8 @@
+use crate::execution_v1::physical_plan::PhysicalOperator;
+use crate::planner::operator::limit::LimitOperator;
+
+#[derive(Debug)]
+pub struct PhysicalLimit {
+    pub(crate) op: LimitOperator,
+    pub(crate) input: Box<PhysicalOperator>
+}

src/execution_v1/physical_plan/physical_plan_builder.rs

@@ -11,11 +11,13 @@ 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_limit::PhysicalLimit;
 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::limit::LimitOperator;
 use crate::planner::operator::project::ProjectOperator;
 use crate::planner::operator::sort::SortOperator;
 use crate::planner::operator::values::ValuesOperator;
@@ -82,6 +84,7 @@ impl PhysicalPlanBuilder {
             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),
+            Operator::Limit(op)=>self.build_physical_limit(plan,op),
             _ => Err(anyhow!(format!(
                 "Unsupported physical plan: {:?}",
                 plan.operator
@@ -119,4 +122,13 @@ impl PhysicalPlanBuilder {
             input: Box::new(input),
         }))
     }
+
+    fn build_physical_limit(&mut self, plan: &LogicalSelectPlan,base : &LimitOperator)->Result<PhysicalOperator>{
+        let input =self.build_select_logical_plan(plan.child(0)?)?;
+
+        Ok(PhysicalOperator::Limit(PhysicalLimit{
+            op:base.clone(),
+            input: Box::new(input),
+        }))
+    }
 }

src/execution_v1/volcano_executor/limit.rs

@@ -0,0 +1,64 @@
+use arrow::record_batch::RecordBatch;
+use futures_async_stream::try_stream;
+use crate::execution_v1::ExecutorError;
+use crate::execution_v1::volcano_executor::BoxedExecutor;
+
+pub struct Limit {}
+
+impl Limit {
+    #[try_stream(boxed, ok = RecordBatch, error = ExecutorError)]
+    pub async fn execute(offset: Option<usize>, limit: Option<usize>, input: BoxedExecutor) {
+        let offset_val = offset.unwrap_or(0);
+        if limit.is_some() && limit.unwrap() == 0 {
+            return Ok(());
+        }
+
+        let mut returned_count = 0;
+
+        #[for_await]
+        for batch in input {
+            let batch = batch?;
+
+            let cardinality = batch.num_rows() as usize;
+            let limit_val = limit.unwrap_or(cardinality);
+
+            let start = returned_count.max(offset_val) - returned_count;
+            let end = {
+                // from total returned rows level, the total_end is end index of whole returned
+                // rows level.
+                let total_end = offset_val + limit_val;
+                let current_batch_end = returned_count + cardinality;
+                // we choose the min of total_end and current_batch_end as the end index of to
+                // match limit semantics.
+                let real_end = total_end.min(current_batch_end);
+                // to calculate the end index of current batch
+                real_end - returned_count
+            };
+            returned_count += cardinality;
+
+            // example: offset=1000, limit=2, cardinality=100
+            // when first loop:
+            // start = 0.max(1000)-0 = 1000
+            // end = (1000+2).min(0+100)-0 = 100
+            // so, start(1000) > end(100), we skip this loop batch.
+            if start >= end {
+                continue;
+            }
+
+            if (start..end) == (0..cardinality) {
+                yield batch;
+            } else {
+                let length = end - start;
+                yield batch.slice(start as usize, length as usize);
+            }
+
+            // dut to returned_count is always += cardinality, and returned_batch maybe slsliced,
+            // so it will larger than real total_end.
+            // example: offset=1, limit=4, cardinality=6, data=[(0..6)]
+            // returned_count=6 > 1+4, meanwhile returned_batch size is 4 ([0..5])
+            if returned_count >= offset_val + limit_val {
+                break;
+            }
+        }
+    }
+}

src/execution_v1/volcano_executor/mod.rs

@@ -5,6 +5,7 @@ mod insert;
 mod values;
 mod filter;
 mod sort;
+mod limit;
 
 use crate::execution_v1::physical_plan::physical_projection::PhysicalProjection;
 use crate::execution_v1::physical_plan::PhysicalOperator;
@@ -16,11 +17,14 @@ use crate::storage::StorageImpl;
 use arrow::record_batch::RecordBatch;
 use futures::stream::BoxStream;
 use futures::TryStreamExt;
+use serde::de::Unexpected::Option;
 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_limit::PhysicalLimit;
 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::limit::Limit;
 use crate::execution_v1::volcano_executor::sort::Sort;
 use crate::execution_v1::volcano_executor::values::Values;
 
@@ -69,6 +73,11 @@ impl VolcanoExecutor {
 
                 Sort::execute(op.sort_fields, op.limit, input)
             }
+            PhysicalOperator::Limit(PhysicalLimit {op,input, ..}) =>{
+                let input = self.build(*input);
+
+                Limit::execute(Some(op.offset), Some(op.limit), input)
+            }
         }
     }
 

src/planner/operator/limit.rs

@@ -7,13 +7,13 @@ use super::Operator;
 #[derive(Debug, PartialEq, Clone)]
 pub struct LimitOperator {
     pub offset: usize,
-    pub count: usize,
+    pub limit: usize,
 }
 
 impl LimitOperator {
-    pub fn new(offset: usize, count: usize, children: LogicalSelectPlan) -> LogicalSelectPlan {
+    pub fn new(offset: usize, limit: usize, children: LogicalSelectPlan) -> LogicalSelectPlan {
         LogicalSelectPlan {
-            operator: Arc::new(Operator::Limit(LimitOperator { offset, count })),
+            operator: Arc::new(Operator::Limit(LimitOperator { offset, limit })),
             children: vec![Arc::new(children)],
         }
     }