datafusion_functions_json/
json_object_keys.rs

1use std::sync::Arc;
2
3use datafusion::arrow::array::{ArrayRef, ListBuilder, StringBuilder};
4use datafusion::arrow::datatypes::{DataType, Field};
5use datafusion::common::{Result as DataFusionResult, ScalarValue};
6use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
7use jiter::Peek;
8
9use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
10use crate::common_macros::make_udf_function;
11
12make_udf_function!(
13    JsonObjectKeys,
14    json_object_keys,
15    json_data path,
16    r"Get the keys of a JSON object as an array."
17);
18
19#[derive(Debug, PartialEq, Eq, Hash)]
20pub(super) struct JsonObjectKeys {
21    signature: Signature,
22    aliases: [String; 2],
23}
24
25impl Default for JsonObjectKeys {
26    fn default() -> Self {
27        Self {
28            signature: Signature::variadic_any(Volatility::Immutable),
29            aliases: ["json_object_keys".to_string(), "json_keys".to_string()],
30        }
31    }
32}
33
34impl ScalarUDFImpl for JsonObjectKeys {
35    fn name(&self) -> &str {
36        self.aliases[0].as_str()
37    }
38
39    fn signature(&self) -> &Signature {
40        &self.signature
41    }
42
43    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
44        return_type_check(
45            arg_types,
46            self.name(),
47            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
48        )
49    }
50
51    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
52        invoke::<BuildListArray>(&args.args, jiter_json_object_keys)
53    }
54
55    fn aliases(&self) -> &[String] {
56        &self.aliases
57    }
58
59    fn placement(
60        &self,
61        args: &[datafusion::logical_expr::ExpressionPlacement],
62    ) -> datafusion::logical_expr::ExpressionPlacement {
63        // If the first argument is a column and the remaining arguments are literals (a path)
64        // then we can push this UDF down to the leaf nodes.
65        if args.len() >= 2
66            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
67            && args[1..]
68                .iter()
69                .all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
70        {
71            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
72        } else {
73            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
74        }
75    }
76}
77
78/// Struct used to build a `ListArray` from the result of `jiter_json_object_keys`.
79#[derive(Debug)]
80struct BuildListArray;
81
82impl InvokeResult for BuildListArray {
83    type Item = Vec<String>;
84
85    type Builder = ListBuilder<StringBuilder>;
86
87    const ACCEPT_DICT_RETURN: bool = true;
88
89    fn builder(capacity: usize) -> Self::Builder {
90        let values_builder = StringBuilder::new();
91        ListBuilder::with_capacity(values_builder, capacity)
92    }
93
94    fn append_value(builder: &mut Self::Builder, value: Option<Self::Item>) {
95        builder.append_option(value.map(|v| v.into_iter().map(Some)));
96    }
97
98    fn finish(mut builder: Self::Builder) -> DataFusionResult<ArrayRef> {
99        Ok(Arc::new(builder.finish()))
100    }
101
102    fn scalar(value: Option<Self::Item>) -> ScalarValue {
103        keys_to_scalar(value)
104    }
105}
106
107fn keys_to_scalar(opt_keys: Option<Vec<String>>) -> ScalarValue {
108    let values_builder = StringBuilder::new();
109    let mut builder = ListBuilder::new(values_builder);
110    if let Some(keys) = opt_keys {
111        for value in keys {
112            builder.values().append_value(value);
113        }
114        builder.append(true);
115    } else {
116        builder.append(false);
117    }
118    let array = builder.finish();
119    ScalarValue::List(Arc::new(array))
120}
121
122fn jiter_json_object_keys(opt_json: Option<&str>, path: &[JsonPath]) -> Result<Vec<String>, GetError> {
123    if let Some((mut jiter, peek)) = jiter_json_find(opt_json, path) {
124        match peek {
125            Peek::Object => {
126                let mut opt_key = jiter.known_object()?;
127
128                let mut keys = Vec::new();
129                while let Some(key) = opt_key {
130                    keys.push(key.to_string());
131                    jiter.next_skip()?;
132                    opt_key = jiter.next_key()?;
133                }
134                Ok(keys)
135            }
136            _ => get_err!(),
137        }
138    } else {
139        get_err!()
140    }
141}
datafusion_functions_json/json_object_keys.rs

datafusion_functions_json/
json_object_keys.rs
