Skip to content

Commit 71e6f9d

Browse files
TCeasonTCeason
committed
migrate string func insert to func-v2
1 parent af7b1cb commit 71e6f9d

File tree

3 files changed

+210
-0
lines changed

3 files changed

+210
-0
lines changed

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,29 @@ pub fn register(registry: &mut FunctionRegistry) {
134134
},
135135
);
136136

137+
registry.register_4_arg::<StringType, NumberType<i64>, NumberType<i64>, StringType, StringType, _, _>(
138+
"insert",
139+
FunctionProperty::default(),
140+
|_, _, _, _| None,
141+
|srcstr, pos, len, substr| {
142+
let mut values: Vec<u8> = vec![];
143+
144+
let sl = srcstr.len() as i64;
145+
if pos < 1 || pos > sl {
146+
values.extend_from_slice(srcstr);
147+
} else {
148+
let p = pos as usize - 1;
149+
values.extend_from_slice(&srcstr[0..p]);
150+
values.extend_from_slice(substr);
151+
if len >= 0 && pos + len < sl {
152+
let l = len as usize;
153+
values.extend_from_slice(&srcstr[p + l..]);
154+
}
155+
}
156+
values
157+
}
158+
);
159+
137160
registry.register_3_arg::<StringType, NumberType<u64>, StringType, StringType, _, _>(
138161
"rpad",
139162
FunctionProperty::default(),

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ fn test_string() {
5555
test_char(file);
5656
test_soundex(file);
5757
test_ord(file);
58+
test_insert(file);
5859
}
5960

6061
fn test_upper(file: &mut impl Write) {
@@ -686,3 +687,58 @@ fn test_ord(file: &mut impl Write) {
686687
run_ast(file, "ord('早ab')", &[]);
687688
run_ast(file, "ord('💖')", &[]);
688689
}
690+
691+
fn test_insert(file: &mut impl Write) {
692+
run_ast(file, "insert('Quadratic', 3, 4, 'What', 4)", &[]);
693+
run_ast(file, "insert('Quadratic', 3, 4)", &[]);
694+
run_ast(file, "insert('Quadratic', 3, 4, 'What')", &[]);
695+
run_ast(file, "insert('Quadratic', -1, 4, 'What')", &[]);
696+
run_ast(file, "insert('Quadratic', 3, 100, 'What')", &[]);
697+
run_ast(file, "insert('Quadratic', 3, 100, NULL)", &[]);
698+
run_ast(file, "insert('Quadratic', 3, NULL, 'NULL')", &[]);
699+
run_ast(file, "insert('Quadratic', NULL, 100, 'NULL')", &[]);
700+
run_ast(file, "insert(NULL, 2, 100, 'NULL')", &[]);
701+
702+
let table = [
703+
(
704+
"a",
705+
DataType::String,
706+
Column::from_data(&["hi", "test", "cc", "q"]),
707+
),
708+
("b", DataType::UInt8, Column::from_data(vec![1, 4, 1, 1])),
709+
("c", DataType::UInt8, Column::from_data(vec![3, 5, 1, 1])),
710+
(
711+
"d",
712+
DataType::String,
713+
Column::from_data(&["xx", "zc", "12", "56"]),
714+
),
715+
];
716+
run_ast(file, "insert(a, b, c, d)", &table);
717+
let columns = [
718+
(
719+
"x",
720+
DataType::Nullable(Box::new(DataType::String)),
721+
Column::from_data_with_validity(&["hi", "test", "cc", "q"], vec![
722+
false, true, true, true,
723+
]),
724+
),
725+
(
726+
"y",
727+
DataType::Nullable(Box::new(DataType::UInt8)),
728+
Column::from_data_with_validity(vec![1, 4, 1, 1], vec![true, true, false, true]),
729+
),
730+
(
731+
"z",
732+
DataType::Nullable(Box::new(DataType::UInt8)),
733+
Column::from_data_with_validity(vec![3, 5, 1, 1], vec![true, false, true, true]),
734+
),
735+
(
736+
"u",
737+
DataType::Nullable(Box::new(DataType::String)),
738+
Column::from_data_with_validity(&["xx", "zc", "12", "56"], vec![
739+
false, true, true, true,
740+
]),
741+
),
742+
];
743+
run_ast(file, "insert(x, y, z, u)", &columns);
744+
}

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2723,3 +2723,134 @@ output domain : Unknown
27232723
output : 4036989590
27242724

27252725

2726+
error:
2727+
--> SQL:1:1
2728+
|
2729+
1 | insert('Quadratic', 3, 4, 'What', 4)
2730+
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `insert(String, UInt8, UInt8, String, UInt8)`
2731+
2732+
2733+
2734+
error:
2735+
--> SQL:1:1
2736+
|
2737+
1 | insert('Quadratic', 3, 4)
2738+
| ^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `insert(String, UInt8, UInt8)`
2739+
2740+
2741+
2742+
ast : insert('Quadratic', 3, 4, 'What')
2743+
raw expr : insert("Quadratic", 3_u8, 4_u8, "What")
2744+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(3_u8 AS Int64), CAST(4_u8 AS Int64), "What")
2745+
optimized expr : "QuWhattic"
2746+
output type : String
2747+
output domain : Unknown
2748+
output : "QuWhattic"
2749+
2750+
2751+
ast : insert('Quadratic', -1, 4, 'What')
2752+
raw expr : insert("Quadratic", minus(1_u8), 4_u8, "What")
2753+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(minus<UInt8>(1_u8) AS Int64), CAST(4_u8 AS Int64), "What")
2754+
optimized expr : "Quadratic"
2755+
output type : String
2756+
output domain : Unknown
2757+
output : "Quadratic"
2758+
2759+
2760+
ast : insert('Quadratic', 3, 100, 'What')
2761+
raw expr : insert("Quadratic", 3_u8, 100_u8, "What")
2762+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(3_u8 AS Int64), CAST(100_u8 AS Int64), "What")
2763+
optimized expr : "QuWhat"
2764+
output type : String
2765+
output domain : Unknown
2766+
output : "QuWhat"
2767+
2768+
2769+
ast : insert('Quadratic', 3, 100, NULL)
2770+
raw expr : insert("Quadratic", 3_u8, 100_u8, NULL)
2771+
checked expr : insert<String NULL, Int64 NULL, Int64 NULL, NULL>(CAST("Quadratic" AS String NULL), CAST(3_u8 AS Int64 NULL), CAST(100_u8 AS Int64 NULL), NULL)
2772+
optimized expr : NULL
2773+
output type : NULL
2774+
output domain : {NULL}
2775+
output : NULL
2776+
2777+
2778+
ast : insert('Quadratic', 3, NULL, 'NULL')
2779+
raw expr : insert("Quadratic", 3_u8, NULL, "NULL")
2780+
checked expr : insert<String NULL, Int64 NULL, NULL, String NULL>(CAST("Quadratic" AS String NULL), CAST(3_u8 AS Int64 NULL), NULL, CAST("NULL" AS String NULL))
2781+
optimized expr : NULL
2782+
output type : NULL
2783+
output domain : {NULL}
2784+
output : NULL
2785+
2786+
2787+
ast : insert('Quadratic', NULL, 100, 'NULL')
2788+
raw expr : insert("Quadratic", NULL, 100_u8, "NULL")
2789+
checked expr : insert<String NULL, NULL, Int64 NULL, String NULL>(CAST("Quadratic" AS String NULL), NULL, CAST(100_u8 AS Int64 NULL), CAST("NULL" AS String NULL))
2790+
optimized expr : NULL
2791+
output type : NULL
2792+
output domain : {NULL}
2793+
output : NULL
2794+
2795+
2796+
ast : insert(NULL, 2, 100, 'NULL')
2797+
raw expr : insert(NULL, 2_u8, 100_u8, "NULL")
2798+
checked expr : insert<NULL, Int64 NULL, Int64 NULL, String NULL>(NULL, CAST(2_u8 AS Int64 NULL), CAST(100_u8 AS Int64 NULL), CAST("NULL" AS String NULL))
2799+
optimized expr : NULL
2800+
output type : NULL
2801+
output domain : {NULL}
2802+
output : NULL
2803+
2804+
2805+
ast : insert(a, b, c, d)
2806+
raw expr : insert(ColumnRef(0)::String, ColumnRef(1)::UInt8, ColumnRef(2)::UInt8, ColumnRef(3)::String)
2807+
checked expr : insert<String, Int64, Int64, String>(ColumnRef(0), CAST(ColumnRef(1) AS Int64), CAST(ColumnRef(2) AS Int64), ColumnRef(3))
2808+
evaluation:
2809+
+--------+-----------------+---------+---------+---------------+---------+
2810+
| | a | b | c | d | Output |
2811+
+--------+-----------------+---------+---------+---------------+---------+
2812+
| Type | String | UInt8 | UInt8 | String | String |
2813+
| Domain | {"cc"..="test"} | {1..=4} | {1..=5} | {"12"..="zc"} | Unknown |
2814+
| Row 0 | "hi" | 1 | 3 | "xx" | "xx" |
2815+
| Row 1 | "test" | 4 | 5 | "zc" | "teszc" |
2816+
| Row 2 | "cc" | 1 | 1 | "12" | "12" |
2817+
| Row 3 | "q" | 1 | 1 | "56" | "56" |
2818+
+--------+-----------------+---------+---------+---------------+---------+
2819+
evaluation (internal):
2820+
+--------+--------------------------------------------------------------------------------------------------------------+
2821+
| Column | Data |
2822+
+--------+--------------------------------------------------------------------------------------------------------------+
2823+
| a | String(StringColumn { data: [104, 105, 116, 101, 115, 116, 99, 99, 113], offsets: [0, 2, 6, 8, 9] }) |
2824+
| b | Int32([1, 4, 1, 1]) |
2825+
| c | Int32([3, 5, 1, 1]) |
2826+
| d | String(StringColumn { data: [120, 120, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 4, 6, 8] }) |
2827+
| Output | String(StringColumn { data: [120, 120, 116, 101, 115, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 7, 9, 11] }) |
2828+
+--------+--------------------------------------------------------------------------------------------------------------+
2829+
2830+
2831+
ast : insert(x, y, z, u)
2832+
raw expr : insert(ColumnRef(0)::String NULL, ColumnRef(1)::UInt8 NULL, ColumnRef(2)::UInt8 NULL, ColumnRef(3)::String NULL)
2833+
checked expr : insert<String NULL, Int64 NULL, Int64 NULL, String NULL>(ColumnRef(0), CAST(ColumnRef(1) AS Int64 NULL), CAST(ColumnRef(2) AS Int64 NULL), ColumnRef(3))
2834+
evaluation:
2835+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2836+
| | x | y | z | u | Output |
2837+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2838+
| Type | String NULL | UInt8 NULL | UInt8 NULL | String NULL | String NULL |
2839+
| Domain | {"cc"..="test"} ∪ {NULL} | {1..=4} ∪ {NULL} | {1..=5} ∪ {NULL} | {"12"..="zc"} ∪ {NULL} | Unknown |
2840+
| Row 0 | NULL | 1 | 3 | NULL | NULL |
2841+
| Row 1 | "test" | 4 | NULL | "zc" | NULL |
2842+
| Row 2 | "cc" | NULL | 1 | "12" | NULL |
2843+
| Row 3 | "q" | 1 | 1 | "56" | "56" |
2844+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2845+
evaluation (internal):
2846+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
2847+
| Column | Data |
2848+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
2849+
| x | Nullable(NullableColumn { column: String(StringColumn { data: [104, 105, 116, 101, 115, 116, 99, 99, 113], offsets: [0, 2, 6, 8, 9] }), validity: [0b____1110] }) |
2850+
| y | Nullable(NullableColumn { column: Int32([1, 4, 1, 1]), validity: [0b____1011] }) |
2851+
| z | Nullable(NullableColumn { column: Int32([3, 5, 1, 1]), validity: [0b____1101] }) |
2852+
| u | Nullable(NullableColumn { column: String(StringColumn { data: [120, 120, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 4, 6, 8] }), validity: [0b____1110] }) |
2853+
| Output | Nullable(NullableColumn { column: String(StringColumn { data: [120, 120, 116, 101, 115, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 7, 9, 11] }), validity: [0b____1000] }) |
2854+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
2855+
2856+

0 commit comments

Comments
 (0)