query/

planner.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::any::Any;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::ops::ControlFlow;
use std::str::FromStr;
use std::sync::Arc;

use arrow_schema::DataType;
use async_trait::async_trait;
use catalog::table_source::DfTableSourceProvider;
use common_error::ext::BoxedError;
use common_query::promql_annotations::promql_annotation_collector;
use common_telemetry::tracing;
use datafusion::common::{DFSchema, plan_err};
use datafusion::execution::SessionStateBuilder;
use datafusion::execution::context::SessionState;
use datafusion::sql::planner::PlannerContext;
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_common::{ScalarValue, ToDFSchema};
use datafusion_expr::expr::{Exists, InSubquery};
use datafusion_expr::{
    Analyze, Explain, ExplainFormat, Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, PlanType,
    ToStringifiedPlan, col,
};
use datafusion_sql::parser::Statement as DfStatement;
use datafusion_sql::planner::{IdentNormalizer, ParserOptions, SqlToRel};
use log_query::LogQuery;
use promql_parser::parser::EvalStmt;
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;
use sqlparser::ast::{AccessExpr, Value, visit_expressions_mut};

use crate::error::{
    CteColumnSchemaMismatchSnafu, PlanSqlSnafu, QueryPlanSnafu, Result, SqlSnafu,
    UnimplementedSnafu,
};
use crate::log_query::planner::LogQueryPlanner;
use crate::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement};
use crate::promql::planner::PromPlanner;
use crate::query_engine::{DefaultPlanDecoder, QueryEngineState};
use crate::range_select::plan_rewrite::RangePlanRewriter;
use crate::{DfContextProviderAdapter, QueryEngineContext};

#[async_trait]
pub trait LogicalPlanner: Send + Sync {
    async fn plan(&self, stmt: &QueryStatement, query_ctx: QueryContextRef) -> Result<LogicalPlan>;

    async fn plan_logs_query(
        &self,
        query: LogQuery,
        query_ctx: QueryContextRef,
    ) -> Result<LogicalPlan>;

    fn optimize(&self, plan: LogicalPlan) -> Result<LogicalPlan>;

    fn as_any(&self) -> &dyn Any;
}

pub struct DfLogicalPlanner {
    engine_state: Arc<QueryEngineState>,
    session_state: SessionState,
}

impl DfLogicalPlanner {
    pub fn new(engine_state: Arc<QueryEngineState>) -> Self {
        let session_state = engine_state.session_state();
        Self {
            engine_state,
            session_state,
        }
    }

    /// Derive a [`SessionState`] whose [`ExecutionProps`] includes
    /// `query_execution_start_time` if a scheduled time extension is present
    /// in the query context.
    fn derive_session_state_with_scheduled_time(
        &self,
        query_ctx: &QueryContextRef,
    ) -> Result<SessionState> {
        let extensions = query_ctx.extensions();
        match crate::options::parse_scheduled_time_datetime(&extensions)? {
            Some(dt) => {
                let execution_props = self
                    .session_state
                    .execution_props()
                    .clone()
                    .with_query_execution_start_time(dt);
                Ok(
                    SessionStateBuilder::new_from_existing(self.session_state.clone())
                        .with_execution_props(execution_props)
                        .build(),
                )
            }
            None => Ok(self.session_state.clone()),
        }
    }

    /// 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,
                format: ExplainFormat::Indent,
                input: plan,
                schema,
                analyze_level: None,
                analyze_categories: None,
            }))
        } else {
            let stringified_plans = vec![plan.to_stringified(PlanType::InitialLogicalPlan)];

            let options = self.session_state.config().options();
            let format = format
                .map(|x| ExplainFormat::from_str(&x))
                .transpose()?
                .unwrap_or_else(|| options.explain.format.clone());

            Ok(LogicalPlan::Explain(Explain {
                verbose,
                explain_format: format,
                plan,
                stringified_plans,
                schema,
                logical_optimization_succeeded: false,
                show_statistics: None,
            }))
        }
    }

    #[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();
            let mut query = match stmt_owned {
                Statement::Query(query) => query.as_ref().clone(),
                _ => unreachable!("has_hybrid_ctes should only return true for Query statements"),
            };
            self.plan_query_with_hybrid_ctes(&query, query_ctx.clone(), &mut planner_context)
                .await?;

            // remove the processed TQL CTEs from the query
            query.hybrid_cte = None;
            stmt = Cow::Owned(Statement::Query(Box::new(query)));
            is_tql_cte = true;
        }

        let mut df_stmt = stmt.as_ref().try_into().context(SqlSnafu)?;
        normalize_field_access_after_subscript(
            &mut df_stmt,
            self.session_state
                .config_options()
                .sql_parser
                .enable_ident_normalization,
        );

        // TODO(LFC): Remove this when Datafusion supports **both** the syntax and implementation of "explain with format".
        if let datafusion::sql::parser::Statement::Statement(stmt) = &mut df_stmt
            && matches!(
                stmt.as_ref(),
                datafusion::sql::sqlparser::ast::Statement::Explain { .. }
            )
        {
            UnimplementedSnafu {
                operation: "EXPLAIN with FORMAT using raw datafusion planner",
            }
            .fail()?;
        }

        let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
        let table_provider = DfTableSourceProvider::new(
            self.engine_state.catalog_manager().clone(),
            self.engine_state.disallow_cross_catalog_query(),
            query_ctx.clone(),
            Arc::new(DefaultPlanDecoder::new(
                scheduled_state.clone(),
                &query_ctx,
            )?),
            scheduled_state
                .config_options()
                .sql_parser
                .enable_ident_normalization,
        );

        let context_provider = DfContextProviderAdapter::try_new(
            self.engine_state.clone(),
            scheduled_state.clone(),
            Some(&df_stmt),
            query_ctx.clone(),
        )
        .await?;

        let config_options = self.session_state.config().options();
        let parser_options = &config_options.sql_parser;
        let parser_options = ParserOptions {
            map_string_types_to_utf8view: false,
            ..parser_options.into()
        };

        let sql_to_rel = SqlToRel::new_with_options(&context_provider, parser_options);

        // this IF is to handle different version of ASTs
        let result = if is_tql_cte {
            let Statement::Query(query) = stmt.into_owned() else {
                unreachable!("is_tql_cte should only be true for Query statements");
            };
            let sqlparser_stmt = sqlparser::ast::Statement::Query(Box::new(query.inner));
            sql_to_rel
                .sql_statement_to_plan_with_context(sqlparser_stmt, &mut planner_context)
                .context(PlanSqlSnafu)?
        } else {
            sql_to_rel
                .statement_to_plan(df_stmt)
                .context(PlanSqlSnafu)?
        };

        common_telemetry::debug!("Logical planner, statement to plan result: {result}");
        let plan = RangePlanRewriter::new(table_provider, query_ctx.clone())
            .rewrite(result)
            .await?;

        // Optimize logical plan by extension rules
        let context = QueryEngineContext::new(scheduled_state, query_ctx);
        let plan = self
            .engine_state
            .optimize_by_extension_rules(plan, &context)?;
        common_telemetry::debug!("Logical planner, optimize result: {plan}");

        Ok(plan)
    }

    /// Generate a relational expression from a SQL expression
    #[tracing::instrument(skip_all)]
    pub(crate) async fn sql_to_expr(
        &self,
        sql: SqlExpr,
        schema: &DFSchema,
        normalize_ident: bool,
        query_ctx: QueryContextRef,
    ) -> Result<DfExpr> {
        let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
        let context_provider = DfContextProviderAdapter::try_new(
            self.engine_state.clone(),
            scheduled_state,
            None,
            query_ctx,
        )
        .await?;

        let config_options = self.session_state.config().options();
        let parser_options = &config_options.sql_parser;
        let parser_options: ParserOptions = ParserOptions {
            map_string_types_to_utf8view: false,
            enable_ident_normalization: normalize_ident,
            ..parser_options.into()
        };

        let sql_to_rel = SqlToRel::new_with_options(&context_provider, parser_options);

        Ok(sql_to_rel.sql_to_expr(sql, schema, &mut PlannerContext::new())?)
    }

    #[tracing::instrument(skip_all)]
    async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
        let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?;
        let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector);
        if let Some(collector) = &promql_annotations {
            scheduled_state
                .config_mut()
                .options_mut()
                .extensions
                .insert(collector.clone());
        }
        let plan_decoder = Arc::new(DefaultPlanDecoder::new(
            scheduled_state.clone(),
            &query_ctx,
        )?);
        let table_provider = DfTableSourceProvider::new(
            self.engine_state.catalog_manager().clone(),
            self.engine_state.disallow_cross_catalog_query(),
            query_ctx.clone(),
            plan_decoder,
            scheduled_state
                .config_options()
                .sql_parser
                .enable_ident_normalization,
        );
        let plan = PromPlanner::stmt_to_plan_with_annotations(
            table_provider,
            stmt,
            &self.engine_state,
            promql_annotations,
        )
        .await
        .map_err(BoxedError::new)
        .context(QueryPlanSnafu)?;

        let context = QueryEngineContext::new(scheduled_state, query_ctx);
        Ok(self
            .engine_state
            .optimize_by_extension_rules(plan, &context)?)
    }

    #[tracing::instrument(skip_all)]
    fn optimize_logical_plan(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
        Ok(self.engine_state.optimize_logical_plan(plan)?)
    }

    /// Check if a statement contains hybrid CTEs (mix of SQL and TQL)
    fn has_hybrid_ctes(&self, stmt: &Statement) -> bool {
        if let Statement::Query(query) = stmt {
            query
                .hybrid_cte
                .as_ref()
                .map(|hybrid_cte| !hybrid_cte.cte_tables.is_empty())
                .unwrap_or(false)
        } else {
            false
        }
    }

    /// Plan a query with hybrid CTEs using DataFusion's native PlannerContext
    async fn plan_query_with_hybrid_ctes(
        &self,
        query: &Query,
        query_ctx: QueryContextRef,
        planner_context: &mut PlannerContext,
    ) -> Result<()> {
        let hybrid_cte = query.hybrid_cte.as_ref().unwrap();

        for cte in &hybrid_cte.cte_tables {
            match &cte.content {
                CteContent::Tql(tql) => {
                    // Plan TQL and register in PlannerContext
                    let mut logical_plan = self.tql_to_logical_plan(tql, query_ctx.clone()).await?;
                    if !cte.columns.is_empty() {
                        let schema = logical_plan.schema();
                        let schema_fields = schema.fields().to_vec();
                        ensure!(
                            schema_fields.len() == cte.columns.len(),
                            CteColumnSchemaMismatchSnafu {
                                cte_name: cte.name.value.clone(),
                                original: schema_fields
                                    .iter()
                                    .map(|field| field.name().clone())
                                    .collect::<Vec<_>>(),
                                expected: cte
                                    .columns
                                    .iter()
                                    .map(|column| column.to_string())
                                    .collect::<Vec<_>>(),
                            }
                        );
                        let aliases = cte
                            .columns
                            .iter()
                            .zip(schema_fields.iter())
                            .map(|(column, field)| col(field.name()).alias(column.to_string()));
                        logical_plan = LogicalPlanBuilder::from(logical_plan)
                            .project(aliases)
                            .context(PlanSqlSnafu)?
                            .build()
                            .context(PlanSqlSnafu)?;
                    }

                    // Wrap in SubqueryAlias to ensure proper table qualification for CTE
                    logical_plan = LogicalPlan::SubqueryAlias(
                        datafusion_expr::SubqueryAlias::try_new(
                            Arc::new(logical_plan),
                            cte.name.value.clone(),
                        )
                        .context(PlanSqlSnafu)?,
                    );

                    planner_context.insert_cte(&cte.name.value, logical_plan);
                }
                CteContent::Sql(_) => {
                    // SQL CTEs should have been moved to the main query's WITH clause
                    // during parsing, so we shouldn't encounter them here
                    unreachable!("SQL CTEs should not be in hybrid_cte.cte_tables");
                }
            }
        }

        Ok(())
    }

    /// Convert TQL to LogicalPlan directly
    async fn tql_to_logical_plan(
        &self,
        tql: &Tql,
        query_ctx: QueryContextRef,
    ) -> Result<LogicalPlan> {
        match tql {
            Tql::Eval(eval) => {
                // Convert TqlEval to PromQuery then to QueryStatement::Promql
                let prom_query = PromQuery {
                    query: eval.query.clone(),
                    start: eval.start.clone(),
                    end: eval.end.clone(),
                    step: eval.step.clone(),
                    lookback: eval
                        .lookback
                        .clone()
                        .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
                    alias: eval.alias.clone(),
                };
                let stmt = QueryLanguageParser::parse_promql(&prom_query, &query_ctx)?;

                self.plan(&stmt, query_ctx).await
            }
            Tql::Explain(_) => UnimplementedSnafu {
                operation: "TQL EXPLAIN in CTEs",
            }
            .fail(),
            Tql::Analyze(_) => UnimplementedSnafu {
                operation: "TQL ANALYZE in CTEs",
            }
            .fail(),
        }
    }

    /// Extracts cast types for all placeholders in a logical plan.
    /// Returns a map where each placeholder ID is mapped to:
    /// - Some(DataType) if the placeholder is cast to a specific type
    /// - None if the placeholder exists but has no cast
    ///
    /// Example: `$1::TEXT` returns `{"$1": Some(DataType::Utf8)}`
    ///
    /// This function walks through all expressions in the logical plan,
    /// including subqueries, to identify placeholders and their cast types.
    fn extract_placeholder_cast_types(
        plan: &LogicalPlan,
    ) -> Result<HashMap<String, Option<DataType>>> {
        let mut placeholder_types = HashMap::new();
        let mut casted_placeholders = HashSet::new();

        Self::extract_from_plan(plan, &mut placeholder_types, &mut casted_placeholders)?;

        Ok(placeholder_types)
    }

    fn extract_from_plan(
        plan: &LogicalPlan,
        placeholder_types: &mut HashMap<String, Option<DataType>>,
        casted_placeholders: &mut HashSet<String>,
    ) -> Result<()> {
        plan.apply(|node| {
            for expr in node.expressions() {
                let _ = expr.apply(|e| {
                    // Handle casted placeholders
                    if let DfExpr::Cast(cast) = e
                        && let DfExpr::Placeholder(ph) = &*cast.expr
                    {
                        placeholder_types
                            .insert(ph.id.clone(), Some(cast.field.data_type().clone()));
                        casted_placeholders.insert(ph.id.clone());
                    }

                    // Handle arrow_cast(Placeholder, 'type_string') generated by SQL rewriter
                    if let DfExpr::ScalarFunction(scalar_func) = e
                        && scalar_func.name() == "arrow_cast"
                        && scalar_func.args.len() == 2
                        && let DfExpr::Placeholder(ph) = &scalar_func.args[0]
                        && let DfExpr::Literal(ScalarValue::Utf8(Some(type_str)), _) =
                            &scalar_func.args[1]
                        && let Ok(data_type) = type_str.parse::<DataType>()
                    {
                        placeholder_types.insert(ph.id.clone(), Some(data_type));
                        casted_placeholders.insert(ph.id.clone());
                    }

                    // Handle bare (non-casted) placeholders
                    if let DfExpr::Placeholder(ph) = e
                        && !casted_placeholders.contains(&ph.id)
                        && !placeholder_types.contains_key(&ph.id)
                    {
                        placeholder_types.insert(ph.id.clone(), None);
                    }

                    // Recurse into subquery plans embedded in expressions
                    match e {
                        DfExpr::Exists(Exists { subquery, .. })
                        | DfExpr::InSubquery(InSubquery { subquery, .. })
                        | DfExpr::ScalarSubquery(subquery) => {
                            Self::extract_from_plan(
                                &subquery.subquery,
                                placeholder_types,
                                casted_placeholders,
                            )?;
                        }
                        _ => {}
                    }

                    Ok(TreeNodeRecursion::Continue)
                });
            }
            Ok(TreeNodeRecursion::Continue)
        })?;
        Ok(())
    }

    fn infer_limit_placeholder_types(
        plan: &LogicalPlan,
        placeholder_types: &mut HashMap<String, Option<DataType>>,
    ) -> Result<()> {
        plan.apply(|node| {
            if let LogicalPlan::Limit(limit) = node {
                for expr in limit.skip.iter().chain(limit.fetch.iter()) {
                    expr.apply(|e| {
                        if let DfExpr::Placeholder(ph) = e {
                            placeholder_types
                                .entry(ph.id.clone())
                                .and_modify(|existing| {
                                    if existing.is_none() {
                                        *existing = Some(DataType::Int64);
                                    }
                                })
                                .or_insert(Some(DataType::Int64));
                        }

                        Ok(TreeNodeRecursion::Continue)
                    })?;
                }
            }

            Ok(TreeNodeRecursion::Continue)
        })?;

        Ok(())
    }

    /// Gets inferred parameter types from a logical plan.
    /// Returns a map where each parameter ID is mapped to:
    /// - Some(DataType) if the parameter type could be inferred
    /// - None if the parameter type could not be inferred
    ///
    /// This function first uses DataFusion's `get_parameter_types()` to infer types.
    /// If any parameters have `None` values (i.e., DataFusion couldn't infer their types),
    /// it falls back to using `extract_placeholder_cast_types()` to detect explicit casts
    /// and applies context-specific inference such as LIMIT/OFFSET placeholders.
    ///
    /// This is because datafusion can only infer types for a limited cases.
    ///
    /// Example: For query `WHERE $1::TEXT AND $2`, DataFusion may not infer `$2`'s type,
    /// but this function will return `{"$1": Some(DataType::Utf8), "$2": None}`.
    pub fn get_inferred_parameter_types(
        plan: &LogicalPlan,
    ) -> Result<HashMap<String, Option<DataType>>> {
        let mut param_types = plan.get_parameter_types().context(PlanSqlSnafu)?;

        let has_none = param_types.values().any(|v| v.is_none());

        if has_none {
            let cast_types = Self::extract_placeholder_cast_types(plan)?;

            for (id, opt_type) in cast_types {
                param_types
                    .entry(id)
                    .and_modify(|existing| {
                        if existing.is_none() {
                            *existing = opt_type.clone();
                        }
                    })
                    .or_insert(opt_type);
            }

            Self::infer_limit_placeholder_types(plan, &mut param_types)?;
        }

        Ok(param_types)
    }
}

/// Normalizes dot field accesses that follow a subscript for DataFusion.
///
/// sqlparser represents `j.o.l[1].inner.l[2]` as a compound field access with
/// the following access chain:
///
/// ```text
/// Dot(Identifier("o")),
/// Dot(Identifier("l")),
/// Subscript(1),
/// Dot(Identifier("inner")),
/// Dot(Identifier("l")),
/// Subscript(2)
/// ```
///
/// DataFusion first resolves the leading `j.o.l` through
/// `JsonExprPlanner::plan_compound_identifier`, which produces an untyped
/// `json_get` with path `$.o.l`. Before invoking `JsonExprPlanner::plan_field_access`,
/// however, DataFusion eagerly converts every remaining access into a
/// `GetFieldAccess`. It accepts string values but not [`SqlExpr::Identifier`]s
/// in [`AccessExpr::Dot`] after a subscript. Without this normalization, that
/// conversion fails at `.inner`, and `plan_field_access` is never called, even
/// for the preceding `[1]`.
///
/// This function converts dot identifiers after the first subscript into
/// `Dot(Value(SingleQuotedString(...)))`, applying DataFusion's identifier
/// normalization before discarding whether each identifier was quoted. It
/// changes neither the SQL text nor the dot accesses into subscript nodes: the
/// resulting AST is conceptually `j.o.l[1].'inner'.'l'[2]`. DataFusion converts
/// the string-valued dot accesses into named field accesses, which
/// `plan_field_access` safely encodes as dot members, quoting names when needed. It can then extend the
/// JSON path to `$.o.l[1].inner.l[2]`.
///
/// This behavior is unchanged in the latest upstream releases checked here:
/// DataFusion 55.0.0 and sqlparser 0.62.0.
///
/// TODO(LFC): Remove this workaround after upstream supports dot identifiers after subscripts.
fn normalize_field_access_after_subscript(stmt: &mut DfStatement, normalize_ident: bool) {
    let DfStatement::Statement(stmt) = stmt else {
        return;
    };
    let normalizer = IdentNormalizer::new(normalize_ident);

    let _ = visit_expressions_mut(stmt.as_mut(), |expr| {
        let SqlExpr::CompoundFieldAccess { access_chain, .. } = expr else {
            return ControlFlow::<()>::Continue(());
        };
        let Some(index) = access_chain
            .iter()
            .position(|x| matches!(x, AccessExpr::Subscript(_)))
        else {
            return ControlFlow::Continue(());
        };

        for access in &mut access_chain[index + 1..] {
            let AccessExpr::Dot(SqlExpr::Identifier(ident)) = access else {
                continue;
            };
            let value = normalizer.normalize(ident.clone());
            *access = AccessExpr::Dot(SqlExpr::Value(
                Value::SingleQuotedString(value).with_span(ident.span),
            ));
        }
        ControlFlow::Continue(())
    });
}

#[async_trait]
impl LogicalPlanner for DfLogicalPlanner {
    #[tracing::instrument(skip_all)]
    async fn plan(&self, stmt: &QueryStatement, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
        match stmt {
            QueryStatement::Sql(stmt) => self.plan_sql(stmt, query_ctx).await,
            QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx).await,
        }
    }

    async fn plan_logs_query(
        &self,
        query: LogQuery,
        query_ctx: QueryContextRef,
    ) -> Result<LogicalPlan> {
        let plan_decoder = Arc::new(DefaultPlanDecoder::new(
            self.session_state.clone(),
            &query_ctx,
        )?);
        let table_provider = DfTableSourceProvider::new(
            self.engine_state.catalog_manager().clone(),
            self.engine_state.disallow_cross_catalog_query(),
            query_ctx,
            plan_decoder,
            self.session_state
                .config_options()
                .sql_parser
                .enable_ident_normalization,
        );

        let mut planner = LogQueryPlanner::new(table_provider, self.session_state.clone());
        planner
            .query_to_plan(query)
            .await
            .map_err(BoxedError::new)
            .context(QueryPlanSnafu)
    }

    fn optimize(&self, plan: LogicalPlan) -> Result<LogicalPlan> {
        self.optimize_logical_plan(plan)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_schema::DataType;
    use catalog::RegisterTableRequest;
    use catalog::memory::MemoryCatalogManager;
    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
    use common_time::Timezone;
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::{ColumnSchema, Schema};
    use session::context::{QueryContext, QueryContextBuilder};
    use store_api::metric_engine_consts::{
        DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
        METRIC_ENGINE_NAME,
    };
    use table::metadata::{TableInfoBuilder, TableMetaBuilder};
    use table::test_util::EmptyTable;

    use super::*;
    use crate::parser::{PromQuery, QueryLanguageParser};
    use crate::{QueryEngineFactory, QueryEngineRef};

    async fn create_test_engine() -> QueryEngineRef {
        let columns = vec![
            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
            ColumnSchema::new("name", ConcreteDataType::string_datatype(), true),
        ];
        let schema = Arc::new(Schema::new(columns));
        let table_meta = TableMetaBuilder::empty()
            .schema(schema)
            .primary_key_indices(vec![0])
            .value_indices(vec![1])
            .next_column_id(1024)
            .build()
            .unwrap();
        let table_info = TableInfoBuilder::new("test", table_meta).build().unwrap();
        let table = EmptyTable::from_table_info(&table_info);

        crate::tests::new_query_engine_with_table(table)
    }

    async fn create_timestamp_test_engine() -> QueryEngineRef {
        let columns = vec![
            ColumnSchema::new(
                "ts",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
            ColumnSchema::new(
                "st",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            ),
            ColumnSchema::new("note", ConcreteDataType::string_datatype(), true),
            ColumnSchema::new(
                "ts_ns",
                ConcreteDataType::timestamp_nanosecond_datatype(),
                true,
            ),
        ];
        let schema = Arc::new(Schema::new(columns));
        let table_meta = TableMetaBuilder::empty()
            .schema(schema)
            .primary_key_indices(vec![])
            .value_indices(vec![0, 1, 2, 3])
            .next_column_id(1024)
            .build()
            .unwrap();
        let table_info = TableInfoBuilder::new("timestamps", table_meta)
            .build()
            .unwrap();
        let table = EmptyTable::from_table_info(&table_info);

        crate::tests::new_query_engine_with_table(table)
    }

    fn create_promql_test_engine() -> QueryEngineRef {
        let catalog_manager = MemoryCatalogManager::with_default_setup();
        let physical_table_name = "phy";
        let physical_table_id = 999u32;

        let physical_schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
                ConcreteDataType::uint32_datatype(),
                false,
            ),
            ColumnSchema::new(
                DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
                ConcreteDataType::uint64_datatype(),
                false,
            ),
            ColumnSchema::new("tag_0", ConcreteDataType::string_datatype(), false),
            ColumnSchema::new("tag_1", ConcreteDataType::string_datatype(), false),
            ColumnSchema::new(
                "timestamp",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
            ColumnSchema::new("field_0", ConcreteDataType::float64_datatype(), true),
        ]));
        let physical_meta = TableMetaBuilder::empty()
            .schema(physical_schema)
            .primary_key_indices(vec![0, 1, 2, 3])
            .value_indices(vec![4, 5])
            .engine(METRIC_ENGINE_NAME.to_string())
            .next_column_id(1024)
            .build()
            .unwrap();
        let physical_info = TableInfoBuilder::default()
            .table_id(physical_table_id)
            .name(physical_table_name)
            .meta(physical_meta)
            .build()
            .unwrap();
        catalog_manager
            .register_table_sync(RegisterTableRequest {
                catalog: DEFAULT_CATALOG_NAME.to_string(),
                schema: DEFAULT_SCHEMA_NAME.to_string(),
                table_name: physical_table_name.to_string(),
                table_id: physical_table_id,
                table: EmptyTable::from_table_info(&physical_info),
            })
            .unwrap();

        let mut options = table::requests::TableOptions::default();
        options.extra_options.insert(
            LOGICAL_TABLE_METADATA_KEY.to_string(),
            physical_table_name.to_string(),
        );
        let logical_schema = Arc::new(Schema::new(vec![
            ColumnSchema::new("tag_0", ConcreteDataType::string_datatype(), false),
            ColumnSchema::new("tag_1", ConcreteDataType::string_datatype(), false),
            ColumnSchema::new(
                "timestamp",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
            ColumnSchema::new("field_0", ConcreteDataType::float64_datatype(), true),
        ]));
        let logical_meta = TableMetaBuilder::empty()
            .schema(logical_schema)
            .primary_key_indices(vec![0, 1])
            .value_indices(vec![3])
            .engine(METRIC_ENGINE_NAME.to_string())
            .options(options)
            .next_column_id(1024)
            .build()
            .unwrap();
        let logical_info = TableInfoBuilder::default()
            .table_id(1024)
            .name("some_metric")
            .meta(logical_meta)
            .build()
            .unwrap();
        catalog_manager
            .register_table_sync(RegisterTableRequest {
                catalog: DEFAULT_CATALOG_NAME.to_string(),
                schema: DEFAULT_SCHEMA_NAME.to_string(),
                table_name: "some_metric".to_string(),
                table_id: 1024,
                table: EmptyTable::from_table_info(&logical_info),
            })
            .unwrap();

        QueryEngineFactory::new(
            catalog_manager,
            None,
            None,
            None,
            None,
            false,
            crate::options::QueryOptions::default(),
        )
        .query_engine()
    }

    async fn parse_sql_to_plan(sql: &str) -> LogicalPlan {
        let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap();
        let engine = create_test_engine().await;
        engine
            .planner()
            .plan(&stmt, QueryContext::arc())
            .await
            .unwrap()
    }

    async fn parse_promql_to_plan(query: &str) -> LogicalPlan {
        let engine = create_promql_test_engine();
        let query_ctx = QueryContext::arc();
        let stmt = QueryLanguageParser::parse_promql(
            &PromQuery {
                query: query.to_string(),
                start: "0".to_string(),
                end: "10".to_string(),
                step: "5s".to_string(),
                lookback: "300s".to_string(),
                alias: None,
            },
            &query_ctx,
        )
        .unwrap();

        engine.planner().plan(&stmt, query_ctx).await.unwrap()
    }

    /// Plans `sql` and runs the DataFusion analyzer, which is where
    /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these
    /// assertions would not see the assignment rewrite at all.
    async fn analyze_insert(
        engine: &QueryEngineRef,
        sql: &str,
        query_ctx: &QueryContextRef,
    ) -> String {
        let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap();
        let plan = engine
            .planner()
            .plan(&stmt, query_ctx.clone())
            .await
            .unwrap();
        let context = engine.engine_context(query_ctx.clone());
        let state = context.state();
        state
            .analyzer()
            .execute_and_check(plan, state.config_options(), |_, _| {})
            .unwrap()
            .display_indent_schema()
            .to_string()
    }

    #[tokio::test]
    async fn test_insert_timestamp_literals_use_query_timezone() {
        let query_ctx = Arc::new(
            QueryContextBuilder::default()
                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
                .build(),
        );
        let engine = create_timestamp_test_engine().await;

        for (sql, expected_timestamps) in [
            (
                "INSERT INTO timestamps (ts, st) \
                 VALUES ('2026-08-02 12:00:00.001', now())",
                &[1_785_643_200_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT '2026-08-03 12:00:00.001', now()",
                &[1_785_729_600_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT '2026-08-04 12:00:00.001', now() LIMIT 1",
                &[1_785_816_000_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT * FROM (\
                     SELECT '2026-08-05 12:00:00.001', now()\
                 ) AS source",
                &[1_785_902_400_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT '2026-08-18 12:00:00.001', max(st) \
                 FROM timestamps GROUP BY note",
                &[1_787_025_600_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT c, s FROM (\
                     SELECT '2026-08-13 12:00:00.001' AS c, now() AS s\
                 ) AS t WHERE c > '2026-01-01'",
                &[1_786_593_600_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT c, s FROM (\
                     SELECT '2026-08-14 12:00:00.001' AS c, now() AS s\
                 ) AS t ORDER BY c",
                &[1_786_680_000_001_i64][..],
            ),
            (
                "INSERT INTO timestamps (ts, st) \
                 SELECT DISTINCT c, s FROM (\
                     SELECT '2026-08-15 12:00:00.001' AS c, now() AS s\
                 ) AS t",
                &[1_786_766_400_001_i64][..],
            ),
        ] {
            let plan = analyze_insert(&engine, sql, &query_ctx).await;

            for expected_timestamp in expected_timestamps {
                assert!(
                    plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")),
                    "{plan}"
                );
            }
        }
    }

    #[tokio::test]
    async fn test_insert_explicit_timestamp_cast_keeps_datafusion_semantics() {
        let query_ctx = Arc::new(
            QueryContextBuilder::default()
                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
                .build(),
        );
        let engine = create_timestamp_test_engine().await;
        let sql = "INSERT INTO timestamps (ts, st) \
                   VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())";
        let plan = analyze_insert(&engine, sql, &query_ctx).await;

        // An explicit cast reaches the analyzer as an `arrow_cast` call rather
        // than an `Expr::Cast`, which is how it stays out of the rewrite.
        assert!(
            plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"),
            "{plan}"
        );
        // 12:00:00.001 read as Shanghai local time; the source query keeps UTC.
        assert!(
            !plan.contains("TimestampMillisecond(1786104000001, None)"),
            "{plan}"
        );
    }

    #[tokio::test]
    async fn test_insert_converts_source_literal_shared_by_several_columns() {
        let query_ctx = Arc::new(
            QueryContextBuilder::default()
                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
                .build(),
        );
        let engine = create_timestamp_test_engine().await;

        // Rewriting `c` in place would also retype `note` and truncate `ts_ns`.
        for (sql, expected) in [
            (
                "INSERT INTO timestamps (ts, note) SELECT a, b FROM (\
                     SELECT c AS a, c AS b FROM (\
                         SELECT '2026-08-12 12:00:00.001' AS c\
                     ) AS t1\
                 ) AS t2",
                &[
                    "TimestampMillisecond(1786507200001, None)",
                    "[a:Utf8, b:Utf8]",
                ][..],
            ),
            (
                "INSERT INTO timestamps (ts, ts_ns) SELECT a, b FROM (\
                     SELECT c AS a, c AS b FROM (\
                         SELECT '2026-08-12 12:00:00.123456789' AS c\
                     ) AS t1\
                 ) AS t2",
                &[
                    "TimestampMillisecond(1786507200123, None)",
                    "TimestampNanosecond(1786507200123456789, None)",
                ][..],
            ),
        ] {
            let plan = analyze_insert(&engine, sql, &query_ctx).await;

            for expected in expected {
                assert!(plan.contains(expected), "{plan}");
            }
        }
    }

    #[tokio::test]
    async fn test_insert_union_converts_via_assignment_cast() {
        let query_ctx = Arc::new(
            QueryContextBuilder::default()
                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
                .build(),
        );
        let engine = create_timestamp_test_engine().await;
        // Branches disagree, so the conversion stays as a cast on the
        // assignment instead of folding. One cast covers every branch, which is
        // why a NULL branch no longer cancels the conversion for the column and
        // why UNION's dedup keys stay on the original strings.
        for sql in [
            "INSERT INTO timestamps (ts, st) \
             SELECT '2026-08-06 12:00:00.001', now() \
             UNION ALL \
             SELECT '2026-08-07 12:00:00.001', NULL",
            "INSERT INTO timestamps (ts, st) \
             SELECT '2026-08-16 12:00:00.001', now() \
             UNION \
             SELECT '2026-08-17 12:00:00.001', now()",
        ] {
            let plan = analyze_insert(&engine, sql, &query_ctx).await;

            assert!(
                plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"),
                "{plan}"
            );
            // The branches themselves are untouched.
            assert!(
                plan.contains("Utf8(\"2026-08-07 12:00:00.001\")")
                    || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"),
                "{plan}"
            );
        }
    }

    #[tokio::test]
    async fn test_insert_mixed_union_keeps_source_coercion() {
        let query_ctx = Arc::new(
            QueryContextBuilder::default()
                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
                .build(),
        );
        let engine = create_timestamp_test_engine().await;
        let sql = "INSERT INTO timestamps (ts, st) \
                   SELECT '2026-08-10 12:00:00.001', now() \
                   UNION ALL \
                   SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()";
        let plan = analyze_insert(&engine, sql, &query_ctx).await;

        assert!(
            !plan.contains("TimestampMillisecond(1786334400001, None)"),
            "{plan}"
        );
        assert!(
            plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"),
            "{plan}"
        );
        // TypeCoercion has already settled this union to timestamp, so the
        // assignment has nothing left to reinterpret. Retargeting the cast here
        // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind,
        // which shifts the value instead of relabelling it.
        assert!(!plan.contains("Asia/Shanghai"), "{plan}");
    }

    #[tokio::test]
    async fn test_extract_placeholder_cast_types_multiple() {
        let plan = parse_sql_to_plan(
            "SELECT $1::INT, $2::TEXT, $3, $4::INTEGER FROM test WHERE $5::FLOAT > 0",
        )
        .await;
        let types = DfLogicalPlanner::extract_placeholder_cast_types(&plan).unwrap();

        assert_eq!(types.len(), 5);
        assert_eq!(types.get("$1"), Some(&Some(DataType::Int32)));
        assert_eq!(types.get("$2"), Some(&Some(DataType::Utf8)));
        assert_eq!(types.get("$3"), Some(&None));
        assert_eq!(types.get("$4"), Some(&Some(DataType::Int32)));
        assert_eq!(types.get("$5"), Some(&Some(DataType::Float32)));
    }

    #[tokio::test]
    async fn test_get_inferred_parameter_types_fallback_for_udf_args() {
        // datafusion is not able to infer type for scalar function arguments
        let plan = parse_sql_to_plan(
            "SELECT parse_ident($1), parse_ident($2::TEXT) FROM test WHERE id > $3",
        )
        .await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.len(), 3);

        let type_1 = types.get("$1").unwrap();
        let type_2 = types.get("$2").unwrap();
        let type_3 = types.get("$3").unwrap();

        assert!(type_1.is_none(), "Expected $1 to be None");
        assert_eq!(type_2, &Some(DataType::Utf8));
        assert_eq!(type_3, &Some(DataType::Int32));
    }

    #[tokio::test]
    async fn test_get_inferred_parameter_types_limit_offset() {
        let plan = parse_sql_to_plan("SELECT id FROM test LIMIT $1 OFFSET $2").await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.get("$1"), Some(&Some(DataType::Int64)));
        assert_eq!(types.get("$2"), Some(&Some(DataType::Int64)));
    }

    #[tokio::test]
    async fn test_plan_pql_applies_extension_rules() {
        for inner_agg in ["count", "sum", "avg", "min", "max", "stddev", "stdvar"] {
            let plan = parse_promql_to_plan(&format!(
                "sum(irate(some_metric[1h])) / scalar(count({inner_agg}(some_metric) by (tag_0)))"
            ))
            .await;
            let plan_str = plan.display_indent_schema().to_string();
            assert!(plan_str.contains("Distinct:"), "{inner_agg}: {plan_str}");
        }
    }

    #[tokio::test]
    async fn test_plan_pql_filters_null_only_groups_for_non_count_inner_aggs() {
        let count_plan = parse_promql_to_plan("scalar(count(count(some_metric) by (tag_0)))").await;
        let count_plan_str = count_plan.display_indent_schema().to_string();
        assert!(
            !count_plan_str.contains("field_0 IS NOT NULL"),
            "{count_plan_str}"
        );

        for inner_agg in ["sum", "avg", "min", "max", "stddev", "stdvar"] {
            let plan = parse_promql_to_plan(&format!(
                "scalar(count({inner_agg}(some_metric) by (tag_0)))"
            ))
            .await;
            let plan_str = plan.display_indent_schema().to_string();
            assert!(
                plan_str.contains("field_0 IS NOT NULL"),
                "{inner_agg}: {plan_str}"
            );
        }
    }

    #[tokio::test]
    async fn test_plan_pql_skips_extension_rules_for_non_direct_or_unsupported_inner_agg() {
        for query in [
            "sum(irate(some_metric[1h])) / scalar(count(sum(irate(some_metric[1h])) by (tag_0)))",
            "sum(irate(some_metric[1h])) / scalar(count(group(some_metric) by (tag_0)))",
        ] {
            let plan = parse_promql_to_plan(query).await;
            let plan_str = plan.display_indent_schema().to_string();
            assert!(!plan_str.contains("Distinct:"), "{query}: {plan_str}");
        }
    }

    #[tokio::test]
    async fn test_plan_sql_does_not_apply_nested_count_rule() {
        let plan = parse_sql_to_plan(
            "SELECT id, count(inner_count) \
             FROM ( \
                 SELECT id, count(name) AS inner_count \
                 FROM test \
                 GROUP BY id \
                 ORDER BY id \
                 LIMIT 1000000 \
             ) t \
             GROUP BY id \
             ORDER BY id",
        )
        .await;

        let plan_str = plan.display_indent_schema().to_string();
        assert!(!plan_str.contains("Distinct:"), "{plan_str}");
    }

    #[tokio::test]
    async fn test_get_inferred_parameter_types_subquery() {
        let plan = parse_sql_to_plan(
            r#"SELECT * FROM test WHERE id = (SELECT id FROM test CROSS JOIN (SELECT parse_ident($1::TEXT) AS parts) p LIMIT 1)"#,
        ).await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.len(), 1);
        let type_1 = types.get("$1").unwrap();
        assert_eq!(type_1, &Some(DataType::Utf8));
    }

    #[tokio::test]
    async fn test_get_inferred_parameter_types_insert() {
        let plan = parse_sql_to_plan("INSERT INTO test (id, name) VALUES ($1, $2), ($3, $4)").await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.len(), 4);
        assert_eq!(types.get("$1"), Some(&Some(DataType::Int32)));
        assert_eq!(types.get("$2"), Some(&Some(DataType::Utf8)));
        assert_eq!(types.get("$3"), Some(&Some(DataType::Int32)));
        assert_eq!(types.get("$4"), Some(&Some(DataType::Utf8)));
    }

    #[tokio::test]
    async fn test_get_inferred_parameter_types_arrow_cast() {
        let plan = parse_sql_to_plan("SELECT $1::INT64, $2::FLOAT64, $3::INT16, $4::INT32, $5::UINT8, $6::UINT16, $7::UINT32").await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.get("$1"), Some(&Some(DataType::Int64)));
        assert_eq!(types.get("$2"), Some(&Some(DataType::Float64)));
        assert_eq!(types.get("$3"), Some(&Some(DataType::Int16)));
        assert_eq!(types.get("$4"), Some(&Some(DataType::Int32)));
        assert_eq!(types.get("$5"), Some(&Some(DataType::UInt8)));
        assert_eq!(types.get("$6"), Some(&Some(DataType::UInt16)));
        assert_eq!(types.get("$7"), Some(&Some(DataType::UInt32)));

        let plan = parse_sql_to_plan("SELECT $1::INT8, $2::FLOAT8, $3::INT2, $4::INT8").await;
        let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap();

        assert_eq!(types.get("$1"), Some(&Some(DataType::Int64)));
        assert_eq!(types.get("$2"), Some(&Some(DataType::Float64)));
        assert_eq!(types.get("$3"), Some(&Some(DataType::Int16)));
    }
}