From d2f3793d156f7d9eb8af7630d8781f4e953d3e1c Mon Sep 17 00:00:00 2001
From: WU Jingdi <taylor-lagrange@qq.com>
Date: Mon, 30 Oct 2023 21:21:25 +0800
Subject: [PATCH] fix: explain range query (#2664)

* fix: explain range query

* chore: add comment
---
 src/query/src/range_select/plan_rewrite.rs    | 40 ++++++++++++++++---
 .../cases/standalone/common/range/nest.result | 31 ++++++++++++++
 tests/cases/standalone/common/range/nest.sql  | 15 +++++++
 3 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/src/query/src/range_select/plan_rewrite.rs b/src/query/src/range_select/plan_rewrite.rs
index 5ac78f2013f0..ca7689af20b7 100644
--- a/src/query/src/range_select/plan_rewrite.rs
+++ b/src/query/src/range_select/plan_rewrite.rs
@@ -26,11 +26,12 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRewriter, VisitRecursion};
 use datafusion_common::{DFSchema, DataFusionError, Result as DFResult};
 use datafusion_expr::expr::ScalarUDF;
 use datafusion_expr::{
-    Aggregate, Expr, ExprSchemable, Extension, LogicalPlan, LogicalPlanBuilder, Projection,
+    Aggregate, Analyze, Explain, Expr, ExprSchemable, Extension, LogicalPlan, LogicalPlanBuilder,
+    Projection,
 };
 use datatypes::prelude::ConcreteDataType;
 use promql_parser::util::parse_duration;
-use snafu::{OptionExt, ResultExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use table::table::adapter::DfTableProviderAdapter;
 
 use super::plan::Fill;
@@ -265,9 +266,38 @@ impl RangePlanRewriter {
                             None => y.clone(),
                         })
                         .collect();
-                    Ok(Some(
-                        plan.with_new_inputs(&inputs).context(DataFusionSnafu)?,
-                    ))
+                    // Due to the limitations of Datafusion, for `LogicalPlan::Analyze` and `LogicalPlan::Explain`,
+                    // directly using the method `with_new_inputs` to rebuild a new `LogicalPlan` will cause an error,
+                    // so here we directly use the `LogicalPlanBuilder` to build a new plan.
+                    let plan = match plan {
+                        LogicalPlan::Analyze(Analyze { verbose, .. }) => {
+                            ensure!(
+                                inputs.len() == 1,
+                                RangeQuerySnafu {
+                                    msg: "Illegal subplan nums when rewrite Analyze logical plan",
+                                }
+                            );
+                            LogicalPlanBuilder::from(inputs[0].clone())
+                                .explain(*verbose, true)
+                                .context(DataFusionSnafu)?
+                                .build()
+                        }
+                        LogicalPlan::Explain(Explain { verbose, .. }) => {
+                            ensure!(
+                                inputs.len() == 1,
+                                RangeQuerySnafu {
+                                    msg: "Illegal subplan nums when rewrite Explain logical plan",
+                                }
+                            );
+                            LogicalPlanBuilder::from(inputs[0].clone())
+                                .explain(*verbose, false)
+                                .context(DataFusionSnafu)?
+                                .build()
+                        }
+                        _ => plan.with_new_inputs(&inputs),
+                    }
+                    .context(DataFusionSnafu)?;
+                    Ok(Some(plan))
                 } else {
                     Ok(None)
                 }
diff --git a/tests/cases/standalone/common/range/nest.result b/tests/cases/standalone/common/range/nest.result
index ae46183cc254..7675b8384c06 100644
--- a/tests/cases/standalone/common/range/nest.result
+++ b/tests/cases/standalone/common/range/nest.result
@@ -45,6 +45,37 @@ SELECT ts, b, min(c) RANGE '5s' FROM (SELECT ts, host AS b, val AS c FROM host W
 | 1970-01-01T00:00:20 | host1 | 2                         |
 +---------------------+-------+---------------------------+
 
+-- Test EXPLAIN and ANALYZE
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (metrics.*) REDACTED
+EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s';
+
++-+-+
+| plan_type_| plan_|
++-+-+
+| logical_plan_| RangeSelect: range_exprs=[MIN(host.val) RANGE 5s FILL NULL], align=5s time_index=ts_|
+|_|_MergeScan [is_placeholder=false]_|
+| physical_plan | RangeSelectExec: range_expr=[RangeFnExec{ MIN(host.val), range: 5000}], align=5000, time_index=ts, by=[host@1] |
+|_|_MergeScanExec: REDACTED
+|_|_|
++-+-+
+
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (metrics.*) REDACTED
+EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s';
+
++-+-+
+| plan_type_| plan_|
++-+-+
+| Plan with Metrics | RangeSelectExec: range_expr=[RangeFnExec{ MIN(host.val), range: 5000}], align=5000, time_index=ts, by=[host@1], REDACTED
+|_|_MergeScanExec: REDACTED
+|_|_|
++-+-+
+
 DROP TABLE host;
 
 Affected Rows: 0
diff --git a/tests/cases/standalone/common/range/nest.sql b/tests/cases/standalone/common/range/nest.sql
index d2478b57311f..50fe3f9305c5 100644
--- a/tests/cases/standalone/common/range/nest.sql
+++ b/tests/cases/standalone/common/range/nest.sql
@@ -22,4 +22,19 @@ SELECT ts, host, foo FROM (SELECT ts, host, min(val) RANGE '5s' AS foo FROM host
 
 SELECT ts, b, min(c) RANGE '5s' FROM (SELECT ts, host AS b, val AS c FROM host WHERE host = 'host1') ALIGN '5s' BY (b) ORDER BY b, ts;
 
+
+-- Test EXPLAIN and ANALYZE
+
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (metrics.*) REDACTED
+EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s';
+
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (metrics.*) REDACTED
+EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s';
+
 DROP TABLE host;