micromegas_analytics/properties/
properties_to_jsonb_udf.rs

1use anyhow::Context;
2use datafusion::arrow::array::{
3    Array, ArrayRef, AsArray, BinaryDictionaryBuilder, DictionaryArray, GenericBinaryArray,
4    GenericListArray, StructArray,
5};
6use datafusion::arrow::datatypes::{DataType, Int32Type};
7use datafusion::common::{Result, internal_err};
8use datafusion::error::DataFusionError;
9use datafusion::logical_expr::{
10    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
11};
12use jsonb::Value;
13use micromegas_tracing::warn;
14use std::borrow::Cow;
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18/// A scalar UDF that converts a list of properties to JSONB binary format with dictionary encoding.
19///
20/// Converts List<Struct<key: String, value: String>> to Dictionary<Int32, Binary> (dictionary-encoded JSONB).
21/// The output uses dictionary encoding to optimize storage of repeated property sets.
22/// Each unique JSONB object like {"key1": "value1", "key2": "value2"} is stored once in the dictionary.
23#[derive(Debug, PartialEq, Eq, Hash)]
24pub struct PropertiesToJsonb {
25    signature: Signature,
26}
27
28impl PropertiesToJsonb {
29    pub fn new() -> Self {
30        Self {
31            signature: Signature::any(1, Volatility::Immutable),
32        }
33    }
34}
35
36impl Default for PropertiesToJsonb {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42fn convert_properties_list_to_jsonb(properties: ArrayRef) -> anyhow::Result<Vec<u8>> {
43    let properties: &StructArray = properties.as_struct();
44    let (key_index, _key_field) = properties
45        .fields()
46        .find("key")
47        .with_context(|| "getting key field")?;
48    let (value_index, _value_field) = properties
49        .fields()
50        .find("value")
51        .with_context(|| "getting value field")?;
52
53    let mut map = BTreeMap::new();
54    let key_column = properties.column(key_index).as_string::<i32>();
55    let value_column = properties.column(value_index).as_string::<i32>();
56
57    for i in 0..properties.len() {
58        if key_column.is_null(i) || value_column.is_null(i) {
59            continue; // Skip null entries
60        }
61        let key = key_column.value(i);
62        let value = value_column.value(i);
63        map.insert(key.to_string(), Value::String(Cow::Borrowed(value)));
64    }
65
66    let jsonb_object = Value::Object(map);
67    let mut buffer = Vec::new();
68    jsonb_object.write_to_vec(&mut buffer);
69    Ok(buffer)
70}
71
72impl ScalarUDFImpl for PropertiesToJsonb {
73    fn name(&self) -> &str {
74        "properties_to_jsonb"
75    }
76
77    fn signature(&self) -> &Signature {
78        &self.signature
79    }
80
81    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
82        Ok(DataType::Dictionary(
83            Box::new(DataType::Int32),
84            Box::new(DataType::Binary),
85        ))
86    }
87
88    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
89        let args = ColumnarValue::values_to_arrays(&args.args)?;
90        if args.len() != 1 {
91            return internal_err!("wrong number of arguments to properties_to_jsonb()");
92        }
93
94        // Handle all input formats and return Dictionary<Int32, Binary>
95        match args[0].data_type() {
96            DataType::List(_) => {
97                // Handle regular list array - convert to dictionary-encoded JSONB
98                let prop_lists = args[0]
99                    .as_any()
100                    .downcast_ref::<GenericListArray<i32>>()
101                    .ok_or_else(|| {
102                        DataFusionError::Internal("error casting property list".into())
103                    })?;
104
105                let mut dict_builder = BinaryDictionaryBuilder::<Int32Type>::new();
106                for i in 0..prop_lists.len() {
107                    if prop_lists.is_null(i) {
108                        dict_builder.append_null();
109                    } else {
110                        match convert_properties_list_to_jsonb(prop_lists.value(i)) {
111                            Ok(jsonb_bytes) => {
112                                dict_builder.append_value(&jsonb_bytes);
113                            }
114                            Err(e) => {
115                                warn!(
116                                    "error converting properties to JSONB at index {}: {:?}",
117                                    i, e
118                                );
119                                dict_builder.append_null();
120                            }
121                        }
122                    }
123                }
124                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
125            }
126            DataType::Binary => {
127                // Pass-through optimization: already JSONB, just need to add dictionary encoding
128                let binary_array = args[0]
129                    .as_any()
130                    .downcast_ref::<GenericBinaryArray<i32>>()
131                    .ok_or_else(|| {
132                        DataFusionError::Internal("error casting to binary array".into())
133                    })?;
134
135                let mut dict_builder = BinaryDictionaryBuilder::<Int32Type>::new();
136                for i in 0..binary_array.len() {
137                    if binary_array.is_null(i) {
138                        dict_builder.append_null();
139                    } else {
140                        let jsonb_bytes = binary_array.value(i);
141                        dict_builder.append_value(jsonb_bytes);
142                    }
143                }
144                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
145            }
146            DataType::Dictionary(_, value_type) => {
147                // Handle dictionary array
148                match value_type.as_ref() {
149                    DataType::List(_) => {
150                        // Convert dictionary-encoded List<Struct> to dictionary-encoded JSONB
151                        let dict_array = args[0]
152                            .as_any()
153                            .downcast_ref::<DictionaryArray<Int32Type>>()
154                            .ok_or_else(|| {
155                                DataFusionError::Internal("error casting dictionary array".into())
156                            })?;
157
158                        let values_array = dict_array.values();
159                        let list_values = values_array
160                            .as_any()
161                            .downcast_ref::<GenericListArray<i32>>()
162                            .ok_or_else(|| {
163                                DataFusionError::Internal(
164                                    "dictionary values are not a list array".into(),
165                                )
166                            })?;
167
168                        let mut dict_builder = BinaryDictionaryBuilder::<Int32Type>::new();
169                        for i in 0..dict_array.len() {
170                            if dict_array.is_null(i) {
171                                dict_builder.append_null();
172                            } else {
173                                let key_index = dict_array.keys().value(i) as usize;
174                                if key_index < list_values.len() {
175                                    let property_list = list_values.value(key_index);
176                                    match convert_properties_list_to_jsonb(property_list) {
177                                        Ok(jsonb_bytes) => {
178                                            dict_builder.append_value(&jsonb_bytes);
179                                        }
180                                        Err(e) => {
181                                            warn!(
182                                                "error converting properties to JSONB at dict index {}: {:?}",
183                                                i, e
184                                            );
185                                            dict_builder.append_null();
186                                        }
187                                    }
188                                } else {
189                                    return internal_err!(
190                                        "Dictionary key index out of bounds in properties_to_jsonb"
191                                    );
192                                }
193                            }
194                        }
195                        Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
196                    }
197                    DataType::Binary => {
198                        // Pass-through optimization: already dictionary-encoded JSONB
199                        Ok(ColumnarValue::Array(args[0].clone()))
200                    }
201                    _ => internal_err!(
202                        "properties_to_jsonb: unsupported dictionary value type, expected List or Binary"
203                    ),
204                }
205            }
206            _ => internal_err!(
207                "properties_to_jsonb: unsupported input type, expected List, Binary, Dictionary<Int32, List>, or Dictionary<Int32, Binary>"
208            ),
209        }
210    }
211}
micromegas_analytics/properties/properties_to_jsonb_udf.rs

micromegas_analytics/properties/
properties_to_jsonb_udf.rs
