From a660c892e998f117b86af85dddbf195f56b422e5 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:46:45 +0800 Subject: [PATCH] feat(promql): hide metric identity in TQL tabular results Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/operator/src/statement/tql.rs | 205 ++++++++++++++++++++++++++- src/query/src/planner.rs | 107 +++++++++++++- src/query/src/promql/planner.rs | 55 ++++--- src/query/src/promql/planner/test.rs | 39 +++++ 4 files changed, 381 insertions(+), 25 deletions(-) diff --git a/src/operator/src/statement/tql.rs b/src/operator/src/statement/tql.rs index fe9baaf4c35..a5a7e3cd4f4 100644 --- a/src/operator/src/statement/tql.rs +++ b/src/operator/src/statement/tql.rs @@ -13,19 +13,24 @@ // limitations under the License. use std::collections::HashMap; +use std::sync::Arc; +use common_error::ext::BoxedError; use common_query::Output; use common_telemetry::tracing; -use datafusion_expr::LogicalPlan; +use datafusion_expr::{Analyze, Explain, LogicalPlan, PlanType, ToStringifiedPlan}; use query::parser::{ ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, DEFAULT_LOOKBACK_STRING, EXPLAIN_NODE_NAME, EXPLAIN_VERBOSE_NODE_NAME, PromQuery, QueryLanguageParser, }; +use query::promql::planner::PromPlanner; use session::context::QueryContextRef; use snafu::ResultExt; use sql::statements::tql::Tql; -use crate::error::{ExecLogicalPlanSnafu, ParseQuerySnafu, PlanStatementSnafu, Result}; +use crate::error::{ + ExecLogicalPlanSnafu, ExternalSnafu, ParseQuerySnafu, PlanStatementSnafu, Result, +}; use crate::statement::StatementExecutor; impl StatementExecutor { @@ -101,11 +106,13 @@ impl StatementExecutor { .context(ParseQuerySnafu)? } }; - self.query_engine + let plan = self + .query_engine .planner() .plan(&stmt, query_ctx.clone()) .await - .context(PlanStatementSnafu) + .context(PlanStatementSnafu)?; + strip_metric_name_output(plan) } /// Execute the given [Tql] query and return the result. @@ -118,3 +125,193 @@ impl StatementExecutor { .context(ExecLogicalPlanSnafu) } } + +/// Removes the PromQL metric-name identity from a TQL plan's result schema. +/// +/// The identity is a PromQL-internal column kept in the plan for the Prometheus HTTP API; a TQL +/// statement is tabular, so its result must not expose it. +/// +/// `EXPLAIN`/`ANALYZE` wrap their input and produce the two fixed planning columns, so the +/// identity is stripped from the wrapped child and the wrapper is rebuilt around it. Rebuilding +/// `EXPLAIN` also regenerates its initial stringified plan so the reported text starts from the +/// stripped child. Stripping adds a projection on top of the child; the identity still shows in +/// the plan text of the nodes below that projection, which is the plan as it was built. +fn strip_metric_name_output(plan: LogicalPlan) -> Result { + let strip = |plan| { + PromPlanner::strip_metric_name_column(plan) + .map_err(BoxedError::new) + .context(ExternalSnafu) + }; + + match plan { + LogicalPlan::Explain(explain) => { + let child = strip(Arc::unwrap_or_clone(explain.plan))?; + let stringified_plans = vec![child.to_stringified(PlanType::InitialLogicalPlan)]; + Ok(LogicalPlan::Explain(Explain { + plan: Arc::new(child), + stringified_plans, + ..explain + })) + } + LogicalPlan::Analyze(analyze) => { + let input = strip(Arc::unwrap_or_clone(analyze.input))?; + Ok(LogicalPlan::Analyze(Analyze { + input: Arc::new(input), + ..analyze + })) + } + plan => strip(plan), + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use common_query::prometheus::{PROMQL_FIELD_ROLE_KEY, PROMQL_METRIC_NAME_ROLE}; + use datafusion::logical_expr::expr::Alias; + use datafusion_expr::{ExplainFormat, Expr as DfExpr, LogicalPlanBuilder, col, lit}; + + use super::*; + + const MARKER: &str = "__promql_metric_name"; + + /// Mirrors the projection `PromPlanner` attaches to a selector result: the identity column + /// is an ordinary `Utf8` column carrying the role metadata. + fn marked_metadata() -> HashMap { + HashMap::from([( + PROMQL_FIELD_ROLE_KEY.to_string(), + PROMQL_METRIC_NAME_ROLE.to_string(), + )]) + } + + fn marker_expr(name: &str) -> DfExpr { + DfExpr::Alias(Alias { + expr: Box::new(lit("metric")), + relation: None, + name: name.to_string(), + metadata: Some(marked_metadata().into()), + }) + } + + fn marked_result_plan() -> LogicalPlan { + let input = LogicalPlanBuilder::empty(false).build().unwrap(); + LogicalPlanBuilder::from(input) + .project(vec![ + lit("0").alias("ts"), + lit(1.0_f64).alias("val"), + marker_expr(MARKER), + ]) + .unwrap() + .build() + .unwrap() + } + + fn schema_column_names(plan: &LogicalPlan) -> Vec { + plan.schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() + } + + fn has_marked_column(plan: &LogicalPlan) -> bool { + plan.schema().fields().iter().any(|field| { + field + .metadata() + .get(PROMQL_FIELD_ROLE_KEY) + .map(String::as_str) + == Some(PROMQL_METRIC_NAME_ROLE) + }) + } + + #[test] + fn strip_metric_name_output_drops_the_identity_from_the_result_schema() { + let plan = strip_metric_name_output(marked_result_plan()).unwrap(); + + assert_eq!(schema_column_names(&plan), vec!["ts", "val"]); + assert!(!has_marked_column(&plan)); + assert!(plan.schema().field_with_name(None, MARKER).is_err()); + } + + #[test] + fn strip_metric_name_output_keeps_an_unmarked_column_with_the_same_name() { + let input = LogicalPlanBuilder::from(marked_result_plan()) + .project(vec![ + col("ts"), + col("val"), + lit("physical").alias(MARKER), + marker_expr(&format!("{MARKER}_")), + ]) + .unwrap() + .build() + .unwrap(); + + let plan = strip_metric_name_output(input).unwrap(); + + // Only the marked column goes; the physically named one stays. + assert_eq!(schema_column_names(&plan), vec!["ts", "val", MARKER]); + assert!(!has_marked_column(&plan)); + } + + #[test] + fn strip_metric_name_output_leaves_an_unmarked_plan_unchanged() { + let input = LogicalPlanBuilder::empty(false).build().unwrap(); + let plan = LogicalPlanBuilder::from(input) + .project(vec![lit("0").alias("ts"), lit(1.0_f64).alias("val")]) + .unwrap() + .build() + .unwrap(); + + assert_eq!(strip_metric_name_output(plan.clone()).unwrap(), plan); + } + + #[test] + fn strip_metric_name_output_rebuilds_explain_around_the_stripped_child() { + let explain = LogicalPlanBuilder::from(marked_result_plan()) + .explain(true, false) + .unwrap() + .build() + .unwrap(); + + let plan = strip_metric_name_output(explain).unwrap(); + + let LogicalPlan::Explain(explain) = &plan else { + panic!("expected an EXPLAIN wrapper: {}", plan.display_indent()); + }; + assert!(explain.verbose); + assert_eq!(explain.explain_format, ExplainFormat::Indent); + assert_eq!(schema_column_names(&explain.plan), vec!["ts", "val"]); + assert!(!has_marked_column(&explain.plan)); + // The initial stringified plan is regenerated from the stripped child: its own + // projection no longer lists the identity, while the child below it still shows the + // identity column it was built with. + let stringified = &explain.stringified_plans; + assert_eq!(stringified.len(), 1); + assert_eq!(stringified[0].plan_type, PlanType::InitialLogicalPlan); + assert_eq!( + stringified[0].plan.lines().next().unwrap(), + "Projection: ts, val" + ); + assert!(stringified[0].plan.contains("AS __promql_metric_name")); + } + + #[test] + fn strip_metric_name_output_rebuilds_analyze_around_the_stripped_input() { + let analyze = LogicalPlanBuilder::from(marked_result_plan()) + .explain(false, true) + .unwrap() + .build() + .unwrap(); + + let plan = strip_metric_name_output(analyze).unwrap(); + + let LogicalPlan::Analyze(analyze) = &plan else { + panic!("expected an ANALYZE wrapper: {}", plan.display_indent()); + }; + assert!(!analyze.verbose); + assert_eq!(analyze.format, ExplainFormat::Indent); + assert_eq!(schema_column_names(&analyze.input), vec!["ts", "val"]); + assert!(!has_marked_column(&analyze.input)); + } +} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index b8501e1a6db..55cbee44f18 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -460,7 +460,13 @@ impl DfLogicalPlanner { }; let stmt = QueryLanguageParser::parse_promql(&prom_query, &query_ctx)?; - self.plan(&stmt, query_ctx).await + let plan = self.plan(&stmt, query_ctx).await?; + // A TQL statement is tabular, so its result must not expose the PromQL-internal + // metric-name identity. Stripping here, before the CTE column arity check and + // alias projection below, keeps both operating on the tabular schema. + PromPlanner::strip_metric_name_column(plan) + .map_err(BoxedError::new) + .context(QueryPlanSnafu) } Tql::Explain(_) => UnimplementedSnafu { operation: "TQL EXPLAIN in CTEs", @@ -746,6 +752,7 @@ mod tests { use catalog::RegisterTableRequest; use catalog::memory::MemoryCatalogManager; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_query::prometheus::{PROMQL_FIELD_ROLE_KEY, PROMQL_METRIC_NAME_ROLE}; use common_time::Timezone; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; @@ -919,6 +926,31 @@ mod tests { .query_engine() } + /// A two-column table (`ts` time index, `val` field), the shape a TQL CTE aliases. + fn create_promql_cte_test_engine() -> QueryEngineRef { + let columns = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![]) + .value_indices(vec![1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::new("metric", table_meta).build().unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + crate::tests::new_query_engine_with_table(table) + } + async fn parse_sql_to_plan(sql: &str) -> LogicalPlan { let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); let engine = create_test_engine().await; @@ -929,6 +961,79 @@ mod tests { .unwrap() } + fn plan_column_names(plan: &LogicalPlan) -> Vec { + plan.schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() + } + + /// The name of the field carrying the PromQL metric-name identity, if any. + fn metric_name_role_column(plan: &LogicalPlan) -> Option { + plan.schema() + .fields() + .iter() + .find(|field| { + field + .metadata() + .get(PROMQL_FIELD_ROLE_KEY) + .map(String::as_str) + == Some(PROMQL_METRIC_NAME_ROLE) + }) + .map(|field| field.name().clone()) + } + + async fn parse_tql_cte_to_plan(engine: &QueryEngineRef, sql: &str) -> LogicalPlan { + let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); + engine + .planner() + .plan(&stmt, QueryContext::arc()) + .await + .unwrap() + } + + #[tokio::test] + async fn test_tql_cte_select_star_keeps_tabular_shape_without_metric_name() { + let engine = create_promql_cte_test_engine(); + let plan = parse_tql_cte_to_plan( + &engine, + "WITH tql AS (TQL EVAL (0, 10, '5s') metric) SELECT * FROM tql", + ) + .await; + + assert_eq!(plan_column_names(&plan), vec!["ts", "val"]); + assert_eq!(metric_name_role_column(&plan), None); + } + + #[tokio::test] + async fn test_tql_cte_column_aliases_cover_the_stripped_result() { + let engine = create_promql_cte_test_engine(); + let plan = parse_tql_cte_to_plan( + &engine, + "WITH tql (the_timestamp, the_value) AS (TQL EVAL (0, 10, '5s') metric) \ + SELECT * FROM tql", + ) + .await; + + // The identity column is stripped before the CTE arity check and the alias + // projection, so the aliases line up with the tabular columns. + assert_eq!(plan_column_names(&plan), vec!["the_timestamp", "the_value"]); + assert_eq!(metric_name_role_column(&plan), None); + } + + #[tokio::test] + async fn test_plan_pql_keeps_metric_name_identity_outside_tql() { + // Plain PromQL still carries the identity for the Prometheus HTTP API; only the + // tabular TQL boundaries strip it. + let plan = parse_promql_to_plan("some_metric").await; + + assert_eq!( + metric_name_role_column(&plan), + Some("__promql_metric_name".to_string()) + ); + } + async fn parse_promql_to_plan(query: &str) -> LogicalPlan { let engine = create_promql_test_engine(); let query_ctx = QueryContext::arc(); diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index cf576fe53e1..d7860c7e910 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1191,26 +1191,13 @@ impl PromPlanner { } fn drop_metric_name(&mut self, input: LogicalPlan) -> Result { - let Some(marker) = Self::metric_name_column(input.schema())? else { - return Ok(input); - }; - self.ctx.tag_columns.retain(|name| name != &marker.name); - self.ctx - .aggregation_field_labels - .retain(|name| name != &marker.name); - let projection = input - .schema() - .iter() - .filter_map(|(qualifier, field)| { - (field.name() != &marker.name) - .then(|| DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))) - }) - .collect::>(); - LogicalPlanBuilder::from(input) - .project(projection) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) + if let Some(marker) = Self::metric_name_column(input.schema())? { + self.ctx.tag_columns.retain(|name| name != &marker.name); + self.ctx + .aggregation_field_labels + .retain(|name| name != &marker.name); + } + Self::strip_metric_name_column(input) } fn negate_field_columns(&mut self, input: LogicalPlan) -> Result { @@ -2103,6 +2090,34 @@ impl PromPlanner { Ok(found) } + /// Strips the metric-name identity from `plan`'s output, if it carries one. + /// + /// The identity is an ordinary column distinguished only by its field metadata; this projects + /// every other column with its qualifier retained and returns `plan` unchanged when no column + /// is marked. Column names are never special-cased, so a physical column that happens to be + /// spelled like the marker is kept. + /// + /// PromQL results keep the identity for the Prometheus HTTP API; tabular boundaries (TQL + /// statements, TQL CTEs) use this to keep it out of user-visible result schemas. + pub fn strip_metric_name_column(plan: LogicalPlan) -> Result { + let Some(marker) = Self::metric_name_column(plan.schema())? else { + return Ok(plan); + }; + let projection = plan + .schema() + .iter() + .filter_map(|(qualifier, field)| { + (field.name() != &marker.name) + .then(|| DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))) + }) + .collect::>(); + LogicalPlanBuilder::from(plan) + .project(projection) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + fn attach_metric_name( &mut self, input: LogicalPlan, diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 12fd6bdfde8..3fb6edf929a 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -9270,6 +9270,45 @@ async fn test_collision_metric_name_label_is_preserved() { ); } +#[tokio::test] +async fn test_dropping_identity_keeps_physical_column_with_the_same_name() { + // Dropping the identity must remove only the marked column, never the physical label + // that happens to be spelled the same way. + let plan = PromPlanner::stmt_to_plan( + build_test_table_provider_with_distinct_tags(&[( + "collision_metric", + &[PROMQL_METRIC_NAME_COLUMN, "job"], + )]) + .await, + // Unary minus drops the metric name of its operand. + &operator_eval_stmt("-collision_metric"), + &build_query_engine_state(), + ) + .await + .unwrap(); + + let schema = plan.schema(); + assert!( + PromPlanner::metric_name_column(schema).unwrap().is_none(), + "the identity must be dropped: {}", + plan.display_indent() + ); + assert!( + schema + .field_with_name(None, PROMQL_METRIC_NAME_COLUMN) + .is_ok(), + "the physical label must survive the drop: {}", + plan.display_indent() + ); + assert!( + schema + .field_with_name(None, &format!("{PROMQL_METRIC_NAME_COLUMN}_")) + .is_err(), + "the marker column must be gone: {}", + plan.display_indent() + ); +} + #[tokio::test] async fn test_unary_sign_controls_metric_name_identity() { let state = build_query_engine_state();