micromegas_datafusion_extensions/properties/
properties_udf.rs

1use datafusion::arrow::array::{
2    Array, AsArray, DictionaryArray, GenericBinaryArray, GenericListArray, Int32Array, StructArray,
3};
4use datafusion::arrow::datatypes::{DataType, Int32Type};
5use datafusion::common::{Result, internal_err};
6use datafusion::error::DataFusionError;
7use datafusion::logical_expr::{
8    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
9};
10use jsonb::RawJsonb;
11use std::sync::Arc;
12
13pub fn extract_properties_as_vec(struct_array: &StructArray) -> Result<Vec<(String, String)>> {
14    let mut properties = Vec::with_capacity(struct_array.len());
15    let key_array = struct_array.column(0).as_string::<i32>();
16    let value_array = struct_array.column(1).as_string::<i32>();
17    for i in 0..struct_array.len() {
18        if struct_array.is_valid(i) {
19            let key = key_array.value(i).to_string();
20            let value = value_array.value(i).to_string();
21            properties.push((key, value));
22        }
23    }
24
25    Ok(properties)
26}
27
28pub fn count_jsonb_properties(jsonb_bytes: &[u8]) -> Result<i32> {
29    let jsonb = RawJsonb::new(jsonb_bytes);
30
31    // Get object keys and count them using array_length
32    match jsonb.object_keys() {
33        Ok(Some(keys_array)) => {
34            // It's an object, get the array length of the keys
35            let keys_raw = keys_array.as_raw();
36            match keys_raw.array_length() {
37                Ok(Some(len)) => Ok(len as i32),
38                Ok(None) => Ok(0), // Empty array
39                Err(e) => Err(DataFusionError::Internal(format!(
40                    "Failed to get keys array length: {e:?}"
41                ))),
42            }
43        }
44        Ok(None) => {
45            // Not an object (array, scalar, null), return 0
46            Ok(0)
47        }
48        Err(e) => Err(DataFusionError::Internal(format!(
49            "Failed to count JSONB properties: {e:?}"
50        ))),
51    }
52}
53
54// Helper UDF to extract properties array from dictionary for use with standard functions
55#[derive(Debug, PartialEq, Eq, Hash)]
56pub struct PropertiesToArray {
57    signature: Signature,
58}
59
60impl PropertiesToArray {
61    pub fn new() -> Self {
62        Self::default()
63    }
64}
65
66impl Default for PropertiesToArray {
67    fn default() -> Self {
68        Self {
69            signature: Signature::any(1, Volatility::Immutable),
70        }
71    }
72}
73
74impl ScalarUDFImpl for PropertiesToArray {
75    fn name(&self) -> &str {
76        "properties_to_array"
77    }
78
79    fn signature(&self) -> &Signature {
80        &self.signature
81    }
82
83    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
84        match &arg_types[0] {
85            DataType::Dictionary(_, value_type) => Ok(value_type.as_ref().clone()),
86            _ => internal_err!("properties_to_array expects a Dictionary input type"),
87        }
88    }
89
90    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
91        let args = args.args;
92        if args.len() != 1 {
93            return internal_err!("properties_to_array expects exactly one argument");
94        }
95
96        match &args[0] {
97            ColumnarValue::Array(array) => {
98                // Reconstruct the full array from dictionary
99                let dict_array = array
100                    .as_any()
101                    .downcast_ref::<DictionaryArray<Int32Type>>()
102                    .ok_or_else(|| {
103                        DataFusionError::Internal(
104                            "properties_to_array requires a dictionary array as input".to_string(),
105                        )
106                    })?;
107
108                // Use Arrow's take function to reconstruct the array
109                use datafusion::arrow::compute::take;
110                let indices = dict_array.keys();
111                let values = dict_array.values();
112
113                let reconstructed = take(values.as_ref(), indices, None)
114                    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
115
116                Ok(ColumnarValue::Array(reconstructed))
117            }
118            ColumnarValue::Scalar(_) => {
119                internal_err!("properties_to_array does not support scalar inputs")
120            }
121        }
122    }
123}
124
125// UDF to get length of properties that works with both regular and dictionary arrays
126#[derive(Debug, PartialEq, Eq, Hash)]
127pub struct PropertiesLength {
128    signature: Signature,
129}
130
131impl PropertiesLength {
132    pub fn new() -> Self {
133        Self::default()
134    }
135}
136
137impl Default for PropertiesLength {
138    fn default() -> Self {
139        Self {
140            signature: Signature::any(1, Volatility::Immutable),
141        }
142    }
143}
144
145impl ScalarUDFImpl for PropertiesLength {
146    fn name(&self) -> &str {
147        "properties_length"
148    }
149
150    fn signature(&self) -> &Signature {
151        &self.signature
152    }
153
154    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
155        Ok(DataType::Int32)
156    }
157
158    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
159        let args = args.args;
160        if args.len() != 1 {
161            return internal_err!("properties_length expects exactly one argument");
162        }
163
164        match &args[0] {
165            ColumnarValue::Array(array) => {
166                match array.data_type() {
167                    DataType::List(_) => {
168                        // Handle regular list array
169                        let list_array = array
170                            .as_any()
171                            .downcast_ref::<GenericListArray<i32>>()
172                            .ok_or_else(|| {
173                                DataFusionError::Internal(
174                                    "properties_length: failed to cast to list array".to_string(),
175                                )
176                            })?;
177
178                        let mut lengths = Vec::with_capacity(list_array.len());
179                        for i in 0..list_array.len() {
180                            if list_array.is_null(i) {
181                                lengths.push(None);
182                            } else {
183                                let start = list_array.value_offsets()[i] as usize;
184                                let end = list_array.value_offsets()[i + 1] as usize;
185                                lengths.push(Some((end - start) as i32));
186                            }
187                        }
188
189                        let length_array = Int32Array::from(lengths);
190                        Ok(ColumnarValue::Array(Arc::new(length_array)))
191                    }
192                    DataType::Binary => {
193                        // Handle JSONB binary array
194                        let binary_array = array
195                            .as_any()
196                            .downcast_ref::<GenericBinaryArray<i32>>()
197                            .ok_or_else(|| {
198                                DataFusionError::Internal(
199                                    "properties_length: failed to cast to binary array".to_string(),
200                                )
201                            })?;
202
203                        let mut lengths = Vec::with_capacity(binary_array.len());
204                        for i in 0..binary_array.len() {
205                            if binary_array.is_null(i) {
206                                lengths.push(None);
207                            } else {
208                                let jsonb_bytes = binary_array.value(i);
209                                match count_jsonb_properties(jsonb_bytes) {
210                                    Ok(len) => lengths.push(Some(len)),
211                                    Err(_) => lengths.push(None), // Error counting, treat as null
212                                }
213                            }
214                        }
215
216                        let length_array = Int32Array::from(lengths);
217                        Ok(ColumnarValue::Array(Arc::new(length_array)))
218                    }
219                    DataType::Dictionary(_, value_type) => {
220                        // Handle dictionary array
221                        match value_type.as_ref() {
222                            DataType::List(_) => {
223                                let dict_array = array
224                                    .as_any()
225                                    .downcast_ref::<DictionaryArray<Int32Type>>()
226                                    .ok_or_else(|| {
227                                        DataFusionError::Internal(
228                                            "properties_length: failed to cast to dictionary array"
229                                                .to_string(),
230                                        )
231                                    })?;
232
233                                let values = dict_array.values();
234                                let list_values = values
235                                    .as_any()
236                                    .downcast_ref::<GenericListArray<i32>>()
237                                    .ok_or_else(|| {
238                                        DataFusionError::Internal(
239                                            "properties_length: dictionary values are not a list array".to_string(),
240                                        )
241                                    })?;
242
243                                // Pre-compute lengths for each unique value in the dictionary
244                                let mut dict_lengths = Vec::with_capacity(list_values.len());
245                                for i in 0..list_values.len() {
246                                    if list_values.is_null(i) {
247                                        dict_lengths.push(None);
248                                    } else {
249                                        let start = list_values.value_offsets()[i] as usize;
250                                        let end = list_values.value_offsets()[i + 1] as usize;
251                                        dict_lengths.push(Some((end - start) as i32));
252                                    }
253                                }
254
255                                // Map dictionary keys to lengths
256                                let keys = dict_array.keys();
257                                let mut lengths = Vec::with_capacity(keys.len());
258                                for i in 0..keys.len() {
259                                    if keys.is_null(i) {
260                                        lengths.push(None);
261                                    } else {
262                                        let key_index = keys.value(i) as usize;
263                                        if key_index < dict_lengths.len() {
264                                            lengths.push(dict_lengths[key_index]);
265                                        } else {
266                                            return internal_err!(
267                                                "Dictionary key index out of bounds"
268                                            );
269                                        }
270                                    }
271                                }
272
273                                let length_array = Int32Array::from(lengths);
274                                Ok(ColumnarValue::Array(Arc::new(length_array)))
275                            }
276                            DataType::Binary => {
277                                // Handle dictionary-encoded JSONB (primary format)
278                                let dict_array = array
279                                    .as_any()
280                                    .downcast_ref::<DictionaryArray<Int32Type>>()
281                                    .ok_or_else(|| {
282                                        DataFusionError::Internal(
283                                            "properties_length: failed to cast to dictionary array"
284                                                .to_string(),
285                                        )
286                                    })?;
287
288                                let values = dict_array.values();
289                                let binary_values = values
290                                    .as_any()
291                                    .downcast_ref::<GenericBinaryArray<i32>>()
292                                    .ok_or_else(|| {
293                                        DataFusionError::Internal(
294                                            "properties_length: dictionary values are not a binary array".to_string(),
295                                        )
296                                    })?;
297
298                                // Pre-compute lengths for each unique JSONB value in the dictionary
299                                let mut dict_lengths = Vec::with_capacity(binary_values.len());
300                                for i in 0..binary_values.len() {
301                                    if binary_values.is_null(i) {
302                                        dict_lengths.push(None);
303                                    } else {
304                                        let jsonb_bytes = binary_values.value(i);
305                                        match count_jsonb_properties(jsonb_bytes) {
306                                            Ok(len) => dict_lengths.push(Some(len)),
307                                            Err(_) => dict_lengths.push(None), // Error counting, treat as null
308                                        }
309                                    }
310                                }
311
312                                // Map dictionary keys to lengths
313                                let keys = dict_array.keys();
314                                let mut lengths = Vec::with_capacity(keys.len());
315                                for i in 0..keys.len() {
316                                    if keys.is_null(i) {
317                                        lengths.push(None);
318                                    } else {
319                                        let key_index = keys.value(i) as usize;
320                                        if key_index < dict_lengths.len() {
321                                            lengths.push(dict_lengths[key_index]);
322                                        } else {
323                                            return internal_err!(
324                                                "Dictionary key index out of bounds"
325                                            );
326                                        }
327                                    }
328                                }
329
330                                let length_array = Int32Array::from(lengths);
331                                Ok(ColumnarValue::Array(Arc::new(length_array)))
332                            }
333                            _ => internal_err!(
334                                "properties_length: unsupported dictionary value type, expected List or Binary"
335                            ),
336                        }
337                    }
338                    _ => internal_err!(
339                        "properties_length: unsupported input type, expected List, Binary, Dictionary<Int32, List>, or Dictionary<Int32, Binary>"
340                    ),
341                }
342            }
343            ColumnarValue::Scalar(_) => {
344                internal_err!("properties_length does not support scalar inputs")
345            }
346        }
347    }
348}
micromegas_datafusion_extensions/properties/properties_udf.rs

micromegas_datafusion_extensions/properties/
properties_udf.rs
