micromegas_datafusion_extensions/jsonb/
keys.rs

1use crate::jsonb::list_udfs::{binary_accessor_or_err, binary_dict, dict_fast_path};
2use datafusion::arrow::array::{ArrayRef, ListBuilder, StringBuilder};
3use datafusion::arrow::datatypes::{DataType, Field};
4use datafusion::common::{Result, exec_err};
5use datafusion::error::DataFusionError;
6use datafusion::logical_expr::{
7    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
8};
9use jsonb::RawJsonb;
10use std::sync::Arc;
11
12/// A scalar UDF that extracts the keys from a JSONB object.
13///
14/// Accepts both `Binary` and `Dictionary<Int32, Binary>` inputs.
15/// Returns `List<Utf8>` containing the object keys, or null if input is not an object.
16/// A `Dictionary<Int32, Binary>` input takes a fast path that extracts each unique blob once
17/// (`take`-expanded to row count), since JSONB values (especially `properties`) are often
18/// repeated; the returned list is never dictionary-encoded — `unnest()` and `SELECT DISTINCT`
19/// need a plain list to work directly, without an `arrow_cast` workaround.
20#[derive(Debug, PartialEq, Eq, Hash)]
21pub struct JsonbObjectKeys {
22    signature: Signature,
23}
24
25impl JsonbObjectKeys {
26    pub fn new() -> Self {
27        Self {
28            signature: Signature::any(1, Volatility::Immutable),
29        }
30    }
31}
32
33impl Default for JsonbObjectKeys {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39fn extract_keys_from_jsonb(jsonb_bytes: &[u8]) -> Result<Option<Vec<String>>> {
40    let jsonb = RawJsonb::new(jsonb_bytes);
41    match jsonb.object_keys() {
42        Ok(Some(keys_jsonb)) => {
43            // keys_jsonb is a JSONB array of string keys
44            let keys_raw = keys_jsonb.as_raw();
45            match keys_raw.array_values() {
46                Ok(Some(values)) => {
47                    let mut keys = Vec::with_capacity(values.len());
48                    for value in values {
49                        let raw = value.as_raw();
50                        match raw.as_str() {
51                            Ok(Some(s)) => keys.push(s.to_string()),
52                            Ok(None) => {
53                                // Key is not a string (shouldn't happen for object keys)
54                                return Ok(None);
55                            }
56                            Err(e) => return Err(DataFusionError::External(e.into())),
57                        }
58                    }
59                    Ok(Some(keys))
60                }
61                Ok(None) => Ok(Some(Vec::new())), // Empty array
62                Err(e) => Err(DataFusionError::External(e.into())),
63            }
64        }
65        Ok(None) => Ok(None), // Input is not an object
66        Err(e) => Err(DataFusionError::External(e.into())),
67    }
68}
69
70impl ScalarUDFImpl for JsonbObjectKeys {
71    fn name(&self) -> &str {
72        "jsonb_object_keys"
73    }
74
75    fn signature(&self) -> &Signature {
76        &self.signature
77    }
78
79    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
80        Ok(DataType::List(Arc::new(Field::new_list_field(
81            DataType::Utf8,
82            true,
83        ))))
84    }
85
86    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
87        let args = ColumnarValue::values_to_arrays(&args.args)?;
88        if args.len() != 1 {
89            return exec_err!("wrong number of arguments to jsonb_object_keys()");
90        }
91
92        Ok(ColumnarValue::Array(build_keys_list(&args[0])?))
93    }
94}
95
96fn append_keys_slot(list_builder: &mut ListBuilder<StringBuilder>, keys: Option<Vec<String>>) {
97    match keys {
98        Some(keys) => {
99            for key in keys {
100                list_builder.values().append_value(key);
101            }
102            list_builder.append(true);
103        }
104        None => list_builder.append(false),
105    }
106}
107
108/// Build `jsonb_object_keys`'s list array: dictionary fast path when the input is
109/// `Dictionary<Int32, Binary>`, row-by-row via `create_binary_accessor` otherwise.
110fn build_keys_list(array: &ArrayRef) -> Result<ArrayRef> {
111    if let Some(dict_array) = binary_dict(array) {
112        let list_builder = ListBuilder::new(StringBuilder::new());
113        return dict_fast_path(dict_array, list_builder, |builder, bytes| {
114            append_keys_slot(builder, extract_keys_from_jsonb(bytes)?);
115            Ok(())
116        });
117    }
118
119    let accessor = binary_accessor_or_err(array, "jsonb_object_keys")?;
120    let mut list_builder = ListBuilder::new(StringBuilder::new());
121    for i in 0..accessor.len() {
122        if accessor.is_null(i) {
123            list_builder.append(false);
124        } else {
125            append_keys_slot(
126                &mut list_builder,
127                extract_keys_from_jsonb(accessor.value(i))?,
128            );
129        }
130    }
131    Ok(Arc::new(list_builder.finish()))
132}
133
134/// Creates a user-defined function to extract the keys from a JSONB object.
135pub fn make_jsonb_object_keys_udf() -> ScalarUDF {
136    ScalarUDF::new_from_impl(JsonbObjectKeys::new())
137}
micromegas_datafusion_extensions/jsonb/keys.rs

micromegas_datafusion_extensions/jsonb/
keys.rs
