Merge pull request gluesql#1441 from PzaThief/feature/implement-Keys-…

…function feat: implement keys function
ChobobDev · Nov 24, 2023 · 58b4d97 · 58b4d97
2 parents ae0809f + c3229a4
commit 58b4d97
Show file tree

Hide file tree

Showing 10 changed files with 119 additions and 28 deletions.
diff --git a/core/src/ast/function.rs b/core/src/ast/function.rs
@@ -69,7 +69,6 @@ pub enum Function {
         filter_chars: Option<Expr>,
         trim_where_field: Option<TrimWhereField>,
     },
-    Entries(Expr),
     Exp(Expr),
     Extract {
         field: DateTimeField,
@@ -201,6 +200,8 @@ pub enum Function {
     },
     IsEmpty(Expr),
     Length(Expr),
+    Entries(Expr),
+    Keys(Expr),
     Values(Expr),
     Splice {
         list_data: Expr,
@@ -467,8 +468,9 @@ impl ToSql for Function {
             }
             Function::IsEmpty(e) => format!("IS_EMPTY({})", e.to_sql()),
             Function::Length(e) => format!("LENGTH({})", e.to_sql()),
-            Function::Values(e) => format!("VALUES({})", e.to_sql()),
             Function::Entries(e) => format!("ENTRIES({})", e.to_sql()),
+            Function::Keys(e) => format!("KEYS({})", e.to_sql()),
+            Function::Values(e) => format!("VALUES({})", e.to_sql()),
             Function::Splice {
                 list_data,
                 begin_index,
@@ -1315,16 +1317,21 @@ mod tests {
         );
 
         assert_eq!(
-            r#"VALUES("map")"#,
-            &Expr::Function(Box::new(Function::Values(Expr::Identifier(
+            r#"ENTRIES("map")"#,
+            &Expr::Function(Box::new(Function::Entries(Expr::Identifier(
                 "map".to_owned()
             ))))
             .to_sql()
         );
 
         assert_eq!(
-            r#"ENTRIES("map")"#,
-            &Expr::Function(Box::new(Function::Entries(Expr::Identifier(
+            r#"KEYS("map")"#,
+            &Expr::Function(Box::new(Function::Keys(Expr::Identifier("map".to_owned())))).to_sql()
+        );
+
+        assert_eq!(
+            r#"VALUES("map")"#,
+            &Expr::Function(Box::new(Function::Values(Expr::Identifier(
                 "map".to_owned()
             ))))
             .to_sql()

diff --git a/core/src/ast_builder/expr/function.rs b/core/src/ast_builder/expr/function.rs
@@ -168,6 +168,7 @@ pub enum FunctionNode<'a> {
     IsEmpty(ExprNode<'a>),
     LastDay(ExprNode<'a>),
     Entries(ExprNode<'a>),
+    Keys(ExprNode<'a>),
     Values(ExprNode<'a>),
 }
 
@@ -385,6 +386,7 @@ impl<'a> TryFrom<FunctionNode<'a>> for Function {
             FunctionNode::IsEmpty(expr) => expr.try_into().map(Function::IsEmpty),
             FunctionNode::LastDay(expr) => expr.try_into().map(Function::LastDay),
             FunctionNode::Entries(expr) => expr.try_into().map(Function::Entries),
+            FunctionNode::Keys(expr) => expr.try_into().map(Function::Keys),
             FunctionNode::Values(expr) => expr.try_into().map(Function::Values),
         }
     }
@@ -558,6 +560,9 @@ impl<'a> ExprNode<'a> {
     pub fn entries(self) -> ExprNode<'a> {
         entries(self)
     }
+    pub fn keys(self) -> ExprNode<'a> {
+        keys(self)
+    }
     pub fn values(self) -> ExprNode<'a> {
         values(self)
     }
@@ -948,6 +953,10 @@ pub fn entries<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
     ExprNode::Function(Box::new(FunctionNode::Entries(expr.into())))
 }
 
+pub fn keys<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
+    ExprNode::Function(Box::new(FunctionNode::Keys(expr.into())))
+}
+
 pub fn values<'a, T: Into<ExprNode<'a>>>(expr: T) -> ExprNode<'a> {
     ExprNode::Function(Box::new(FunctionNode::Values(expr.into())))
 }
@@ -1757,4 +1766,15 @@ mod tests {
         let expected = "VALUES(map)";
         test_expr(actual, expected);
     }
+
+    #[test]
+    fn function_keys() {
+        let actual = f::keys(col("map"));
+        let expected = "KEYS(map)";
+        test_expr(actual, expected);
+
+        let actual = col("map").keys();
+        let expected = "KEYS(map)";
+        test_expr(actual, expected);
+    }
 }
diff --git a/core/src/executor/evaluate/function.rs b/core/src/executor/evaluate/function.rs
@@ -695,14 +695,6 @@ pub fn is_empty<'a>(expr: Evaluated<'_>) -> ControlFlow<Evaluated<'a>> {
     Continue(Evaluated::Value(Value::Bool(length == 0)))
 }
 
-pub fn values<'a>(expr: Evaluated<'_>) -> ControlFlow<Evaluated<'a>> {
-    match expr.try_into().break_if_null()? {
-        Value::Map(m) => Ok(Evaluated::Value(Value::List(m.into_values().collect()))),
-        _ => Err(EvaluateError::MapTypeRequired.into()),
-    }
-    .into_control_flow()
-}
-
 // --- etc ---
 
 pub fn unwrap<'a>(
@@ -1002,6 +994,24 @@ pub fn entries<'a>(name: String, expr: Evaluated<'_>) -> ControlFlow<Evaluated<'
     .into_control_flow()
 }
 
+pub fn keys<'a>(expr: Evaluated<'_>) -> ControlFlow<Evaluated<'a>> {
+    match expr.try_into().break_if_null()? {
+        Value::Map(m) => Ok(Evaluated::Value(Value::List(
+            m.into_keys().map(Value::Str).collect(),
+        ))),
+        _ => Err(EvaluateError::MapTypeRequired.into()),
+    }
+    .into_control_flow()
+}
+
+pub fn values<'a>(expr: Evaluated<'_>) -> ControlFlow<Evaluated<'a>> {
+    match expr.try_into().break_if_null()? {
+        Value::Map(m) => Ok(Evaluated::Value(Value::List(m.into_values().collect()))),
+        _ => Err(EvaluateError::MapTypeRequired.into()),
+    }
+    .into_control_flow()
+}
+
 pub fn splice<'a>(
     name: String,
     list_data: Evaluated<'_>,

diff --git a/core/src/executor/evaluate/mod.rs b/core/src/executor/evaluate/mod.rs
@@ -672,6 +672,7 @@ async fn evaluate_function<'a, 'b: 'a, 'c: 'a, T: GStore>(
         }
         Function::Length(expr) => f::length(name, eval(expr).await?),
         Function::Entries(expr) => f::entries(name, eval(expr).await?),
+        Function::Keys(expr) => f::keys(eval(expr).await?),
         Function::Values(expr) => {
             let expr = eval(expr).await?;
             f::values(expr)

diff --git a/core/src/plan/expr/function.rs b/core/src/plan/expr/function.rs
@@ -61,8 +61,10 @@ impl Function {
             | Self::GetY(expr)
             | Self::IsEmpty(expr)
             | Self::Sort { expr, order: None }
-            | Self::Values(expr)
-            | Self::Dedup(expr) => Exprs::Single([expr].into_iter()),
+            | Self::Dedup(expr)
+            | Self::Entries(expr)
+            | Self::Keys(expr)
+            | Self::Values(expr) => Exprs::Single([expr].into_iter()),
             Self::Left { expr, size: expr2 }
             | Self::Right { expr, size: expr2 }
             | Self::Lpad {
@@ -203,7 +205,6 @@ impl Function {
                 Exprs::VariableArgsWithSingle(once(separator).chain(exprs.iter()))
             }
             Self::Greatest(exprs) => Exprs::VariableArgs(exprs.iter()),
-            Self::Entries(expr) => Exprs::Single([expr].into_iter()),
             Self::Splice {
                 list_data: expr,
                 begin_index: expr2,

diff --git a/core/src/translate/function.rs b/core/src/translate/function.rs
@@ -406,13 +406,6 @@ pub fn translate_function(sql_function: &SqlFunction) -> Result<Expr> {
                 divisor,
             })))
         }
-        "ENTRIES" => {
-            check_len(name, args.len(), 1)?;
-
-            let expr = translate_expr(args[0])?;
-
-            Ok(Expr::Function(Box::new(Function::Entries(expr))))
-        }
         "MOD" => {
             check_len(name, args.len(), 2)?;
 
@@ -633,6 +626,18 @@ pub fn translate_function(sql_function: &SqlFunction) -> Result<Expr> {
                 .collect::<Result<Vec<_>>>()?;
             Ok(Expr::Function(Box::new(Function::Greatest(exprs))))
         }
+        "ENTRIES" => {
+            check_len(name, args.len(), 1)?;
+
+            let expr = translate_expr(args[0])?;
+            Ok(Expr::Function(Box::new(Function::Entries(expr))))
+        }
+        "KEYS" => {
+            check_len(name, args.len(), 1)?;
+
+            let expr = translate_expr(args[0])?;
+            Ok(Expr::Function(Box::new(Function::Keys(expr))))
+        }
         "VALUES" => {
             check_len(name, args.len(), 1)?;
 

diff --git a/test-suite/src/function/keys.rs b/test-suite/src/function/keys.rs
@@ -0,0 +1,45 @@
+use {
+    crate::*,
+    gluesql_core::{error::EvaluateError, prelude::Value},
+};
+
+test_case!(keys, {
+    let g = get_tester!();
+
+    g.run("CREATE TABLE USER (id INTEGER, data MAP);").await;
+    g.run(
+        r#"
+            INSERT INTO USER VALUES 
+            (1, '{"id": 1, "name": "alice", "is_male": false}'),
+            (2, '{"name": "bob"}'),
+            (3, '{}');
+        "#,
+    )
+    .await;
+
+    g.named_test(
+         "return all keys from map by ascending order",
+        r#"SELECT SORT(KEYS(data), 'ASC') as result FROM USER WHERE id=1"#,
+        {
+            Ok(select!(result; Value::List; vec![Value::Str("id".to_owned()), Value::Str("is_male".to_owned()), Value::Str("name".to_owned())]))
+        }
+    ).await;
+    g.named_test(
+        "return one key from map",
+        r#"SELECT KEYS(data) as result FROM USER WHERE id=2"#,
+        Ok(select!(result; Value::List; vec![Value::Str("name".to_owned())])),
+    )
+    .await;
+    g.named_test(
+        "return null from empty map",
+        r#"SELECT KEYS(data) as result FROM USER WHERE id=3"#,
+        Ok(select!(result; Value::List; vec![])),
+    )
+    .await;
+    g.named_test(
+        "return argument type error",
+        r#"SELECT KEYS(id) FROM USER WHERE id=1"#,
+        Err(EvaluateError::MapTypeRequired.into()),
+    )
+    .await;
+});
diff --git a/test-suite/src/function/mod.rs b/test-suite/src/function/mod.rs
@@ -24,6 +24,7 @@ pub mod greatest;
 pub mod ifnull;
 pub mod initcap;
 pub mod is_empty;
+pub mod keys;
 pub mod last_day;
 pub mod left_right;
 pub mod length;

diff --git a/test-suite/src/function/values.rs b/test-suite/src/function/values.rs
@@ -18,14 +18,14 @@ test_case!(values, {
     .await;
 
     g.named_test(
-         "return all values from map by ascending order",
+         "return all values from map by descending order",
         r#"SELECT SORT(VALUES(data), 'DESC') as result FROM USER WHERE id=1"#,
         {
             Ok(select!(result; Value::List; vec![Value::I64(1), Value::Bool(false), Value::Str("alice".to_owned())]))
         }
     ).await;
     g.named_test(
-         "return all values from map by descending order",
+         "return all values from map by ascending order",
         r#"SELECT SORT(VALUES(data), 'ASC') as result FROM USER WHERE id=1"#,
         {
             Ok(select!(result; Value::List; vec![Value::Str("alice".to_owned()), Value::Bool(false), Value::I64(1)]))
@@ -44,7 +44,7 @@ test_case!(values, {
     )
     .await;
     g.named_test(
-        "return arguemnt type error",
+        "return argument type error",
         r#"SELECT VALUES(id) FROM USER WHERE id=1"#,
         Err(EvaluateError::MapTypeRequired.into()),
     )

diff --git a/test-suite/src/lib.rs b/test-suite/src/lib.rs
@@ -102,7 +102,6 @@ macro_rules! generate_store_tests {
         glue!(inline_view, inline_view::inline_view);
         glue!(values, values::values);
         glue!(unary_operator, unary_operator::unary_operator);
-        glue!(function_values, function::values::values);
         glue!(function_upper_lower, function::upper_lower::upper_lower);
         glue!(function_initcap, function::initcap::initcap);
         glue!(function_gcd_lcm, function::gcd_lcm::gcd_lcm);
@@ -158,6 +157,8 @@ macro_rules! generate_store_tests {
         glue!(function_add_month, function::add_month::add_month);
         glue!(function_slice, function::slice::slice);
         glue!(function_entries, function::entries::entries);
+        glue!(function_keys, function::keys::keys);
+        glue!(function_values, function::values::values);
         glue!(join, join::join);
         glue!(join_project, join::project);
         glue!(migrate, migrate::migrate);