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

@@ -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<LogicalPlan> {
+    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<String, String> {
+        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<String> {
+        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));
+    }
+}

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<String> {
+        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<String> {
+        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();

src/query/src/promql/planner.rs

@@ -1191,26 +1191,13 @@ impl PromPlanner {
     }
 
     fn drop_metric_name(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
-        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::<Vec<_>>();
-        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<LogicalPlan> {
@@ -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<LogicalPlan> {
+        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::<Vec<_>>();
+        LogicalPlanBuilder::from(plan)
+            .project(projection)
+            .context(DataFusionPlanningSnafu)?
+            .build()
+            .context(DataFusionPlanningSnafu)
+    }
+
     fn attach_metric_name(
         &mut self,
         input: LogicalPlan,

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();