feat(promql): checkpoint internal selector metric identity

Attach field-marked name columns at selector lowering and remove identity on unary minus. Raw selector execution and collision tests passed; unary fixture lookup corrected but rerun pending. Intermediate stack checkpoint: remaining operator lifecycle and tabular output boundaries must be completed before activation is merge-ready.

Signed-off-by: Greptime Agent <agent@localhost>
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/query/src/promql/planner.rs

@@ -33,6 +33,7 @@ use common_query::prelude::{
     GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_native_histogram,
     greptime_value,
 };
+use common_query::prometheus::{PROMQL_FIELD_ROLE_KEY, PROMQL_METRIC_NAME_ROLE};
 use common_query::promql_annotations::PromqlAnnotationCollector;
 use datafusion::common::DFSchemaRef;
 use datafusion::datasource::DefaultTableSource;
@@ -141,6 +142,8 @@ const DEFAULT_TIME_INDEX_COLUMN: &str = "time";
 
 /// default value column name for empty metric
 const DEFAULT_FIELD_COLUMN: &str = "value";
+/// Base name of the private column carrying a selector's metric name.
+const PROMQL_METRIC_NAME_COLUMN: &str = "__promql_metric_name";
 
 /// Special modifier to project field columns under multi-field mode
 const FIELD_COLUMN_MATCHER: &str = "__field__";
@@ -1107,7 +1110,31 @@ impl PromPlanner {
         let UnaryExpr { expr } = unary_expr;
         // Unary Expr in PromQL implys the `-` operator
         let input = self.prom_expr_to_plan(expr, query_engine_state).await?;
-        self.negate_field_columns(input)
+        let plan = self.negate_field_columns(input)?;
+        self.drop_metric_name(plan)
+    }
+
+    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)
     }
 
     fn negate_field_columns(&mut self, input: LogicalPlan) -> Result<LogicalPlan> {
@@ -1932,6 +1959,83 @@ impl PromPlanner {
         }
     }
 
+    fn metric_name_column(schema: &DFSchemaRef) -> Result<Option<Column>> {
+        let mut found = None;
+        for (qualifier, field) in schema.iter() {
+            if field
+                .metadata()
+                .get(PROMQL_FIELD_ROLE_KEY)
+                .map(String::as_str)
+                == Some(PROMQL_METRIC_NAME_ROLE)
+            {
+                ensure!(
+                    field.data_type() == &ArrowDataType::Utf8,
+                    UnexpectedPlanExprSnafu {
+                        desc: "PromQL metric-name column must be Utf8"
+                    }
+                );
+                ensure!(
+                    found.is_none(),
+                    UnexpectedPlanExprSnafu {
+                        desc: "multiple PromQL metric-name columns"
+                    }
+                );
+                found = Some(Column::new(qualifier.cloned(), field.name().clone()));
+            }
+        }
+        Ok(found)
+    }
+
+    fn attach_metric_name(
+        &mut self,
+        input: LogicalPlan,
+        metric_name: Option<&str>,
+    ) -> Result<LogicalPlan> {
+        let Some(metric_name) = metric_name else {
+            return Ok(input);
+        };
+        ensure!(
+            Self::metric_name_column(input.schema())?.is_none(),
+            UnexpectedPlanExprSnafu {
+                desc: "metric-name identity already attached"
+            }
+        );
+        let occupied = input
+            .schema()
+            .fields()
+            .iter()
+            .map(|field| field.name().as_str())
+            .collect::<HashSet<_>>();
+        let mut name = PROMQL_METRIC_NAME_COLUMN.to_string();
+        while occupied.contains(name.as_str()) {
+            name.push('_');
+        }
+        let mut exprs = input
+            .schema()
+            .iter()
+            .map(|(qualifier, field)| {
+                DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone()))
+            })
+            .collect::<Vec<_>>();
+        let mut metadata = std::collections::HashMap::new();
+        metadata.insert(
+            PROMQL_FIELD_ROLE_KEY.to_string(),
+            PROMQL_METRIC_NAME_ROLE.to_string(),
+        );
+        exprs.push(DfExpr::Alias(Alias {
+            expr: Box::new(lit(metric_name.to_string())),
+            relation: None,
+            name: name.clone(),
+            metadata: Some(metadata.into()),
+        }));
+        self.ctx.tag_columns.push(name);
+        LogicalPlanBuilder::from(input)
+            .project(exprs)
+            .context(DataFusionPlanningSnafu)?
+            .build()
+            .context(DataFusionPlanningSnafu)
+    }
+
     async fn prom_vector_selector_to_plan(
         &mut self,
         vector_selector: &VectorSelector,
@@ -1944,8 +2048,9 @@ impl PromPlanner {
             at,
         } = vector_selector;
         let matchers = self.preprocess_label_matchers(matchers, name)?;
+        let metric_name = self.ctx.table_name.clone();
         if let Some(empty_plan) = self.setup_context().await? {
-            return Ok(empty_plan);
+            return self.attach_metric_name(empty_plan, metric_name.as_deref());
         }
         let offset_ms = Self::offset_millis(offset);
         // `@` anchors the sample window at a fixed timestamp: the selector selects its samples
@@ -2058,7 +2163,7 @@ impl PromPlanner {
         if let Some(timestamp_value_column) = timestamp_value_column {
             self.create_timestamp_func_plan(manipulate, ident(timestamp_value_column))
         } else {
-            Ok(manipulate)
+            self.attach_metric_name(manipulate, metric_name.as_deref())
         }
     }
 
@@ -2128,6 +2233,7 @@ impl PromPlanner {
             at,
         } = vs;
         let matchers = self.preprocess_label_matchers(matchers, name)?;
+        let metric_name = self.ctx.table_name.clone();
         ensure!(!range.is_zero(), ZeroRangeSelectorSnafu);
         let range_ms = range.as_millis() as _;
         self.ctx.range = Some(range_ms);
@@ -2226,7 +2332,7 @@ impl PromPlanner {
             }
         };
 
-        Ok(manipulate)
+        self.attach_metric_name(manipulate, metric_name.as_deref())
     }
 
     async fn prom_call_expr_to_plan(

src/query/src/promql/planner/test.rs

@@ -23,7 +23,7 @@ use common_query::native_histogram::{
     CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array,
 };
 use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value};
-use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS;
+use common_query::prometheus::{PROMETHEUS_STALE_NAN_BITS, PROMQL_FIELD_ROLE_KEY};
 use common_query::test_util::DummyDecoder;
 use common_recordbatch::RecordBatch as GreptimeRecordBatch;
 use datafusion::arrow::array::{
@@ -8900,3 +8900,178 @@ async fn eval_times_before_epoch() {
         assert!(plan(too_early, too_early).await.is_err());
     }
 }
+
+/// Collects `(series tag, metric-name marker, sample value)` triples of every row in `batches`,
+/// one row per emitted sample of the `cv_metric` fixture, sorted by series tag.
+fn metric_name_rows(batches: &[RecordBatch], marker: &str) -> Vec<(String, String, f64)> {
+    let mut rows = batches
+        .iter()
+        .flat_map(|batch| {
+            let tag = batch
+                .column_by_name("k")
+                .expect("no series tag column")
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .expect("the series tag must be a string column");
+            let name = batch
+                .column_by_name(marker)
+                .expect("no metric-name marker column")
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .expect("the metric-name marker must be a string column");
+            let value = batch
+                .column_by_name(greptime_value())
+                .expect("no sample value column")
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .expect("the sample value must be a float column");
+            (0..batch.num_rows())
+                .map(|row| {
+                    (
+                        tag.value(row).to_string(),
+                        name.value(row).to_string(),
+                        value.value(row),
+                    )
+                })
+                .collect::<Vec<_>>()
+        })
+        .collect::<Vec<_>>();
+    rows.sort_by(|left, right| left.0.cmp(&right.0));
+    rows
+}
+
+#[tokio::test]
+async fn test_raw_selector_attaches_metric_name_identity() {
+    let state = build_query_engine_state();
+    for query in [
+        "cv_metric",
+        r#"cv_metric{k=~"k.*"}"#,
+        r#"{__name__="cv_metric"}"#,
+    ] {
+        let plan = PromPlanner::stmt_to_plan(
+            build_count_values_table_provider().await,
+            &operator_eval_stmt(query),
+            &state,
+        )
+        .await
+        .unwrap();
+
+        let marker = PromPlanner::metric_name_column(plan.schema())
+            .unwrap()
+            .unwrap_or_else(|| panic!("{query}: the selector must attach its metric name"));
+        assert_eq!(marker.name, PROMQL_METRIC_NAME_COLUMN, "{query}");
+
+        let (_, batches) = execute(plan, &state).await;
+        assert_eq!(
+            metric_name_rows(&batches, &marker.name),
+            vec![
+                ("k0".to_string(), "cv_metric".to_string(), 1.0),
+                ("k1".to_string(), "cv_metric".to_string(), 2.0),
+                ("k2".to_string(), "cv_metric".to_string(), 1.0),
+            ],
+            "{query}"
+        );
+    }
+}
+
+#[tokio::test]
+async fn test_collision_metric_name_label_is_preserved() {
+    let plan = PromPlanner::stmt_to_plan(
+        build_test_table_provider_with_distinct_tags(&[(
+            "collision_metric",
+            &[PROMQL_METRIC_NAME_COLUMN, "job"],
+        )])
+        .await,
+        &operator_eval_stmt("collision_metric"),
+        &build_query_engine_state(),
+    )
+    .await
+    .unwrap();
+
+    // The label physically named like the marker keeps its name and metadata; the marker
+    // moves to a suffixed column instead of overwriting it.
+    let physical = plan
+        .schema()
+        .fields()
+        .iter()
+        .find(|field| field.name() == PROMQL_METRIC_NAME_COLUMN)
+        .expect("the physical label must be preserved");
+    assert!(
+        physical.metadata().get(PROMQL_FIELD_ROLE_KEY).is_none(),
+        "the physical label must not carry the marker role"
+    );
+    let marker = PromPlanner::metric_name_column(plan.schema())
+        .unwrap()
+        .expect("the selector must attach its metric name");
+    assert_eq!(
+        marker.name,
+        format!("{PROMQL_METRIC_NAME_COLUMN}_"),
+        "{}",
+        plan.display_indent()
+    );
+}
+
+#[tokio::test]
+async fn test_unary_sign_controls_metric_name_identity() {
+    let state = build_query_engine_state();
+    // The parser returns the operand of unary plus unchanged and wraps unary minus in a
+    // negation, which drops the metric name.
+    for (sign, negated) in [("-", true), ("+", false)] {
+        let query = format!("{sign}cv_metric");
+        let plan = PromPlanner::stmt_to_plan(
+            build_count_values_table_provider().await,
+            &operator_eval_stmt(&query),
+            &state,
+        )
+        .await
+        .unwrap();
+
+        let marker = PromPlanner::metric_name_column(plan.schema()).unwrap();
+        assert_eq!(
+            marker.is_some(),
+            !negated,
+            "{query}: {}",
+            plan.display_indent()
+        );
+
+        // The negated sample column is aliased after the negation expression and no longer
+        // carries the source field name, so identify it by type. The fixture has a single
+        // `Float64` column; the tag, time index and marker are not `Float64`.
+        let marker_name = marker.as_ref().map(|marker| marker.name.clone());
+        let sample_columns = plan
+            .schema()
+            .fields()
+            .iter()
+            .filter(|field| {
+                field.data_type() == &ArrowDataType::Float64
+                    && marker_name.as_ref() != Some(field.name())
+            })
+            .map(|field| field.name().clone())
+            .collect::<Vec<_>>();
+        assert_eq!(
+            sample_columns.len(),
+            1,
+            "{query}: expected exactly one float sample column: {}",
+            plan.display_indent()
+        );
+        let sample_column = sample_columns.into_iter().next().unwrap();
+
+        let (_, batches) = execute(plan, &state).await;
+        match marker {
+            Some(marker) => assert_eq!(
+                metric_name_rows(&batches, &marker.name),
+                vec![
+                    ("k0".to_string(), "cv_metric".to_string(), 1.0),
+                    ("k1".to_string(), "cv_metric".to_string(), 2.0),
+                    ("k2".to_string(), "cv_metric".to_string(), 1.0),
+                ],
+                "{query}"
+            ),
+            None => {
+                let mut values = values(&batches, &sample_column);
+                values.sort_by(f64::total_cmp);
+                assert_eq!(values, vec![-2.0, -1.0, -1.0], "{query}");
+            }
+        }
+    }
+}