micromegas_datafusion_extensions/properties/
property_get.rs

1use anyhow::Context;
2use datafusion::arrow::array::{Array, StringDictionaryBuilder};
3use datafusion::arrow::array::{
4    ArrayRef, DictionaryArray, GenericBinaryArray, GenericListArray, StringArray,
5};
6use datafusion::arrow::array::{AsArray, StructArray};
7use datafusion::arrow::datatypes::{DataType, Int32Type};
8use datafusion::common::{Result, internal_err};
9use datafusion::error::DataFusionError;
10use datafusion::logical_expr::{
11    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
12};
13use jsonb::RawJsonb;
14use std::sync::Arc;
15
16/// A scalar UDF that retrieves a property from a list of properties.
17#[derive(Debug, PartialEq, Eq, Hash)]
18pub struct PropertyGet {
19    signature: Signature,
20}
21
22impl PropertyGet {
23    pub fn new() -> Self {
24        Self {
25            signature: Signature::any(2, Volatility::Immutable),
26        }
27    }
28}
29
30impl Default for PropertyGet {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36fn find_property_in_list(properties: ArrayRef, name: &str) -> anyhow::Result<Option<String>> {
37    let properties: &StructArray = properties.as_struct();
38    let (key_index, _key_field) = properties
39        .fields()
40        .find("key")
41        .with_context(|| "getting key field")?;
42    let (value_index, _value_field) = properties
43        .fields()
44        .find("value")
45        .with_context(|| "getting value field")?;
46    for i in 0..properties.len() {
47        let key = properties.column(key_index).as_string::<i32>().value(i);
48        if key.eq_ignore_ascii_case(name) {
49            let value = properties.column(value_index).as_string::<i32>().value(i);
50            return Ok(Some(value.into()));
51        }
52    }
53    Ok(None)
54}
55
56fn extract_from_jsonb(jsonb_bytes: &[u8], name: &str) -> anyhow::Result<Option<String>> {
57    let jsonb = RawJsonb::new(jsonb_bytes);
58    if let Some(value_jsonb) = jsonb
59        .get_by_name(name, true)
60        .with_context(|| "getting JSONB property by name")?
61    {
62        // The value_jsonb is an OwnedJsonb, convert it to RawJsonb to access its value
63        let raw_value = value_jsonb.as_raw();
64
65        // Try to get the value as a string (handles unescaping properly)
66        if let Some(str_value) = raw_value
67            .as_str()
68            .with_context(|| "extracting string value from JSONB")?
69        {
70            Ok(Some(str_value.to_string()))
71        } else {
72            // If it's not a string, convert it to JSON representation
73            Ok(Some(raw_value.to_string()))
74        }
75    } else {
76        Ok(None)
77    }
78}
79
80impl ScalarUDFImpl for PropertyGet {
81    fn name(&self) -> &str {
82        "property_get"
83    }
84    fn signature(&self) -> &Signature {
85        &self.signature
86    }
87    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
88        Ok(DataType::Dictionary(
89            Box::new(DataType::Int32),
90            Box::new(DataType::Utf8),
91        ))
92    }
93    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
94        let args = ColumnarValue::values_to_arrays(&args.args)?;
95        if args.len() != 2 {
96            return internal_err!("wrong number of arguments to property_get()");
97        }
98
99        let names = args[1]
100            .as_any()
101            .downcast_ref::<StringArray>()
102            .ok_or_else(|| DataFusionError::Execution("downcasting names in PropertyGet".into()))?;
103
104        // Handle both regular arrays and dictionary arrays
105        match args[0].data_type() {
106            DataType::List(_) => {
107                // Handle regular list array
108                let prop_lists = args[0]
109                    .as_any()
110                    .downcast_ref::<GenericListArray<i32>>()
111                    .ok_or_else(|| {
112                        DataFusionError::Internal("error casting property list".into())
113                    })?;
114
115                if prop_lists.len() != names.len() {
116                    return internal_err!("arrays of different lengths in property_get()");
117                }
118
119                let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
120                for i in 0..prop_lists.len() {
121                    let name = names.value(i);
122                    if let Some(value) = find_property_in_list(prop_lists.value(i), name)
123                        .map_err(|e| DataFusionError::Internal(format!("{e:?}")))?
124                    {
125                        dict_builder.append_value(value);
126                    } else {
127                        dict_builder.append_null();
128                    }
129                }
130                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
131            }
132            DataType::Binary => {
133                // Handle non-dictionary JSONB binary array
134                let binary_array = args[0]
135                    .as_any()
136                    .downcast_ref::<GenericBinaryArray<i32>>()
137                    .ok_or_else(|| {
138                        DataFusionError::Internal("error casting to binary array".into())
139                    })?;
140
141                if binary_array.len() != names.len() {
142                    return internal_err!("arrays of different lengths in property_get()");
143                }
144
145                let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
146                for i in 0..binary_array.len() {
147                    if binary_array.is_null(i) {
148                        dict_builder.append_null();
149                    } else {
150                        let jsonb_bytes = binary_array.value(i);
151                        let name = names.value(i);
152                        if let Some(value) = extract_from_jsonb(jsonb_bytes, name).map_err(|e| {
153                            DataFusionError::Internal(format!("JSONB extraction error: {e:?}"))
154                        })? {
155                            dict_builder.append_value(value);
156                        } else {
157                            dict_builder.append_null();
158                        }
159                    }
160                }
161                Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
162            }
163            DataType::Dictionary(_, value_type) => {
164                // Handle dictionary array
165                match value_type.as_ref() {
166                    DataType::Binary => {
167                        // Handle dictionary-encoded JSONB (new primary format)
168                        let dict_array = args[0]
169                            .as_any()
170                            .downcast_ref::<DictionaryArray<Int32Type>>()
171                            .ok_or_else(|| {
172                                DataFusionError::Internal("error casting dictionary array".into())
173                            })?;
174
175                        if dict_array.len() != names.len() {
176                            return internal_err!("arrays of different lengths in property_get()");
177                        }
178
179                        let values_array = dict_array.values();
180                        let binary_values = values_array
181                            .as_any()
182                            .downcast_ref::<GenericBinaryArray<i32>>()
183                            .ok_or_else(|| {
184                                DataFusionError::Internal(
185                                    "dictionary values are not a binary array".into(),
186                                )
187                            })?;
188
189                        let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
190                        for i in 0..dict_array.len() {
191                            if dict_array.is_null(i) {
192                                dict_builder.append_null();
193                            } else {
194                                let key_index = dict_array.keys().value(i) as usize;
195                                if key_index < binary_values.len() {
196                                    let jsonb_bytes = binary_values.value(key_index);
197                                    let name = names.value(i);
198                                    if let Some(value) = extract_from_jsonb(jsonb_bytes, name)
199                                        .map_err(|e| {
200                                            DataFusionError::Internal(format!(
201                                                "JSONB extraction error: {e:?}"
202                                            ))
203                                        })?
204                                    {
205                                        dict_builder.append_value(value);
206                                    } else {
207                                        dict_builder.append_null();
208                                    }
209                                } else {
210                                    return internal_err!(
211                                        "Dictionary key index out of bounds in property_get"
212                                    );
213                                }
214                            }
215                        }
216                        Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
217                    }
218                    DataType::List(_) => {
219                        let dict_array = args[0]
220                            .as_any()
221                            .downcast_ref::<DictionaryArray<Int32Type>>()
222                            .ok_or_else(|| {
223                                DataFusionError::Internal("error casting dictionary array".into())
224                            })?;
225
226                        if dict_array.len() != names.len() {
227                            return internal_err!("arrays of different lengths in property_get()");
228                        }
229
230                        let values_array = dict_array.values();
231                        let list_values = values_array
232                            .as_any()
233                            .downcast_ref::<GenericListArray<i32>>()
234                            .ok_or_else(|| {
235                                DataFusionError::Internal(
236                                    "dictionary values are not a list array".into(),
237                                )
238                            })?;
239
240                        let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
241                        for i in 0..dict_array.len() {
242                            let name = names.value(i);
243
244                            if dict_array.is_null(i) {
245                                dict_builder.append_null();
246                            } else {
247                                let key_index = dict_array.keys().value(i) as usize;
248                                if key_index < list_values.len() {
249                                    let property_list = list_values.value(key_index);
250                                    if let Some(value) = find_property_in_list(property_list, name)
251                                        .map_err(|e| DataFusionError::Internal(format!("{e:?}")))?
252                                    {
253                                        dict_builder.append_value(value);
254                                    } else {
255                                        dict_builder.append_null();
256                                    }
257                                } else {
258                                    return internal_err!(
259                                        "Dictionary key index out of bounds in property_get"
260                                    );
261                                }
262                            }
263                        }
264                        Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
265                    }
266                    _ => internal_err!(
267                        "property_get: unsupported dictionary value type, expected List or Binary"
268                    ),
269                }
270            }
271            _ => internal_err!(
272                "property_get: unsupported input type, expected List, Binary, Dictionary<Int32, List>, or Dictionary<Int32, Binary>"
273            ),
274        }
275    }
276}
micromegas_datafusion_extensions/properties/property_get.rs

micromegas_datafusion_extensions/properties/
property_get.rs
