rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-31 04:10:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: explain custom statement (#7058)

Browse Source 
* feat: explain tql cte

Signed-off-by: discord9 <discord9@163.com>

* chore: unused

Signed-off-by: discord9 <discord9@163.com>

* fix: analyze format

Signed-off-by: discord9 <discord9@163.com>

* Update src/sql/src/statements/statement.rs

Co-authored-by: Yingwen <realevenyag@gmail.com>
Signed-off-by: discord9 <discord9@163.com>

* test: sqlness

Signed-off-by: discord9 <discord9@163.com>

* pcr

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
Co-authored-by: Yingwen <realevenyag@gmail.com>
This commit is contained in:
discord9
2025-10-11 14:27:51 +08:00
committed by GitHub
parent e46ce7c6da
commit ba034c5a9e
6 changed files with 935 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
src/query/src/planner.rs
View File

@@ -20,10 +20,14 @@ use async_trait::async_trait;
use catalog::table_source::DfTableSourceProvider;
use common_error::ext::BoxedError;
use common_telemetry::tracing;
use datafusion::common::DFSchema;
use datafusion::common::{DFSchema, plan_err};
use datafusion::execution::context::SessionState;
use datafusion::sql::planner::PlannerContext;
use datafusion_expr::{Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, col};
use datafusion_common::ToDFSchema;
use datafusion_expr::{
Analyze, Explain, ExplainFormat, Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, PlanType,
ToStringifiedPlan, col,
};
use datafusion_sql::planner::{ParserOptions, SqlToRel};
use log_query::LogQuery;
use promql_parser::parser::EvalStmt;
@@ -31,6 +35,7 @@ use session::context::QueryContextRef;
use snafu::{ResultExt, ensure};
use sql::CteContent;
use sql::ast::Expr as SqlExpr;
use sql::statements::explain::ExplainStatement;
use sql::statements::query::Query;
use sql::statements::statement::Statement;
use sql::statements::tql::Tql;
@@ -75,12 +80,69 @@ impl DfLogicalPlanner {
}
}
/// Basically the same with `explain_to_plan` in DataFusion, but adapted to Greptime's
/// `plan_sql` to support Greptime Statements.
async fn explain_to_plan(
&self,
explain: &ExplainStatement,
query_ctx: QueryContextRef,
) -> Result<LogicalPlan> {
let plan = self.plan_sql(&explain.statement, query_ctx).await?;
if matches!(plan, LogicalPlan::Explain(_)) {
return plan_err!("Nested EXPLAINs are not supported").context(PlanSqlSnafu);
}
let verbose = explain.verbose;
let analyze = explain.analyze;
let format = explain.format.map(|f| f.to_string());
let plan = Arc::new(plan);
let schema = LogicalPlan::explain_schema();
let schema = ToDFSchema::to_dfschema_ref(schema)?;
if verbose && format.is_some() {
return plan_err!("EXPLAIN VERBOSE with FORMAT is not supported").context(PlanSqlSnafu);
}
if analyze {
// notice format is already set in query context, so can be ignore here
Ok(LogicalPlan::Analyze(Analyze {
verbose,
input: plan,
schema,
}))
} else {
let stringified_plans = vec![plan.to_stringified(PlanType::InitialLogicalPlan)];
// default to configuration value
let options = self.session_state.config().options();
let format = format.as_ref().unwrap_or(&options.explain.format);
let format: ExplainFormat = format.parse()?;
Ok(LogicalPlan::Explain(Explain {
verbose,
explain_format: format,
plan,
stringified_plans,
schema,
logical_optimization_succeeded: false,
}))
}
}
#[tracing::instrument(skip_all)]
#[async_recursion::async_recursion]
async fn plan_sql(&self, stmt: &Statement, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
let mut planner_context = PlannerContext::new();
let mut stmt = Cow::Borrowed(stmt);
let mut is_tql_cte = false;
// handle explain before normal processing so we can explain Greptime Statements
if let Statement::Explain(explain) = stmt.as_ref() {
return self.explain_to_plan(explain, query_ctx).await;
}
// Check for hybrid CTEs before normal processing
if self.has_hybrid_ctes(stmt.as_ref()) {
let stmt_owned = stmt.into_owned();
@@ -101,10 +163,13 @@ impl DfLogicalPlanner {
// TODO(LFC): Remove this when Datafusion supports **both** the syntax and implementation of "explain with format".
if let datafusion::sql::parser::Statement::Statement(
box datafusion::sql::sqlparser::ast::Statement::Explain { format, .. },
box datafusion::sql::sqlparser::ast::Statement::Explain { .. },
) = &mut df_stmt
{
format.take();
UnimplementedSnafu {
operation: "EXPLAIN with FORMAT using raw datafusion planner",
}
.fail()?;
}
let table_provider = DfTableSourceProvider::new(

83
src/sql/src/parsers/explain_parser.rs
View File

@@ -13,37 +13,47 @@
// limitations under the License.
use snafu::ResultExt;
use sqlparser::ast::DescribeAlias;
use sqlparser::keywords::Keyword;
use crate::error::{self, Result};
use crate::parser::ParserContext;
use crate::statements::explain::Explain;
use crate::statements::explain::ExplainStatement;
use crate::statements::statement::Statement;
/// EXPLAIN statement parser implementation
impl ParserContext<'_> {
/// explain that support use our own parser to parse the inner statement
pub(crate) fn parse_explain(&mut self) -> Result<Statement> {
let explain_statement = self
.parser
.parse_explain(DescribeAlias::Explain)
.with_context(|_| error::UnexpectedSnafu {
expected: "a query statement",
actual: self.peek_token_as_string(),
})?;
let analyze = self.parser.parse_keyword(Keyword::ANALYZE);
let verbose = self.parser.parse_keyword(Keyword::VERBOSE);
let format =
if self.parser.parse_keyword(Keyword::FORMAT) {
Some(self.parser.parse_analyze_format().with_context(|_| {
error::UnexpectedSnafu {
expected: "analyze format",
actual: self.peek_token_as_string(),
}
})?)
} else {
None
};
Ok(Statement::Explain(Box::new(Explain::try_from(
explain_statement,
)?)))
let statement = self.parse_statement()?;
let explain = ExplainStatement {
analyze,
verbose,
format,
statement: Box::new(statement),
};
Ok(Statement::Explain(Box::new(explain)))
}
}
#[cfg(test)]
mod tests {
use sqlparser::ast::helpers::attached_token::AttachedToken;
use sqlparser::ast::{
GroupByExpr, Query as SpQuery, SelectFlavor, Statement as SpStatement,
WildcardAdditionalOptions,
};
use sqlparser::ast::{GroupByExpr, Query as SpQuery, SelectFlavor, WildcardAdditionalOptions};
use super::*;
use crate::dialect::GreptimeDbDialect;
@@ -97,31 +107,30 @@ mod tests {
flavor: SelectFlavor::Standard,
};
let sp_statement = SpStatement::Query(Box::new(SpQuery {
with: None,
body: Box::new(sqlparser::ast::SetExpr::Select(Box::new(select))),
order_by: None,
limit: None,
limit_by: vec![],
offset: None,
fetch: None,
locks: vec![],
for_clause: None,
settings: None,
format_clause: None,
}));
let sp_query = Box::new(
SpQuery {
with: None,
body: Box::new(sqlparser::ast::SetExpr::Select(Box::new(select))),
order_by: None,
limit: None,
limit_by: vec![],
offset: None,
fetch: None,
locks: vec![],
for_clause: None,
settings: None,
format_clause: None,
}
.try_into()
.unwrap(),
);
let explain = Explain::try_from(SpStatement::Explain {
describe_alias: DescribeAlias::Explain,
let explain = ExplainStatement {
analyze: false,
verbose: false,
statement: Box::new(sp_statement),
format: None,
query_plan: false,
options: None,
estimate: false,
})
.unwrap();
statement: Box::new(Statement::Query(sp_query)),
};
assert_eq!(stmts[0], Statement::Explain(Box::new(explain)))
}

46
src/sql/src/statements/explain.rs
View File

@@ -15,36 +15,44 @@
use std::fmt::{Display, Formatter};
use serde::Serialize;
use sqlparser::ast::{AnalyzeFormat, Statement as SpStatement};
use sqlparser::ast::AnalyzeFormat;
use sqlparser_derive::{Visit, VisitMut};
use crate::error::Error;
use crate::statements::statement::Statement;
/// Explain statement.
#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
pub struct Explain {
pub inner: SpStatement,
pub struct ExplainStatement {
/// `EXPLAIN ANALYZE ..`
pub analyze: bool,
/// `EXPLAIN .. VERBOSE ..`
pub verbose: bool,
/// `EXPLAIN .. FORMAT `
pub format: Option<AnalyzeFormat>,
/// The statement to analyze. Note this is a Greptime [`Statement`] (not a
/// [`sqlparser::ast::Statement`] so that we can use
/// Greptime specific statements
pub statement: Box<Statement>,
}
impl Explain {
impl ExplainStatement {
pub fn format(&self) -> Option<AnalyzeFormat> {
match self.inner {
SpStatement::Explain { format, .. } => format,
_ => None,
}
self.format
}
}
impl TryFrom<SpStatement> for Explain {
type Error = Error;
fn try_from(value: SpStatement) -> Result<Self, Self::Error> {
Ok(Explain { inner: value })
}
}
impl Display for Explain {
impl Display for ExplainStatement {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.inner)
write!(f, "EXPLAIN")?;
if self.analyze {
write!(f, " ANALYZE")?;
}
if self.verbose {
write!(f, " VERBOSE")?;
}
if let Some(format) = &self.format {
write!(f, " FORMAT {}", format)?;
}
write!(f, " {}", self.statement)
}
}

5
src/sql/src/statements/statement.rs
View File

@@ -30,7 +30,7 @@ use crate::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
use crate::statements::delete::Delete;
use crate::statements::describe::DescribeTable;
use crate::statements::drop::{DropDatabase, DropFlow, DropTable, DropView};
use crate::statements::explain::Explain;
use crate::statements::explain::ExplainStatement;
use crate::statements::insert::Insert;
use crate::statements::kill::Kill;
use crate::statements::query::Query;
@@ -126,7 +126,7 @@ pub enum Statement {
// DESCRIBE TABLE
DescribeTable(DescribeTable),
// EXPLAIN QUERY
Explain(Box<Explain>),
Explain(Box<ExplainStatement>),
// COPY
Copy(Copy),
// Telemetry Query Language
@@ -301,7 +301,6 @@ impl TryFrom<&Statement> for DfStatement {
fn try_from(s: &Statement) -> Result<Self, Self::Error> {
let s = match s {
Statement::Query(query) => SpStatement::Query(Box::new(query.inner.clone())),
Statement::Explain(explain) => explain.inner.clone(),
Statement::Insert(insert) => insert.inner.clone(),
Statement::Delete(delete) => delete.inner.clone(),
_ => {