Skip to content

Commit 654e40d

Browse files
JLDLaughlinJLDLaughlin
authored
sql: add support for various map operators (MaterializeInc#4844)
Adds support for the following map operators: &?, &|, <@, @>.
1 parent 6d4fc63 commit 654e40d

File tree

3 files changed

+203
-7
lines changed

3 files changed

+203
-7
lines changed

‎src/expr/src/scalar/func.rs

Lines changed: 46 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1202,6 +1202,38 @@ fn map_contains_key<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a> {
12021202
map.iter().any(|(k2, _v)| k == k2).into()
12031203
}
12041204

1205+
fn map_contains_all_keys<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a> {
1206+
let map = a.unwrap_map();
1207+
let keys = b.unwrap_array();
1208+
1209+
keys.elements()
1210+
.iter()
1211+
.all(|key| !key.is_null() && map.iter().any(|(k, _v)| k == key.unwrap_str()))
1212+
.into()
1213+
}
1214+
1215+
fn map_contains_any_keys<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a> {
1216+
let map = a.unwrap_map();
1217+
let keys = b.unwrap_array();
1218+
1219+
keys.elements()
1220+
.iter()
1221+
.any(|key| !key.is_null() && map.iter().any(|(k, _v)| k == key.unwrap_str()))
1222+
.into()
1223+
}
1224+
1225+
fn map_contains_map<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a> {
1226+
let map_a = a.unwrap_map();
1227+
b.unwrap_map()
1228+
.iter()
1229+
.all(|(b_key, b_val)| {
1230+
map_a
1231+
.iter()
1232+
.any(|(a_key, a_val)| (a_key == b_key) && (a_val == b_val))
1233+
})
1234+
.into()
1235+
}
1236+
12051237
fn map_get_value<'a>(a: Datum<'a>, b: Datum<'a>) -> Datum<'a> {
12061238
let target_key = b.unwrap_str();
12071239
match a.unwrap_map().iter().find(|(key, _v)| target_key == *key) {
@@ -1866,6 +1898,9 @@ pub enum BinaryFunc {
18661898
MapContainsKey,
18671899
MapGetValue,
18681900
MapGetValues,
1901+
MapContainsAllKeys,
1902+
MapContainsAnyKeys,
1903+
MapContainsMap,
18691904
ConvertFrom,
18701905
Trim,
18711906
TrimLeading,
@@ -1988,6 +2023,9 @@ impl BinaryFunc {
19882023
BinaryFunc::MapContainsKey => Ok(eager!(map_contains_key)),
19892024
BinaryFunc::MapGetValue => Ok(eager!(map_get_value)),
19902025
BinaryFunc::MapGetValues => Ok(eager!(map_get_values, temp_storage)),
2026+
BinaryFunc::MapContainsAllKeys => Ok(eager!(map_contains_all_keys)),
2027+
BinaryFunc::MapContainsAnyKeys => Ok(eager!(map_contains_any_keys)),
2028+
BinaryFunc::MapContainsMap => Ok(eager!(map_contains_map)),
19912029
BinaryFunc::RoundDecimal(scale) => Ok(eager!(round_decimal_binary, *scale)),
19922030
BinaryFunc::ConvertFrom => eager!(convert_from),
19932031
BinaryFunc::Trim => Ok(eager!(trim)),
@@ -2129,9 +2167,8 @@ impl BinaryFunc {
21292167
| JsonbDeleteInt64
21302168
| JsonbDeleteString => ScalarType::Jsonb.nullable(true),
21312169

2132-
JsonbContainsString | JsonbContainsJsonb | MapContainsKey => {
2133-
ScalarType::Bool.nullable(in_nullable)
2134-
}
2170+
JsonbContainsString | JsonbContainsJsonb | MapContainsKey | MapContainsAllKeys
2171+
| MapContainsAnyKeys | MapContainsMap => ScalarType::Bool.nullable(in_nullable),
21352172

21362173
MapGetValue => input1_type
21372174
.scalar_type
@@ -2294,6 +2331,9 @@ impl BinaryFunc {
22942331
| MapContainsKey
22952332
| MapGetValue
22962333
| MapGetValues
2334+
| MapContainsAllKeys
2335+
| MapContainsAnyKeys
2336+
| MapContainsMap
22972337
| TextConcat
22982338
| ListIndex
22992339
| IsRegexpMatch { .. }
@@ -2397,10 +2437,12 @@ impl fmt::Display for BinaryFunc {
23972437
BinaryFunc::JsonbGetString { .. } => f.write_str("->"),
23982438
BinaryFunc::JsonbContainsString | BinaryFunc::MapContainsKey => f.write_str("?"),
23992439
BinaryFunc::JsonbConcat => f.write_str("||"),
2400-
BinaryFunc::JsonbContainsJsonb => f.write_str("@>"),
2440+
BinaryFunc::JsonbContainsJsonb | BinaryFunc::MapContainsMap => f.write_str("@>"),
24012441
BinaryFunc::JsonbDeleteInt64 => f.write_str("-"),
24022442
BinaryFunc::JsonbDeleteString => f.write_str("-"),
24032443
BinaryFunc::MapGetValue | BinaryFunc::MapGetValues => f.write_str("->"),
2444+
BinaryFunc::MapContainsAllKeys => f.write_str("?&"),
2445+
BinaryFunc::MapContainsAnyKeys => f.write_str("?|"),
24042446
BinaryFunc::RoundDecimal(_) => f.write_str("round"),
24052447
BinaryFunc::ConvertFrom => f.write_str("convert_from"),
24062448
BinaryFunc::Trim => f.write_str("btrim"),

‎src/sql/src/plan/func.rs

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1822,7 +1822,7 @@ lazy_static! {
18221822
params!(ListElementAny, ListAny) => ElementListConcat
18231823
},
18241824

1825-
//JSON
1825+
//JSON and MAP
18261826
"->" => Scalar {
18271827
params!(Jsonb, Int64) => JsonbGetInt64 { stringify: false },
18281828
params!(Jsonb, String) => JsonbGetString { stringify: false },
@@ -1844,7 +1844,8 @@ lazy_static! {
18441844
params!(String, Jsonb) => Operation::binary(|_ecx, lhs, rhs| {
18451845
Ok(lhs.call_unary(UnaryFunc::CastStringToJsonb)
18461846
.call_binary(rhs, JsonbContainsJsonb))
1847-
})
1847+
}),
1848+
params!(MapAny, MapAny) => MapContainsMap
18481849
},
18491850
"<@" => Scalar {
18501851
params!(Jsonb, Jsonb) => Operation::binary(|_ecx, lhs, rhs| {
@@ -1862,12 +1863,21 @@ lazy_static! {
18621863
lhs.call_unary(UnaryFunc::CastStringToJsonb),
18631864
BinaryFunc::JsonbContainsJsonb,
18641865
))
1866+
}),
1867+
params!(MapAny, MapAny) => Operation::binary(|_ecx, lhs, rhs| {
1868+
Ok(rhs.call_binary(lhs, MapContainsMap))
18651869
})
18661870
},
18671871
"?" => Scalar {
18681872
params!(Jsonb, String) => JsonbContainsString,
18691873
params!(MapAny, String) => MapContainsKey
18701874
},
1875+
"?&" => Scalar {
1876+
params!(MapAny, Plain(Array(Box::new(String)))) => MapContainsAllKeys
1877+
},
1878+
"?|" => Scalar {
1879+
params!(MapAny, Plain(Array(Box::new(String)))) => MapContainsAnyKeys
1880+
},
18711881
// COMPARISON OPS
18721882
// n.b. Decimal impls are separated from other types because they
18731883
// require a function pointer, which you cannot dynamically generate.

‎test/sqllogictest/map.slt

Lines changed: 145 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ SELECT '{a=>{b=>c}}'::map[text=>map[text=>text]]
8787

8888
# Test map operators.
8989

90-
## Test ?
90+
## ?
9191
query T
9292
SELECT '{a=>1, b=>2}'::map[text=>int] ? 'a'
9393
----
@@ -116,6 +116,150 @@ SELECT '{""=>1}'::map[text=>int] ? ''
116116
----
117117
true
118118

119+
## ?&
120+
query error string literal does not support casting from string to string\[\]
121+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& 'a'
122+
123+
query error arguments cannot be implicitly cast to any implementation's parameters
124+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY[1]
125+
126+
query error cannot determine type of empty array
127+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY[]
128+
129+
query T
130+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY[NULL]
131+
----
132+
false
133+
134+
query T
135+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY['a']
136+
----
137+
true
138+
139+
query T
140+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY['b', 'a']
141+
----
142+
true
143+
144+
query T
145+
SELECT '{a=>1, b=>2}'::map[text=>int] ?& ARRAY['c', 'b']
146+
----
147+
false
148+
149+
query error arguments cannot be implicitly cast to any implementation's parameters
150+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?& ARRAY[1]
151+
152+
query T
153+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?& ARRAY['1']
154+
----
155+
true
156+
157+
query T
158+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?& ARRAY['']
159+
----
160+
false
161+
162+
query T
163+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?& ARRAY['']
164+
----
165+
false
166+
167+
## ?|
168+
query error string literal does not support casting from string to string\[\]
169+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| 'a'
170+
171+
query error arguments cannot be implicitly cast to any implementation's parameters
172+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| ARRAY[1]
173+
174+
query T
175+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| ARRAY[NULL]
176+
----
177+
false
178+
179+
query T
180+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| ARRAY['a']
181+
----
182+
true
183+
184+
query T
185+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| ARRAY['c', 'b']
186+
----
187+
true
188+
189+
query T
190+
SELECT '{a=>1, b=>2}'::map[text=>int] ?| ARRAY['c', 'd', '1']
191+
----
192+
false
193+
194+
query error arguments cannot be implicitly cast to any implementation's parameters
195+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?| ARRAY[1]
196+
197+
query T
198+
SELECT '{1=>t, 2=>f}'::map[text=>bool] ?| ARRAY['1']
199+
----
200+
true
201+
202+
## @>
203+
query error could not determine polymorphic type because input has type unknown
204+
SELECT '{a=>1, b=>2}'::map[text=>int] @> 'a'
205+
206+
query error arguments cannot be implicitly cast to any implementation's parameters
207+
SELECT '{a=>1, b=>2}'::map[text=>int] @> 'a'::text
208+
209+
query error arguments cannot be implicitly cast to any implementation's parameters
210+
SELECT '{a=>1, b=>2}'::map[text=>int] @> ARRAY[1]
211+
212+
query T
213+
SELECT '{a=>1, b=>2}'::map[text=>int] @> '{a=>t}'::map[text=>bool]
214+
----
215+
false
216+
217+
query T
218+
SELECT '{a=>1, b=>2}'::map[text=>int] @> '{a=>1}'::map[text=>int]
219+
----
220+
true
221+
222+
query T
223+
SELECT '{a=>1, b=>2}'::map[text=>int] @> '{a=>1, b=>2}'::map[text=>int]
224+
----
225+
true
226+
227+
query T
228+
SELECT '{a=>1, b=>2}'::map[text=>int] @> '{a=>10, b=>20}'::map[text=>int]
229+
----
230+
false
231+
232+
query T
233+
SELECT '{a=>1, b=>2}'::map[text=>int] @> '{a=>1, b=>2, c=>3}'::map[text=>int]
234+
----
235+
false
236+
237+
## <@
238+
query T
239+
SELECT '{a=>1, b=>2}'::map[text=>int] <@ '{a=>t}'::map[text=>bool]
240+
----
241+
false
242+
243+
query T
244+
SELECT '{a=>1, b=>2}'::map[text=>int] <@ '{a=>1}'::map[text=>int]
245+
----
246+
false
247+
248+
query T
249+
SELECT '{a=>1, b=>2}'::map[text=>int] <@ '{a=>1, b=>2}'::map[text=>int]
250+
----
251+
true
252+
253+
query T
254+
SELECT '{a=>1, b=>2}'::map[text=>int] <@ '{a=>10, b=>20}'::map[text=>int]
255+
----
256+
false
257+
258+
query T
259+
SELECT '{a=>1, b=>2}'::map[text=>int] <@ '{a=>1, b=>2, c=>3}'::map[text=>int]
260+
----
261+
true
262+
119263
## ->
120264
query T
121265
SELECT '{a=>1, b=>2}'::map[text=>int] -> ''

0 commit comments

Comments
 (0)