Skip to content

Commit 564e0cf

Browse files
TCeasonTCeason
committed
modify with suggestion
1 parent 4da5044 commit 564e0cf

File tree

2 files changed

+86
-80
lines changed

2 files changed

+86
-80
lines changed

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

Lines changed: 67 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use common_expression::types::NumberType;
2525
use common_expression::types::StringType;
2626
use common_expression::vectorize_with_builder_1_arg;
2727
use common_expression::vectorize_with_builder_2_arg;
28+
use common_expression::vectorize_with_builder_3_arg;
2829
use common_expression::vectorize_with_builder_4_arg;
2930
use common_expression::FunctionProperty;
3031
use common_expression::FunctionRegistry;
@@ -111,106 +112,111 @@ pub fn register(registry: &mut FunctionRegistry) {
111112
);
112113
registry.register_aliases("char_length", &["character_length"]);
113114

114-
registry.register_3_arg::<StringType, NumberType<u64>, StringType, StringType, _, _>(
115+
registry.register_passthrough_nullable_3_arg::<StringType, NumberType<u64>, StringType, StringType, _, _>(
115116
"lpad",
116117
FunctionProperty::default(),
117118
|_, _, _| None,
118-
|str: &[u8], l: u64, pad: &[u8]| {
119-
let mut buff: Vec<u8> = vec![];
120-
if l != 0 {
121-
if l > str.len() as u64 {
122-
let l = l - str.len() as u64;
123-
while buff.len() < l as usize {
124-
if buff.len() + pad.len() <= l as usize {
125-
buff.extend_from_slice(pad);
119+
vectorize_with_builder_3_arg::<StringType, NumberType<u64>, StringType, StringType>(
120+
|s, pad_len, pad, output| {
121+
let pad_len = pad_len as usize;
122+
if pad_len <= s.len() {
123+
output.put_slice(&s[..pad_len])
124+
} else {
125+
let mut remain_pad_len = pad_len - s.len();
126+
while remain_pad_len > 0 {
127+
if remain_pad_len < pad.len() {
128+
output.put_slice(&pad[..remain_pad_len]);
129+
remain_pad_len = 0;
126130
} else {
127-
buff.extend_from_slice(&pad[0..l as usize - buff.len()])
131+
output.put_slice(pad);
132+
remain_pad_len -= pad.len();
128133
}
129134
}
130-
buff.extend_from_slice(str);
131-
} else {
132-
buff.extend_from_slice(&str[0..l as usize]);
135+
output.put_slice(s);
133136
}
137+
output.commit_row();
138+
Ok(())
134139
}
135-
buff
136-
},
140+
),
137141
);
138142

139143
registry.register_passthrough_nullable_4_arg::<StringType, NumberType<i64>, NumberType<i64>, StringType, StringType, _, _>(
140-
"insert",
141-
FunctionProperty::default(),
142-
|_, _, _, _| None,
143-
vectorize_with_builder_4_arg::<StringType, NumberType<i64>, NumberType<i64>, StringType, StringType>(
144-
|srcstr, pos, len, substr, builder| {
145-
let mut values: Vec<u8> = vec![];
146-
147-
let sl = srcstr.len() as i64;
148-
if pos < 1 || pos > sl {
149-
values.extend_from_slice(srcstr);
144+
"insert",
145+
FunctionProperty::default(),
146+
|_, _, _, _| None,
147+
vectorize_with_builder_4_arg::<StringType, NumberType<i64>, NumberType<i64>, StringType, StringType>(
148+
|srcstr, pos, len, substr, output| {
149+
let pos = pos as usize;
150+
let len = len as usize;
151+
if pos < 1 || pos > srcstr.len() {
152+
output.put_slice(srcstr);
150153
} else {
151-
let p = pos as usize - 1;
152-
values.extend_from_slice(&srcstr[0..p]);
153-
values.extend_from_slice(substr);
154-
if len >= 0 && pos + len < sl {
155-
let l = len as usize;
156-
values.extend_from_slice(&srcstr[p + l..]);
154+
let pos = pos - 1;
155+
output.put_slice(&srcstr[0..pos]);
156+
output.put_slice(substr);
157+
if pos + len < srcstr.len() {
158+
output.put_slice(&srcstr[(pos + len)..]);
157159
}
158160
}
159-
builder.put_slice(&values);
160-
builder.commit_row();
161-
Ok(())
161+
output.commit_row();
162+
Ok(())
162163
}),
163-
);
164+
);
164165

165-
registry.register_3_arg::<StringType, NumberType<u64>, StringType, StringType, _, _>(
166+
registry.register_passthrough_nullable_3_arg::<StringType, NumberType<u64>, StringType, StringType, _, _>(
166167
"rpad",
167168
FunctionProperty::default(),
168169
|_, _, _| None,
169-
|str: &[u8], l: u64, pad: &[u8]| {
170-
let mut buff: Vec<u8> = vec![];
171-
if l != 0 {
172-
if l > str.len() as u64 {
173-
buff.extend_from_slice(str);
174-
while buff.len() < l as usize {
175-
if buff.len() + pad.len() <= l as usize {
176-
buff.extend_from_slice(pad);
177-
} else {
178-
buff.extend_from_slice(&pad[0..l as usize - buff.len()])
179-
}
170+
vectorize_with_builder_3_arg::<StringType, NumberType<u64>, StringType, StringType>(
171+
|s: &[u8], pad_len: u64, pad: &[u8], output| {
172+
let pad_len = pad_len as usize;
173+
if pad_len <= s.len() {
174+
output.put_slice(&s[..pad_len])
175+
} else {
176+
output.put_slice(s);
177+
let mut remain_pad_len = pad_len - s.len();
178+
while remain_pad_len > 0 {
179+
if remain_pad_len < pad.len() {
180+
output.put_slice(&pad[..remain_pad_len]);
181+
remain_pad_len = 0;
182+
} else {
183+
output.put_slice(pad);
184+
remain_pad_len -= pad.len();
180185
}
181-
} else {
182-
buff.extend_from_slice(&str[0..l as usize]);
183186
}
184187
}
185-
buff
186-
},
188+
output.commit_row();
189+
Ok(())
190+
}),
187191
);
188192

189-
registry.register_3_arg::<StringType, StringType, StringType, StringType, _, _>(
193+
registry.register_passthrough_nullable_3_arg::<StringType, StringType, StringType, StringType, _, _>(
190194
"replace",
191195
FunctionProperty::default(),
192196
|_, _, _| None,
193-
|str, from, to| {
194-
let mut buf: Vec<u8> = vec![];
197+
vectorize_with_builder_3_arg::<StringType, StringType, StringType, StringType>(
198+
|str, from, to, output| {
195199
if from.is_empty() || from == to {
196-
buf.extend_from_slice(str);
197-
return buf;
200+
output.put_slice(str);
201+
output.commit_row();
202+
return Ok(());
198203
}
199204
let mut skip = 0;
200205
for (p, w) in str.windows(from.len()).enumerate() {
201206
if w == from {
202-
buf.extend_from_slice(to);
207+
output.put_slice(to);
203208
skip = from.len();
204209
} else if p + w.len() == str.len() {
205-
buf.extend_from_slice(w);
210+
output.put_slice(w);
206211
} else if skip > 1 {
207212
skip -= 1;
208213
} else {
209-
buf.extend_from_slice(&w[0..1]);
214+
output.put_slice(&w[0..1]);
210215
}
211216
}
212-
buf
213-
},
217+
output.commit_row();
218+
Ok(())
219+
}),
214220
);
215221

216222
registry.register_2_arg::<StringType, StringType, NumberType<i8>, _, _>(

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

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2862,19 +2862,19 @@ evaluation:
28622862
| Domain | {"cc"..="test"} | {1..=4} | {1..=5} | {"12"..="zc"} | Unknown |
28632863
| Row 0 | "hi" | 1 | 3 | "xx" | "xx" |
28642864
| Row 1 | "test" | 4 | 5 | "zc" | "teszc" |
2865-
| Row 2 | "cc" | 1 | 1 | "12" | "12" |
2865+
| Row 2 | "cc" | 1 | 1 | "12" | "12c" |
28662866
| Row 3 | "q" | 1 | 1 | "56" | "56" |
28672867
+--------+-----------------+---------+---------+---------------+---------+
28682868
evaluation (internal):
2869-
+--------+----------------------------------------------------------------------------+
2870-
| Column | Data |
2871-
+--------+----------------------------------------------------------------------------+
2872-
| a | StringColumn { data: 0x686974657374636371, offsets: [0, 2, 6, 8, 9] } |
2873-
| b | Int32([1, 4, 1, 1]) |
2874-
| c | Int32([3, 5, 1, 1]) |
2875-
| d | StringColumn { data: 0x78787a6331323536, offsets: [0, 2, 4, 6, 8] } |
2876-
| Output | StringColumn { data: 0x78787465737a6331323536, offsets: [0, 2, 7, 9, 11] } |
2877-
+--------+----------------------------------------------------------------------------+
2869+
+--------+-------------------------------------------------------------------------------+
2870+
| Column | Data |
2871+
+--------+-------------------------------------------------------------------------------+
2872+
| a | StringColumn { data: 0x686974657374636371, offsets: [0, 2, 6, 8, 9] } |
2873+
| b | Int32([1, 4, 1, 1]) |
2874+
| c | Int32([3, 5, 1, 1]) |
2875+
| d | StringColumn { data: 0x78787a6331323536, offsets: [0, 2, 4, 6, 8] } |
2876+
| Output | StringColumn { data: 0x78787465737a633132633536, offsets: [0, 2, 7, 10, 12] } |
2877+
+--------+-------------------------------------------------------------------------------+
28782878

28792879

28802880
ast : insert(x, y, z, u)
@@ -2892,14 +2892,14 @@ evaluation:
28922892
| Row 3 | "q" | 1 | 1 | "56" | "56" |
28932893
+--------+--------------------------+------------------+------------------+------------------------+-------------+
28942894
evaluation (internal):
2895-
+--------+-------------------------------------------------------------------------------------------------------------------------------+
2896-
| Column | Data |
2897-
+--------+-------------------------------------------------------------------------------------------------------------------------------+
2898-
| x | NullableColumn { column: StringColumn { data: 0x686974657374636371, offsets: [0, 2, 6, 8, 9] }, validity: [0b____1110] } |
2899-
| y | NullableColumn { column: Int32([1, 4, 1, 1]), validity: [0b____1011] } |
2900-
| z | NullableColumn { column: Int32([3, 5, 1, 1]), validity: [0b____1101] } |
2901-
| u | NullableColumn { column: StringColumn { data: 0x78787a6331323536, offsets: [0, 2, 4, 6, 8] }, validity: [0b____1110] } |
2902-
| Output | NullableColumn { column: StringColumn { data: 0x78787465737a6331323536, offsets: [0, 2, 7, 9, 11] }, validity: [0b____1000] } |
2903-
+--------+-------------------------------------------------------------------------------------------------------------------------------+
2895+
+--------+----------------------------------------------------------------------------------------------------------------------------------+
2896+
| Column | Data |
2897+
+--------+----------------------------------------------------------------------------------------------------------------------------------+
2898+
| x | NullableColumn { column: StringColumn { data: 0x686974657374636371, offsets: [0, 2, 6, 8, 9] }, validity: [0b____1110] } |
2899+
| y | NullableColumn { column: Int32([1, 4, 1, 1]), validity: [0b____1011] } |
2900+
| z | NullableColumn { column: Int32([3, 5, 1, 1]), validity: [0b____1101] } |
2901+
| u | NullableColumn { column: StringColumn { data: 0x78787a6331323536, offsets: [0, 2, 4, 6, 8] }, validity: [0b____1110] } |
2902+
| Output | NullableColumn { column: StringColumn { data: 0x78787465737a633132633536, offsets: [0, 2, 7, 10, 12] }, validity: [0b____1000] } |
2903+
+--------+----------------------------------------------------------------------------------------------------------------------------------+
29042904

29052905

0 commit comments

Comments
 (0)