Implement take function in ast_builder (gluesql#1346)

ChobobDev · Jul 29, 2023 · 4cbe8a4 · 4cbe8a4
1 parent 596a150
commit 4cbe8a4
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 3 deletions.
diff --git a/core/src/ast_builder/expr/function.rs b/core/src/ast_builder/expr/function.rs
@@ -89,6 +89,10 @@ pub enum FunctionNode<'a> {
         separator: ExprNode<'a>,
         exprs: ExprList<'a>,
     },
+    Take {
+        expr: ExprNode<'a>,
+        size: ExprNode<'a>,
+    },
     Substr {
         expr: ExprNode<'a>,
         start: ExprNode<'a>,
@@ -262,6 +266,11 @@ impl<'a> TryFrom<FunctionNode<'a>> for Function {
                 let exprs = exprs.try_into()?;
                 Ok(Function::ConcatWs { separator, exprs })
             }
+            FunctionNode::Take { expr, size } => {
+                let expr = expr.try_into()?;
+                let size = size.try_into()?;
+                Ok(Function::Take { expr, size })
+            }
             FunctionNode::Degrees(expr) => expr.try_into().map(Function::Degrees),
             FunctionNode::Radians(expr) => expr.try_into().map(Function::Radians),
             FunctionNode::Exp(expr) => expr.try_into().map(Function::Exp),
@@ -485,6 +494,9 @@ impl<'a> ExprNode<'a> {
     pub fn rpad<T: Into<ExprNode<'a>>>(self, size: T, fill: Option<ExprNode<'a>>) -> ExprNode<'a> {
         rpad(self, size, fill)
     }
+    pub fn take<T: Into<ExprNode<'a>>>(self, size: T) -> ExprNode<'a> {
+        take(self, size)
+    }
     pub fn exp(self) -> ExprNode<'a> {
         exp(self)
     }
@@ -727,6 +739,13 @@ pub fn radians<'a, V: Into<ExprNode<'a>>>(expr: V) -> ExprNode<'a> {
     ExprNode::Function(Box::new(FunctionNode::Radians(expr.into())))
 }
 
+pub fn take<'a, T: Into<ExprNode<'a>>, U: Into<ExprNode<'a>>>(expr: T, size: U) -> ExprNode<'a> {
+    ExprNode::Function(Box::new(FunctionNode::Take {
+        expr: expr.into(),
+        size: size.into(),
+    }))
+}
+
 pub fn exp<'a, V: Into<ExprNode<'a>>>(expr: V) -> ExprNode<'a> {
     ExprNode::Function(Box::new(FunctionNode::Exp(expr.into())))
 }
@@ -914,7 +933,7 @@ mod tests {
             generate_uuid, get_x, get_y, ifnull, initcap, is_empty, last_day, lcm, left, length,
             ln, log, log10, log2, lower, lpad, ltrim, md5, modulo, now, num, pi, point, position,
             power, radians, rand, repeat, replace, reverse, right, round, rpad, rtrim, sign, sin,
-            skip, sqrt, substr, tan, test_expr, text, time, timestamp, to_date, to_time,
+            skip, sqrt, substr, take, tan, test_expr, text, time, timestamp, to_date, to_time,
             to_timestamp, upper,
         },
         prelude::DataType,
@@ -1312,6 +1331,17 @@ mod tests {
         test_expr(actual, expected);
     }
 
+    #[test]
+    fn function_take() {
+        let actual = take(col("list"), num(3));
+        let expected = "TAKE(list,3)";
+        test_expr(actual, expected);
+
+        let actual = expr("list").take(num(3));
+        let expected = "TAKE(list,3)";
+        test_expr(actual, expected);
+    }
+
     #[test]
     fn function_exp() {
         let actual = exp(num(2));

diff --git a/core/src/ast_builder/mod.rs b/core/src/ast_builder/mod.rs
@@ -75,8 +75,8 @@ pub use expr::{
         degrees, divide, exp, extract, find_idx, floor, format, gcd, generate_uuid, get_x, get_y,
         ifnull, initcap, is_empty, last_day, lcm, left, length, ln, log, log10, log2, lower, lpad,
         ltrim, md5, modulo, now, pi, point, position, power, radians, rand, repeat, replace,
-        reverse, right, round, rpad, rtrim, sign, sin, skip, sqrt, substr, tan, to_date, to_time,
-        to_timestamp, upper, FunctionNode,
+        reverse, right, round, rpad, rtrim, sign, sin, skip, sqrt, substr, take, tan, to_date,
+        to_time, to_timestamp, upper, FunctionNode,
     },
 };