feat/header: support set chunk header (#818)

* feat/header: support set chunk header Signed-off-by: caicancai <2356672992@qq.com> * feat/header: support set chunk header Signed-off-by: caicancai <2356672992@qq.com> * feat/header: support set chunk header Signed-off-by: caicancai <2356672992@qq.com> * feat/header: support set chunk header Signed-off-by: caicancai <2356672992@qq.com> * feat/header: support set chunk header Signed-off-by: caicancai <2356672992@qq.com> --------- Signed-off-by: caicancai <2356672992@qq.com>
risinglightdb · Dec 5, 2023 · 57bedd5 · 57bedd5
1 parent 803f041
commit 57bedd5
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/src/binder/mod.rs b/src/binder/mod.rs
@@ -8,6 +8,7 @@ use std::vec::Vec;
 use egg::{Id, Language};
 use itertools::Itertools;
 
+use crate::array;
 use crate::catalog::{RootCatalog, TableRefId, DEFAULT_SCHEMA_NAME};
 use crate::parser::*;
 use crate::planner::{Expr as Node, RecExpr, TypeError, TypeSchemaAnalysis};
@@ -101,6 +102,23 @@ pub struct Binder {
     table_occurrences: HashMap<TableRefId, u32>,
 }
 
+pub fn bind_header(mut chunk: array::Chunk, stmt: &Statement) -> array::Chunk {
+    let header_values = match stmt {
+        Statement::CreateTable { .. } => vec!["$create".to_string()],
+        Statement::Drop { .. } => vec!["$drop".to_string()],
+        Statement::Insert { .. } => vec!["$insert.row_counts".to_string()],
+        Statement::Explain { .. } => vec!["$explain".to_string()],
+        Statement::Delete { .. } => vec!["$delete.row_counts".to_string()],
+        _ => Vec::new(),
+    };
+
+    if !header_values.is_empty() {
+        chunk.set_header(header_values);
+    }
+
+    chunk
+}
+
 /// The context of binder execution.
 #[derive(Debug, Default)]
 struct Context {

diff --git a/src/db.rs b/src/db.rs
@@ -8,6 +8,7 @@ use risinglight_proto::rowset::block_statistics::BlockStatisticsType;
 use crate::array::{
     ArrayBuilder, ArrayBuilderImpl, Chunk, DataChunk, I32ArrayBuilder, StringArrayBuilder,
 };
+use crate::binder::bind_header;
 use crate::catalog::{RootCatalogRef, TableRefId, INTERNAL_SCHEMA_NAME};
 use crate::parser::{parse, ParserError, Statement};
 use crate::planner::Statistics;
@@ -210,8 +211,9 @@ impl Database {
             if self.handle_set(&stmt)? {
                 continue;
             }
+
             let mut binder = crate::binder::Binder::new(self.catalog.clone());
-            let bound = binder.bind(stmt)?;
+            let bound = binder.bind(stmt.clone())?;
             let optimized = optimizer.optimize(&bound);
             let executor = match self.storage.clone() {
                 StorageImpl::InMemoryStorage(s) => {
@@ -222,8 +224,8 @@ impl Database {
                 }
             };
             let output = executor.try_collect().await?;
-            let chunk = Chunk::new(output);
-            // TODO: set name
+            let mut chunk = Chunk::new(output);
+            chunk = bind_header(chunk, &stmt);
             outputs.push(chunk);
         }
         Ok(outputs)