micromegas_datafusion_extensions/jsonb/
get.rs

1use datafusion::arrow::array::{
2    Array, BinaryDictionaryBuilder, DictionaryArray, GenericBinaryArray, StringArray,
3};
4use datafusion::arrow::datatypes::{DataType, Int32Type};
5use datafusion::common::{Result, internal_err};
6use datafusion::error::DataFusionError;
7use datafusion::logical_expr::{
8    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
9};
10use jsonb::RawJsonb;
11use std::sync::Arc;
12
13/// A scalar UDF that retrieves a value from a JSONB object by name.
14///
15/// Accepts both Binary and Dictionary<Int32, Binary> inputs.
16/// Returns Dictionary<Int32, Binary> for memory efficiency.
17#[derive(Debug, PartialEq, Eq, Hash)]
18pub struct JsonbGet {
19    signature: Signature,
20}
21
22impl JsonbGet {
23    pub fn new() -> Self {
24        Self {
25            signature: Signature::any(2, Volatility::Immutable),
26        }
27    }
28}
29
30impl Default for JsonbGet {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36fn extract_jsonb_value(jsonb_bytes: &[u8], name: &str) -> Result<Option<Vec<u8>>> {
37    let jsonb = RawJsonb::new(jsonb_bytes);
38    match jsonb.get_by_name(name, true) {
39        Ok(Some(value)) => Ok(Some(value.to_vec())),
40        Ok(None) => Ok(None),
41        Err(e) => Err(DataFusionError::External(e.into())),
42    }
43}
44
45impl ScalarUDFImpl for JsonbGet {
46    fn name(&self) -> &str {
47        "jsonb_get"
48    }
49
50    fn signature(&self) -> &Signature {
51        &self.signature
52    }
53
54    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
55        Ok(DataType::Dictionary(
56            Box::new(DataType::Int32),
57            Box::new(DataType::Binary),
58        ))
59    }
60
61    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
62        let args = ColumnarValue::values_to_arrays(&args.args)?;
63        if args.len() != 2 {
64            return internal_err!("wrong number of arguments to jsonb_get()");
65        }
66
67        let names = args[1]
68            .as_any()
69            .downcast_ref::<StringArray>()
70            .ok_or_else(|| {
71                DataFusionError::Execution("second argument must be a string array".into())
72            })?;
73
74        match args[0].data_type() {
75            DataType::Binary => {
76                // Handle plain Binary JSONB array
77                let binary_array = args[0]
78                    .as_any()
79                    .downcast_ref::<GenericBinaryArray<i32>>()
80                    .ok_or_else(|| {
81                        DataFusionError::Internal("error casting to binary array".into())
82                    })?;
83
84                if binary_array.len() != names.len() {
85                    return internal_err!("arrays of different lengths in jsonb_get()");
86                }
87
88                let mut dict_builder = BinaryDictionaryBuilder::<Int32Type>::new();
89                for i in 0..binary_array.len() {
90                    if binary_array.is_null(i) {
91                        dict_builder.append_null();
92                    } else {
93                        let jsonb_bytes = binary_array.value(i);
94                        let name = names.value(i);
95                        if let Some(value) = extract_jsonb_value(jsonb_bytes, name)? {
96                            dict_builder.append_value(&value);
97                        } else {
98                            dict_builder.append_null();
99                        }
100                    }
101                }
102                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
103            }
104            DataType::Dictionary(_, value_type)
105                if matches!(value_type.as_ref(), DataType::Binary) =>
106            {
107                // Handle dictionary-encoded JSONB array
108                let dict_array = args[0]
109                    .as_any()
110                    .downcast_ref::<DictionaryArray<Int32Type>>()
111                    .ok_or_else(|| {
112                        DataFusionError::Internal("error casting dictionary array".into())
113                    })?;
114
115                if dict_array.len() != names.len() {
116                    return internal_err!("arrays of different lengths in jsonb_get()");
117                }
118
119                let binary_values = dict_array
120                    .values()
121                    .as_any()
122                    .downcast_ref::<GenericBinaryArray<i32>>()
123                    .ok_or_else(|| {
124                        DataFusionError::Internal("dictionary values are not a binary array".into())
125                    })?;
126
127                let mut dict_builder = BinaryDictionaryBuilder::<Int32Type>::new();
128                for i in 0..dict_array.len() {
129                    if dict_array.is_null(i) {
130                        dict_builder.append_null();
131                    } else {
132                        let key_index = dict_array.keys().value(i) as usize;
133                        if key_index < binary_values.len() {
134                            let jsonb_bytes = binary_values.value(key_index);
135                            let name = names.value(i);
136                            if let Some(value) = extract_jsonb_value(jsonb_bytes, name)? {
137                                dict_builder.append_value(&value);
138                            } else {
139                                dict_builder.append_null();
140                            }
141                        } else {
142                            return internal_err!(
143                                "Dictionary key index out of bounds in jsonb_get"
144                            );
145                        }
146                    }
147                }
148                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
149            }
150            _ => internal_err!(
151                "jsonb_get: unsupported input type, expected Binary or Dictionary<Int32, Binary>"
152            ),
153        }
154    }
155}
156
157/// Creates a user-defined function to get a value from a JSONB object by name.
158pub fn make_jsonb_get_udf() -> ScalarUDF {
159    ScalarUDF::new_from_impl(JsonbGet::new())
160}
micromegas_datafusion_extensions/jsonb/get.rs

micromegas_datafusion_extensions/jsonb/
get.rs
