feat(query): expand trim/ltrim/rtrim

databendlabs · Nov 11, 2024 · cf7752c · cf7752c
1 parent f4e599f
commit cf7752c
Show file tree

Hide file tree

Showing 5 changed files with 157 additions and 102 deletions.
diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs
@@ -961,12 +961,22 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         rule! {
             TRIM
             ~ "("
-            ~ #subexpr(0)
+            ~ #subexpr(0) ~ ("," ~ #subexpr(0))?
             ~ ^")"
         },
-        |(_, _, expr, _)| ExprElement::Trim {
-            expr: Box::new(expr),
-            trim_where: None,
+        |(_, _, expr, trim_str, _)| {
+            if let Some(trim_str) = trim_str {
+                let trim_str = trim_str.1;
+                ExprElement::Trim {
+                    expr: Box::new(expr),
+                    trim_where: Some((TrimWhere::Both, Box::new(trim_str))),
+                }
+            } else {
+                ExprElement::Trim {
+                    expr: Box::new(expr),
+                    trim_where: None,
+                }
+            }
         },
     );
     let trim_from = map(

diff --git a/src/query/functions/src/scalars/string.rs b/src/query/functions/src/scalars/string.rs
@@ -430,6 +430,22 @@ pub fn register(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_passthrough_nullable_2_arg::<StringType, StringType, StringType, _, _>(
+        "ltrim",
+        |_, _, _| FunctionDomain::Full,
+        vectorize_with_builder_2_arg::<StringType, StringType, StringType>(
+            |val, trim_str, output, _| {
+                if trim_str.is_empty() {
+                    output.put_and_commit(val);
+                    return;
+                }
+
+                output.put_and_commit(val.trim_start_matches(trim_str));
+            },
+        ),
+    );
+
+
     registry.register_passthrough_nullable_2_arg::<StringType, StringType, StringType, _, _>(
         "trim_leading",
         |_, _, _| FunctionDomain::Full,
@@ -445,6 +461,21 @@ pub fn register(registry: &mut FunctionRegistry) {
         ),
     );
 
+    registry.register_passthrough_nullable_2_arg::<StringType, StringType, StringType, _, _>(
+        "rtrim",
+        |_, _, _| FunctionDomain::Full,
+        vectorize_with_builder_2_arg::<StringType, StringType, StringType>(
+            |val, trim_str, output, _| {
+                if trim_str.is_empty() {
+                    output.put_and_commit(val);
+                    return;
+                }
+
+                output.put_and_commit(val.trim_end_matches(trim_str));
+            },
+        ),
+    );
+
     registry.register_passthrough_nullable_2_arg::<StringType, StringType, StringType, _, _>(
         "trim_trailing",
         |_, _, _| FunctionDomain::Full,

diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -2534,6 +2534,8 @@ Functions overloads:
 35 lte FACTORY
 0 ltrim(String) :: String
 1 ltrim(String NULL) :: String NULL
+2 ltrim(String, String) :: String
+3 ltrim(String NULL, String NULL) :: String NULL
 0 map(Array(Nothing), Array(Nothing)) :: Map(Nothing)
 1 map(Array(Nothing) NULL, Array(Nothing) NULL) :: Map(Nothing) NULL
 2 map(Array(T0), Array(T1)) :: Map(T0, T1)
@@ -3532,6 +3534,8 @@ Functions overloads:
 1 rpad(String NULL, UInt64 NULL, String NULL) :: String NULL
 0 rtrim(String) :: String
 1 rtrim(String NULL) :: String NULL
+2 rtrim(String, String) :: String
+3 rtrim(String NULL, String NULL) :: String NULL
 0 running_difference(Int64) :: Int64
 1 running_difference(Int64 NULL) :: Int64 NULL
 2 running_difference(Date) :: Int32

diff --git a/tests/sqllogictests/suites/base/20+_others/20_0001_planner.test b/tests/sqllogictests/suites/base/20+_others/20_0001_planner.test
@@ -779,104 +779,6 @@ SELECT DISTINCT count(number %3) as c FROM numbers(10)  group by number % 3 ORDE
 3
 4
 
-query T
-select trim(leading ' ' from '      abc')
-----
-abc
-
-
-statement ok
-select trim(leading ' ' from '')
-
-
-statement ok
-select trim(leading 'ab' from 'abab')
-
-
-query T
-select trim(leading 'ab' from 'abc')
-----
-c
-
-
-query T
-select trim(leading ' ' from NULL)
-----
-NULL
-
-
-query T
-select trim(leading NULL from 'aaa')
-----
-NULL
-
-
-query T
-select trim(trailing ' ' from 'abc    ')
-----
-abc
-
-
-statement ok
-select trim(trailing ' ' from '')
-
-
-statement ok
-select trim(trailing 'ab' from 'abab')
-
-
-query T
-select trim(trailing 'ab' from 'cab')
-----
-c
-
-
-query T
-select trim(trailing ' ' from NULL)
-----
-NULL
-
-
-query T
-select trim(trailing NULL from 'aaa')
-----
-NULL
-
-
-statement ok
-select trim(both 'ab' from 'abab')
-
-
-query T
-select trim(both 'ab' from 'abcab')
-----
-c
-
-
-query T
-select trim(both ' ' from NULL)
-----
-NULL
-
-
-query T
-select trim(both NULL from 'aaa')
-----
-NULL
-
-
-query T
-select trim(' abc ')
-----
-abc
-
-
-query T
-select trim(NULL)
-----
-NULL
-
-
 query I
 select [1, 2, 3]
 ----

diff --git a/tests/sqllogictests/suites/query/functions/02_0055_function_strings_trim.test b/tests/sqllogictests/suites/query/functions/02_0055_function_strings_trim.test
@@ -31,12 +31,120 @@ select trim_leading('aaabbaaa', '')
 ----
 aaabbaaa
 
+query T
+select ltrim('aaabbaaa', 'aa')
+----
+abbaaa
+
 query T
 select trim_trailing('aaabbaaa', 'aa')
 ----
 aaabba
 
+query T
+select rtrim('aaabbaaa', 'aa')
+----
+aaabba
+
 query T
 select trim_trailing('aaabbaaa', '')
 ----
 aaabbaaa
+
+query T
+select trim(leading ' ' from '      abc')
+----
+abc
+
+
+query T
+select trim(leading ' ' from '')
+----
+(empty)
+
+statement ok
+select trim(leading 'ab' from 'abab')
+
+
+query T
+select trim(leading 'ab' from 'abc')
+----
+c
+
+
+query T
+select trim(leading ' ' from NULL)
+----
+NULL
+
+
+query T
+select trim(leading NULL from 'aaa')
+----
+NULL
+
+
+query T
+select trim(trailing ' ' from 'abc    ')
+----
+abc
+
+
+query T
+select trim(trailing ' ' from '')
+----
+(empty)
+
+query T
+select trim(trailing 'ab' from 'abab')
+----
+(empty)
+
+query T
+select trim(trailing 'ab' from 'cab')
+----
+c
+
+
+query T
+select trim(trailing ' ' from NULL)
+----
+NULL
+
+
+query T
+select trim(trailing NULL from 'aaa')
+----
+NULL
+
+
+query T
+select trim(both 'ab' from 'abab')
+----
+(empty)
+
+query T
+select trim(both 'ab' from 'abcab')
+----
+c
+
+query T
+select trim('abcab', 'ab')
+----
+c
+
+query T
+select trim(both ' ' from NULL)
+----
+NULL
+
+
+query T
+select trim(both NULL from 'aaa')
+----
+NULL
+
+query T
+select trim(NULL)
+----
+NULL