diff --git a/e2e_test/batch/func.slt b/e2e_test/batch/func.slt index 28a68697479ee..05b723f0c8240 100644 --- a/e2e_test/batch/func.slt +++ b/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) diff --git a/proto/expr.proto b/proto/expr.proto index 193e9cf102a8b..d97561ac99d0a 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -68,6 +68,7 @@ message ExprNode { TO_CHAR = 223; MD5 = 224; CHAR_LENGTH = 225; + REPEAT = 226; // Boolean comparison IS_TRUE = 301; diff --git a/src/expr/src/expr/build_expr_from_prost.rs b/src/expr/src/expr/build_expr_from_prost.rs index fec596d503ea6..b53a4773ebcb1 100644 --- a/src/expr/src/expr/build_expr_from_prost.rs +++ b/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 Result { + 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 { let (children, ret_type) = get_children_and_return_type(prost)?; let child = expr_build_from_prost(&children[0])?; diff --git a/src/expr/src/expr/expr_binary_bytes.rs b/src/expr/src/expr/expr_binary_bytes.rs index 2c4295dfe48d8..4a43d2a877519 100644 --- a/src/expr/src/expr/expr_binary_bytes.rs +++ b/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::::new(expr_ia1, expr_ia2, return_type, repeat) + .boxed() +} + #[cfg(test)] mod tests { use risingwave_common::array::{DataChunk, Row}; diff --git a/src/expr/src/expr/mod.rs b/src/expr/src/expr/mod.rs index df3c2f255dd59..a3371ff89a372 100644 --- a/src/expr/src/expr/mod.rs +++ b/src/expr/src/expr/mod.rs @@ -99,6 +99,7 @@ pub fn build_from_prost(prost: &ExprNode) -> Result { 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), diff --git a/src/expr/src/vector_op/mod.rs b/src/expr/src/vector_op/mod.rs index 62b08f09d9bee..6c32b91b59f3c 100644 --- a/src/expr/src/vector_op/mod.rs +++ b/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; diff --git a/src/expr/src/vector_op/repeat.rs b/src/expr/src/vector_op/repeat.rs new file mode 100644 index 0000000000000..0b108347ca8cb --- /dev/null +++ b/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 { + 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(()) + } +} diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index d3e5899b8b574..daa4a00a82542 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/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), diff --git a/src/frontend/src/expr/type_inference.rs b/src/frontend/src/expr/type_inference.rs index 92d659174d90b..cb948cd0a07e9 100644 --- a/src/frontend/src/expr/type_inference.rs +++ b/src/frontend/src/expr/type_inference.rs @@ -326,10 +326,9 @@ fn build_type_derive_map() -> HashMap { 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,