From 994a4c564fef932ca8ff90c85547cd199eb2bbf6 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:55:36 +0000 Subject: [PATCH] fix: support multi-measure per-entity schemas safely Rebased onto current main from PR #412 (fix/per-entity-avg-rewrite). Co-Authored-By: Claude Opus 5.5 --- .../asap-aware-mapping/src/analytical_cost.rs | 2 + .../src/query_physical_lowering.rs | 39 ++++- crates/asap-aware-mapping/src/rewrite.rs | 29 +++- .../tests/avg_over_time_rewrite.rs | 28 ++++ .../fixtures/avg_over_time_overflow.test.yml | 34 +++++ crates/types/src/pre_asap/query_expr.rs | 143 +++++++++++++++++- crates/types/src/pre_asap/resolve.rs | 60 ++++++++ 7 files changed, 324 insertions(+), 11 deletions(-) create mode 100644 crates/integration-tests/tests/avg_over_time_rewrite.rs create mode 100644 crates/integration-tests/tests/fixtures/avg_over_time_overflow.test.yml diff --git a/crates/asap-aware-mapping/src/analytical_cost.rs b/crates/asap-aware-mapping/src/analytical_cost.rs index ed2a316db..bdbfb795e 100644 --- a/crates/asap-aware-mapping/src/analytical_cost.rs +++ b/crates/asap-aware-mapping/src/analytical_cost.rs @@ -2193,6 +2193,8 @@ pub enum AnalyticalCostError { UnsupportedCandidate, #[error("query operator has no physical implementation in the analytical model")] UnsupportedQueryOperator, + #[error("multi-measure per-entity aggregates have no physical implementation; lower each measure separately")] + UnsupportedMultiMeasurePerEntity, #[error("inconsistent operator statistics: {0}")] InconsistentOperatorStatistics(&'static str), #[error("summary operation {0} has no lifecycle-aware cost formula")] diff --git a/crates/asap-aware-mapping/src/query_physical_lowering.rs b/crates/asap-aware-mapping/src/query_physical_lowering.rs index aa934e8c9..ffe51394a 100644 --- a/crates/asap-aware-mapping/src/query_physical_lowering.rs +++ b/crates/asap-aware-mapping/src/query_physical_lowering.rs @@ -296,7 +296,7 @@ pub fn lower_query_physical_dag( } if matches!(reduction, asap_types::pre_asap::Reduction::PerEntity) { if measures.len() != 1 { - return Err(AnalyticalCostError::UnsupportedQueryOperator); + return Err(AnalyticalCostError::UnsupportedMultiMeasurePerEntity); } let accumulator_count = u64::try_from(measures.len()) .map_err(|_| AnalyticalCostError::Overflow)?; @@ -2545,6 +2545,43 @@ mod tests { )); } + /// Multi-measure schemas must not silently lower to a single accumulator. + #[test] + fn multi_measure_per_entity_lowering_is_explicitly_unsupported() { + use asap_types::pre_asap::{ + AggIntent, Column, DataType, QueryExpr, Reduction, Schema, Source, + }; + let source = Source::TimeSeries { + metric: "requests".into(), + }; + let root = Rc::new(QueryExpr::Aggregate { + reduction: Reduction::PerEntity, + measures: vec![ + AggIntent::Sum { col: None }, + AggIntent::Count { + accuracy: asap_types::types::AccuracyTarget::Exact, + }, + ], + output_names: vec![], + filters: vec![], + having: None, + child: Rc::new(QueryExpr::Scan { + source: source.clone(), + predicates: vec![], + schema: Schema::new(vec![Column::new("value", DataType::Float64, false)]), + }), + }); + let provided = HashMap::new(); + assert!(matches!( + lower_query_physical_dag( + &root, + &scope(vec![coverage(source, vec![])]), + &scripted(&provided) + ), + Err(AnalyticalCostError::UnsupportedMultiMeasurePerEntity) + )); + } + #[test] fn promql_relabel_sample_and_per_series_lower_as_a_complete_chain() { use asap_types::pre_asap::{ diff --git a/crates/asap-aware-mapping/src/rewrite.rs b/crates/asap-aware-mapping/src/rewrite.rs index 06d1724fd..9ea84f433 100644 --- a/crates/asap-aware-mapping/src/rewrite.rs +++ b/crates/asap-aware-mapping/src/rewrite.rs @@ -545,19 +545,44 @@ mod tests { assert!(AvgToSumOverCountStrategy.replacements(&target).is_empty()); } + /// Matching schemas do not make range SUM/COUNT safe for unbounded samples. #[test] - fn does_not_match_a_per_entity_avg_aggregate() { + fn per_entity_avg_rewrite_is_rejected_without_arithmetic_proof() { let q = Rc::new(QueryExpr::Aggregate { reduction: Reduction::PerEntity, measures: vec![AggIntent::Avg { col: None }], output_names: vec![], filters: vec![], having: None, - child: Rc::new(metric_scan(&[])), + child: Rc::new(QueryExpr::TimeRange { + range: Duration::from_secs(300), + child: Rc::new(metric_scan(&["job", "instance"])), + }), }); let target = TargetSubDAG::new(&q); assert!(!AvgToSumOverCountStrategy.matches(&target)); assert!(AvgToSumOverCountStrategy.replacements(&target).is_empty()); + assert!(build_rewrite(&q).is_none()); + } + + /// COUNT(*) cannot replace the denominator of a nullable sample average. + #[test] + fn per_entity_nullable_or_non_sample_average_is_not_rewritten() { + for (nullable, column) in [(true, None), (false, Some(2)), (false, Some(99))] { + let mut scan = metric_scan(&["job"]); + if let QueryExpr::Scan { schema, .. } = &mut scan { + schema.columns[1].nullable = nullable; + } + let root = Rc::new(QueryExpr::Aggregate { + reduction: Reduction::PerEntity, + measures: vec![AggIntent::Avg { col: column }], + output_names: vec![], + filters: vec![], + having: None, + child: Rc::new(scan), + }); + assert!(build_rewrite(&root).is_none()); + } } #[test] diff --git a/crates/integration-tests/tests/avg_over_time_rewrite.rs b/crates/integration-tests/tests/avg_over_time_rewrite.rs new file mode 100644 index 000000000..9832e0832 --- /dev/null +++ b/crates/integration-tests/tests/avg_over_time_rewrite.rs @@ -0,0 +1,28 @@ +use std::rc::Rc; + +use asap_aware_mapping::replacement::{ReplacementStrategy, TargetSubDAG}; +use asap_aware_mapping::rewrite::AvgToSumOverCountStrategy; +use asap_integration_tests::fixtures::lower_promql; +use asap_types::types::AccuracyTarget; + +/// Unbounded Float64 ranges must retain AVG: finite samples can overflow SUM. +/// The independent Prometheus oracle is fixtures/avg_over_time_overflow.test.yml. +#[test] +fn range_average_does_not_offer_unconditional_sum_count() { + for query in [ + "avg_over_time(latency[5m])", + "avg_over_time(latency{job=\"api\"}[5m])", + "avg_over_time(latency[5m:1m])", + ] { + let root = Rc::new(lower_promql(query, AccuracyTarget::Exact).unwrap()); + let target = TargetSubDAG::new(&root); + assert!( + !AvgToSumOverCountStrategy.matches(&target), + "unbounded average must not match: {query}" + ); + assert!( + AvgToSumOverCountStrategy.replacements(&target).is_empty(), + "unbounded average must not produce a sum/count rewrite: {query}" + ); + } +} diff --git a/crates/integration-tests/tests/fixtures/avg_over_time_overflow.test.yml b/crates/integration-tests/tests/fixtures/avg_over_time_overflow.test.yml new file mode 100644 index 000000000..48918d740 --- /dev/null +++ b/crates/integration-tests/tests/fixtures/avg_over_time_overflow.test.yml @@ -0,0 +1,34 @@ +# Official Prometheus 3.5.0 oracle: AVG remains finite when SUM overflows. +# Run: promtool test rules crates/integration-tests/tests/fixtures/avg_over_time_overflow.test.yml +evaluation_interval: 1m +tests: +- interval: 1m + input_series: + - series: 'latency{job="api"}' + values: '1e308 1e308' + promql_expr_test: + - expr: 'avg_over_time(latency[5m])' + eval_time: 1m + exp_samples: + - labels: '{job="api"}' + value: 1e308 + - expr: 'sum_over_time(latency[5m]) / count_over_time(latency[5m])' + eval_time: 1m + exp_samples: + - labels: '{job="api"}' + value: .inf +- interval: 1m + input_series: + - series: 'latency{job="api"}' + values: '-1e308 -1e308' + promql_expr_test: + - expr: 'avg_over_time(latency[5m])' + eval_time: 1m + exp_samples: + - labels: '{job="api"}' + value: -1e308 + - expr: 'sum_over_time(latency[5m]) / count_over_time(latency[5m])' + eval_time: 1m + exp_samples: + - labels: '{job="api"}' + value: -.inf diff --git a/crates/types/src/pre_asap/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index 9f0e800c1..0da7421f7 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -52,6 +52,8 @@ impl ColState for ColumnRef { /// Errors from schema derivation over a canonical tree. #[derive(Debug, Error)] pub enum QueryExprError { + #[error("invalid per-entity aggregate: {0}")] + InvalidPerEntityAggregate(String), #[error("invalid scalar function signature: {0}")] InvalidScalarSignature(String), #[error("by-column id {0} out of range (input has {1} columns)")] @@ -1540,8 +1542,14 @@ fn per_series_reduction_schema(input: &Schema, agg: &AggIntent) -> Result Result { let by = match reduction { Reduction::PerEntity => { - debug_assert_eq!( - measures.len(), - 1, - "a per-entity reduction is single-aggregate" - ); - return per_series_reduction_schema(in_schema, &measures[0]); + if measures.is_empty() { + return Err(QueryExprError::InvalidPerEntityAggregate( + "at least one measure is required".into(), + )); + } + if measures.len() == 1 { + return per_series_reduction_schema(in_schema, &measures[0]); + } + let invalid = |message: &str| QueryExprError::InvalidPerEntityAggregate(message.into()); + let vi = in_schema + .column_id("value") + .ok_or_else(|| invalid("multiple measures require an input value column"))?; + if in_schema.time_index == Some(vi) { + return Err(invalid("the value column cannot be the timestamp")); + } + if output_names.len() > measures.len() { + return Err(invalid("more output names than measures")); + } + let mut output = in_schema.clone(); + let mut names: std::collections::HashSet = in_schema + .columns + .iter() + .enumerate() + .filter(|(i, _)| *i != vi) + .map(|(_, c)| c.name.clone()) + .collect(); + for (i, measure) in measures.iter().enumerate() { + if matches!(measure, AggIntent::CountValues { .. }) { + return Err(invalid("count_values changes series identity and cannot be combined with other measures")); + } + let input = in_schema + .columns + .get(measure.input_cols().first().copied().unwrap_or(vi)) + .ok_or_else(|| invalid("measure input column is out of range"))?; + let mut column = measure.output_column(input); + if let Some(name) = output_names.get(i).filter(|n| !n.is_empty()) { + column.name = name.clone(); + } + if !names.insert(column.name.clone()) { + return Err(invalid("measure names must be unique and must not collide with labels or timestamp")); + } + column.dtype = DataType::Float64; + if i == 0 { + output.columns[vi] = column; + } else { + output.columns.push(column); + } + } + // A computed sample cannot retain a uniqueness proof about raw values. + output.unique_keys.retain(|key| !key.contains(&vi)); + return Ok(output); } Reduction::Reduce(by) => by, }; @@ -2393,6 +2446,80 @@ mod tests { } } + /// Multiple measures retain series metadata and expose separate named float samples. + #[test] + fn per_entity_multiple_measures_preserve_schema() { + let input = Schema::with_time_index( + vec![ + col("ts", DataType::Timestamp, false), + col("value", DataType::Float64, false), + col("job", DataType::Utf8, true), + ], + 0, + vec![vec![0, 2]], + ); + let measures = vec![ + AggIntent::Sum { col: None }, + AggIntent::Count { + accuracy: crate::types::AccuracyTarget::Exact, + }, + ]; + let output = + aggregate_output_schema(&input, &Reduction::PerEntity, &measures, &[]).unwrap(); + assert_eq!( + output + .columns + .iter() + .map(|c| c.name.as_str()) + .collect::>(), + vec!["ts", "sum", "job", "count"] + ); + assert_eq!(output.time_index, input.time_index); + assert_eq!(output.unique_keys, input.unique_keys); + assert_eq!(output.closed, input.closed); + assert_eq!(output.columns[1].dtype, DataType::Float64); + assert_eq!(output.columns[3].dtype, DataType::Float64); + } + + /// Invalid or ambiguous multi-measure shapes fail instead of losing outputs. + #[test] + fn per_entity_measure_names_are_validated() { + let input = Schema::new(vec![ + col("value", DataType::Float64, false), + col("job", DataType::Utf8, true), + ]); + let measures = vec![AggIntent::Sum { col: None }, AggIntent::Sum { col: None }]; + for names in [ + vec![], + vec!["a".into(), "a".into()], + vec!["job".into(), "b".into()], + ] { + assert!( + aggregate_output_schema(&input, &Reduction::PerEntity, &measures, &names).is_err() + ); + } + let output = aggregate_output_schema( + &input, + &Reduction::PerEntity, + &measures, + &["first".into(), "second".into()], + ) + .unwrap(); + assert_eq!(output.column_id("first"), Some(0)); + assert_eq!(output.column_id("second"), Some(2)); + assert!(aggregate_output_schema(&input, &Reduction::PerEntity, &[], &[]).is_err()); + assert!(aggregate_output_schema( + &input, + &Reduction::PerEntity, + &[ + AggIntent::Sum { col: Some(99) }, + AggIntent::Sum { col: None } + ], + &[] + ) + .is_err()); + } + #[test] fn per_series_rate_preserves_labels() { // A per-series range reduction (`rate`) is label-preserving: it produces diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index 95b482e54..3ddfb0e0a 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -764,6 +764,66 @@ mod tests { assert!(resolve_agg_intent(&missing, &schema).is_err()); } + /// Parent expressions can resolve every named per-entity measure after binding. + #[test] + fn resolves_multiple_per_entity_measure_outputs() { + use crate::pre_asap::query_expr::ProjectItem; + let aggregate = QueryExpr::Aggregate { + reduction: Reduction::PerEntity, + measures: vec![ + AggIntent::Sum { + col: Some(ColumnRef::SampleValue), + }, + AggIntent::Count { + accuracy: crate::types::AccuracyTarget::Exact, + }, + ], + output_names: vec!["total".into(), "samples".into()], + filters: vec![], + having: None, + child: Rc::new(QueryExpr::Scan { + source: Source::TimeSeries { + metric: "latency".into(), + }, + predicates: vec![], + schema: Some(crate::pre_asap::schema::Schema::new(vec![ + crate::pre_asap::schema::Column::new( + "value", + crate::pre_asap::schema::DataType::Float64, + false, + ), + ])), + }), + }; + let query = QueryExpr::Project { + cols: vec![ + ProjectItem { + alias: None, + expr: QueryExpr::Column(ColumnRef::Named("total".into())), + }, + ProjectItem { + alias: None, + expr: QueryExpr::Column(ColumnRef::Named("samples".into())), + }, + ], + qualifier: None, + child: Rc::new(aggregate), + }; + let output = resolve_root(&query).unwrap().output_schema().unwrap(); + assert_eq!( + output + .columns + .iter() + .map(|c| c.name.as_str()) + .collect::>(), + vec!["total", "samples"] + ); + assert!(output + .columns + .iter() + .all(|c| c.dtype == crate::pre_asap::schema::DataType::Float64)); + } + /// `resolve_root` over a `BinaryOp { , PromqlScalarBridge, vector_match }` /// (issue #220): the bridged scalar operand resolves through the same /// generic walk as every other node (its `Literal` child has no