feat(expr): support repeat function (#3148)

feat(expr): support repeat function

e2e_test/batch/func.slt

@@ -148,3 +148,13 @@ query I
 select char_length('abcdefghijklmnopqrstuvwxyz');
 ----
 26
+
+query I
+select repeat('hello', 3);
+----
+hellohellohello
+
+query I
+select repeat('hello', -1);
+----
+(empty)

proto/expr.proto

@@ -68,6 +68,7 @@ message ExprNode {
     TO_CHAR = 223;
     MD5 = 224;
     CHAR_LENGTH = 225;
+    REPEAT = 226;
 
     // Boolean comparison
     IS_TRUE = 301;

src/expr/src/expr/build_expr_from_prost.rs

@@ -17,7 +17,7 @@ use risingwave_common::types::{DataType, ToOwnedDatum};
 use risingwave_pb::expr::expr_node::RexNode;
 use risingwave_pb::expr::ExprNode;
 
-use crate::expr::expr_binary_bytes::{new_substr_start, new_to_char};
+use crate::expr::expr_binary_bytes::{new_repeat, new_substr_start, new_to_char};
 use crate::expr::expr_binary_nonnull::{new_binary_expr, new_like_default};
 use crate::expr::expr_binary_nullable::new_nullable_binary_expr;
 use crate::expr::expr_case::{CaseExpression, WhenClause};
@@ -73,6 +73,14 @@ pub fn build_nullable_binary_expr_prost(prost: &ExprNode) -> Result<BoxedExpress
     ))
 }
 
+pub fn build_repeat_expr(prost: &ExprNode) -> Result<BoxedExpression> {
+    let (children, ret_type) = get_children_and_return_type(prost)?;
+    ensure!(children.len() == 2);
+    let left_expr = expr_build_from_prost(&children[0])?;
+    let right_expr = expr_build_from_prost(&children[1])?;
+    Ok(new_repeat(left_expr, right_expr, ret_type))
+}
+
 pub fn build_substr_expr(prost: &ExprNode) -> Result<BoxedExpression> {
     let (children, ret_type) = get_children_and_return_type(prost)?;
     let child = expr_build_from_prost(&children[0])?;

src/expr/src/expr/expr_binary_bytes.rs

@@ -20,6 +20,7 @@ use risingwave_common::types::DataType;
 use super::Expression;
 use crate::expr::template::BinaryBytesExpression;
 use crate::expr::BoxedExpression;
+use crate::vector_op::repeat::repeat;
 use crate::vector_op::substr::*;
 use crate::vector_op::to_char::to_char_timestamp;
 
@@ -67,6 +68,15 @@ pub fn new_to_char(
     .boxed()
 }
 
+pub fn new_repeat(
+    expr_ia1: BoxedExpression,
+    expr_ia2: BoxedExpression,
+    return_type: DataType,
+) -> BoxedExpression {
+    BinaryBytesExpression::<Utf8Array, I32Array, _>::new(expr_ia1, expr_ia2, return_type, repeat)
+        .boxed()
+}
+
 #[cfg(test)]
 mod tests {
     use risingwave_common::array::{DataChunk, Row};

src/expr/src/expr/mod.rs

@@ -99,6 +99,7 @@ pub fn build_from_prost(prost: &ExprNode) -> Result<BoxedExpression> {
         Trim => build_trim_expr(prost),
         Ltrim => build_ltrim_expr(prost),
         Rtrim => build_rtrim_expr(prost),
+        Repeat => build_repeat_expr(prost),
         ConcatWs => ConcatWsExpression::try_from(prost).map(Expression::boxed),
         SplitPart => build_split_part_expr(prost),
         ConstantValue => LiteralExpression::try_from(prost).map(Expression::boxed),

src/expr/src/vector_op/mod.rs

@@ -27,6 +27,7 @@ pub mod lower;
 pub mod ltrim;
 pub mod md5;
 pub mod position;
+pub mod repeat;
 pub mod replace;
 pub mod round;
 pub mod rtrim;

src/expr/src/vector_op/repeat.rs

@@ -0,0 +1,54 @@
+// Copyright 2022 Singularity Data
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::iter;
+
+use risingwave_common::array::{BytesGuard, BytesWriter};
+
+use crate::Result;
+
+#[inline(always)]
+pub fn repeat(s: &str, count: i32, writer: BytesWriter) -> Result<BytesGuard> {
+    let repeated = iter::repeat(s)
+        .take(count.try_into().unwrap_or(0))
+        .flat_map(|s| s.chars());
+    writer.write_from_char_iter(repeated).map_err(Into::into)
+}
+
+#[cfg(test)]
+mod tests {
+    use risingwave_common::array::{Array, ArrayBuilder, Utf8ArrayBuilder};
+
+    use super::*;
+
+    #[test]
+    fn test_repeat() -> Result<()> {
+        let cases = vec![
+            ("hello, world", 1, "hello, world"),
+            ("114514", 3, "114514114514114514"),
+            ("ssss", 0, ""),
+            ("ssss", -114514, ""),
+        ];
+
+        for (s, count, expected) in cases {
+            let builder = Utf8ArrayBuilder::new(1).unwrap();
+            let writer = builder.writer();
+            let guard = repeat(s, count, writer).unwrap();
+            let array = guard.into_inner().finish().unwrap();
+            let v = array.value_at(0).unwrap();
+            assert_eq!(v, expected);
+        }
+        Ok(())
+    }
+}

src/frontend/src/binder/expr/function.rs

@@ -101,6 +101,7 @@ impl Binder {
                 }
                 "char_length" => ExprType::CharLength,
                 "character_length" => ExprType::CharLength,
+                "repeat" => ExprType::Repeat,
                 _ => {
                     return Err(ErrorCode::NotImplemented(
                         format!("unsupported function: {:?}", function_name),

src/frontend/src/expr/type_inference.rs

@@ -326,10 +326,9 @@ fn build_type_derive_map() -> HashMap<FuncSign, DataTypeName> {
     for e in [E::Trim, E::Ltrim, E::Rtrim] {
         map.insert(FuncSign::new(e, vec![T::Varchar, T::Varchar]), T::Varchar);
     }
-    map.insert(
-        FuncSign::new(E::Substr, vec![T::Varchar, T::Int32]),
-        T::Varchar,
-    );
+    for e in [E::Repeat, E::Substr] {
+        map.insert(FuncSign::new(e, vec![T::Varchar, T::Int32]), T::Varchar);
+    }
     map.insert(
         FuncSign::new(E::Substr, vec![T::Varchar, T::Int32, T::Int32]),
         T::Varchar,