micromegas_datafusion_extensions/jsonb/
array_length.rs

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

micromegas_datafusion_extensions/jsonb/
array_length.rs
