laminar_sql/datafusion/json_udf/
predicates.rs

1//! JSONB existence and containment predicates.
2
3use std::hash::{Hash, Hasher};
4use std::sync::Arc;
5
6use arrow::datatypes::DataType;
7use arrow_array::{Array, BooleanArray, LargeBinaryArray, ListArray, StringArray};
8use datafusion_common::Result;
9use datafusion_expr::{
10    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
11};
12
13use super::{expand_args, json_types};
14
15/// Tests whether a JSONB object contains a key.
16#[derive(Debug)]
17pub struct JsonbExists {
18    signature: Signature,
19}
20
21impl JsonbExists {
22    /// Creates a new `jsonb_exists` UDF.
23    #[must_use]
24    pub fn new() -> Self {
25        Self {
26            signature: Signature::new(
27                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
28                Volatility::Immutable,
29            ),
30        }
31    }
32}
33
34impl Default for JsonbExists {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl PartialEq for JsonbExists {
41    fn eq(&self, _other: &Self) -> bool {
42        true
43    }
44}
45
46impl Eq for JsonbExists {}
47
48impl Hash for JsonbExists {
49    fn hash<H: Hasher>(&self, state: &mut H) {
50        "jsonb_exists".hash(state);
51    }
52}
53
54impl ScalarUDFImpl for JsonbExists {
55    fn as_any(&self) -> &dyn std::any::Any {
56        self
57    }
58
59    fn name(&self) -> &'static str {
60        "jsonb_exists"
61    }
62
63    fn signature(&self) -> &Signature {
64        &self.signature
65    }
66
67    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
68        Ok(DataType::Boolean)
69    }
70
71    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
72        let expanded = expand_args(&args.args)?;
73        let jsonb_arr = expanded[0]
74            .as_any()
75            .downcast_ref::<LargeBinaryArray>()
76            .ok_or_else(|| {
77                datafusion_common::DataFusionError::Internal(
78                    "jsonb_exists: first arg must be LargeBinary".into(),
79                )
80            })?;
81        let key_arr = expanded[1]
82            .as_any()
83            .downcast_ref::<StringArray>()
84            .ok_or_else(|| {
85                datafusion_common::DataFusionError::Internal(
86                    "jsonb_exists: second arg must be Utf8".into(),
87                )
88            })?;
89
90        let result: BooleanArray = (0..jsonb_arr.len())
91            .map(|i| {
92                if jsonb_arr.is_null(i) || key_arr.is_null(i) {
93                    None
94                } else {
95                    Some(json_types::jsonb_has_key(
96                        jsonb_arr.value(i),
97                        key_arr.value(i),
98                    ))
99                }
100            })
101            .collect();
102        Ok(ColumnarValue::Array(Arc::new(result)))
103    }
104}
105
106// ══════════════════════════════════════════════════════════════════
107// jsonb_exists_any(jsonb, text[]) -> bool  (SQL: ?|)
108// ══════════════════════════════════════════════════════════════════
109
110/// `jsonb_exists_any(jsonb, text[]) -> bool`
111///
112/// Returns true if the JSONB object contains any of the given keys.
113/// Maps to the `?|` operator.
114#[derive(Debug)]
115pub struct JsonbExistsAny {
116    signature: Signature,
117}
118
119impl JsonbExistsAny {
120    /// Creates a new `jsonb_exists_any` UDF.
121    #[must_use]
122    pub fn new() -> Self {
123        Self {
124            signature: Signature::new(
125                TypeSignature::Exact(vec![
126                    DataType::LargeBinary,
127                    DataType::List(Arc::new(arrow_schema::Field::new(
128                        "item",
129                        DataType::Utf8,
130                        true,
131                    ))),
132                ]),
133                Volatility::Immutable,
134            ),
135        }
136    }
137}
138
139impl Default for JsonbExistsAny {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145impl PartialEq for JsonbExistsAny {
146    fn eq(&self, _other: &Self) -> bool {
147        true
148    }
149}
150
151impl Eq for JsonbExistsAny {}
152
153impl Hash for JsonbExistsAny {
154    fn hash<H: Hasher>(&self, state: &mut H) {
155        "jsonb_exists_any".hash(state);
156    }
157}
158
159impl ScalarUDFImpl for JsonbExistsAny {
160    fn as_any(&self) -> &dyn std::any::Any {
161        self
162    }
163
164    fn name(&self) -> &'static str {
165        "jsonb_exists_any"
166    }
167
168    fn signature(&self) -> &Signature {
169        &self.signature
170    }
171
172    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
173        Ok(DataType::Boolean)
174    }
175
176    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
177        let expanded = expand_args(&args.args)?;
178        let jsonb_arr = expanded[0]
179            .as_any()
180            .downcast_ref::<LargeBinaryArray>()
181            .ok_or_else(|| {
182                datafusion_common::DataFusionError::Internal(
183                    "jsonb_exists_any: first arg must be LargeBinary".into(),
184                )
185            })?;
186        let keys_arr = expanded[1]
187            .as_any()
188            .downcast_ref::<ListArray>()
189            .ok_or_else(|| {
190                datafusion_common::DataFusionError::Internal(
191                    "jsonb_exists_any: second arg must be List<Utf8>".into(),
192                )
193            })?;
194
195        let result: BooleanArray =
196            (0..jsonb_arr.len())
197                .map(|i| {
198                    if jsonb_arr.is_null(i) || keys_arr.is_null(i) {
199                        return None;
200                    }
201                    let jsonb = jsonb_arr.value(i);
202                    let keys_list = keys_arr.value(i);
203                    let keys = keys_list.as_any().downcast_ref::<StringArray>()?;
204                    Some((0..keys.len()).any(|k| {
205                        !keys.is_null(k) && json_types::jsonb_has_key(jsonb, keys.value(k))
206                    }))
207                })
208                .collect();
209        Ok(ColumnarValue::Array(Arc::new(result)))
210    }
211}
212
213// ══════════════════════════════════════════════════════════════════
214// jsonb_exists_all(jsonb, text[]) -> bool  (SQL: ?&)
215// ══════════════════════════════════════════════════════════════════
216
217/// `jsonb_exists_all(jsonb, text[]) -> bool`
218///
219/// Returns true if the JSONB object contains all of the given keys.
220/// Maps to the `?&` operator.
221#[derive(Debug)]
222pub struct JsonbExistsAll {
223    signature: Signature,
224}
225
226impl JsonbExistsAll {
227    /// Creates a new `jsonb_exists_all` UDF.
228    #[must_use]
229    pub fn new() -> Self {
230        Self {
231            signature: Signature::new(
232                TypeSignature::Exact(vec![
233                    DataType::LargeBinary,
234                    DataType::List(Arc::new(arrow_schema::Field::new(
235                        "item",
236                        DataType::Utf8,
237                        true,
238                    ))),
239                ]),
240                Volatility::Immutable,
241            ),
242        }
243    }
244}
245
246impl Default for JsonbExistsAll {
247    fn default() -> Self {
248        Self::new()
249    }
250}
251
252impl PartialEq for JsonbExistsAll {
253    fn eq(&self, _other: &Self) -> bool {
254        true
255    }
256}
257
258impl Eq for JsonbExistsAll {}
259
260impl Hash for JsonbExistsAll {
261    fn hash<H: Hasher>(&self, state: &mut H) {
262        "jsonb_exists_all".hash(state);
263    }
264}
265
266impl ScalarUDFImpl for JsonbExistsAll {
267    fn as_any(&self) -> &dyn std::any::Any {
268        self
269    }
270
271    fn name(&self) -> &'static str {
272        "jsonb_exists_all"
273    }
274
275    fn signature(&self) -> &Signature {
276        &self.signature
277    }
278
279    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
280        Ok(DataType::Boolean)
281    }
282
283    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
284        let expanded = expand_args(&args.args)?;
285        let jsonb_arr = expanded[0]
286            .as_any()
287            .downcast_ref::<LargeBinaryArray>()
288            .ok_or_else(|| {
289                datafusion_common::DataFusionError::Internal(
290                    "jsonb_exists_all: first arg must be LargeBinary".into(),
291                )
292            })?;
293        let keys_arr = expanded[1]
294            .as_any()
295            .downcast_ref::<ListArray>()
296            .ok_or_else(|| {
297                datafusion_common::DataFusionError::Internal(
298                    "jsonb_exists_all: second arg must be List<Utf8>".into(),
299                )
300            })?;
301
302        let result: BooleanArray =
303            (0..jsonb_arr.len())
304                .map(|i| {
305                    if jsonb_arr.is_null(i) || keys_arr.is_null(i) {
306                        return None;
307                    }
308                    let jsonb = jsonb_arr.value(i);
309                    let keys_list = keys_arr.value(i);
310                    let keys = keys_list.as_any().downcast_ref::<StringArray>()?;
311                    Some((0..keys.len()).all(|k| {
312                        !keys.is_null(k) && json_types::jsonb_has_key(jsonb, keys.value(k))
313                    }))
314                })
315                .collect();
316        Ok(ColumnarValue::Array(Arc::new(result)))
317    }
318}
319
320// ══════════════════════════════════════════════════════════════════
321// jsonb_contains(jsonb, jsonb) -> bool  (SQL: @>)
322// ══════════════════════════════════════════════════════════════════
323
324/// `jsonb_contains(jsonb, jsonb) -> bool`
325///
326/// Returns true if the left JSONB value contains the right.
327/// Maps to the `@>` operator.
328#[derive(Debug)]
329pub struct JsonbContains {
330    signature: Signature,
331}
332
333impl JsonbContains {
334    /// Creates a new `jsonb_contains` UDF.
335    #[must_use]
336    pub fn new() -> Self {
337        Self {
338            signature: Signature::new(
339                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::LargeBinary]),
340                Volatility::Immutable,
341            ),
342        }
343    }
344}
345
346impl Default for JsonbContains {
347    fn default() -> Self {
348        Self::new()
349    }
350}
351
352impl PartialEq for JsonbContains {
353    fn eq(&self, _other: &Self) -> bool {
354        true
355    }
356}
357
358impl Eq for JsonbContains {}
359
360impl Hash for JsonbContains {
361    fn hash<H: Hasher>(&self, state: &mut H) {
362        "jsonb_contains".hash(state);
363    }
364}
365
366impl ScalarUDFImpl for JsonbContains {
367    fn as_any(&self) -> &dyn std::any::Any {
368        self
369    }
370
371    fn name(&self) -> &'static str {
372        "jsonb_contains"
373    }
374
375    fn signature(&self) -> &Signature {
376        &self.signature
377    }
378
379    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
380        Ok(DataType::Boolean)
381    }
382
383    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
384        let expanded = expand_args(&args.args)?;
385        let left_arr = expanded[0]
386            .as_any()
387            .downcast_ref::<LargeBinaryArray>()
388            .ok_or_else(|| {
389                datafusion_common::DataFusionError::Internal(
390                    "jsonb_contains: first arg must be LargeBinary".into(),
391                )
392            })?;
393        let right_arr = expanded[1]
394            .as_any()
395            .downcast_ref::<LargeBinaryArray>()
396            .ok_or_else(|| {
397                datafusion_common::DataFusionError::Internal(
398                    "jsonb_contains: second arg must be LargeBinary".into(),
399                )
400            })?;
401
402        let result: BooleanArray = (0..left_arr.len())
403            .map(|i| {
404                if left_arr.is_null(i) || right_arr.is_null(i) {
405                    None
406                } else {
407                    json_types::jsonb_contains(left_arr.value(i), right_arr.value(i))
408                }
409            })
410            .collect();
411        Ok(ColumnarValue::Array(Arc::new(result)))
412    }
413}
414
415// ══════════════════════════════════════════════════════════════════
416// jsonb_contained_by(jsonb, jsonb) -> bool  (SQL: <@)
417// ══════════════════════════════════════════════════════════════════
418
419/// `jsonb_contained_by(jsonb, jsonb) -> bool`
420///
421/// Returns true if the left JSONB value is contained by the right.
422/// Maps to the `<@` operator.
423#[derive(Debug)]
424pub struct JsonbContainedBy {
425    signature: Signature,
426}
427
428impl JsonbContainedBy {
429    /// Creates a new `jsonb_contained_by` UDF.
430    #[must_use]
431    pub fn new() -> Self {
432        Self {
433            signature: Signature::new(
434                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::LargeBinary]),
435                Volatility::Immutable,
436            ),
437        }
438    }
439}
440
441impl Default for JsonbContainedBy {
442    fn default() -> Self {
443        Self::new()
444    }
445}
446
447impl PartialEq for JsonbContainedBy {
448    fn eq(&self, _other: &Self) -> bool {
449        true
450    }
451}
452
453impl Eq for JsonbContainedBy {}
454
455impl Hash for JsonbContainedBy {
456    fn hash<H: Hasher>(&self, state: &mut H) {
457        "jsonb_contained_by".hash(state);
458    }
459}
460
461impl ScalarUDFImpl for JsonbContainedBy {
462    fn as_any(&self) -> &dyn std::any::Any {
463        self
464    }
465
466    fn name(&self) -> &'static str {
467        "jsonb_contained_by"
468    }
469
470    fn signature(&self) -> &Signature {
471        &self.signature
472    }
473
474    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
475        Ok(DataType::Boolean)
476    }
477
478    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
479        let expanded = expand_args(&args.args)?;
480        let left_arr = expanded[0]
481            .as_any()
482            .downcast_ref::<LargeBinaryArray>()
483            .ok_or_else(|| {
484                datafusion_common::DataFusionError::Internal(
485                    "jsonb_contained_by: first arg must be LargeBinary".into(),
486                )
487            })?;
488        let right_arr = expanded[1]
489            .as_any()
490            .downcast_ref::<LargeBinaryArray>()
491            .ok_or_else(|| {
492                datafusion_common::DataFusionError::Internal(
493                    "jsonb_contained_by: second arg must be LargeBinary".into(),
494                )
495            })?;
496
497        // <@ is just @> with swapped args
498        let result: BooleanArray = (0..left_arr.len())
499            .map(|i| {
500                if left_arr.is_null(i) || right_arr.is_null(i) {
501                    None
502                } else {
503                    json_types::jsonb_contains(right_arr.value(i), left_arr.value(i))
504                }
505            })
506            .collect();
507        Ok(ColumnarValue::Array(Arc::new(result)))
508    }
509}
laminar_sql/datafusion/json_udf/predicates.rs

laminar_sql/datafusion/json_udf/
predicates.rs
