micromegas_analytics/properties/
properties_to_dict_udf.rs

1use datafusion::arrow::array::{
2    Array, AsArray, DictionaryArray, GenericListArray, Int32Array, ListBuilder, StringBuilder,
3    StructArray, StructBuilder,
4};
5use datafusion::arrow::datatypes::{DataType, Field, Fields, Int32Type};
6use datafusion::common::{Result, internal_err};
7use datafusion::error::DataFusionError;
8use datafusion::logical_expr::{
9    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
10};
11use micromegas_datafusion_extensions::properties::properties_udf::extract_properties_as_vec;
12use std::collections::HashMap;
13use std::sync::Arc;
14
15#[derive(Debug, PartialEq, Eq, Hash)]
16pub struct PropertiesToDict {
17    signature: Signature,
18}
19
20impl PropertiesToDict {
21    pub fn new() -> Self {
22        Self::default()
23    }
24}
25
26impl Default for PropertiesToDict {
27    fn default() -> Self {
28        Self {
29            signature: Signature::exact(
30                vec![DataType::List(Arc::new(Field::new(
31                    "Property",
32                    DataType::Struct(Fields::from(vec![
33                        Field::new("key", DataType::Utf8, false),
34                        Field::new("value", DataType::Utf8, false),
35                    ])),
36                    false,
37                )))],
38                Volatility::Immutable,
39            ),
40        }
41    }
42}
43
44impl ScalarUDFImpl for PropertiesToDict {
45    fn name(&self) -> &str {
46        "properties_to_dict"
47    }
48
49    fn signature(&self) -> &Signature {
50        &self.signature
51    }
52
53    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
54        Ok(DataType::Dictionary(
55            Box::new(DataType::Int32),
56            Box::new(DataType::List(Arc::new(Field::new(
57                "Property",
58                DataType::Struct(Fields::from(vec![
59                    Field::new("key", DataType::Utf8, false),
60                    Field::new("value", DataType::Utf8, false),
61                ])),
62                false,
63            )))),
64        ))
65    }
66
67    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
68        let args = args.args;
69        if args.len() != 1 {
70            return internal_err!("properties_to_dict expects exactly one argument");
71        }
72
73        match &args[0] {
74            ColumnarValue::Array(array) => {
75                let list_array = array
76                    .as_any()
77                    .downcast_ref::<GenericListArray<i32>>()
78                    .ok_or_else(|| {
79                        DataFusionError::Internal(
80                            "properties_to_dict requires a list array as input".to_string(),
81                        )
82                    })?;
83
84                let dict_array = build_dictionary_from_properties(list_array)?;
85                Ok(ColumnarValue::Array(Arc::new(dict_array)))
86            }
87            ColumnarValue::Scalar(_) => {
88                internal_err!("properties_to_dict does not support scalar inputs")
89            }
90        }
91    }
92}
93
94struct PropertiesDictionaryBuilder {
95    map: HashMap<Vec<(String, String)>, usize>,
96    values_builder: ListBuilder<StructBuilder>,
97    keys: Vec<Option<i32>>,
98}
99
100impl PropertiesDictionaryBuilder {
101    fn new(capacity: usize) -> Self {
102        let prop_struct_fields = vec![
103            Field::new("key", DataType::Utf8, false),
104            Field::new("value", DataType::Utf8, false),
105        ];
106        let prop_field = Arc::new(Field::new(
107            "Property",
108            DataType::Struct(Fields::from(prop_struct_fields.clone())),
109            false,
110        ));
111        let values_builder =
112            ListBuilder::new(StructBuilder::from_fields(prop_struct_fields, capacity))
113                .with_field(prop_field);
114
115        Self {
116            map: HashMap::new(),
117            values_builder,
118            keys: Vec::with_capacity(capacity),
119        }
120    }
121
122    fn append_property_list(&mut self, struct_array: &StructArray) -> Result<()> {
123        let prop_vec = extract_properties_as_vec(struct_array)?;
124
125        match self.map.get(&prop_vec) {
126            Some(&index) => {
127                self.keys.push(Some(index as i32));
128            }
129            None => {
130                let new_index = self.map.len();
131                self.add_to_values(&prop_vec)?;
132                self.map.insert(prop_vec, new_index);
133                self.keys.push(Some(new_index as i32));
134            }
135        }
136        Ok(())
137    }
138
139    fn append_null(&mut self) {
140        self.keys.push(None);
141    }
142
143    fn add_to_values(&mut self, properties: &[(String, String)]) -> Result<()> {
144        let struct_builder = self.values_builder.values();
145        for (key, value) in properties {
146            struct_builder
147                .field_builder::<StringBuilder>(0)
148                .ok_or_else(|| DataFusionError::Internal("Failed to get key builder".to_string()))?
149                .append_value(key);
150            struct_builder
151                .field_builder::<StringBuilder>(1)
152                .ok_or_else(|| {
153                    DataFusionError::Internal("Failed to get value builder".to_string())
154                })?
155                .append_value(value);
156            struct_builder.append(true);
157        }
158        self.values_builder.append(true);
159        Ok(())
160    }
161
162    fn finish(mut self) -> Result<DictionaryArray<Int32Type>> {
163        let keys = Int32Array::from(self.keys);
164        let values = Arc::new(self.values_builder.finish());
165        DictionaryArray::try_new(keys, values)
166            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
167    }
168}
169
170pub fn build_dictionary_from_properties(
171    list_array: &GenericListArray<i32>,
172) -> Result<DictionaryArray<Int32Type>> {
173    let mut builder = PropertiesDictionaryBuilder::new(list_array.len());
174    for i in 0..list_array.len() {
175        if list_array.is_null(i) {
176            builder.append_null();
177        } else {
178            let start = list_array.value_offsets()[i] as usize;
179            let end = list_array.value_offsets()[i + 1] as usize;
180            let sliced_values = list_array.values().slice(start, end - start);
181            let struct_array = sliced_values.as_struct();
182            builder.append_property_list(struct_array)?;
183        }
184    }
185
186    builder.finish()
187}
micromegas_analytics/properties/properties_to_dict_udf.rs

micromegas_analytics/properties/
properties_to_dict_udf.rs
