micromegas_datafusion_extensions/jsonb/
keys.rs

1use datafusion::arrow::array::{
2    Array, ArrayRef, DictionaryArray, GenericBinaryArray, Int32Array, ListBuilder, StringBuilder,
3};
4use datafusion::arrow::datatypes::{DataType, Field, Int32Type};
5use datafusion::common::{Result, internal_err};
6use datafusion::error::DataFusionError;
7use datafusion::logical_expr::{
8    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
9};
10use jsonb::RawJsonb;
11use std::collections::HashMap;
12use std::sync::Arc;
13
14/// A scalar UDF that extracts the keys from a JSONB object.
15///
16/// Accepts both `Binary` and `Dictionary<Int32, Binary>` inputs.
17/// Returns `Dictionary<Int32, List<Utf8>>` containing the object keys, or null if input is not an object.
18/// Dictionary encoding is used because JSONB values (especially properties) are often repeated.
19#[derive(Debug, PartialEq, Eq, Hash)]
20pub struct JsonbObjectKeys {
21    signature: Signature,
22}
23
24impl JsonbObjectKeys {
25    pub fn new() -> Self {
26        Self {
27            signature: Signature::any(1, Volatility::Immutable),
28        }
29    }
30}
31
32impl Default for JsonbObjectKeys {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38fn extract_keys_from_jsonb(jsonb_bytes: &[u8]) -> Result<Option<Vec<String>>> {
39    let jsonb = RawJsonb::new(jsonb_bytes);
40    match jsonb.object_keys() {
41        Ok(Some(keys_jsonb)) => {
42            // keys_jsonb is a JSONB array of string keys
43            let keys_raw = keys_jsonb.as_raw();
44            match keys_raw.array_values() {
45                Ok(Some(values)) => {
46                    let mut keys = Vec::with_capacity(values.len());
47                    for value in values {
48                        let raw = value.as_raw();
49                        match raw.as_str() {
50                            Ok(Some(s)) => keys.push(s.to_string()),
51                            Ok(None) => {
52                                // Key is not a string (shouldn't happen for object keys)
53                                return Ok(None);
54                            }
55                            Err(e) => return Err(DataFusionError::External(e.into())),
56                        }
57                    }
58                    Ok(Some(keys))
59                }
60                Ok(None) => Ok(Some(Vec::new())), // Empty array
61                Err(e) => Err(DataFusionError::External(e.into())),
62            }
63        }
64        Ok(None) => Ok(None), // Input is not an object
65        Err(e) => Err(DataFusionError::External(e.into())),
66    }
67}
68
69impl ScalarUDFImpl for JsonbObjectKeys {
70    fn name(&self) -> &str {
71        "jsonb_object_keys"
72    }
73
74    fn signature(&self) -> &Signature {
75        &self.signature
76    }
77
78    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
79        Ok(DataType::Dictionary(
80            Box::new(DataType::Int32),
81            Box::new(DataType::List(Arc::new(Field::new_list_field(
82                DataType::Utf8,
83                true,
84            )))),
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 jsonb_object_keys()");
92        }
93
94        match args[0].data_type() {
95            DataType::Binary => {
96                let binary_array = args[0]
97                    .as_any()
98                    .downcast_ref::<GenericBinaryArray<i32>>()
99                    .ok_or_else(|| {
100                        DataFusionError::Internal("error casting to binary array".into())
101                    })?;
102
103                let result = build_dict_list_array(binary_array.len(), |i| {
104                    if binary_array.is_null(i) {
105                        Ok(None)
106                    } else {
107                        extract_keys_from_jsonb(binary_array.value(i))
108                    }
109                })?;
110                Ok(ColumnarValue::Array(result))
111            }
112            DataType::Dictionary(_, value_type)
113                if matches!(value_type.as_ref(), DataType::Binary) =>
114            {
115                let dict_array = args[0]
116                    .as_any()
117                    .downcast_ref::<DictionaryArray<Int32Type>>()
118                    .ok_or_else(|| {
119                        DataFusionError::Internal("error casting dictionary array".into())
120                    })?;
121
122                let binary_values = dict_array
123                    .values()
124                    .as_any()
125                    .downcast_ref::<GenericBinaryArray<i32>>()
126                    .ok_or_else(|| {
127                        DataFusionError::Internal("dictionary values are not a binary array".into())
128                    })?;
129
130                let result = build_dict_list_array(dict_array.len(), |i| {
131                    if dict_array.is_null(i) {
132                        Ok(None)
133                    } else {
134                        let key_index = dict_array.keys().value(i) as usize;
135                        if key_index < binary_values.len() {
136                            extract_keys_from_jsonb(binary_values.value(key_index))
137                        } else {
138                            internal_err!("Dictionary key index out of bounds in jsonb_object_keys")
139                        }
140                    }
141                })?;
142                Ok(ColumnarValue::Array(result))
143            }
144            _ => internal_err!(
145                "jsonb_object_keys: unsupported input type, expected Binary or Dictionary<Int32, Binary>"
146            ),
147        }
148    }
149}
150
151/// Build a Dictionary<Int32, List<Utf8>> array from a function that returns keys for each index.
152/// Uses a HashMap to deduplicate identical key lists for memory efficiency.
153/// Returns None from get_keys to indicate a null output (distinct from Some(empty vec) for empty objects).
154fn build_dict_list_array<F>(len: usize, mut get_keys: F) -> Result<ArrayRef>
155where
156    F: FnMut(usize) -> Result<Option<Vec<String>>>,
157{
158    // Map from key list to dictionary index (only for non-null results)
159    let mut unique_lists: HashMap<Vec<String>, i32> = HashMap::new();
160    let mut key_indices: Vec<Option<i32>> = Vec::with_capacity(len);
161    let mut ordered_lists: Vec<Vec<String>> = Vec::new();
162
163    // First pass: collect all values and deduplicate
164    for i in 0..len {
165        let keys = get_keys(i)?;
166        match keys {
167            Some(key_list) => {
168                if let Some(idx) = unique_lists.get(&key_list) {
169                    key_indices.push(Some(*idx));
170                } else {
171                    let idx = ordered_lists.len() as i32;
172                    unique_lists.insert(key_list.clone(), idx);
173                    key_indices.push(Some(idx));
174                    ordered_lists.push(key_list);
175                }
176            }
177            None => {
178                // Null input produces null dictionary entry (null key)
179                key_indices.push(None);
180            }
181        }
182    }
183
184    // Build the values array (List<Utf8>) from unique lists
185    let mut list_builder = ListBuilder::new(StringBuilder::new());
186    for keys in &ordered_lists {
187        for key in keys {
188            list_builder.values().append_value(key);
189        }
190        list_builder.append(true);
191    }
192    let values_array = Arc::new(list_builder.finish());
193
194    // Build the keys array (None values become null keys)
195    let keys_array = Int32Array::from(key_indices);
196
197    // Construct the dictionary array
198    let dict_array =
199        DictionaryArray::<Int32Type>::try_new(keys_array, values_array).map_err(|e| {
200            DataFusionError::Internal(format!("Failed to create dictionary array: {e}"))
201        })?;
202
203    Ok(Arc::new(dict_array))
204}
205
206/// Creates a user-defined function to extract the keys from a JSONB object.
207pub fn make_jsonb_object_keys_udf() -> ScalarUDF {
208    ScalarUDF::new_from_impl(JsonbObjectKeys::new())
209}
micromegas_datafusion_extensions/jsonb/keys.rs

micromegas_datafusion_extensions/jsonb/
keys.rs
