From 6f08fc5590a152425fea9a2a843b3783d8d53843 Mon Sep 17 00:00:00 2001
From: Fedomn <fedomn.ma@gmail.com>
Date: Sun, 11 Dec 2022 17:14:43 +0800
Subject: [PATCH] feat(planner): expand star expression

Signed-off-by: Fedomn <fedomn.ma@gmail.com>
---
 src/planner_v2/binder/bind_context.rs         | 43 +++++++++++++++++++
 .../binder/query_node/bind_select_node.rs     | 37 ++++++++++++++--
 .../binder/statement/bind_select.rs           |  1 +
 tests/slt/select.slt                          | 21 +++++++++
 4 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/src/planner_v2/binder/bind_context.rs b/src/planner_v2/binder/bind_context.rs
index 4c269b0..8e5e3df 100644
--- a/src/planner_v2/binder/bind_context.rs
+++ b/src/planner_v2/binder/bind_context.rs
@@ -99,4 +99,47 @@ impl BindContext {
             )))
         }
     }
+
+    pub fn generate_all_column_expressions(
+        &mut self,
+        table_name: Option<String>,
+    ) -> Result<Vec<sqlparser::ast::SelectItem>, BindError> {
+        use sqlparser::ast;
+        let select_items = if let Some(table_name) = table_name {
+            if let Some(binding) = self.get_binding(table_name.as_str()) {
+                binding
+                    .names
+                    .iter()
+                    .map(|col_name| {
+                        ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(vec![
+                            ast::Ident::new(binding.alias.clone()),
+                            ast::Ident::new(col_name.clone()),
+                        ]))
+                    })
+                    .collect::<Vec<_>>()
+            } else {
+                return Err(BindError::Internal(format!(
+                    "Table {} not found in context",
+                    table_name
+                )));
+            }
+        } else {
+            self.binding_list
+                .iter()
+                .flat_map(|binding| {
+                    binding
+                        .names
+                        .iter()
+                        .map(|col_name| {
+                            ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(vec![
+                                ast::Ident::new(binding.alias.clone()),
+                                ast::Ident::new(col_name.clone()),
+                            ]))
+                        })
+                        .collect::<Vec<_>>()
+                })
+                .collect::<Vec<_>>()
+        };
+        Ok(select_items)
+    }
 }
diff --git a/src/planner_v2/binder/query_node/bind_select_node.rs b/src/planner_v2/binder/query_node/bind_select_node.rs
index ab47f2c..c9863c0 100644
--- a/src/planner_v2/binder/query_node/bind_select_node.rs
+++ b/src/planner_v2/binder/query_node/bind_select_node.rs
@@ -2,7 +2,8 @@ use derive_new::new;
 use sqlparser::ast::{Ident, Query};
 
 use crate::planner_v2::{
-    BindError, Binder, BoundExpression, BoundTableRef, ExpressionBinder, VALUES_LIST_ALIAS,
+    BindError, Binder, BoundExpression, BoundTableRef, ExpressionBinder, SqlparserResolver,
+    VALUES_LIST_ALIAS,
 };
 use crate::types_v2::LogicalType;
 
@@ -68,8 +69,13 @@ impl Binder {
 
         let mut result_names = vec![];
         let mut result_types = vec![];
-        let select_list = select
-            .projection
+        // expand any "*" statements
+        let new_select_list = self.expand_star_expressions(select.projection.clone())?;
+        if new_select_list.is_empty() {
+            return Err(BindError::Internal("empty select list".to_string()));
+        }
+
+        let select_list = new_select_list
             .iter()
             .map(|item| self.bind_select_item(item, &mut result_names, &mut result_types))
             .collect::<Result<Vec<_>, _>>()?;
@@ -82,6 +88,31 @@ impl Binder {
         ))
     }
 
+    fn expand_star_expressions(
+        &mut self,
+        select_list: Vec<sqlparser::ast::SelectItem>,
+    ) -> Result<Vec<sqlparser::ast::SelectItem>, BindError> {
+        let mut new_select_list = vec![];
+        for item in select_list {
+            match item {
+                sqlparser::ast::SelectItem::Wildcard(_) => {
+                    let col_exprs = self.bind_context.generate_all_column_expressions(None)?;
+                    new_select_list.extend(col_exprs);
+                }
+                sqlparser::ast::SelectItem::QualifiedWildcard(object_name, _) => {
+                    let (_schema_name, table_name) =
+                        SqlparserResolver::object_name_to_schema_table(&object_name)?;
+                    let col_exprs = self
+                        .bind_context
+                        .generate_all_column_expressions(Some(table_name))?;
+                    new_select_list.extend(col_exprs);
+                }
+                other => new_select_list.push(other),
+            }
+        }
+        Ok(new_select_list)
+    }
+
     fn bind_select_item(
         &mut self,
         item: &sqlparser::ast::SelectItem,
diff --git a/src/planner_v2/binder/statement/bind_select.rs b/src/planner_v2/binder/statement/bind_select.rs
index 0d53ddd..dea07f3 100644
--- a/src/planner_v2/binder/statement/bind_select.rs
+++ b/src/planner_v2/binder/statement/bind_select.rs
@@ -8,6 +8,7 @@ impl Binder {
         match stmt {
             Statement::Query(query) => {
                 let node = self.bind_select_node(query)?;
+                // println!("bind context: {:#?}", self.bind_context);
                 self.create_plan_for_select_node(node)
             }
             _ => Err(BindError::UnsupportedStmt(format!("{:?}", stmt))),
diff --git a/tests/slt/select.slt b/tests/slt/select.slt
index 58734cb..0967db3 100644
--- a/tests/slt/select.slt
+++ b/tests/slt/select.slt
@@ -7,3 +7,24 @@ Gregg CO 2 10000
 John CO 3 11500
 Von (empty) 4 NULL
 
+
+onlyif sqlrs_v2
+statement ok
+create table t1(v1 int, v2 int, v3 int);
+insert into t1(v3, v2, v1) values (0, 4, 1), (1, 5, 2);
+
+
+onlyif sqlrs_v2
+query III
+select t1.v1, v2 from t1;
+----
+1 4
+2 5
+
+
+onlyif sqlrs_v2
+query III
+select *, t1.* from t1;
+----
+1 4 0 1 4 0
+2 5 1 2 5 1