diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5867301f972..5cc8a91e4cf 100644 --- a/src/query/src/promql/planner.rs +++ b/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 { + 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) } fn negate_field_columns(&mut self, input: LogicalPlan) -> Result { @@ -1932,6 +1959,83 @@ impl PromPlanner { } } + fn metric_name_column(schema: &DFSchemaRef) -> Result> { + 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 { + 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::>(); + 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::>(); + 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( diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 19de1098288..1f8c116d265 100644 --- a/src/query/src/promql/planner/test.rs +++ b/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::() + .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::() + .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::() + .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::>() + }) + .collect::>(); + 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::>(); + 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}"); + } + } + } +}