From 4dd92c774e3cceca0b3e505f23334f4d55c18444 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 14 Aug 2026 08:24:50 +0000 Subject: [PATCH] feat: add json_object function and use it in the entity-graph derivation (#8870) * feat: add json_object scalar function Builds a JSONB object from interleaved (key, value, ...) arguments, like MySQL's JSON_OBJECT. Values are written into the binary directly, so JSON-hostile characters (quotes, backslashes, control characters) need no text-level escaping. Keys must be non-NULL strings; values may be strings, numbers, booleans, or NULL (JSON null). Signed-off-by: Dennis Zhuang * fix: build entity-graph JSON objects with json_object The derivation assembled entity_id_attrs and descriptive by concatenating a JSON text and parsing it, escaping only backslash and double quote in runtime values. A label containing a control character (e.g. a newline) produced unparseable text and failed the whole semantic_entities scan instead of one attribute. json_object assembles the JSONB binary directly from the value columns, so no text escaping is involved; NULL-to-'' stays at the call site. Signed-off-by: Dennis Zhuang * chore: trim comments and fold duplicate test coverage Signed-off-by: Dennis Zhuang * fix: json_object() returns an empty object; narrow values to integers and floats MySQL's JSON_OBJECT allows an empty pair list, so the signature accepts zero arguments and the row count falls back to number_rows. Decimals stay rejected instead of casting to Float64: JSONB numbers (i64/u64/f64) cannot represent them exactly and a silent precision loss is worse than an explicit cast. Signed-off-by: Dennis Zhuang * chore: document key-to-string conversion and align test naming Keys follow MySQL JSON_OBJECT: any castable type is converted to string. Rustdoc and the cast-failure message now say so, with a numeric-key test. Test names take the module-conventional test_ prefix. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/scalars/json.rs | 2 + .../function/src/scalars/json/json_object.rs | 281 ++++++++++++++++++ src/operator/src/statement/semantic_graph.rs | 66 ++-- .../common/function/json/json.result | 70 +++++ .../standalone/common/function/json/json.sql | 25 ++ 5 files changed, 406 insertions(+), 38 deletions(-) create mode 100644 src/common/function/src/scalars/json/json_object.rs diff --git a/src/common/function/src/scalars/json.rs b/src/common/function/src/scalars/json.rs index 8a06c972c9..d0e7f412ac 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; mod json_object_keys; mod json_path_exists; mod json_path_match; @@ -54,6 +55,7 @@ impl JsonFunction { registry.register_scalar(JsonIsArray::default()); registry.register_scalar(JsonIsObject::default()); + registry.register_scalar(json_object::JsonObjectFunction::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.rs b/src/common/function/src/scalars/json/json_object.rs new file mode 100644 index 0000000000..a9602c3abc --- /dev/null +++ b/src/common/function/src/scalars/json/json_object.rs @@ -0,0 +1,281 @@ +// 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 datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, BinaryViewBuilder}; +use datafusion_common::arrow::compute; +use datafusion_common::arrow::datatypes::{DataType, Float64Type, Int64Type, UInt64Type}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, TypeSignature, Volatility}; + +use crate::function::Function; + +const NAME: &str = "json_object"; + +/// Builds a `JSONB` object from interleaved `(key, value, key, value, ...)` +/// arguments, like MySQL's `JSON_OBJECT`; called with no arguments it returns +/// `{}`. Values are written into the binary directly, so they need no JSON +/// text escaping. Keys must be non-NULL and are converted to strings (so a +/// numeric key like `1` becomes `"1"`, as in MySQL); values may be strings, +/// integers, floats, booleans, or NULL (rendered as JSON null). Other types — +/// including decimals, which JSONB numbers cannot represent exactly — are +/// rejected; cast them explicitly. A duplicate key keeps the last value. +#[derive(Clone, Debug)] +pub(crate) struct JsonObjectFunction { + signature: Signature, +} + +impl Default for JsonObjectFunction { + fn default() -> Self { + Self { + signature: Signature::one_of( + vec![TypeSignature::Nullary, TypeSignature::VariadicAny], + Volatility::Immutable, + ), + } + } +} + +/// A value column normalized to the canonical arrow type its JSON rendering +/// reads from. +enum ValueColumn { + Null, + Bool(ArrayRef), + Int(ArrayRef), + UInt(ArrayRef), + Float(ArrayRef), + String(ArrayRef), +} + +impl ValueColumn { + fn try_new(array: &ArrayRef) -> datafusion_common::Result { + let normalized = match array.data_type() { + DataType::Null => return Ok(ValueColumn::Null), + DataType::Boolean => ValueColumn::Bool(array.clone()), + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { + ValueColumn::Int(compute::cast(array, &DataType::Int64)?) + } + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => { + ValueColumn::UInt(compute::cast(array, &DataType::UInt64)?) + } + DataType::Float16 | DataType::Float32 | DataType::Float64 => { + ValueColumn::Float(compute::cast(array, &DataType::Float64)?) + } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + ValueColumn::String(compute::cast(array, &DataType::Utf8View)?) + } + other => { + return Err(DataFusionError::Execution(format!( + "{NAME} does not support values of type {other}; cast the value to a string" + ))); + } + }; + Ok(normalized) + } + + fn value(&self, row: usize) -> jsonb::Value<'_> { + let array = match self { + ValueColumn::Null => return jsonb::Value::Null, + ValueColumn::Bool(array) + | ValueColumn::Int(array) + | ValueColumn::UInt(array) + | ValueColumn::Float(array) + | ValueColumn::String(array) => array, + }; + if !array.is_valid(row) { + return jsonb::Value::Null; + } + match self { + ValueColumn::Null => unreachable!(), + ValueColumn::Bool(array) => array.as_boolean().value(row).into(), + ValueColumn::Int(array) => array.as_primitive::().value(row).into(), + ValueColumn::UInt(array) => array.as_primitive::().value(row).into(), + ValueColumn::Float(array) => array.as_primitive::().value(row).into(), + ValueColumn::String(array) => array.as_string_view().value(row).into(), + } + } +} + +impl Function for JsonObjectFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::BinaryView) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + if arrays.len() % 2 != 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} expects (key, value) argument pairs, got {} arguments", + arrays.len() + ))); + } + let pairs = arrays + .chunks(2) + .map(|pair| { + let keys = compute::cast(&pair[0], &DataType::Utf8View).map_err(|_| { + DataFusionError::Execution(format!( + "{NAME} cannot convert keys of type {} to string", + pair[0].data_type() + )) + })?; + if keys.null_count() > 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} does not allow NULL keys" + ))); + } + Ok((keys, ValueColumn::try_new(&pair[1])?)) + }) + .collect::>>()?; + + let rows = arrays.first().map_or(args.number_rows, |a| a.len()); + let mut builder = BinaryViewBuilder::with_capacity(rows); + let mut buf = Vec::new(); + for row in 0..rows { + let mut object = jsonb::Object::new(); + for (keys, values) in &pairs { + object.insert( + keys.as_string_view().value(row).to_string(), + values.value(row), + ); + } + buf.clear(); + jsonb::Value::Object(object).write_to_vec(&mut buf); + builder.append_value(&buf); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +impl Display for JsonObjectFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "JSON_OBJECT") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{ + Int64Array, NullArray, StringArray, TimestampMillisecondArray, + }; + + use super::*; + + fn invoke(args: Vec, rows: usize) -> datafusion_common::Result> { + let function = JsonObjectFunction::default(); + let result = function.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("x", DataType::BinaryView, true)), + config_options: Arc::new(Default::default()), + })?; + let array = result.to_array(rows)?; + let array = array.as_binary_view(); + Ok((0..array.len()) + .map(|i| jsonb::from_slice(array.value(i)).unwrap().to_string()) + .collect()) + } + + fn key(name: &str) -> ColumnarValue { + ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(name.to_string()))) + } + + #[test] + fn test_builds_objects_from_mixed_types_without_escaping() { + let texts = invoke( + vec![ + key("host"), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("we\"ird\\\nhost"), + None, + ]))), + key("pid"), + ColumnarValue::Array(Arc::new(Int64Array::from(vec![42, 7]))), + key("up"), + ColumnarValue::Scalar(datafusion_common::ScalarValue::Boolean(Some(true))), + key("v"), + ColumnarValue::Array(Arc::new(NullArray::new(2))), + ], + 2, + ) + .unwrap(); + assert_eq!( + texts, + vec![ + r#"{"host":"we\"ird\\\nhost","pid":42,"up":true,"v":null}"#, + r#"{"host":null,"pid":7,"up":true,"v":null}"#, + ] + ); + } + + #[test] + fn test_empty_call_builds_one_empty_object_per_row() { + assert_eq!(invoke(vec![], 3).unwrap(), vec!["{}", "{}", "{}"]); + } + + #[test] + fn test_numeric_key_converts_to_string() { + let texts = invoke( + vec![ + ColumnarValue::Scalar(datafusion_common::ScalarValue::Int64(Some(7))), + key("v"), + ], + 1, + ) + .unwrap(); + assert_eq!(texts, vec![r#"{"7":"v"}"#]); + } + + #[test] + fn test_rejects_odd_arguments_null_keys_and_unsupported_values() { + let err = invoke(vec![key("a")], 1).unwrap_err(); + assert!(err.to_string().contains("(key, value) argument pairs")); + + let err = invoke( + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec![None::<&str>]))), + key("v"), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("NULL keys")); + + let err = invoke( + vec![ + key("ts"), + ColumnarValue::Array(Arc::new(TimestampMillisecondArray::from(vec![1_000]))), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not support values")); + } +} diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 777c6f677e..7014a94f8a 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -465,43 +465,34 @@ fn null_json() -> Expr { lit(ScalarValue::Binary(None)) } -/// Renders a compile-time-known string as JSON text (quoted, fully escaped — -/// including control characters, unlike the runtime value escaping). +/// Renders a compile-time-known string as JSON text (quoted, fully escaped). fn json_quote(value: &str) -> String { serde_json::Value::from(value).to_string() } -/// Wraps a column reference in `replace` calls so its runtime value is -/// JSON-escaped (`\` then `"`); NULL becomes `''` so one NULL column does not -/// invalidate the whole JSON text (descriptive columns are nullable). -fn json_escaped_value_expr(column: &str) -> Expr { - let escaped_backslash = - string_fns::replace().call(vec![cast_string_or_empty(column), lit("\\"), lit("\\\\")]); - string_fns::replace().call(vec![escaped_backslash, lit("\""), lit("\\\"")]) -} +/// The `json_object` UDF, resolved like [`PARSE_JSON_UDF`]. It assembles the +/// JSONB binary directly from the value columns, so runtime values need no +/// JSON text escaping. +static JSON_OBJECT_UDF: LazyLock> = LazyLock::new(|| { + Arc::new( + FUNCTION_REGISTRY + .get_function("json_object") + .expect("json_object must be registered") + .provide(FunctionContext::default()), + ) +}); -/// Builds a JSONB object from `columns` by concatenating a JSON text and parsing -/// it — GreptimeDB has no struct→json function. Keys are JSON-escaped in Rust; -/// values are JSON-escaped at runtime via [`json_escaped_value_expr`]. -/// -/// TODO(entity-graph): replace the text round-trip with a UDF that assembles -/// JSONB directly from the value columns (`jsonb::ObjectBuilder`, keys baked -/// in), dropping the escaping helpers and the per-row parse cost. +/// Builds a JSONB object with one entry per column: key = the column name, +/// value = the column rendered as a string, NULL coalesced to `""` so one NULL +/// column does not null the entry (descriptive columns are nullable). Keys come +/// out sorted — JSONB objects are key-ordered regardless of input order. fn json_object_expr(columns: &[String]) -> Expr { - if columns.is_empty() { - return parse_json_expr(lit("{}")); + let mut args = Vec::with_capacity(columns.len() * 2); + for column in columns { + args.push(lit(column.as_str())); + args.push(cast_string_or_empty(column)); } - let mut parts = vec![lit("{")]; - for (i, column) in columns.iter().enumerate() { - if i > 0 { - parts.push(lit(",")); - } - parts.push(lit(format!("{}:\"", json_quote(column)))); - parts.push(json_escaped_value_expr(column)); - parts.push(lit("\"")); - } - parts.push(lit("}")); - parse_json_expr(concat_expr(parts)) + cast(JSON_OBJECT_UDF.call(args), DataType::Binary) } /// Renders pre-sorted columns as a `k=v,k=v` concatenation. `nullable` @@ -599,9 +590,7 @@ fn registry_source( let entity_id_attrs = if decl.id_columns.len() == 1 { null_json() } else { - let mut cols = decl.id_columns.clone(); - cols.sort(); - json_object_expr(&cols) + json_object_expr(&decl.id_columns) }; let scope = match decl.scope_columns.as_slice() { @@ -799,7 +788,7 @@ mod tests { Arc::new(StringArray::from(vec!["cart", "cart", "cart"])), Arc::new(Int64Array::from(vec![42, 42, 42])), Arc::new(StringArray::from(vec![ - Some(r#"we"ird\host"#), + Some("we\"ird\\\nhost"), None, Some("h2"), ])), @@ -912,7 +901,8 @@ mod tests { rows.sort(); // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON escapes `\` and `"` in runtime values, NULL -> "". + // descriptive JSON keeps `\`, `"` and control characters intact in + // runtime values, NULL -> "". assert_eq!( rows, vec![ @@ -929,7 +919,7 @@ mod tests { ( "pid=42,service_name=cart".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), - Some(r#"{"host":"we\"ird\\host"}"#.to_string()), + Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), ] ); @@ -955,7 +945,7 @@ mod tests { let batches = collect(&ctx, plan).await; let mut scopes: Vec = batches.iter().flat_map(|b| strings(b, 7)).collect(); scopes.sort(); - assert_eq!(scopes, vec!["", "h2", r#"we"ird\host"#]); + assert_eq!(scopes, vec!["", "h2", "we\"ird\\\nhost"]); // Multiple scope columns: sorted `k=v,k=v`. let mut multi = decl("service", &["service_name"]); @@ -978,7 +968,7 @@ mod tests { vec![ "host=,pid=42", "host=h2,pid=42", - r#"host=we"ird\host,pid=42"# + "host=we\"ird\\\nhost,pid=42" ] ); } diff --git a/tests/cases/standalone/common/function/json/json.result b/tests/cases/standalone/common/function/json/json.result index ab32c9c013..c035db72e4 100644 --- a/tests/cases/standalone/common/function/json/json.result +++ b/tests/cases/standalone/common/function/json/json.result @@ -145,3 +145,73 @@ SELECT json_path_match(parse_json('null'), '$.a == 1'); | | +------------------------------------------------------------+ +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + ++-----------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("a"),Int64(1),Utf8("b"),Utf8("text"),Utf8("c"),Boolean(true),Utf8("d"),Float64(1.5))) | ++-----------------------------------------------------------------------------------------------------------------------+ +| {"a":1,"b":"text","c":true,"d":1.5} | ++-----------------------------------------------------------------------------------------------------------------------+ + +SELECT json_to_string(json_object()); + ++-------------------------------+ +| json_to_string(json_object()) | ++-------------------------------+ +| {} | ++-------------------------------+ + +SELECT json_to_string(json_object('nul', NULL)); + ++-----------------------------------------------+ +| json_to_string(json_object(Utf8("nul"),NULL)) | ++-----------------------------------------------+ +| {"nul":null} | ++-----------------------------------------------+ + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + ++-------------------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("nl"),concat(Utf8("line1"),chr(Int64(10)),Utf8("line2")),Utf8("q"),Utf8("quote"back\slash"))) | ++-------------------------------------------------------------------------------------------------------------------------------+ +| {"nl":"line1\nline2","q":"quote\"back\\slash"} | ++-------------------------------------------------------------------------------------------------------------------------------+ + +SELECT json_object('a'); + +Error: 3001(EngineExecuteQuery), Execution error: json_object expects (key, value) argument pairs, got 1 arguments + +SELECT json_object(NULL, 1); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not allow NULL keys + +SELECT json_object('ts', to_timestamp(0)); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Timestamp(ns); cast the value to a string + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Decimal128(10, 2); cast the value to a string + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +Affected Rows: 0 + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +Affected Rows: 2 + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + ++------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("host"),json_object_src.host,Utf8("pid"),json_object_src.pid)) | ++------------------------------------------------------------------------------------------------+ +| {"host":"h1","pid":42} | +| {"host":null,"pid":7} | ++------------------------------------------------------------------------------------------------+ + +DROP TABLE json_object_src; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/json/json.sql b/tests/cases/standalone/common/function/json/json.sql index e09f206df9..16517b54f6 100644 --- a/tests/cases/standalone/common/function/json/json.sql +++ b/tests/cases/standalone/common/function/json/json.sql @@ -37,3 +37,28 @@ SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), '$.b[1 to last] >= 2') SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), 'null'); SELECT json_path_match(parse_json('null'), '$.a == 1'); + +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + +SELECT json_to_string(json_object()); + +SELECT json_to_string(json_object('nul', NULL)); + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + +SELECT json_object('a'); + +SELECT json_object(NULL, 1); + +SELECT json_object('ts', to_timestamp(0)); + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + +DROP TABLE json_object_src;