Skip to content

Commit d2b3fd9

Browse files
hanxuanlianghanxuanliang
authored
feat(functions): add map_contains_key scalar function (#15465)
* feat(functions): add map_contains_key scalar function * feat(functions): update function_list.txt and add more test_map case * feat(functions): add sqllogictests in map_contains_key * feat(functions): fix function_list.txt * feat(functions): update logic tests for map_contains_key
1 parent e90dd66 commit d2b3fd9

File tree

6 files changed

+200
-3
lines changed

6 files changed

+200
-3
lines changed

‎src/query/functions/README.md

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,26 @@
44

55
- [Aggregate function](https://docs.databend.com/guides/community/contributor/how-to-write-aggregate-functions)
66

7+
## Don't forget to add tests
78

8-
# Don't forget to add tests:
99
- run:
1010
`env REGENERATE_GOLDENFILES=1 cargo test` after adding new functions
11-
- Add logic tests in "tests/sqllogictests/suites/query/02_function" and run:
12-
`cargo run -p databend-sqllogictests -- --handlers http --run_file <file_name>`
11+
12+
- build and run databend, prepare for logical tests:
13+
```shell
14+
make build
15+
make run-debug
16+
```
17+
18+
- Add logic tests file in "tests/sqllogictests/suites/query/functions/*.test" and run:
19+
```shell
20+
cargo run \
21+
-p databend-sqllogictests \
22+
--bin databend-sqllogictests \
23+
-- \
24+
--handlers "http" \
25+
--run_dir functions \
26+
--run_file <test_file_name> \
27+
--debug \
28+
--parallel 2
29+
```

‎src/query/functions/src/scalars/map.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use std::hash::Hash;
1818
use databend_common_expression::types::nullable::NullableDomain;
1919
use databend_common_expression::types::ArgType;
2020
use databend_common_expression::types::ArrayType;
21+
use databend_common_expression::types::BooleanType;
2122
use databend_common_expression::types::EmptyArrayType;
2223
use databend_common_expression::types::EmptyMapType;
2324
use databend_common_expression::types::GenericType;
@@ -227,4 +228,20 @@ pub fn register(registry: &mut FunctionRegistry) {
227228
|_, _| FunctionDomain::Full,
228229
|map, _| map.len() as u64,
229230
);
231+
232+
registry.register_2_arg_core::<EmptyMapType, GenericType<0>, BooleanType, _, _>(
233+
"map_contains_key",
234+
|_, _, _| FunctionDomain::Full,
235+
|_, _, _| Value::Scalar(false),
236+
);
237+
238+
registry
239+
.register_2_arg::<MapType<GenericType<0>, GenericType<1>>, GenericType<0>, BooleanType, _, _>(
240+
"map_contains_key",
241+
|_, _, _| FunctionDomain::Full,
242+
|map, key, _| {
243+
map.iter()
244+
.any(|(k, _)| k == key)
245+
},
246+
);
230247
}

‎src/query/functions/tests/it/scalars/map.rs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ fn test_map() {
3131
test_map_values(file);
3232
test_map_size(file);
3333
test_map_cat(file);
34+
test_map_contains_key(file);
3435
}
3536

3637
fn test_map_cat(file: &mut impl Write) {
@@ -182,6 +183,43 @@ fn test_map_keys(file: &mut impl Write) {
182183
);
183184
}
184185

186+
fn test_map_contains_key(file: &mut impl Write) {
187+
run_ast(file, "map_contains_key({'a':1,'b':2,'c':3}, 'a')", &[]);
188+
run_ast(file, "map_contains_key({}, 'a')", &[]);
189+
run_ast(file, "map_contains_key({'a':1,'b':2,'c':3}, 'd')", &[]);
190+
run_ast(file, "map_contains_key({'a':NULL,'b':2,'c':NULL}, 'a')", &[
191+
]);
192+
193+
// Nested Maps:: examines recursive key checking capabilities
194+
let columns = [
195+
("a_col", StringType::from_data(vec!["a", "b", "c"])),
196+
("b_col", StringType::from_data(vec!["d", "e", "f"])),
197+
("c_col", StringType::from_data(vec!["x", "y", "z"])),
198+
(
199+
"d_col",
200+
StringType::from_data_with_validity(vec!["v1", "v2", "v3"], vec![true, true, true]),
201+
),
202+
(
203+
"e_col",
204+
StringType::from_data_with_validity(vec!["v4", "v5", ""], vec![true, true, false]),
205+
),
206+
(
207+
"f_col",
208+
StringType::from_data_with_validity(vec!["v6", "", "v7"], vec![true, false, true]),
209+
),
210+
];
211+
run_ast(
212+
file,
213+
"map_contains_key(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'a')",
214+
&columns,
215+
);
216+
run_ast(
217+
file,
218+
"map_contains_key(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'd')",
219+
&columns,
220+
);
221+
}
222+
185223
fn test_map_values(file: &mut impl Write) {
186224
run_ast(file, "map_values({})", &[]);
187225
run_ast(file, "map_values({'a':1,'b':2,'c':3})", &[]);

‎src/query/functions/tests/it/scalars/testdata/function_list.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2440,6 +2440,9 @@ Functions overloads:
24402440
1 map_cat(Map(Nothing) NULL, Map(Nothing) NULL) :: Map(Nothing) NULL
24412441
2 map_cat(Map(T0, T1), Map(T0, T1)) :: Map(T0, T1)
24422442
3 map_cat(Map(T0, T1) NULL, Map(T0, T1) NULL) :: Map(T0, T1) NULL
2443+
0 map_contains_key(Map(Nothing), T0) :: Boolean
2444+
1 map_contains_key(Map(T0, T1), T0) :: Boolean
2445+
2 map_contains_key(Map(T0, T1) NULL, T0 NULL) :: Boolean NULL
24432446
0 map_keys(Map(Nothing)) :: Array(Nothing)
24442447
1 map_keys(Map(T0, T1)) :: Array(T0)
24452448
2 map_keys(Map(T0, T1) NULL) :: Array(T0) NULL

‎src/query/functions/tests/it/scalars/testdata/map.txt

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,3 +527,93 @@ output domain : {[{"k1"..="k2"}], [{[{"nk1"..="nk3"}], [{"new_nv1"..="nv3"}]}]}
527527
output : {'k1':{'nk1':'new_nv1'}, 'k2':{'nk3':'nv3'}}
528528

529529

530+
ast : map_contains_key({'a':1,'b':2,'c':3}, 'a')
531+
raw expr : map_contains_key(map(array('a', 'b', 'c'), array(1, 2, 3)), 'a')
532+
checked expr : map_contains_key<T0=String, T1=UInt8><Map(T0, T1), T0>(map<T0=String, T1=UInt8><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>("a", "b", "c"), array<T0=UInt8><T0, T0, T0>(1_u8, 2_u8, 3_u8)), "a")
533+
optimized expr : true
534+
output type : Boolean
535+
output domain : {TRUE}
536+
output : true
537+
538+
539+
ast : map_contains_key({}, 'a')
540+
raw expr : map_contains_key(map(array(), array()), 'a')
541+
checked expr : map_contains_key<T0=String><Map(Nothing), T0>(map<Array(Nothing), Array(Nothing)>(array<>(), array<>()), "a")
542+
optimized expr : false
543+
output type : Boolean
544+
output domain : {FALSE}
545+
output : false
546+
547+
548+
ast : map_contains_key({'a':1,'b':2,'c':3}, 'd')
549+
raw expr : map_contains_key(map(array('a', 'b', 'c'), array(1, 2, 3)), 'd')
550+
checked expr : map_contains_key<T0=String, T1=UInt8><Map(T0, T1), T0>(map<T0=String, T1=UInt8><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>("a", "b", "c"), array<T0=UInt8><T0, T0, T0>(1_u8, 2_u8, 3_u8)), "d")
551+
optimized expr : false
552+
output type : Boolean
553+
output domain : {FALSE}
554+
output : false
555+
556+
557+
ast : map_contains_key({'a':NULL,'b':2,'c':NULL}, 'a')
558+
raw expr : map_contains_key(map(array('a', 'b', 'c'), array(NULL, 2, NULL)), 'a')
559+
checked expr : map_contains_key<T0=String, T1=UInt8 NULL><Map(T0, T1), T0>(map<T0=String, T1=UInt8 NULL><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>("a", "b", "c"), array<T0=UInt8 NULL><T0, T0, T0>(CAST(NULL AS UInt8 NULL), CAST(2_u8 AS UInt8 NULL), CAST(NULL AS UInt8 NULL))), "a")
560+
optimized expr : true
561+
output type : Boolean
562+
output domain : {TRUE}
563+
output : true
564+
565+
566+
ast : map_contains_key(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'a')
567+
raw expr : map_contains_key(map(array(a_col::String, b_col::String, c_col::String), array(d_col::String NULL, e_col::String NULL, f_col::String NULL)), 'a')
568+
checked expr : map_contains_key<T0=String, T1=String NULL><Map(T0, T1), T0>(map<T0=String, T1=String NULL><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>(a_col, b_col, c_col), array<T0=String NULL><T0, T0, T0>(d_col, e_col, f_col)), "a")
569+
evaluation:
570+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
571+
| | a_col | b_col | c_col | d_col | e_col | f_col | Output |
572+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
573+
| Type | String | String | String | String NULL | String NULL | String NULL | Boolean |
574+
| Domain | {"a"..="c"} | {"d"..="f"} | {"x"..="z"} | {"v1"..="v3"} | {""..="v5"} ∪ {NULL} | {""..="v7"} ∪ {NULL} | Unknown |
575+
| Row 0 | 'a' | 'd' | 'x' | 'v1' | 'v4' | 'v6' | true |
576+
| Row 1 | 'b' | 'e' | 'y' | 'v2' | 'v5' | NULL | false |
577+
| Row 2 | 'c' | 'f' | 'z' | 'v3' | NULL | 'v7' | false |
578+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
579+
evaluation (internal):
580+
+--------+-----------------------------------------------------------------------------------------------------------------+
581+
| Column | Data |
582+
+--------+-----------------------------------------------------------------------------------------------------------------+
583+
| a_col | StringColumn { data: 0x616263, offsets: [0, 1, 2, 3] } |
584+
| b_col | StringColumn { data: 0x646566, offsets: [0, 1, 2, 3] } |
585+
| c_col | StringColumn { data: 0x78797a, offsets: [0, 1, 2, 3] } |
586+
| d_col | NullableColumn { column: StringColumn { data: 0x763176327633, offsets: [0, 2, 4, 6] }, validity: [0b_____111] } |
587+
| e_col | NullableColumn { column: StringColumn { data: 0x76347635, offsets: [0, 2, 4, 4] }, validity: [0b_____011] } |
588+
| f_col | NullableColumn { column: StringColumn { data: 0x76367637, offsets: [0, 2, 2, 4] }, validity: [0b_____101] } |
589+
| Output | Boolean([0b_____001]) |
590+
+--------+-----------------------------------------------------------------------------------------------------------------+
591+
592+
593+
ast : map_contains_key(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'd')
594+
raw expr : map_contains_key(map(array(a_col::String, b_col::String, c_col::String), array(d_col::String NULL, e_col::String NULL, f_col::String NULL)), 'd')
595+
checked expr : map_contains_key<T0=String, T1=String NULL><Map(T0, T1), T0>(map<T0=String, T1=String NULL><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>(a_col, b_col, c_col), array<T0=String NULL><T0, T0, T0>(d_col, e_col, f_col)), "d")
596+
evaluation:
597+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
598+
| | a_col | b_col | c_col | d_col | e_col | f_col | Output |
599+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
600+
| Type | String | String | String | String NULL | String NULL | String NULL | Boolean |
601+
| Domain | {"a"..="c"} | {"d"..="f"} | {"x"..="z"} | {"v1"..="v3"} | {""..="v5"} ∪ {NULL} | {""..="v7"} ∪ {NULL} | Unknown |
602+
| Row 0 | 'a' | 'd' | 'x' | 'v1' | 'v4' | 'v6' | true |
603+
| Row 1 | 'b' | 'e' | 'y' | 'v2' | 'v5' | NULL | false |
604+
| Row 2 | 'c' | 'f' | 'z' | 'v3' | NULL | 'v7' | false |
605+
+--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
606+
evaluation (internal):
607+
+--------+-----------------------------------------------------------------------------------------------------------------+
608+
| Column | Data |
609+
+--------+-----------------------------------------------------------------------------------------------------------------+
610+
| a_col | StringColumn { data: 0x616263, offsets: [0, 1, 2, 3] } |
611+
| b_col | StringColumn { data: 0x646566, offsets: [0, 1, 2, 3] } |
612+
| c_col | StringColumn { data: 0x78797a, offsets: [0, 1, 2, 3] } |
613+
| d_col | NullableColumn { column: StringColumn { data: 0x763176327633, offsets: [0, 2, 4, 6] }, validity: [0b_____111] } |
614+
| e_col | NullableColumn { column: StringColumn { data: 0x76347635, offsets: [0, 2, 4, 4] }, validity: [0b_____011] } |
615+
| f_col | NullableColumn { column: StringColumn { data: 0x76367637, offsets: [0, 2, 2, 4] }, validity: [0b_____101] } |
616+
| Output | Boolean([0b_____001]) |
617+
+--------+-----------------------------------------------------------------------------------------------------------------+
618+
619+

‎tests/sqllogictests/suites/query/functions/02_0074_function_map.test

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,5 +135,37 @@ SELECT MAP_CAT(
135135
----
136136
{'k1':{'nk1':'nv1'},'k2':{'nk2':'nv2'},'k3':{'nk1':'new_nv1'},'k4':{'nk3':'nv3'}}
137137

138+
# Test map_contains_key function
139+
statement ok
140+
CREATE TABLE map_contains_test(col_str Map(String, String Null) Not Null, col_int Map(String, Int Null) Null)
141+
142+
statement ok
143+
INSERT INTO map_contains_test VALUES ({'k1':'v1','k2':'v2','k3':null},{'a':10,'b':20}), ({'k5':'v5','k6':'v6'}, {'d':40,'e':null,'f':50}), ({}, null)
144+
145+
query TT
146+
select map_contains_key(col_str, 'k1'), map_contains_key(col_str, 'k2'), map_contains_key(col_int, 'a') from map_contains_test
147+
----
148+
1 1 1
149+
0 0 0
150+
0 0 NULL
151+
152+
# Test empty map
153+
query
154+
SELECT map_contains_key({}, 'k1')
155+
----
156+
0
157+
158+
# Test non-existent key
159+
query
160+
SELECT map_contains_key({'k1': 'v1', 'k2': 'v2'}, 'k3')
161+
----
162+
0
163+
164+
# Test NULL value
165+
query
166+
SELECT map_contains_key({'k1': 'v1', 'k2': NULL}, 'k2')
167+
----
168+
1
169+
138170
statement ok
139171
DROP DATABASE map_func_test

0 commit comments

Comments
 (0)