laminar_sql/datafusion/json_udf/
extraction.rs

1//! JSONB field, index, and path extraction UDFs.
2
3use std::hash::{Hash, Hasher};
4use std::sync::Arc;
5
6use arrow::datatypes::DataType;
7use arrow_array::{
8    builder::{LargeBinaryBuilder, StringBuilder},
9    Array, LargeBinaryArray, ListArray, StringArray,
10};
11use datafusion_common::Result;
12use datafusion_expr::{
13    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
14};
15
16use super::{expand_args, json_types};
17
18/// Extracts a JSON object field by key, returning JSONB.
19#[derive(Debug)]
20pub struct JsonbGet {
21    signature: Signature,
22}
23
24impl JsonbGet {
25    /// Creates a new `jsonb_get` UDF.
26    #[must_use]
27    pub fn new() -> Self {
28        Self {
29            signature: Signature::new(
30                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
31                Volatility::Immutable,
32            ),
33        }
34    }
35}
36
37impl Default for JsonbGet {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl PartialEq for JsonbGet {
44    fn eq(&self, _other: &Self) -> bool {
45        true
46    }
47}
48
49impl Eq for JsonbGet {}
50
51impl Hash for JsonbGet {
52    fn hash<H: Hasher>(&self, state: &mut H) {
53        "jsonb_get".hash(state);
54    }
55}
56
57impl ScalarUDFImpl for JsonbGet {
58    fn as_any(&self) -> &dyn std::any::Any {
59        self
60    }
61
62    fn name(&self) -> &'static str {
63        "jsonb_get"
64    }
65
66    fn signature(&self) -> &Signature {
67        &self.signature
68    }
69
70    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
71        Ok(DataType::LargeBinary)
72    }
73
74    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
75        let expanded = expand_args(&args.args)?;
76        let jsonb_arr = expanded[0]
77            .as_any()
78            .downcast_ref::<LargeBinaryArray>()
79            .ok_or_else(|| {
80                datafusion_common::DataFusionError::Internal(
81                    "jsonb_get: first arg must be LargeBinary".into(),
82                )
83            })?;
84        let key_arr = expanded[1]
85            .as_any()
86            .downcast_ref::<StringArray>()
87            .ok_or_else(|| {
88                datafusion_common::DataFusionError::Internal(
89                    "jsonb_get: second arg must be Utf8".into(),
90                )
91            })?;
92
93        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
94        for i in 0..jsonb_arr.len() {
95            if jsonb_arr.is_null(i) || key_arr.is_null(i) {
96                builder.append_null();
97            } else {
98                let jsonb = jsonb_arr.value(i);
99                let key = key_arr.value(i);
100                match json_types::jsonb_get_field(jsonb, key) {
101                    Some(val) => builder.append_value(val),
102                    None => builder.append_null(),
103                }
104            }
105        }
106        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
107    }
108}
109
110// ══════════════════════════════════════════════════════════════════
111// jsonb_get_idx(jsonb, int32) -> jsonb  (SQL: -> with int)
112// ══════════════════════════════════════════════════════════════════
113
114/// `jsonb_get_idx(jsonb, int32) -> jsonb`
115///
116/// Extracts a JSON array element by index, returning JSONB.
117/// Maps to the `->` operator with an integer index.
118#[derive(Debug)]
119pub struct JsonbGetIdx {
120    signature: Signature,
121}
122
123impl JsonbGetIdx {
124    /// Creates a new `jsonb_get_idx` UDF.
125    #[must_use]
126    pub fn new() -> Self {
127        Self {
128            signature: Signature::new(
129                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Int32]),
130                Volatility::Immutable,
131            ),
132        }
133    }
134}
135
136impl Default for JsonbGetIdx {
137    fn default() -> Self {
138        Self::new()
139    }
140}
141
142impl PartialEq for JsonbGetIdx {
143    fn eq(&self, _other: &Self) -> bool {
144        true
145    }
146}
147
148impl Eq for JsonbGetIdx {}
149
150impl Hash for JsonbGetIdx {
151    fn hash<H: Hasher>(&self, state: &mut H) {
152        "jsonb_get_idx".hash(state);
153    }
154}
155
156impl ScalarUDFImpl for JsonbGetIdx {
157    fn as_any(&self) -> &dyn std::any::Any {
158        self
159    }
160
161    fn name(&self) -> &'static str {
162        "jsonb_get_idx"
163    }
164
165    fn signature(&self) -> &Signature {
166        &self.signature
167    }
168
169    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
170        Ok(DataType::LargeBinary)
171    }
172
173    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
174        let expanded = expand_args(&args.args)?;
175        let jsonb_arr = expanded[0]
176            .as_any()
177            .downcast_ref::<LargeBinaryArray>()
178            .ok_or_else(|| {
179                datafusion_common::DataFusionError::Internal(
180                    "jsonb_get_idx: first arg must be LargeBinary".into(),
181                )
182            })?;
183        let idx_arr = expanded[1]
184            .as_any()
185            .downcast_ref::<arrow_array::Int32Array>()
186            .ok_or_else(|| {
187                datafusion_common::DataFusionError::Internal(
188                    "jsonb_get_idx: second arg must be Int32".into(),
189                )
190            })?;
191
192        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
193        for i in 0..jsonb_arr.len() {
194            if jsonb_arr.is_null(i) || idx_arr.is_null(i) {
195                builder.append_null();
196            } else {
197                let jsonb = jsonb_arr.value(i);
198                let idx = idx_arr.value(i);
199                match usize::try_from(idx)
200                    .ok()
201                    .and_then(|u| json_types::jsonb_array_get(jsonb, u))
202                {
203                    Some(val) => builder.append_value(val),
204                    None => builder.append_null(),
205                }
206            }
207        }
208        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
209    }
210}
211
212// ══════════════════════════════════════════════════════════════════
213// jsonb_get_text(jsonb, text) -> text  (SQL: ->>)
214// ══════════════════════════════════════════════════════════════════
215
216/// `jsonb_get_text(jsonb, text) -> text`
217///
218/// Extracts a JSON object field by key, returning TEXT.
219/// Maps to the `->>` operator with a text key.
220#[derive(Debug)]
221pub struct JsonbGetText {
222    signature: Signature,
223}
224
225impl JsonbGetText {
226    /// Creates a new `jsonb_get_text` UDF.
227    #[must_use]
228    pub fn new() -> Self {
229        Self {
230            signature: Signature::new(
231                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
232                Volatility::Immutable,
233            ),
234        }
235    }
236}
237
238impl Default for JsonbGetText {
239    fn default() -> Self {
240        Self::new()
241    }
242}
243
244impl PartialEq for JsonbGetText {
245    fn eq(&self, _other: &Self) -> bool {
246        true
247    }
248}
249
250impl Eq for JsonbGetText {}
251
252impl Hash for JsonbGetText {
253    fn hash<H: Hasher>(&self, state: &mut H) {
254        "jsonb_get_text".hash(state);
255    }
256}
257
258impl ScalarUDFImpl for JsonbGetText {
259    fn as_any(&self) -> &dyn std::any::Any {
260        self
261    }
262
263    fn name(&self) -> &'static str {
264        "jsonb_get_text"
265    }
266
267    fn signature(&self) -> &Signature {
268        &self.signature
269    }
270
271    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
272        Ok(DataType::Utf8)
273    }
274
275    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
276        let expanded = expand_args(&args.args)?;
277        let jsonb_arr = expanded[0]
278            .as_any()
279            .downcast_ref::<LargeBinaryArray>()
280            .ok_or_else(|| {
281                datafusion_common::DataFusionError::Internal(
282                    "jsonb_get_text: first arg must be LargeBinary".into(),
283                )
284            })?;
285        let key_arr = expanded[1]
286            .as_any()
287            .downcast_ref::<StringArray>()
288            .ok_or_else(|| {
289                datafusion_common::DataFusionError::Internal(
290                    "jsonb_get_text: second arg must be Utf8".into(),
291                )
292            })?;
293
294        let mut builder = StringBuilder::with_capacity(jsonb_arr.len(), 256);
295        for i in 0..jsonb_arr.len() {
296            if jsonb_arr.is_null(i) || key_arr.is_null(i) {
297                builder.append_null();
298            } else {
299                let jsonb = jsonb_arr.value(i);
300                let key = key_arr.value(i);
301                match json_types::jsonb_get_field(jsonb, key) {
302                    Some(val) => match json_types::jsonb_to_text(val) {
303                        Some(text) => builder.append_value(&text),
304                        None => builder.append_null(),
305                    },
306                    None => builder.append_null(),
307                }
308            }
309        }
310        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
311    }
312}
313
314// ══════════════════════════════════════════════════════════════════
315// jsonb_get_text_idx(jsonb, int32) -> text  (SQL: ->> with int)
316// ══════════════════════════════════════════════════════════════════
317
318/// `jsonb_get_text_idx(jsonb, int32) -> text`
319///
320/// Extracts a JSON array element by index, returning TEXT.
321/// Maps to the `->>` operator with an integer index.
322#[derive(Debug)]
323pub struct JsonbGetTextIdx {
324    signature: Signature,
325}
326
327impl JsonbGetTextIdx {
328    /// Creates a new `jsonb_get_text_idx` UDF.
329    #[must_use]
330    pub fn new() -> Self {
331        Self {
332            signature: Signature::new(
333                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Int32]),
334                Volatility::Immutable,
335            ),
336        }
337    }
338}
339
340impl Default for JsonbGetTextIdx {
341    fn default() -> Self {
342        Self::new()
343    }
344}
345
346impl PartialEq for JsonbGetTextIdx {
347    fn eq(&self, _other: &Self) -> bool {
348        true
349    }
350}
351
352impl Eq for JsonbGetTextIdx {}
353
354impl Hash for JsonbGetTextIdx {
355    fn hash<H: Hasher>(&self, state: &mut H) {
356        "jsonb_get_text_idx".hash(state);
357    }
358}
359
360impl ScalarUDFImpl for JsonbGetTextIdx {
361    fn as_any(&self) -> &dyn std::any::Any {
362        self
363    }
364
365    fn name(&self) -> &'static str {
366        "jsonb_get_text_idx"
367    }
368
369    fn signature(&self) -> &Signature {
370        &self.signature
371    }
372
373    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
374        Ok(DataType::Utf8)
375    }
376
377    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
378        let expanded = expand_args(&args.args)?;
379        let jsonb_arr = expanded[0]
380            .as_any()
381            .downcast_ref::<LargeBinaryArray>()
382            .ok_or_else(|| {
383                datafusion_common::DataFusionError::Internal(
384                    "jsonb_get_text_idx: first arg must be LargeBinary".into(),
385                )
386            })?;
387        let idx_arr = expanded[1]
388            .as_any()
389            .downcast_ref::<arrow_array::Int32Array>()
390            .ok_or_else(|| {
391                datafusion_common::DataFusionError::Internal(
392                    "jsonb_get_text_idx: second arg must be Int32".into(),
393                )
394            })?;
395
396        let mut builder = StringBuilder::with_capacity(jsonb_arr.len(), 256);
397        for i in 0..jsonb_arr.len() {
398            if jsonb_arr.is_null(i) || idx_arr.is_null(i) {
399                builder.append_null();
400            } else {
401                let jsonb = jsonb_arr.value(i);
402                let idx = idx_arr.value(i);
403                match usize::try_from(idx)
404                    .ok()
405                    .and_then(|u| json_types::jsonb_array_get(jsonb, u))
406                {
407                    Some(val) => match json_types::jsonb_to_text(val) {
408                        Some(text) => builder.append_value(&text),
409                        None => builder.append_null(),
410                    },
411                    None => builder.append_null(),
412                }
413            }
414        }
415        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
416    }
417}
418
419// ══════════════════════════════════════════════════════════════════
420// jsonb_get_path(jsonb, text[]) -> jsonb  (SQL: #>)
421// ══════════════════════════════════════════════════════════════════
422
423/// `jsonb_get_path(jsonb, text[]) -> jsonb`
424///
425/// Extracts a JSONB value at a nested path given as a text array.
426/// Maps to the `#>` operator.
427#[derive(Debug)]
428pub struct JsonbGetPath {
429    signature: Signature,
430}
431
432impl JsonbGetPath {
433    /// Creates a new `jsonb_get_path` UDF.
434    #[must_use]
435    pub fn new() -> Self {
436        Self {
437            signature: Signature::new(
438                TypeSignature::Exact(vec![
439                    DataType::LargeBinary,
440                    DataType::List(Arc::new(arrow_schema::Field::new(
441                        "item",
442                        DataType::Utf8,
443                        true,
444                    ))),
445                ]),
446                Volatility::Immutable,
447            ),
448        }
449    }
450}
451
452impl Default for JsonbGetPath {
453    fn default() -> Self {
454        Self::new()
455    }
456}
457
458impl PartialEq for JsonbGetPath {
459    fn eq(&self, _other: &Self) -> bool {
460        true
461    }
462}
463
464impl Eq for JsonbGetPath {}
465
466impl Hash for JsonbGetPath {
467    fn hash<H: Hasher>(&self, state: &mut H) {
468        "jsonb_get_path".hash(state);
469    }
470}
471
472impl ScalarUDFImpl for JsonbGetPath {
473    fn as_any(&self) -> &dyn std::any::Any {
474        self
475    }
476
477    fn name(&self) -> &'static str {
478        "jsonb_get_path"
479    }
480
481    fn signature(&self) -> &Signature {
482        &self.signature
483    }
484
485    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
486        Ok(DataType::LargeBinary)
487    }
488
489    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
490        let expanded = expand_args(&args.args)?;
491        let jsonb_arr = expanded[0]
492            .as_any()
493            .downcast_ref::<LargeBinaryArray>()
494            .ok_or_else(|| {
495                datafusion_common::DataFusionError::Internal(
496                    "jsonb_get_path: first arg must be LargeBinary".into(),
497                )
498            })?;
499        let path_arr = expanded[1]
500            .as_any()
501            .downcast_ref::<ListArray>()
502            .ok_or_else(|| {
503                datafusion_common::DataFusionError::Internal(
504                    "jsonb_get_path: second arg must be List<Utf8>".into(),
505                )
506            })?;
507
508        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
509        for i in 0..jsonb_arr.len() {
510            if jsonb_arr.is_null(i) || path_arr.is_null(i) {
511                builder.append_null();
512            } else {
513                let jsonb = jsonb_arr.value(i);
514                let path_list = path_arr.value(i);
515                let path_strings = path_list
516                    .as_any()
517                    .downcast_ref::<StringArray>()
518                    .ok_or_else(|| {
519                        datafusion_common::DataFusionError::Internal(
520                            "jsonb_get_path: path elements must be Utf8".into(),
521                        )
522                    })?;
523                match walk_path(jsonb, path_strings) {
524                    Some(val) => builder.append_value(val),
525                    None => builder.append_null(),
526                }
527            }
528        }
529        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
530    }
531}
532
533/// Walk a JSONB value through a sequence of string keys.
534fn walk_path<'a>(mut jsonb: &'a [u8], path: &StringArray) -> Option<&'a [u8]> {
535    for i in 0..path.len() {
536        if path.is_null(i) {
537            return None;
538        }
539        let key = path.value(i);
540        // Try object field access first
541        if let Some(next) = json_types::jsonb_get_field(jsonb, key) {
542            jsonb = next;
543        } else if let Ok(idx) = key.parse::<usize>() {
544            // Fall back to array index
545            jsonb = json_types::jsonb_array_get(jsonb, idx)?;
546        } else {
547            return None;
548        }
549    }
550    Some(jsonb)
551}
552
553// ══════════════════════════════════════════════════════════════════
554// jsonb_get_path_text(jsonb, text[]) -> text  (SQL: #>>)
555// ══════════════════════════════════════════════════════════════════
556
557/// `jsonb_get_path_text(jsonb, text[]) -> text`
558///
559/// Extracts a text value at a nested path given as a text array.
560/// Maps to the `#>>` operator.
561#[derive(Debug)]
562pub struct JsonbGetPathText {
563    signature: Signature,
564}
565
566impl JsonbGetPathText {
567    /// Creates a new `jsonb_get_path_text` UDF.
568    #[must_use]
569    pub fn new() -> Self {
570        Self {
571            signature: Signature::new(
572                TypeSignature::Exact(vec![
573                    DataType::LargeBinary,
574                    DataType::List(Arc::new(arrow_schema::Field::new(
575                        "item",
576                        DataType::Utf8,
577                        true,
578                    ))),
579                ]),
580                Volatility::Immutable,
581            ),
582        }
583    }
584}
585
586impl Default for JsonbGetPathText {
587    fn default() -> Self {
588        Self::new()
589    }
590}
591
592impl PartialEq for JsonbGetPathText {
593    fn eq(&self, _other: &Self) -> bool {
594        true
595    }
596}
597
598impl Eq for JsonbGetPathText {}
599
600impl Hash for JsonbGetPathText {
601    fn hash<H: Hasher>(&self, state: &mut H) {
602        "jsonb_get_path_text".hash(state);
603    }
604}
605
606impl ScalarUDFImpl for JsonbGetPathText {
607    fn as_any(&self) -> &dyn std::any::Any {
608        self
609    }
610
611    fn name(&self) -> &'static str {
612        "jsonb_get_path_text"
613    }
614
615    fn signature(&self) -> &Signature {
616        &self.signature
617    }
618
619    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
620        Ok(DataType::Utf8)
621    }
622
623    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
624        let expanded = expand_args(&args.args)?;
625        let jsonb_arr = expanded[0]
626            .as_any()
627            .downcast_ref::<LargeBinaryArray>()
628            .ok_or_else(|| {
629                datafusion_common::DataFusionError::Internal(
630                    "jsonb_get_path_text: first arg must be LargeBinary".into(),
631                )
632            })?;
633        let path_arr = expanded[1]
634            .as_any()
635            .downcast_ref::<ListArray>()
636            .ok_or_else(|| {
637                datafusion_common::DataFusionError::Internal(
638                    "jsonb_get_path_text: second arg must be List<Utf8>".into(),
639                )
640            })?;
641
642        let mut builder = StringBuilder::with_capacity(jsonb_arr.len(), 256);
643        for i in 0..jsonb_arr.len() {
644            if jsonb_arr.is_null(i) || path_arr.is_null(i) {
645                builder.append_null();
646            } else {
647                let jsonb = jsonb_arr.value(i);
648                let path_list = path_arr.value(i);
649                let path_strings = path_list
650                    .as_any()
651                    .downcast_ref::<StringArray>()
652                    .ok_or_else(|| {
653                        datafusion_common::DataFusionError::Internal(
654                            "jsonb_get_path_text: path elements must be Utf8".into(),
655                        )
656                    })?;
657                match walk_path(jsonb, path_strings) {
658                    Some(val) => match json_types::jsonb_to_text(val) {
659                        Some(text) => builder.append_value(&text),
660                        None => builder.append_null(),
661                    },
662                    None => builder.append_null(),
663                }
664            }
665        }
666        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
667    }
668}
laminar_sql/datafusion/json_udf/extraction.rs

laminar_sql/datafusion/json_udf/
extraction.rs
