feat(frontend): support explain format json (#19041)

e2e_test/batch/explain.slt

@@ -11,4 +11,23 @@ statement ok
 drop table t;
 
 statement error Not supported: EXPLAIN CREATE VIEW
-explain create view v as select 1;
+explain create view v as select 1;
+
+query error
+explain (trace, format json) select 1;
+----
+    db error: ERROR: Failed to run the query
+
+Caused by:
+  Not supported: EXPLAIN (TRACE, JSON FORMAT)
+HINT: Only EXPLAIN (LOGICAL | PHYSICAL, JSON FORMAT) is supported.
+
+
+query error
+explain (distsql, format json) select 1;
+----
+    db error: ERROR: Failed to run the query
+
+Caused by:
+  Not supported: EXPLAIN (TRACE, JSON FORMAT)
+HINT: Only EXPLAIN (LOGICAL | PHYSICAL, JSON FORMAT) is supported.

src/frontend/planner_test/tests/testdata/input/explain_json_format.yaml

@@ -0,0 +1,18 @@
+- name: test json output format (logical)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (logical, format json) SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  expected_outputs:
+    - explain_output
+- name: test json output format (batch)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (physical, format json) SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  expected_outputs:
+    - explain_output
+- name: test json output format (stream)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (physical, format json) create materialized view m1 as SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  expected_outputs:
+    - explain_output

src/frontend/planner_test/tests/testdata/output/explain_json_format.yaml

@@ -0,0 +1,147 @@
+# This file is automatically generated. See `src/frontend/planner_test/README.md` for more information.
+- name: test json output format (logical)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (logical, format json) SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  explain_output: |
+    {
+      "name": "LogicalAgg",
+      "fields": {
+        "aggs": [
+          "approx_percentile($expr1)"
+        ]
+      },
+      "children": [
+        {
+          "name": "LogicalProject",
+          "fields": {
+            "exprs": [
+              "t.v1::Float64 as $expr1"
+            ]
+          },
+          "children": [
+            {
+              "name": "LogicalScan",
+              "fields": {
+                "columns": [
+                  "v1"
+                ],
+                "table": "t"
+              },
+              "children": []
+            }
+          ]
+        }
+      ]
+    }
+- name: test json output format (batch)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (physical, format json) SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  explain_output: |
+    {
+      "name": "BatchSimpleAgg",
+      "fields": {
+        "aggs": [
+          "approx_percentile($expr1)"
+        ]
+      },
+      "children": [
+        {
+          "name": "BatchExchange",
+          "fields": {
+            "dist": "Single",
+            "order": []
+          },
+          "children": [
+            {
+              "name": "BatchProject",
+              "fields": {
+                "exprs": [
+                  "t.v1::Float64 as $expr1"
+                ]
+              },
+              "children": [
+                {
+                  "name": "BatchScan",
+                  "fields": {
+                    "columns": [
+                      "v1"
+                    ],
+                    "table": "t"
+                  },
+                  "children": []
+                }
+              ]
+            }
+          ]
+        }
+      ]
+    }
+- name: test json output format (stream)
+  sql: |
+    CREATE TABLE t (v1 int);
+    explain (physical, format json) create materialized view m1 as SELECT approx_percentile(0.5) WITHIN GROUP (order by v1) from t;
+  explain_output: |
+    {
+      "name": "StreamMaterialize",
+      "fields": {
+        "columns": [
+          "approx_percentile"
+        ],
+        "pk_columns": [],
+        "pk_conflict": "NoCheck",
+        "stream_key": []
+      },
+      "children": [
+        {
+          "name": "StreamGlobalApproxPercentile",
+          "fields": {
+            "quantile": "0.5:Float64",
+            "relative_error": "0.01:Float64"
+          },
+          "children": [
+            {
+              "name": "StreamExchange",
+              "fields": {
+                "dist": "Single"
+              },
+              "children": [
+                {
+                  "name": "StreamLocalApproxPercentile",
+                  "fields": {
+                    "percentile_col": "$expr1",
+                    "quantile": "0.5:Float64",
+                    "relative_error": "0.01:Float64"
+                  },
+                  "children": [
+                    {
+                      "name": "StreamProject",
+                      "fields": {
+                        "exprs": [
+                          "t.v1::Float64 as $expr1",
+                          "t._row_id"
+                        ]
+                      },
+                      "children": [
+                        {
+                          "name": "StreamTableScan",
+                          "fields": {
+                            "columns": [
+                              "v1",
+                              "_row_id"
+                            ],
+                            "table": "t"
+                          },
+                          "children": []
+                        }
+                      ]
+                    }
+                  ]
+                }
+              ]
+            }
+          ]
+        }
+      ]
+    }

src/frontend/src/handler/explain.rs

@@ -16,7 +16,7 @@ use pgwire::pg_response::{PgResponse, StatementType};
 use risingwave_batch::worker_manager::worker_node_manager::WorkerNodeSelector;
 use risingwave_common::bail_not_implemented;
 use risingwave_common::types::Fields;
-use risingwave_sqlparser::ast::{ExplainOptions, ExplainType, Statement};
+use risingwave_sqlparser::ast::{ExplainFormat, ExplainOptions, ExplainType, Statement};
 use thiserror_ext::AsReport;
 
 use super::create_index::{gen_create_index_plan, resolve_index_schema};
@@ -179,6 +179,7 @@ async fn do_handle_explain(
         let explain_trace = context.is_explain_trace();
         let explain_verbose = context.is_explain_verbose();
         let explain_type = context.explain_type();
+        let explain_format = context.explain_format();
 
         if explain_trace {
             let trace = context.take_trace();
@@ -212,7 +213,10 @@ async fn do_handle_explain(
             ExplainType::Physical => {
                 // if explain trace is on, the plan has been in the rows
                 if !explain_trace && let Ok(plan) = &plan {
-                    blocks.push(plan.explain_to_string());
+                    match explain_format {
+                        ExplainFormat::Text => blocks.push(plan.explain_to_string()),
+                        ExplainFormat::Json => blocks.push(plan.explain_to_json()),
+                    }
                 }
             }
             ExplainType::Logical => {
@@ -248,6 +252,21 @@ pub async fn handle_explain(
     if analyze {
         bail_not_implemented!(issue = 4856, "explain analyze");
     }
+    if options.trace && options.explain_format == ExplainFormat::Json {
+        return Err(ErrorCode::NotSupported(
+            "EXPLAIN (TRACE, JSON FORMAT)".to_string(),
+            "Only EXPLAIN (LOGICAL | PHYSICAL, JSON FORMAT) is supported.".to_string(),
+        )
+        .into());
+    }
+    if options.explain_type == ExplainType::DistSql && options.explain_format == ExplainFormat::Json
+    {
+        return Err(ErrorCode::NotSupported(
+            "EXPLAIN (TRACE, JSON FORMAT)".to_string(),
+            "Only EXPLAIN (LOGICAL | PHYSICAL, JSON FORMAT) is supported.".to_string(),
+        )
+        .into());
+    }
 
     let mut blocks = Vec::new();
     let result = do_handle_explain(handler_args, options.clone(), stmt, &mut blocks).await;

src/frontend/src/optimizer/logical_optimization.rs

@@ -108,6 +108,8 @@ impl OptimizationStage {
 
 use std::sync::LazyLock;
 
+use risingwave_sqlparser::ast::ExplainFormat;
+
 pub struct LogicalOptimizer {}
 
 static DAG_TO_TREE: LazyLock<OptimizationStage> = LazyLock::new(|| {
@@ -684,7 +686,14 @@ impl LogicalOptimizer {
         InputRefValidator.validate(plan.clone());
 
         if ctx.is_explain_logical() {
-            ctx.store_logical(plan.explain_to_string());
+            match ctx.explain_format() {
+                ExplainFormat::Text => {
+                    ctx.store_logical(plan.explain_to_string());
+                }
+                ExplainFormat::Json => {
+                    ctx.store_logical(plan.explain_to_json());
+                }
+            }
         }
 
         Ok(plan)
@@ -789,7 +798,14 @@ impl LogicalOptimizer {
         InputRefValidator.validate(plan.clone());
 
         if ctx.is_explain_logical() {
-            ctx.store_logical(plan.explain_to_string());
+            match ctx.explain_format() {
+                ExplainFormat::Text => {
+                    ctx.store_logical(plan.explain_to_string());
+                }
+                ExplainFormat::Json => {
+                    ctx.store_logical(plan.explain_to_json());
+                }
+            }
         }
 
         Ok(plan)

src/frontend/src/optimizer/optimizer_context.rs

@@ -19,7 +19,7 @@ use std::marker::PhantomData;
 use std::rc::Rc;
 use std::sync::Arc;
 
-use risingwave_sqlparser::ast::{ExplainOptions, ExplainType};
+use risingwave_sqlparser::ast::{ExplainFormat, ExplainOptions, ExplainType};
 
 use crate::binder::ShareId;
 use crate::expr::{CorrelatedId, SessionTimezone};
@@ -176,6 +176,10 @@ impl OptimizerContext {
         self.explain_options.explain_type.clone()
     }
 
+    pub fn explain_format(&self) -> ExplainFormat {
+        self.explain_options.explain_format.clone()
+    }
+
     pub fn is_explain_logical(&self) -> bool {
         self.explain_type() == ExplainType::Logical
     }

src/frontend/src/optimizer/plan_node/mod.rs

@@ -53,6 +53,7 @@ use super::property::{Distribution, FunctionalDependencySet, MonotonicityMap, Or
 use crate::error::{ErrorCode, Result};
 use crate::optimizer::ExpressionSimplifyRewriter;
 use crate::session::current::notice_to_user;
+use crate::utils::PrettySerde;
 
 /// A marker trait for different conventions, used for enforcing type safety.
 ///
@@ -643,6 +644,9 @@ pub trait Explain {
 
     /// Explain the plan node and return a string.
     fn explain_to_string(&self) -> String;
+
+    /// Explain the plan node and return a json string.
+    fn explain_to_json(&self) -> String;
 }
 
 impl Explain for PlanRef {
@@ -665,6 +669,14 @@ impl Explain for PlanRef {
         config.unicode(&mut output, &plan.explain());
         output
     }
+
+    /// Explain the plan node and return a json string.
+    fn explain_to_json(&self) -> String {
+        let plan = reorganize_elements_id(self.clone());
+        let explain_ir = plan.explain();
+        serde_json::to_string_pretty(&PrettySerde(explain_ir))
+            .expect("failed to serialize plan to json")
+    }
 }
 
 pub(crate) fn pretty_config() -> PrettyConfig {

src/frontend/src/utils/mod.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod pretty_serde;
+pub use pretty_serde::PrettySerde;
 mod column_index_mapping;
 use std::any::Any;
 use std::hash::{Hash, Hasher};