micromegas_datafusion_extensions/jsonb/
cast.rs

1use datafusion::arrow::array::{
2    Array, DictionaryArray, Float64Array, GenericBinaryArray, Int64Array, StringDictionaryBuilder,
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 casts a JSONB value to a string.
14///
15/// Accepts both Binary and Dictionary<Int32, Binary> inputs.
16/// Returns Dictionary<Int32, Utf8> for memory efficiency.
17#[derive(Debug, PartialEq, Eq, Hash)]
18pub struct JsonbAsString {
19    signature: Signature,
20}
21
22impl JsonbAsString {
23    pub fn new() -> Self {
24        Self {
25            signature: Signature::any(1, Volatility::Immutable),
26        }
27    }
28}
29
30impl Default for JsonbAsString {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36fn extract_string_from_jsonb(jsonb_bytes: &[u8]) -> Result<Option<String>> {
37    let jsonb = RawJsonb::new(jsonb_bytes);
38    match jsonb.as_str() {
39        Ok(Some(value)) => Ok(Some(value.to_string())),
40        Ok(None) => Ok(None),
41        Err(e) => Err(DataFusionError::External(e.into())),
42    }
43}
44
45impl ScalarUDFImpl for JsonbAsString {
46    fn name(&self) -> &str {
47        "jsonb_as_string"
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::Utf8),
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() != 1 {
64            return internal_err!("wrong number of arguments to jsonb_as_string()");
65        }
66
67        match args[0].data_type() {
68            DataType::Binary => {
69                let binary_array = args[0]
70                    .as_any()
71                    .downcast_ref::<GenericBinaryArray<i32>>()
72                    .ok_or_else(|| {
73                        DataFusionError::Internal("error casting to binary array".into())
74                    })?;
75
76                let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
77                for i in 0..binary_array.len() {
78                    if binary_array.is_null(i) {
79                        dict_builder.append_null();
80                    } else {
81                        let jsonb_bytes = binary_array.value(i);
82                        if let Some(value) = extract_string_from_jsonb(jsonb_bytes)? {
83                            dict_builder.append_value(value);
84                        } else {
85                            dict_builder.append_null();
86                        }
87                    }
88                }
89                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
90            }
91            DataType::Dictionary(_, value_type)
92                if matches!(value_type.as_ref(), DataType::Binary) =>
93            {
94                let dict_array = args[0]
95                    .as_any()
96                    .downcast_ref::<DictionaryArray<Int32Type>>()
97                    .ok_or_else(|| {
98                        DataFusionError::Internal("error casting dictionary array".into())
99                    })?;
100
101                let binary_values = dict_array
102                    .values()
103                    .as_any()
104                    .downcast_ref::<GenericBinaryArray<i32>>()
105                    .ok_or_else(|| {
106                        DataFusionError::Internal("dictionary values are not a binary array".into())
107                    })?;
108
109                let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
110                for i in 0..dict_array.len() {
111                    if dict_array.is_null(i) {
112                        dict_builder.append_null();
113                    } else {
114                        let key_index = dict_array.keys().value(i) as usize;
115                        if key_index < binary_values.len() {
116                            let jsonb_bytes = binary_values.value(key_index);
117                            if let Some(value) = extract_string_from_jsonb(jsonb_bytes)? {
118                                dict_builder.append_value(value);
119                            } else {
120                                dict_builder.append_null();
121                            }
122                        } else {
123                            return internal_err!(
124                                "Dictionary key index out of bounds in jsonb_as_string"
125                            );
126                        }
127                    }
128                }
129                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
130            }
131            _ => internal_err!(
132                "jsonb_as_string: unsupported input type, expected Binary or Dictionary<Int32, Binary>"
133            ),
134        }
135    }
136}
137
138/// Creates a user-defined function to cast a JSONB value to a string.
139pub fn make_jsonb_as_string_udf() -> ScalarUDF {
140    ScalarUDF::new_from_impl(JsonbAsString::new())
141}
142
143/// A scalar UDF that casts a JSONB value to a f64.
144///
145/// Accepts both Binary and Dictionary<Int32, Binary> inputs.
146/// Returns Float64.
147#[derive(Debug, PartialEq, Eq, Hash)]
148pub struct JsonbAsF64 {
149    signature: Signature,
150}
151
152impl JsonbAsF64 {
153    pub fn new() -> Self {
154        Self {
155            signature: Signature::any(1, Volatility::Immutable),
156        }
157    }
158}
159
160impl Default for JsonbAsF64 {
161    fn default() -> Self {
162        Self::new()
163    }
164}
165
166fn extract_f64_from_jsonb(jsonb_bytes: &[u8]) -> Result<Option<f64>> {
167    let jsonb = RawJsonb::new(jsonb_bytes);
168    match jsonb.as_f64() {
169        Ok(value) => Ok(value),
170        Err(e) => Err(DataFusionError::External(e.into())),
171    }
172}
173
174impl ScalarUDFImpl for JsonbAsF64 {
175    fn name(&self) -> &str {
176        "jsonb_as_f64"
177    }
178
179    fn signature(&self) -> &Signature {
180        &self.signature
181    }
182
183    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
184        Ok(DataType::Float64)
185    }
186
187    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
188        let args = ColumnarValue::values_to_arrays(&args.args)?;
189        if args.len() != 1 {
190            return internal_err!("wrong number of arguments to jsonb_as_f64()");
191        }
192
193        match args[0].data_type() {
194            DataType::Binary => {
195                let binary_array = args[0]
196                    .as_any()
197                    .downcast_ref::<GenericBinaryArray<i32>>()
198                    .ok_or_else(|| {
199                        DataFusionError::Internal("error casting to binary array".into())
200                    })?;
201
202                let mut builder = Float64Array::builder(binary_array.len());
203                for i in 0..binary_array.len() {
204                    if binary_array.is_null(i) {
205                        builder.append_null();
206                    } else {
207                        let jsonb_bytes = binary_array.value(i);
208                        if let Some(value) = extract_f64_from_jsonb(jsonb_bytes)? {
209                            builder.append_value(value);
210                        } else {
211                            builder.append_null();
212                        }
213                    }
214                }
215                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
216            }
217            DataType::Dictionary(_, value_type)
218                if matches!(value_type.as_ref(), DataType::Binary) =>
219            {
220                let dict_array = args[0]
221                    .as_any()
222                    .downcast_ref::<DictionaryArray<Int32Type>>()
223                    .ok_or_else(|| {
224                        DataFusionError::Internal("error casting dictionary array".into())
225                    })?;
226
227                let binary_values = dict_array
228                    .values()
229                    .as_any()
230                    .downcast_ref::<GenericBinaryArray<i32>>()
231                    .ok_or_else(|| {
232                        DataFusionError::Internal("dictionary values are not a binary array".into())
233                    })?;
234
235                let mut builder = Float64Array::builder(dict_array.len());
236                for i in 0..dict_array.len() {
237                    if dict_array.is_null(i) {
238                        builder.append_null();
239                    } else {
240                        let key_index = dict_array.keys().value(i) as usize;
241                        if key_index < binary_values.len() {
242                            let jsonb_bytes = binary_values.value(key_index);
243                            if let Some(value) = extract_f64_from_jsonb(jsonb_bytes)? {
244                                builder.append_value(value);
245                            } else {
246                                builder.append_null();
247                            }
248                        } else {
249                            return internal_err!(
250                                "Dictionary key index out of bounds in jsonb_as_f64"
251                            );
252                        }
253                    }
254                }
255                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
256            }
257            _ => internal_err!(
258                "jsonb_as_f64: unsupported input type, expected Binary or Dictionary<Int32, Binary>"
259            ),
260        }
261    }
262}
263
264/// Creates a user-defined function to cast a JSONB value to a f64.
265pub fn make_jsonb_as_f64_udf() -> ScalarUDF {
266    ScalarUDF::new_from_impl(JsonbAsF64::new())
267}
268
269/// A scalar UDF that casts a JSONB value to an i64.
270///
271/// Accepts both Binary and Dictionary<Int32, Binary> inputs.
272/// Returns Int64.
273#[derive(Debug, PartialEq, Eq, Hash)]
274pub struct JsonbAsI64 {
275    signature: Signature,
276}
277
278impl JsonbAsI64 {
279    pub fn new() -> Self {
280        Self {
281            signature: Signature::any(1, Volatility::Immutable),
282        }
283    }
284}
285
286impl Default for JsonbAsI64 {
287    fn default() -> Self {
288        Self::new()
289    }
290}
291
292fn extract_i64_from_jsonb(jsonb_bytes: &[u8]) -> Result<Option<i64>> {
293    let jsonb = RawJsonb::new(jsonb_bytes);
294    match jsonb.as_i64() {
295        Ok(value) => Ok(value),
296        Err(e) => Err(DataFusionError::External(e.into())),
297    }
298}
299
300impl ScalarUDFImpl for JsonbAsI64 {
301    fn name(&self) -> &str {
302        "jsonb_as_i64"
303    }
304
305    fn signature(&self) -> &Signature {
306        &self.signature
307    }
308
309    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
310        Ok(DataType::Int64)
311    }
312
313    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
314        let args = ColumnarValue::values_to_arrays(&args.args)?;
315        if args.len() != 1 {
316            return internal_err!("wrong number of arguments to jsonb_as_i64()");
317        }
318
319        match args[0].data_type() {
320            DataType::Binary => {
321                let binary_array = args[0]
322                    .as_any()
323                    .downcast_ref::<GenericBinaryArray<i32>>()
324                    .ok_or_else(|| {
325                        DataFusionError::Internal("error casting to binary array".into())
326                    })?;
327
328                let mut builder = Int64Array::builder(binary_array.len());
329                for i in 0..binary_array.len() {
330                    if binary_array.is_null(i) {
331                        builder.append_null();
332                    } else {
333                        let jsonb_bytes = binary_array.value(i);
334                        if let Some(value) = extract_i64_from_jsonb(jsonb_bytes)? {
335                            builder.append_value(value);
336                        } else {
337                            builder.append_null();
338                        }
339                    }
340                }
341                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
342            }
343            DataType::Dictionary(_, value_type)
344                if matches!(value_type.as_ref(), DataType::Binary) =>
345            {
346                let dict_array = args[0]
347                    .as_any()
348                    .downcast_ref::<DictionaryArray<Int32Type>>()
349                    .ok_or_else(|| {
350                        DataFusionError::Internal("error casting dictionary array".into())
351                    })?;
352
353                let binary_values = dict_array
354                    .values()
355                    .as_any()
356                    .downcast_ref::<GenericBinaryArray<i32>>()
357                    .ok_or_else(|| {
358                        DataFusionError::Internal("dictionary values are not a binary array".into())
359                    })?;
360
361                let mut builder = Int64Array::builder(dict_array.len());
362                for i in 0..dict_array.len() {
363                    if dict_array.is_null(i) {
364                        builder.append_null();
365                    } else {
366                        let key_index = dict_array.keys().value(i) as usize;
367                        if key_index < binary_values.len() {
368                            let jsonb_bytes = binary_values.value(key_index);
369                            if let Some(value) = extract_i64_from_jsonb(jsonb_bytes)? {
370                                builder.append_value(value);
371                            } else {
372                                builder.append_null();
373                            }
374                        } else {
375                            return internal_err!(
376                                "Dictionary key index out of bounds in jsonb_as_i64"
377                            );
378                        }
379                    }
380                }
381                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
382            }
383            _ => internal_err!(
384                "jsonb_as_i64: unsupported input type, expected Binary or Dictionary<Int32, Binary>"
385            ),
386        }
387    }
388}
389
390/// Creates a user-defined function to cast a JSONB value to an i64.
391pub fn make_jsonb_as_i64_udf() -> ScalarUDF {
392    ScalarUDF::new_from_impl(JsonbAsI64::new())
393}
micromegas_datafusion_extensions/jsonb/cast.rs

micromegas_datafusion_extensions/jsonb/
cast.rs
