fix: support multi-measure per-entity schemas safely

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")]

crates/asap-aware-mapping/src/query_physical_lowering.rs

@@ -292,7 +292,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)?;
@@ -2486,6 +2486,42 @@ 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![],
+            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::{

crates/asap-aware-mapping/src/rewrite.rs

@@ -26,33 +26,15 @@
 //! that reshaping — see "Non-goals" below for why it does not also decide
 //! whether the reshaping is worth it.
 //!
-//! ## Scope: `by(...)` grouping only (issue #253's own scope note)
+//! ## Scope
 //!
-//! [`AvgToSumOverCountStrategy::matches`] additionally requires
-//! `Reduction::Reduce(by)` with `by` an ordinary (non-`without`) grouping —
-//! narrower than [`SketchAlgorithmStrategy`]'s `bindable_intent`, which is
-//! `Reduction`-agnostic. Two concrete reasons, not stylistic ones:
-//!
-//! - **`Reduction::PerEntity`** (`rate`/`increase`/`*_over_time`) is
-//!   single-measure by construction —
-//!   [`aggregate_output_schema`](asap_types::pre_asap::query_expr::aggregate_output_schema)
-//!   `debug_assert!`s exactly one measure for it. This rewrite's entire
-//!   point is introducing a *second* measure (`Count` alongside `Sum`)
-//!   under the same node, which would violate that invariant outright, not
-//!   just drift a schema detail.
-//! - **`without(...)` grouping** leaves an `Aggregate`'s own output schema
-//!   *open* (`closed: false`, see `without_output_schema`), while the
-//!   `Project` this strategy always wraps the rewrite in forces
-//!   `closed: true` (see `QueryExpr::output_schema`'s `Project` arm). Under
-//!   `without(...)` the rewritten form's `closed` flag would silently flip
-//!   relative to the original — exactly the kind of schema drift this
-//!   module exists to avoid.
-//!
-//! Both are follow-ups (issue #253 itself scopes to "the concrete case in
-//! Peilin's comment"), not correctness bugs in what ships here — a node
-//! outside this scope simply doesn't `match`, the same "safe but
-//! uninformative" fallback [`SketchAlgorithmStrategy`]/[`SharedSubtreeStrategy`]
-//! already use for shapes they don't have an opinion on.
+//! Ordinary `by(...)` averages with non-null inputs are eligible.
+//! Per-entity range averages are excluded: even finite Float64 samples can
+//! overflow SUM while AVG remains finite. A future decomposition needs a
+//! proven arithmetic domain or an overflow-safe execution path; schema
+//! compatibility alone cannot establish semantic equivalence.
+//! `without(...)` remains excluded because the grouped cast projection closes
+//! its input schema. Nullable inputs are excluded because Count means COUNT(*).
 //!
 //! ## Non-goals (mirrors [`replacement`]'s own discipline)
 //!
@@ -76,9 +58,8 @@ use asap_types::types::AccuracyTarget;
 use crate::replacement::{Replacement, ReplacementStrategy, ReplacementSubDAG, TargetSubDAG};
 
 /// The shape [`AvgToSumOverCountStrategy`] rewrites: a single `Avg{col}`
-/// measure, no `HAVING`, grouped with an ordinary `by(...)` reduction (see
-/// the module docs' "Scope" for why `without(...)`/`PerEntity` are
-/// excluded). Returns the grouping key count and the summed column so
+/// measure, no `HAVING`, with ordinary grouping.
+/// Returns the grouping key count and summed column so
 /// [`build_rewrite`] doesn't have to re-match.
 fn avg_rewrite_target(node: &QueryExpr) -> Option<(usize, Option<ColumnId>)> {
     let QueryExpr::Aggregate {
@@ -91,12 +72,10 @@ fn avg_rewrite_target(node: &QueryExpr) -> Option<(usize, Option<ColumnId>)> {
     else {
         return None;
     };
-    let Reduction::Reduce(by) = reduction else {
-        return None;
+    let by = match reduction {
+        Reduction::Reduce(by) if !by.is_without() => by,
+        _ => return None,
     };
-    if by.is_without() {
-        return None;
-    }
     let [AggIntent::Avg { col }] = measures.as_slice() else {
         return None;
     };
@@ -114,12 +93,12 @@ fn avg_rewrite_target(node: &QueryExpr) -> Option<(usize, Option<ColumnId>)> {
     Some((by.keys().len(), *col))
 }
 
-/// Build the rewritten `Project{ cast(sum) } / Aggregate{ Count }` tree for
+/// Build `Sum / Count` (with a cast projection for grouped aggregates) for
 /// `root`, or `None` if `root` isn't [`avg_rewrite_target`]'s shape. `Sum` and
 /// `Count` deliberately live in separate, single-measure aggregates so the
 /// replacement fixpoint discovers each as an independently bindable target.
 ///
-/// The `Project`'s leading `by.len()` items are bare `Column(i)`
+/// For grouped aggregates, the `Project`'s leading `by.len()` items are bare `Column(i)`
 /// pass-throughs of the grouping keys — identical in name/type to the
 /// original `Avg` aggregate's own leading columns, since both aggregates
 /// share the same `reduction`/`child` and only differ in `measures`
@@ -463,18 +442,42 @@ 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![],
             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![],
+                having: None,
+                child: Rc::new(scan),
+            });
+            assert!(build_rewrite(&root).is_none());
+        }
     }
 
     #[test]

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_frontend_promql::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}"
+        );
+    }
+}

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