Merge pull request #7571 from sundy-li/function-v2-repeat

feat(query): add function v2 repeat

Author: mergify[bot]

src/query/functions-v2/src/scalars/string.rs

@@ -24,6 +24,7 @@ use common_expression::types::GenericMap;
 use common_expression::types::NumberType;
 use common_expression::types::StringType;
 use common_expression::vectorize_with_builder_1_arg;
+use common_expression::vectorize_with_builder_2_arg;
 use common_expression::FunctionProperty;
 use common_expression::FunctionRegistry;
 use common_expression::Value;
@@ -534,6 +535,26 @@ pub fn register(registry: &mut FunctionRegistry) {
         ),
     );
 
+    const MAX_REPEAT_TIMES: u64 = 1000000;
+    registry.register_passthrough_nullable_2_arg::<StringType, NumberType<u64>, StringType, _, _>(
+        "repeat",
+        FunctionProperty::default(),
+        |_, _| None,
+        vectorize_with_builder_2_arg::<StringType, NumberType<u64>, StringType>(
+            |a, times, builder| {
+                if times > MAX_REPEAT_TIMES {
+                    return Err(format!(
+                        "Too many times to repeat: ({}), maximum is: {}",
+                        times, MAX_REPEAT_TIMES
+                    ));
+                }
+                (0..times).for_each(|_| builder.put_slice(a));
+                builder.commit_row();
+                Ok(())
+            },
+        ),
+    );
+
     registry.register_passthrough_nullable_1_arg::<StringType, StringType, _, _>(
         "unhex",
         FunctionProperty::default(),

src/query/functions-v2/tests/it/scalars/string.rs

@@ -55,6 +55,7 @@ fn test_string() {
     test_char(file);
     test_soundex(file);
     test_ord(file);
+    test_repeat(file);
     test_insert(file);
 }
 
@@ -688,6 +689,14 @@ fn test_ord(file: &mut impl Write) {
     run_ast(file, "ord('💖')", &[]);
 }
 
+fn test_repeat(file: &mut impl Write) {
+    run_ast(file, "repeat('3', NULL)", &[]);
+    run_ast(file, "repeat('3', 5)", &[]);
+    run_ast(file, "repeat('3', 1000001)", &[]);
+    let table = [("a", DataType::String, Column::from_data(&["a", "b", "c"]))];
+    run_ast(file, "repeat(a, 3)", &table);
+}
+
 fn test_insert(file: &mut impl Write) {
     run_ast(file, "insert('Quadratic', 3, 4, 'What', 4)", &[]);
     run_ast(file, "insert('Quadratic', 3, 4)", &[]);

src/query/functions-v2/tests/it/scalars/testdata/string.txt

@@ -2723,6 +2723,55 @@ output domain  : Unknown
 output         : 4036989590
 
 
+ast            : repeat('3', NULL)
+raw expr       : repeat("3", NULL)
+checked expr   : repeat<String NULL, NULL>(CAST("3" AS String NULL), NULL)
+optimized expr : NULL
+output type    : NULL
+output domain  : {NULL}
+output         : NULL
+
+
+ast            : repeat('3', 5)
+raw expr       : repeat("3", 5_u8)
+checked expr   : repeat<String, UInt64>("3", CAST(5_u8 AS UInt64))
+optimized expr : "33333"
+output type    : String
+output domain  : Unknown
+output         : "33333"
+
+
+error: 
+  --> SQL:1:1
+  |
+1 | repeat('3', 1000001)
+  | ^^^^^^^^^^^^^^^^^^^^ Too many times to repeat: (1000001), maximum is: 1000000
+
+
+
+ast            : repeat(a, 3)
+raw expr       : repeat(ColumnRef(0)::String, 3_u8)
+checked expr   : repeat<String, UInt64>(ColumnRef(0), CAST(3_u8 AS UInt64))
+optimized expr : repeat<String, UInt64>(ColumnRef(0), 3_u64)
+evaluation:
++--------+-------------+---------+
+|        | a           | Output  |
++--------+-------------+---------+
+| Type   | String      | String  |
+| Domain | {"a"..="c"} | Unknown |
+| Row 0  | "a"         | "aaa"   |
+| Row 1  | "b"         | "bbb"   |
+| Row 2  | "c"         | "ccc"   |
++--------+-------------+---------+
+evaluation (internal):
++--------+------------------------------------------------------------------------------------+
+| Column | Data                                                                               |
++--------+------------------------------------------------------------------------------------+
+| a      | StringColumn { data: [97, 98, 99], offsets: [0, 1, 2, 3] }                         |
+| Output | StringColumn { data: [97, 97, 97, 98, 98, 98, 99, 99, 99], offsets: [0, 3, 6, 9] } |
++--------+------------------------------------------------------------------------------------+
+
+
 error: 
   --> SQL:1:1
   |