feat: Add range table function (risingwavelabs#6646)

feat: Add table function Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
cyliu0 · Nov 29, 2022 · 04514b8 · 04514b8
1 parent 1b4cfdb
commit 04514b8
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 2 deletions.
diff --git a/e2e_test/batch/basic/generate_series.slt → ...test/batch/basic/generate_series.slt.part b/e2e_test/batch/basic/generate_series.slt → ...test/batch/basic/generate_series.slt.part
diff --git a/e2e_test/batch/basic/range.slt.part b/e2e_test/batch/basic/range.slt.part
@@ -0,0 +1,84 @@
+statement ok
+SET RW_IMPLICIT_FLUSH TO true;
+
+query I
+select * from range('2008-03-01 00:00:00'::TIMESTAMP,'2008-03-04 12:00:00'::TIMESTAMP, interval '12' hour)
+----
+2008-03-01 00:00:00
+2008-03-01 12:00:00
+2008-03-02 00:00:00
+2008-03-02 12:00:00
+2008-03-03 00:00:00
+2008-03-03 12:00:00
+2008-03-04 00:00:00
+
+query I
+select * from range(timestamp '2020-01-01 01:02:03', timestamp '2021-01-01 03:04:05', interval '1' month);
+----
+2020-01-01 01:02:03
+2020-02-01 01:02:03
+2020-03-01 01:02:03
+2020-04-01 01:02:03
+2020-05-01 01:02:03
+2020-06-01 01:02:03
+2020-07-01 01:02:03
+2020-08-01 01:02:03
+2020-09-01 01:02:03
+2020-10-01 01:02:03
+2020-11-01 01:02:03
+2020-12-01 01:02:03
+2021-01-01 01:02:03
+
+query I
+select * from range(timestamp '2021-01-01 03:04:05', timestamp '2020-01-01 01:02:03', interval '1' month-interval '60' day);
+----
+2021-01-01 03:04:05
+2020-12-03 03:04:05
+2020-11-04 03:04:05
+2020-10-05 03:04:05
+2020-09-06 03:04:05
+2020-08-07 03:04:05
+2020-07-09 03:04:05
+2020-06-10 03:04:05
+2020-05-11 03:04:05
+2020-04-12 03:04:05
+2020-03-13 03:04:05
+2020-02-13 03:04:05
+2020-01-13 03:04:05
+
+query I
+SELECT * FROM range('2'::INT,'10'::INT,'2'::INT)
+----
+2
+4
+6
+8
+
+query I
+SELECT * FROM range('2'::INT + '2'::INT,'10'::INT,'2'::INT);
+----
+4
+6
+8
+
+query I
+SELECT * FROM range('2'::INT + '2'::INT,'10'::INT,'2'::INT);
+----
+4
+6
+8
+
+query I
+SELECT * FROM range('10'::INT,'2'::INT,'-2'::INT);
+----
+10
+8
+6
+4
+
+statement error
+SELECT * FROM range('2'::INT,'10'::INT,'0'::INT);
+
+query I
+SELECT * FROM range('2'::INT,'10'::INT,'-2'::INT);
+----
diff --git a/src/frontend/src/expr/table_function.rs b/src/frontend/src/expr/table_function.rs
@@ -37,6 +37,7 @@ pub struct TableFunction {
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum TableFunctionType {
     Generate,
+    Range,
     Unnest,
     RegexpMatches,
 }
@@ -45,6 +46,7 @@ impl TableFunctionType {
     fn to_protobuf(self) -> Type {
         match self {
             TableFunctionType::Generate => Type::Generate,
+            TableFunctionType::Range => Type::Range,
             TableFunctionType::Unnest => Type::Unnest,
             TableFunctionType::RegexpMatches => Type::RegexpMatches,
         }
@@ -55,6 +57,7 @@ impl TableFunctionType {
     pub fn name(&self) -> &str {
         match self {
             TableFunctionType::Generate => "generate_series",
+            TableFunctionType::Range => "range",
             TableFunctionType::Unnest => "unnest",
             TableFunctionType::RegexpMatches => "regexp_matches",
         }
@@ -67,6 +70,8 @@ impl FromStr for TableFunctionType {
     fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
         if s.eq_ignore_ascii_case("generate_series") {
             Ok(TableFunctionType::Generate)
+        } else if s.eq_ignore_ascii_case("range") {
+            Ok(TableFunctionType::Range)
         } else if s.eq_ignore_ascii_case("unnest") {
             Ok(TableFunctionType::Unnest)
         } else if s.eq_ignore_ascii_case("regexp_matches") {
@@ -85,7 +90,7 @@ impl TableFunction {
         // Current implementation is copied from legacy code.
 
         match func_type {
-            TableFunctionType::Generate => {
+            function_type @ (TableFunctionType::Generate | TableFunctionType::Range) => {
                 // generate_series ( start timestamp, stop timestamp, step interval ) or
                 // generate_series ( start i32, stop i32, step i32 )
 
@@ -116,7 +121,7 @@ impl TableFunction {
                 Ok(TableFunction {
                     args,
                     return_type: data_type,
-                    function_type: TableFunctionType::Generate,
+                    function_type,
                 })
             }
             TableFunctionType::Unnest => {