feat: support dot access for JSON columns (#9775)

Signed-off-by: luofucong <luofc@foxmail.com>

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<PlannerResult<Vec<Expr>>> {
-        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(())

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
+

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;