Skip to content

Commit fc546d1

Browse files
sundy-lisundy-li
committed
Merge branch 'main' into function-v2-repeat
2 parents e2c3f09 + 39112ff commit fc546d1

File tree

4 files changed

+241
-7
lines changed

4 files changed

+241
-7
lines changed

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

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

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

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ fn test_string() {
5656
test_soundex(file);
5757
test_ord(file);
5858
test_repeat(file);
59+
test_insert(file);
5960
}
6061

6162
fn test_upper(file: &mut impl Write) {
@@ -694,3 +695,66 @@ fn test_repeat(file: &mut impl Write) {
694695
let table = [("a", DataType::String, Column::from_data(&["a", "b", "c"]))];
695696
run_ast(file, "repeat(a, 3)", &table);
696697
}
698+
699+
fn test_insert(file: &mut impl Write) {
700+
run_ast(file, "insert('Quadratic', 3, 4, 'What', 4)", &[]);
701+
run_ast(file, "insert('Quadratic', 3, 4)", &[]);
702+
run_ast(file, "insert('Quadratic', 3, 4, 'What')", &[]);
703+
run_ast(file, "insert('Quadratic', -1, 4, 'What')", &[]);
704+
run_ast(file, "insert('Quadratic', 3, 100, 'What')", &[]);
705+
run_ast(file, "insert('Quadratic', 3, 100, NULL)", &[]);
706+
run_ast(file, "insert('Quadratic', 3, NULL, 'NULL')", &[]);
707+
run_ast(file, "insert('Quadratic', NULL, 100, 'NULL')", &[]);
708+
run_ast(file, "insert(NULL, 2, 100, 'NULL')", &[]);
709+
710+
let table = [
711+
(
712+
"a",
713+
DataType::String,
714+
Column::from_data(&["hi", "test", "cc", "q"]),
715+
),
716+
(
717+
"b",
718+
DataType::Number(NumberDataType::UInt8),
719+
Column::from_data(vec![1, 4, 1, 1]),
720+
),
721+
(
722+
"c",
723+
DataType::Number(NumberDataType::UInt8),
724+
Column::from_data(vec![3, 5, 1, 1]),
725+
),
726+
(
727+
"d",
728+
DataType::String,
729+
Column::from_data(&["xx", "zc", "12", "56"]),
730+
),
731+
];
732+
run_ast(file, "insert(a, b, c, d)", &table);
733+
let columns = [
734+
(
735+
"x",
736+
DataType::Nullable(Box::new(DataType::String)),
737+
Column::from_data_with_validity(&["hi", "test", "cc", "q"], vec![
738+
false, true, true, true,
739+
]),
740+
),
741+
(
742+
"y",
743+
DataType::Nullable(Box::new(DataType::Number(NumberDataType::UInt8))),
744+
Column::from_data_with_validity(vec![1, 4, 1, 1], vec![true, true, false, true]),
745+
),
746+
(
747+
"z",
748+
DataType::Nullable(Box::new(DataType::Number(NumberDataType::UInt8))),
749+
Column::from_data_with_validity(vec![3, 5, 1, 1], vec![true, false, true, true]),
750+
),
751+
(
752+
"u",
753+
DataType::Nullable(Box::new(DataType::String)),
754+
Column::from_data_with_validity(&["xx", "zc", "12", "56"], vec![
755+
false, true, true, true,
756+
]),
757+
),
758+
];
759+
run_ast(file, "insert(x, y, z, u)", &columns);
760+
}

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2764,3 +2764,134 @@ evaluation (internal):
27642764
+--------+------------------------------------------------------------------------------------+
27652765

27662766

2767+
error:
2768+
--> SQL:1:1
2769+
|
2770+
1 | insert('Quadratic', 3, 4, 'What', 4)
2771+
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `insert(String, UInt8, UInt8, String, UInt8)`
2772+
2773+
2774+
2775+
error:
2776+
--> SQL:1:1
2777+
|
2778+
1 | insert('Quadratic', 3, 4)
2779+
| ^^^^^^^^^^^^^^^^^^^^^^^^^ no overload satisfies `insert(String, UInt8, UInt8)`
2780+
2781+
2782+
2783+
ast : insert('Quadratic', 3, 4, 'What')
2784+
raw expr : insert("Quadratic", 3_u8, 4_u8, "What")
2785+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(3_u8 AS Int64), CAST(4_u8 AS Int64), "What")
2786+
optimized expr : "QuWhattic"
2787+
output type : String
2788+
output domain : Unknown
2789+
output : "QuWhattic"
2790+
2791+
2792+
ast : insert('Quadratic', -1, 4, 'What')
2793+
raw expr : insert("Quadratic", minus(1_u8), 4_u8, "What")
2794+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(minus<UInt8>(1_u8) AS Int64), CAST(4_u8 AS Int64), "What")
2795+
optimized expr : "Quadratic"
2796+
output type : String
2797+
output domain : Unknown
2798+
output : "Quadratic"
2799+
2800+
2801+
ast : insert('Quadratic', 3, 100, 'What')
2802+
raw expr : insert("Quadratic", 3_u8, 100_u8, "What")
2803+
checked expr : insert<String, Int64, Int64, String>("Quadratic", CAST(3_u8 AS Int64), CAST(100_u8 AS Int64), "What")
2804+
optimized expr : "QuWhat"
2805+
output type : String
2806+
output domain : Unknown
2807+
output : "QuWhat"
2808+
2809+
2810+
ast : insert('Quadratic', 3, 100, NULL)
2811+
raw expr : insert("Quadratic", 3_u8, 100_u8, NULL)
2812+
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)
2813+
optimized expr : NULL
2814+
output type : NULL
2815+
output domain : {NULL}
2816+
output : NULL
2817+
2818+
2819+
ast : insert('Quadratic', 3, NULL, 'NULL')
2820+
raw expr : insert("Quadratic", 3_u8, NULL, "NULL")
2821+
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))
2822+
optimized expr : NULL
2823+
output type : NULL
2824+
output domain : {NULL}
2825+
output : NULL
2826+
2827+
2828+
ast : insert('Quadratic', NULL, 100, 'NULL')
2829+
raw expr : insert("Quadratic", NULL, 100_u8, "NULL")
2830+
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))
2831+
optimized expr : NULL
2832+
output type : NULL
2833+
output domain : {NULL}
2834+
output : NULL
2835+
2836+
2837+
ast : insert(NULL, 2, 100, 'NULL')
2838+
raw expr : insert(NULL, 2_u8, 100_u8, "NULL")
2839+
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))
2840+
optimized expr : NULL
2841+
output type : NULL
2842+
output domain : {NULL}
2843+
output : NULL
2844+
2845+
2846+
ast : insert(a, b, c, d)
2847+
raw expr : insert(ColumnRef(0)::String, ColumnRef(1)::UInt8, ColumnRef(2)::UInt8, ColumnRef(3)::String)
2848+
checked expr : insert<String, Int64, Int64, String>(ColumnRef(0), CAST(ColumnRef(1) AS Int64), CAST(ColumnRef(2) AS Int64), ColumnRef(3))
2849+
evaluation:
2850+
+--------+-----------------+---------+---------+---------------+---------+
2851+
| | a | b | c | d | Output |
2852+
+--------+-----------------+---------+---------+---------------+---------+
2853+
| Type | String | UInt8 | UInt8 | String | String |
2854+
| Domain | {"cc"..="test"} | {1..=4} | {1..=5} | {"12"..="zc"} | Unknown |
2855+
| Row 0 | "hi" | 1 | 3 | "xx" | "xx" |
2856+
| Row 1 | "test" | 4 | 5 | "zc" | "teszc" |
2857+
| Row 2 | "cc" | 1 | 1 | "12" | "12" |
2858+
| Row 3 | "q" | 1 | 1 | "56" | "56" |
2859+
+--------+-----------------+---------+---------+---------------+---------+
2860+
evaluation (internal):
2861+
+--------+------------------------------------------------------------------------------------------------------+
2862+
| Column | Data |
2863+
+--------+------------------------------------------------------------------------------------------------------+
2864+
| a | StringColumn { data: [104, 105, 116, 101, 115, 116, 99, 99, 113], offsets: [0, 2, 6, 8, 9] } |
2865+
| b | Int32([1, 4, 1, 1]) |
2866+
| c | Int32([3, 5, 1, 1]) |
2867+
| d | StringColumn { data: [120, 120, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 4, 6, 8] } |
2868+
| Output | StringColumn { data: [120, 120, 116, 101, 115, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 7, 9, 11] } |
2869+
+--------+------------------------------------------------------------------------------------------------------+
2870+
2871+
2872+
ast : insert(x, y, z, u)
2873+
raw expr : insert(ColumnRef(0)::String NULL, ColumnRef(1)::UInt8 NULL, ColumnRef(2)::UInt8 NULL, ColumnRef(3)::String NULL)
2874+
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))
2875+
evaluation:
2876+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2877+
| | x | y | z | u | Output |
2878+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2879+
| Type | String NULL | UInt8 NULL | UInt8 NULL | String NULL | String NULL |
2880+
| Domain | {"cc"..="test"} ∪ {NULL} | {1..=4} ∪ {NULL} | {1..=5} ∪ {NULL} | {"12"..="zc"} ∪ {NULL} | Unknown |
2881+
| Row 0 | NULL | 1 | 3 | NULL | NULL |
2882+
| Row 1 | "test" | 4 | NULL | "zc" | NULL |
2883+
| Row 2 | "cc" | NULL | 1 | "12" | NULL |
2884+
| Row 3 | "q" | 1 | 1 | "56" | "56" |
2885+
+--------+--------------------------+------------------+------------------+------------------------+-------------+
2886+
evaluation (internal):
2887+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
2888+
| Column | Data |
2889+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
2890+
| x | NullableColumn { column: StringColumn { data: [104, 105, 116, 101, 115, 116, 99, 99, 113], offsets: [0, 2, 6, 8, 9] }, validity: [0b____1110] } |
2891+
| y | NullableColumn { column: Int32([1, 4, 1, 1]), validity: [0b____1011] } |
2892+
| z | NullableColumn { column: Int32([3, 5, 1, 1]), validity: [0b____1101] } |
2893+
| u | NullableColumn { column: StringColumn { data: [120, 120, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 4, 6, 8] }, validity: [0b____1110] } |
2894+
| Output | NullableColumn { column: StringColumn { data: [120, 120, 116, 101, 115, 122, 99, 49, 50, 53, 54], offsets: [0, 2, 7, 9, 11] }, validity: [0b____1000] } |
2895+
+--------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
2896+
2897+

‎src/query/storages/fuse/src/operations/read.rs

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -222,12 +222,6 @@ impl FuseTableSource {
222222
fn generate_one_block(&mut self, block: DataBlock) -> Result<()> {
223223
let mut partitions = self.ctx.try_get_partitions(1)?;
224224

225-
let progress_values = ProgressValues {
226-
rows: block.num_rows(),
227-
bytes: block.memory_size(),
228-
};
229-
self.scan_progress.incr(&progress_values);
230-
231225
self.state = match partitions.is_empty() {
232226
true => State::Generated(None, block),
233227
false => State::Generated(Some(partitions.remove(0)), block),
@@ -322,9 +316,21 @@ impl Processor for FuseTableSource {
322316
return Err(ErrorCode::LogicalError("It's a bug. Need remain reader"));
323317
};
324318
// the last step of prewhere
319+
let progress_values = ProgressValues {
320+
rows: block.num_rows(),
321+
bytes: block.memory_size(),
322+
};
323+
self.scan_progress.incr(&progress_values);
325324
DataBlock::filter_block(block, &filter)?
326325
} else {
327-
self.output_reader.deserialize(part, chunks)?
326+
let block = self.output_reader.deserialize(part, chunks)?;
327+
let progress_values = ProgressValues {
328+
rows: block.num_rows(),
329+
bytes: block.memory_size(),
330+
};
331+
self.scan_progress.incr(&progress_values);
332+
333+
block
328334
};
329335

330336
self.generate_one_block(data_block)?;
@@ -341,11 +347,21 @@ impl Processor for FuseTableSource {
341347
if !DataBlock::filter_exists(&filter)? {
342348
// all rows in this block are filtered out
343349
// turn to read next part
350+
let progress_values = ProgressValues {
351+
rows: data_block.num_rows(),
352+
bytes: data_block.memory_size(),
353+
};
354+
self.scan_progress.incr(&progress_values);
344355
self.generate_one_empty_block()?;
345356
return Ok(());
346357
}
347358
if self.remain_reader.is_none() {
348359
// shortcut, we don't need to read remain data
360+
let progress_values = ProgressValues {
361+
rows: data_block.num_rows(),
362+
bytes: data_block.memory_size(),
363+
};
364+
self.scan_progress.incr(&progress_values);
349365
let block = DataBlock::filter_block(data_block, &filter)?;
350366
self.generate_one_block(block)?;
351367
} else {

0 commit comments

Comments
 (0)