diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index 4837403c73f..a60e6155445 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -33,7 +33,8 @@ use datafusion_expr::{ Expr, ExprSchemable, GetFieldAccess, Operator, ScalarUDF, WindowFunctionDefinition, }; use datatypes::extension::json::{ - Json2ExtensionType, is_json2_extension_type, parse_legacy_json2_settings, + Json2ExtensionType, is_any_json_extension_type, is_json2_extension_type, + parse_legacy_json2_settings, }; use datatypes::types::json_type::JsonNativeType; use jsonb::jsonpath::Path; @@ -138,7 +139,13 @@ impl ExprPlanner for JsonExprPlanner { qualifier: Option<&TableReference>, nested_names: &[String], ) -> Result>> { - if !is_json2_extension_type(field) { + if !is_any_json_extension_type(field) + // SQL-created JSONB columns use Binary with greptime:type=Json, without an Arrow extension. + && field + .metadata() + .get(datatypes::schema::TYPE_KEY) + .is_none_or(|value| value != "Json") + { return Ok(PlannerResult::Original(Vec::new())); } @@ -514,9 +521,10 @@ mod tests { use datafusion_expr::WindowFrame; use datafusion_functions::core::coalesce; use datafusion_functions::math::{abs, power}; - use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType, JsonMetadata}; use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::prelude::ConcreteDataType; + use std::collections::HashMap; use super::*; @@ -654,78 +662,85 @@ mod tests { let qualifier = TableReference::bare("events"); let nested_names = vec!["payload".to_string(), "cpu".to_string()]; - let planned = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) + for field in [ + Field::new("labels", DataType::Struct(Fields::empty()), true) .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &nested_names, - )?; + Field::new("labels", DataType::Binary, true).with_extension_type(JsonExtensionType), + Field::new("labels", DataType::Binary, true).with_metadata(HashMap::from([( + datatypes::schema::TYPE_KEY.to_string(), + "Json".to_string(), + )])), + ] { + let planned = + planner.plan_compound_identifier(&field, Some(&qualifier), &nested_names)?; - match planned { - PlannerResult::Planned(Expr::ScalarFunction(func)) => { - assert_eq!(func.func.name(), JsonGetWithType::NAME); - assert_eq!(func.args.len(), 2); - assert_eq!( - func.args[0], - Expr::Column(Column::new(Some(qualifier.clone()), "labels")) - ); + match planned { + PlannerResult::Planned(Expr::ScalarFunction(func)) => { + assert_eq!(func.func.name(), JsonGetWithType::NAME); + assert_eq!(func.args.len(), 2); + assert_eq!( + func.args[0], + Expr::Column(Column::new(Some(qualifier.clone()), "labels")) + ); + assert_eq!( + func.args[1], + Expr::Literal(ScalarValue::Utf8(Some("$.payload.cpu".to_string())), None) + ); + } + other => panic!("expected json_get scalar function, got {other:?}"), + } + + for (key, path) in [ + ("http.status_code", r#"$."http.status_code""#), + ("a\"b", r#"$."a\"b""#), + ("a\\b", r#"$."a\\b""#), + ] { + let planned = planner.plan_compound_identifier( + &field, + Some(&qualifier), + &[key.to_string()], + )?; + let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else { + unreachable!() + }; assert_eq!( func.args[1], - Expr::Literal(ScalarValue::Utf8(Some("$.payload.cpu".to_string())), None) + Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None) ); } - other => panic!("expected json_get scalar function, got {other:?}"), - } - for (key, path) in [ - ("http.status_code", r#"$."http.status_code""#), - ("a\"b", r#"$."a\"b""#), - ("a\\b", r#"$."a\\b""#), - ] { - let planned = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) - .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &[key.to_string()], - )?; - let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else { + let PlannerResult::Planned(Expr::ScalarFunction(func)) = planner + .plan_compound_identifier( + &field, + Some(&qualifier), + &["resource".to_string(), "http.status_code".to_string()], + )? + else { unreachable!() }; assert_eq!( func.args[1], - Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None) + Expr::Literal( + ScalarValue::Utf8(Some(r#"$.resource."http.status_code""#.to_string())), + None + ) ); } - let PlannerResult::Planned(Expr::ScalarFunction(func)) = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) - .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &["resource".to_string(), "http.status_code".to_string()], - )? - else { - unreachable!() - }; - assert_eq!( - func.args[1], - Expr::Literal( - ScalarValue::Utf8(Some(r#"$.resource."http.status_code""#.to_string())), - None - ) - ); + for data_type in [DataType::Utf8, DataType::Binary] { + let original = planner.plan_compound_identifier( + &Field::new("plain", data_type, true), + Some(&qualifier), + &nested_names, + )?; - let original = planner.plan_compound_identifier( - &Field::new("plain", DataType::Utf8, true), - Some(&qualifier), - &nested_names, - )?; - - match original { - PlannerResult::Original(exprs) => assert!(exprs.is_empty()), - other => panic!( - "expected original empty result for non-json field, got {:?}", - other, - ), + match original { + PlannerResult::Original(exprs) => assert!(exprs.is_empty()), + other => panic!( + "expected original empty result for non-json field, got {:?}", + other, + ), + } } Ok(()) diff --git a/tests/cases/standalone/common/types/json/json_dot_access.result b/tests/cases/standalone/common/types/json/json_dot_access.result new file mode 100644 index 00000000000..fe14dde72eb --- /dev/null +++ b/tests/cases/standalone/common/types/json/json_dot_access.result @@ -0,0 +1,97 @@ +CREATE TABLE json_dot_access (ts TIMESTAMP TIME INDEX, j JSON); + +Affected Rows: 0 + +INSERT INTO json_dot_access VALUES + (1, '{"a":{"b":2},"items":[{"name":"first"},{"name":"second"}],"a.b":"literal","UPPER":"quoted","enabled":true}'), + (2, '{"a":{"b":4},"items":[],"enabled":false}'), + (3, '{"a":null}'), + (4, '{}'), + (5, NULL); + +Affected Rows: 5 + +SELECT ts, j.a.b AS nested, j."a.b" AS dotted_key, j."UPPER" AS quoted_key, + j.missing AS missing, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + ++-------------------------+--------+------------+------------+---------+--------+ +| ts | nested | dotted_key | quoted_key | missing | item | ++-------------------------+--------+------------+------------+---------+--------+ +| 1970-01-01T00:00:00.001 | 2 | literal | quoted | | second | +| 1970-01-01T00:00:00.002 | 4 | | | | | +| 1970-01-01T00:00:00.003 | | | | | | +| 1970-01-01T00:00:00.004 | | | | | | +| 1970-01-01T00:00:00.005 | | | | | | ++-------------------------+--------+------------+------------+---------+--------+ + +SELECT ts, j.a.b = json_get(j, '$.a.b') AS same_value, + j.a.b + 1 AS incremented, j.a.b::BIGINT AS cast_value, + abs(j.a.b) AS absolute +FROM json_dot_access ORDER BY ts; + ++-------------------------+------------+-------------+------------+----------+ +| ts | same_value | incremented | cast_value | absolute | ++-------------------------+------------+-------------+------------+----------+ +| 1970-01-01T00:00:00.001 | true | 3 | 2 | 2.0 | +| 1970-01-01T00:00:00.002 | true | 5 | 4 | 4.0 | +| 1970-01-01T00:00:00.003 | | | | | +| 1970-01-01T00:00:00.004 | | | | | +| 1970-01-01T00:00:00.005 | | | | | ++-------------------------+------------+-------------+------------+----------+ + +SELECT t.ts, t.j.a.b AS qualified FROM json_dot_access AS t +WHERE t.j.a.b > 2 AND t.j.enabled = false ORDER BY t.ts; + ++-------------------------+-----------+ +| ts | qualified | ++-------------------------+-----------+ +| 1970-01-01T00:00:00.002 | 4 | ++-------------------------+-----------+ + +SELECT sum(j.a.b) AS total FROM json_dot_access; + ++-------+ +| total | ++-------+ +| 6.0 | ++-------+ + +SELECT ts, payload.a.b AS nested +FROM (SELECT ts, j AS payload FROM json_dot_access) ORDER BY ts; + ++-------------------------+--------+ +| ts | nested | ++-------------------------+--------+ +| 1970-01-01T00:00:00.001 | 2 | +| 1970-01-01T00:00:00.002 | 4 | +| 1970-01-01T00:00:00.003 | | +| 1970-01-01T00:00:00.004 | | +| 1970-01-01T00:00:00.005 | | ++-------------------------+--------+ + +ADMIN FLUSH_TABLE('json_dot_access'); + ++--------------------------------------+ +| ADMIN FLUSH_TABLE('json_dot_access') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +SELECT ts, j.a.b AS nested, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + ++-------------------------+--------+--------+ +| ts | nested | item | ++-------------------------+--------+--------+ +| 1970-01-01T00:00:00.001 | 2 | second | +| 1970-01-01T00:00:00.002 | 4 | | +| 1970-01-01T00:00:00.003 | | | +| 1970-01-01T00:00:00.004 | | | +| 1970-01-01T00:00:00.005 | | | ++-------------------------+--------+--------+ + +DROP TABLE json_dot_access; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json_dot_access.sql b/tests/cases/standalone/common/types/json/json_dot_access.sql new file mode 100644 index 00000000000..b8cbc86dd23 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json_dot_access.sql @@ -0,0 +1,32 @@ +CREATE TABLE json_dot_access (ts TIMESTAMP TIME INDEX, j JSON); + +INSERT INTO json_dot_access VALUES + (1, '{"a":{"b":2},"items":[{"name":"first"},{"name":"second"}],"a.b":"literal","UPPER":"quoted","enabled":true}'), + (2, '{"a":{"b":4},"items":[],"enabled":false}'), + (3, '{"a":null}'), + (4, '{}'), + (5, NULL); + +SELECT ts, j.a.b AS nested, j."a.b" AS dotted_key, j."UPPER" AS quoted_key, + j.missing AS missing, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + +SELECT ts, j.a.b = json_get(j, '$.a.b') AS same_value, + j.a.b + 1 AS incremented, j.a.b::BIGINT AS cast_value, + abs(j.a.b) AS absolute +FROM json_dot_access ORDER BY ts; + +SELECT t.ts, t.j.a.b AS qualified FROM json_dot_access AS t +WHERE t.j.a.b > 2 AND t.j.enabled = false ORDER BY t.ts; + +SELECT sum(j.a.b) AS total FROM json_dot_access; + +SELECT ts, payload.a.b AS nested +FROM (SELECT ts, j AS payload FROM json_dot_access) ORDER BY ts; + +ADMIN FLUSH_TABLE('json_dot_access'); + +SELECT ts, j.a.b AS nested, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + +DROP TABLE json_dot_access;