From f639a2099d9a9a82004b54918872db1392ca8a03 Mon Sep 17 00:00:00 2001 From: Whis Liao Date: Tue, 4 Aug 2026 10:31:30 +0800 Subject: [PATCH] feat(function): add json_object_keys scalar function (#8722) Expose JSON object key listing for outermost objects, with sqlness coverage. --- src/common/function/src/scalars/json.rs | 2 + .../src/scalars/json/json_object_keys.rs | 211 ++++++++++++++++++ .../common/function/json/json.result | 41 ++++ .../standalone/common/function/json/json.sql | 11 + 4 files changed, 265 insertions(+) create mode 100644 src/common/function/src/scalars/json/json_object_keys.rs diff --git a/src/common/function/src/scalars/json.rs b/src/common/function/src/scalars/json.rs index e9d638c6cb..8a06c972c9 100644 --- a/src/common/function/src/scalars/json.rs +++ b/src/common/function/src/scalars/json.rs @@ -15,6 +15,7 @@ pub mod json_get; mod json_get_rewriter; mod json_is; +mod json_object_keys; mod json_path_exists; mod json_path_match; mod json_to_string; @@ -53,6 +54,7 @@ impl JsonFunction { registry.register_scalar(JsonIsArray::default()); registry.register_scalar(JsonIsObject::default()); + registry.register_scalar(json_object_keys::JsonObjectKeysFunction::default()); registry.register_scalar(json_path_exists::JsonPathExistsFunction::default()); registry.register_scalar(json_path_match::JsonPathMatchFunction::default()); diff --git a/src/common/function/src/scalars/json/json_object_keys.rs b/src/common/function/src/scalars/json/json_object_keys.rs new file mode 100644 index 0000000000..6883ef7d33 --- /dev/null +++ b/src/common/function/src/scalars/json/json_object_keys.rs @@ -0,0 +1,211 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{self, Display}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, ListBuilder, StringViewBuilder}; +use arrow::compute; +use arrow::datatypes::{DataType, Field}; +use datafusion_common::DataFusionError; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; + +use crate::function::{Function, extract_args}; + +/// Returns the keys of the outermost JSON object as a list of strings. +#[derive(Clone, Debug)] +pub(crate) struct JsonObjectKeysFunction { + signature: Signature, +} + +impl Default for JsonObjectKeysFunction { + fn default() -> Self { + Self { + signature: Signature::uniform( + 1, + vec![ + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::Null, + ], + Volatility::Immutable, + ), + } + } +} + +const NAME: &str = "json_object_keys"; + +impl Function for JsonObjectKeysFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8View, + true, + )))) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [jsons] = extract_args(self.name(), &args)?; + let jsons = compute::cast(&jsons, &DataType::BinaryView)?; + let jsons = jsons.as_binary_view(); + + let size = jsons.len(); + let mut builder = ListBuilder::with_capacity(StringViewBuilder::new(), size); + + for i in 0..size { + let Some(json) = jsons.is_valid(i).then(|| jsons.value(i)) else { + builder.append_null(); + continue; + }; + + match jsonb::from_slice(json) { + Ok(jsonb::Value::Object(object)) => { + for key in object.keys() { + builder.values().append_value(key); + } + builder.append(true); + } + Ok(_) => builder.append_null(), + Err(e) => { + return Err(DataFusionError::Execution(format!( + "invalid json binary: {e}" + ))); + } + } + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +impl Display for JsonObjectKeysFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "JSON_OBJECT_KEYS") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{BinaryArray, NullArray}; + use arrow_schema::Field; + + use super::*; + + #[test] + fn test_json_object_keys_function() { + let json_object_keys = JsonObjectKeysFunction::default(); + let return_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8View, true))); + + assert_eq!("json_object_keys", json_object_keys.name()); + assert_eq!( + return_type, + json_object_keys.return_type(&[DataType::Binary]).unwrap() + ); + + let json_strings = [ + Some(r#"{"b": 2, "a": 1}"#), + Some("{}"), + Some(r#"{"outer": {"inner": 1}, "value": 2}"#), + Some("[1, 2]"), + Some("42"), + Some("null"), + None, + ]; + + let results = [ + Some(vec!["a", "b"]), + Some(vec![]), + Some(vec!["outer", "value"]), + None, + None, + None, + None, + ]; + + let jsonbs = json_strings + .into_iter() + .map(|s| s.map(|json| jsonb::parse_value(json.as_bytes()).unwrap().to_vec())) + .collect::>(); + + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(BinaryArray::from_iter( + jsonbs, + )))], + arg_fields: vec![], + number_rows: 7, + return_field: Arc::new(Field::new("x", return_type.clone(), true)), + config_options: Arc::new(Default::default()), + }; + let result = json_object_keys + .invoke_with_args(args) + .and_then(|x| x.to_array(7)) + .unwrap(); + let vector = result.as_list::(); + + assert_eq!(7, vector.len()); + for (i, expected) in results.iter().enumerate() { + match expected { + Some(expected) => { + let values = vector.value(i); + let values = values.as_string_view(); + let keys = values.iter().flatten().collect::>(); + assert_eq!(expected, &keys); + } + None => assert!(vector.is_null(i)), + } + } + + let invalid_jsonb = vec![b"invalid json"]; + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new( + BinaryArray::from_iter_values(invalid_jsonb), + ))], + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(Field::new("x", return_type.clone(), true)), + config_options: Arc::new(Default::default()), + }; + let result = json_object_keys.invoke_with_args(args); + assert!(result.is_err()); + + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(NullArray::new(1)))], + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(Field::new("x", return_type, true)), + config_options: Arc::new(Default::default()), + }; + let result = json_object_keys + .invoke_with_args(args) + .and_then(|x| x.to_array(1)) + .unwrap(); + let vector = result.as_list::(); + assert!(vector.is_null(0)); + } +} diff --git a/tests/cases/standalone/common/function/json/json.result b/tests/cases/standalone/common/function/json/json.result index 62c562cb9f..ab32c9c013 100644 --- a/tests/cases/standalone/common/function/json/json.result +++ b/tests/cases/standalone/common/function/json/json.result @@ -1,3 +1,44 @@ +--- json_object_keys --- +SELECT json_object_keys(parse_json('{"a": 1, "b": {"c": 2}}')); + ++---------------------------------------------------------------+ +| json_object_keys(parse_json(Utf8("{"a": 1, "b": {"c": 2}}"))) | ++---------------------------------------------------------------+ +| [a, b] | ++---------------------------------------------------------------+ + +SELECT json_object_keys(parse_json('{}')); + ++------------------------------------------+ +| json_object_keys(parse_json(Utf8("{}"))) | ++------------------------------------------+ +| [] | ++------------------------------------------+ + +SELECT json_object_keys(parse_json('[1, 2]')); + ++----------------------------------------------+ +| json_object_keys(parse_json(Utf8("[1, 2]"))) | ++----------------------------------------------+ +| | ++----------------------------------------------+ + +SELECT json_object_keys(parse_json('null')); + ++--------------------------------------------+ +| json_object_keys(parse_json(Utf8("null"))) | ++--------------------------------------------+ +| | ++--------------------------------------------+ + +SELECT json_object_keys(NULL); + ++------------------------+ +| json_object_keys(NULL) | ++------------------------+ +| | ++------------------------+ + --- json_path_exists --- SELECT json_path_exists(parse_json('{"a": 1, "b": 2}'), '$.a'); diff --git a/tests/cases/standalone/common/function/json/json.sql b/tests/cases/standalone/common/function/json/json.sql index f8d6527ecc..e09f206df9 100644 --- a/tests/cases/standalone/common/function/json/json.sql +++ b/tests/cases/standalone/common/function/json/json.sql @@ -1,3 +1,14 @@ +--- json_object_keys --- +SELECT json_object_keys(parse_json('{"a": 1, "b": {"c": 2}}')); + +SELECT json_object_keys(parse_json('{}')); + +SELECT json_object_keys(parse_json('[1, 2]')); + +SELECT json_object_keys(parse_json('null')); + +SELECT json_object_keys(NULL); + --- json_path_exists --- SELECT json_path_exists(parse_json('{"a": 1, "b": 2}'), '$.a');