From 9b6f557d5f401a8deaed17376fd3d62b81fa3599 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 29 Sep 2026 14:39:33 +0800 Subject: [PATCH] feat(promql): apply function metric-name lifecycle Preserve identity for sorting and last_over_time, drop it after other function results and histogram folding, and resolve semantic name sort keys exclusively through metadata. Seven focused tests pass on Rust 1.96.1. Remaining aggregation, binary/set and output-boundary activation are separate work. Signed-off-by: Greptime Agent Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/planner.rs | 51 +++- src/query/src/promql/planner/test.rs | 423 ++++++++++++++++++++++++++- 2 files changed, 465 insertions(+), 9 deletions(-) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5cc8a91e4cf..6e337c90849 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1114,6 +1114,21 @@ impl PromPlanner { self.drop_metric_name(plan) } + /// Whether a PromQL function keeps the metric-name identity of its input. The functions + /// absent from the list drop it in [`Self::drop_metric_name`] once their result is built. + fn keeps_metric_name(function_name: &str) -> bool { + matches!( + function_name, + "label_join" + | "label_replace" + | "sort" + | "sort_desc" + | "sort_by_label" + | "sort_by_label_desc" + | "last_over_time" + ) + } + fn drop_metric_name(&mut self, input: LogicalPlan) -> Result { let Some(marker) = Self::metric_name_column(input.schema())? else { return Ok(input); @@ -2344,9 +2359,11 @@ impl PromPlanner { // some special functions that are not expression but a plan match func.name { SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => { - return self + let plan = self .create_histogram_plan(func.name, args, query_engine_state) - .await; + .await?; + // Keep the input identity through folding, then drop it from the result. + return self.drop_metric_name(plan); } SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await, SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await, @@ -2434,6 +2451,7 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)? .filter(self.create_empty_values_filter_expr(true)?) .context(DataFusionPlanningSnafu)?; + let result_schema = builder.schema().clone(); let builder = match func.name { "sort" => builder @@ -2446,6 +2464,7 @@ impl PromPlanner { .sort(Self::create_sort_exprs_by_tags( func.name, args.literals, + &result_schema, true, )?) .context(DataFusionPlanningSnafu)?, @@ -2453,6 +2472,7 @@ impl PromPlanner { .sort(Self::create_sort_exprs_by_tags( func.name, args.literals, + &result_schema, false, )?) .context(DataFusionPlanningSnafu)?, @@ -2487,7 +2507,13 @@ impl PromPlanner { } common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}"); - Ok(plan) + if Self::keeps_metric_name(func.name) { + Ok(plan) + } else { + // Ordinary functions drop the metric-name identity of their input. Dropping it on + // the completed result leaves the identity visible to the function itself. + self.drop_metric_name(plan) + } } async fn prom_ext_expr_to_plan( @@ -4637,6 +4663,7 @@ impl PromPlanner { fn create_sort_exprs_by_tags( func: &str, tags: Vec, + input_schema: &DFSchemaRef, asc: bool, ) -> Result> { ensure!( @@ -4644,10 +4671,26 @@ impl PromPlanner { FunctionInvalidArgumentSnafu { fn_name: func } ); + let metric_name_column = Self::metric_name_column(input_schema)?; tags.iter() .map(|col| match col { DfExpr::Literal(ScalarValue::Utf8(Some(label)), _) => { - Ok(DfExpr::Column(Column::from_name(label)).sort(asc, false)) + // `__name__` is not a physical column: the metric identity lives in the + // marked column of the input. Without one, the label keeps its ordinary + // missing-label behavior. + let column = match (label.as_str(), &metric_name_column) { + (METRIC_NAME, Some(metric_name_column)) => metric_name_column.clone(), + (METRIC_NAME, None) => { + return Err(datafusion::common::field_not_found( + None::, + label, + input_schema.as_ref(), + )) + .context(DataFusionPlanningSnafu); + } + _ => Column::from_name(label), + }; + Ok(DfExpr::Column(column).sort(asc, false)) } other => UnexpectedPlanExprSnafu { desc: format!("expected label string literal, but found {:?}", other), diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 1f8c116d265..06c444d3134 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -23,7 +23,9 @@ 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, PROMQL_FIELD_ROLE_KEY}; +use common_query::prometheus::{ + PROMETHEUS_STALE_NAN_BITS, PROMQL_FIELD_ROLE_KEY, PROMQL_METRIC_NAME_ROLE, +}; use common_query::test_util::DummyDecoder; use common_recordbatch::RecordBatch as GreptimeRecordBatch; use datafusion::arrow::array::{ @@ -36,7 +38,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, provider_as_source}; use datafusion::execution::context::SessionContext; -use datafusion::logical_expr::{Extension, col}; +use datafusion::logical_expr::{Extension, UserDefinedLogicalNode, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql::extension_plan::HistogramFold; @@ -66,6 +68,121 @@ fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { plan.inputs().into_iter().find_map(find_instant_manipulate) } +fn find_histogram_fold(plan: &LogicalPlan) -> Option<&HistogramFold> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(histogram_fold) = node.as_any().downcast_ref::() + { + return Some(histogram_fold); + } + + plan.inputs().into_iter().find_map(find_histogram_fold) +} + +/// The name of the single `Float64` sample column of `plan`, ignoring the metric-name marker. +fn float_sample_column(plan: &LogicalPlan) -> String { + let marker = PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .map(|marker| marker.name); + let mut sample_columns = plan + .schema() + .fields() + .iter() + .filter(|field| { + field.data_type() == &ArrowDataType::Float64 && marker.as_ref() != Some(field.name()) + }) + .map(|field| field.name().clone()); + let sample_column = sample_columns + .next() + .unwrap_or_else(|| panic!("expected a float sample column: {}", plan.display_indent())); + assert!( + sample_columns.next().is_none(), + "expected exactly one float sample column: {}", + plan.display_indent() + ); + sample_column +} + +/// `(series tag, timestamp in milliseconds, sample value)` of every emitted `cv_metric` sample, +/// sorted by series tag. +fn cv_rows(batches: &[RecordBatch], value_column: &str) -> Vec<(String, i64, 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 timestamp = batch + .column_by_name("timestamp") + .expect("no timestamp column") + .as_any() + .downcast_ref::() + .expect("the time index must be a millisecond timestamp column"); + let value = batch + .column_by_name(value_column) + .unwrap_or_else(|| { + panic!( + "no sample value column {value_column} in {}", + batch.schema() + ) + }) + .as_any() + .downcast_ref::() + .expect("the sample value must be a float column"); + (0..batch.num_rows()) + .map(|row| { + ( + tag.value(row).to_string(), + timestamp.value(row), + value.value(row), + ) + }) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.0.cmp(&right.0)); + rows +} + +/// The dropped metric-name identity must be gone from the physical batches as well. +fn assert_metric_name_not_in_batches(batches: &[RecordBatch]) { + for batch in batches { + let schema = batch.schema(); + assert!( + schema.field_with_name(PROMQL_METRIC_NAME_COLUMN).is_err(), + "the metric-name marker must not reach the result batches: {schema:?}" + ); + assert!( + schema + .fields() + .iter() + .all(|field| field.metadata().get(PROMQL_FIELD_ROLE_KEY).is_none()), + "no result batch field may carry the metric-name role: {schema:?}" + ); + } +} + +/// The kept metric-name identity must still be visible in the result batches' Arrow metadata. +fn assert_metric_name_in_batches(batches: &[RecordBatch], marker: &str) { + assert!(!batches.is_empty(), "expected at least one result batch"); + for batch in batches { + let schema = batch.schema(); + let field = schema + .field_with_name(marker) + .unwrap_or_else(|_| panic!("the marker must reach the result batches: {schema:?}")); + assert_eq!( + field + .metadata() + .get(PROMQL_FIELD_ROLE_KEY) + .map(String::as_str), + Some(PROMQL_METRIC_NAME_ROLE), + "the result batch must mark the metric-name field: {schema:?}" + ); + } +} + fn build_query_engine_state() -> QueryEngineState { QueryEngineState::new( new_memory_catalog_manager().unwrap(), @@ -7299,7 +7416,11 @@ async fn test_and_preserves_left_context_when_le_is_missing() { ) .await .unwrap(); - assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); + assert!( + PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .is_none() + ); assert!(!plan.schema().fields().is_empty()); assert!(!contains_histogram_fold(&plan), "{plan:?}"); } @@ -8904,6 +9025,16 @@ async fn eval_times_before_epoch() { /// 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)> { + metric_name_rows_of(batches, marker, greptime_value()) +} + +/// Like [`metric_name_rows`], but reading the sample from `value_column` (range functions name +/// their output column after the function instead of `greptime_value`). +fn metric_name_rows_of( + batches: &[RecordBatch], + marker: &str, + value_column: &str, +) -> Vec<(String, String, f64)> { let mut rows = batches .iter() .flat_map(|batch| { @@ -8920,8 +9051,13 @@ fn metric_name_rows(batches: &[RecordBatch], marker: &str) -> Vec<(String, Strin .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") + .column_by_name(value_column) + .unwrap_or_else(|| { + panic!( + "no sample value column {value_column} in {}", + batch.schema() + ) + }) .as_any() .downcast_ref::() .expect("the sample value must be a float column"); @@ -9075,3 +9211,280 @@ async fn test_unary_sign_controls_metric_name_identity() { } } } + +#[tokio::test] +async fn test_ordinary_functions_drop_metric_name_from_their_result() { + let state = build_query_engine_state(); + for query in ["abs(cv_metric)", "ceil(cv_metric)"] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + assert!( + PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .is_none(), + "{query}: an ordinary function must drop the metric name: {}", + plan.display_indent() + ); + + let sample_column = float_sample_column(&plan); + let (_, batches) = execute(plan, &state).await; + assert_metric_name_not_in_batches(&batches); + assert_eq!( + cv_rows(&batches, &sample_column), + vec![ + ("k0".to_string(), 1_000, 1.0), + ("k1".to_string(), 1_000, 2.0), + ("k2".to_string(), 1_000, 1.0), + ], + "{query}" + ); + } +} + +#[tokio::test] +async fn test_sort_functions_keep_metric_name_identity() { + let state = build_query_engine_state(); + for (query, expected_values) in [ + ("sort(cv_metric)", vec![1.0, 1.0, 2.0]), + ("sort_desc(cv_metric)", vec![2.0, 1.0, 1.0]), + ] { + 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}: sort keeps the metric name")); + assert_eq!(marker.name, PROMQL_METRIC_NAME_COLUMN, "{query}"); + + let sample_column = float_sample_column(&plan); + let (_, batches) = execute(plan, &state).await; + assert_metric_name_in_batches(&batches, &marker.name); + 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}" + ); + assert_eq!(values(&batches, &sample_column), expected_values, "{query}"); + } +} + +#[tokio::test] +async fn test_sort_by_label_resolves_semantic_name_to_the_marker() { + let state = build_query_engine_state(); + // `__name__` is not a physical column: it resolves to the marked identity column, while + // ordinary labels keep their existing lookup. + for (query, sort_column, direction) in [ + ( + r#"sort_by_label(cv_metric, "__name__")"#, + "__promql_metric_name", + "ASC", + ), + (r#"sort_by_label_desc(cv_metric, "k")"#, "k", "DESC"), + ] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let plan_text = plan.display_indent().to_string(); + assert!( + plan_text.contains(&format!("{sort_column} {direction}")), + "{query}: {plan_text}" + ); + let marker = PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .unwrap_or_else(|| panic!("{query}: sort_by_label keeps the metric name")); + assert_eq!(marker.name, PROMQL_METRIC_NAME_COLUMN, "{query}"); + + let (_, batches) = execute(plan, &state).await; + assert_metric_name_in_batches(&batches, &marker.name); + 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_range_functions_control_metric_name_identity() { + let state = build_query_engine_state(); + // The fixture's only samples sit at `ts=1000`; the evaluation point is the same second. + for (query, keeps_name) in [ + ("last_over_time(cv_metric[5m])", true), + ("sum_over_time(cv_metric[5m])", false), + ] { + 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(), + keeps_name, + "{query}: {}", + plan.display_indent() + ); + + let sample_column = float_sample_column(&plan); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + cv_rows(&batches, &sample_column), + vec![ + ("k0".to_string(), 1_000, 1.0), + ("k1".to_string(), 1_000, 2.0), + ("k2".to_string(), 1_000, 1.0), + ], + "{query}" + ); + match marker { + Some(marker) => { + assert_metric_name_in_batches(&batches, &marker.name); + assert_eq!( + metric_name_rows_of(&batches, &marker.name, &sample_column), + 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 => assert_metric_name_not_in_batches(&batches), + } + } +} + +#[tokio::test] +async fn test_histogram_quantile_folds_with_identity_then_drops_it() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + ) + .await + .unwrap(); + + // The completed result drops the metric name ... + assert!( + PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .is_none(), + "{}", + plan.display_indent() + ); + // ... while the fold still groups by the input's identity. + let fold = find_histogram_fold(&plan).expect("the plan must contain a HistogramFold"); + assert!( + PromPlanner::metric_name_column(fold.inputs()[0].schema()) + .unwrap() + .is_some(), + "the fold must see the input metric identity: {}", + plan.display_indent() + ); + + let value_field = float_sample_column(&plan); + let (_, batches) = execute(plan, &state).await; + assert_metric_name_not_in_batches(&batches); + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("tag") + .expect("no series tag column") + .as_any() + .downcast_ref::() + .expect("the series tag must be a string column"); + let values = batch + .column_by_name(&value_field) + .expect("no sample value column") + .as_any() + .downcast_ref::() + .expect("the sample value must be a float column"); + (0..batch.num_rows()) + .map(|row| (tags.value(row).to_string(), values.value(row))) + .collect::>() + }) + .collect::>(); + actual.sort_by(|left, right| left.0.cmp(&right.0)); + assert_eq!( + actual, + vec![("classic".to_string(), 1.0), ("native".to_string(), 0.0)] + ); +} + +#[tokio::test] +async fn test_sort_by_label_semantic_name_without_marker_keeps_missing_label_error() { + let state = build_query_engine_state(); + let provider = || async { + build_test_table_provider_with_distinct_tags(&[("collision_metric", &["__name__", "job"])]) + .await + }; + let plan = PromPlanner::stmt_to_plan( + provider().await, + &operator_eval_stmt("abs(collision_metric)"), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .field_with_unqualified_name("__name__") + .is_ok() + ); + assert!( + PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .is_none() + ); + for function in ["sort_by_label", "sort_by_label_desc"] { + let query = format!(r#"{function}(abs(collision_metric), "__name__")"#); + let error = + PromPlanner::stmt_to_plan(provider().await, &operator_eval_stmt(&query), &state) + .await + .unwrap_err(); + assert!( + matches!(error, crate::promql::error::Error::DataFusionPlanning { .. }), + "{error}" + ); + let query = format!(r#"{function}(collision_metric, "__name__")"#); + let plan = PromPlanner::stmt_to_plan(provider().await, &operator_eval_stmt(&query), &state) + .await + .unwrap(); + let marker = PromPlanner::metric_name_column(plan.schema()) + .unwrap() + .unwrap(); + assert!( + plan.display_indent() + .to_string() + .contains(&format!("Sort: {}", marker.name)) + ); + } +}