feat(lower): arbitrary PromQL function nesting in aggregates (#27)

crates/core/src/intent_algebra/column_resolution.rs

@@ -76,6 +76,22 @@ pub fn resolve_column_ref(col: &ColumnRef, schema: &Schema) -> Result<ColumnId,
                     .collect();
                 (non_ts.len() == 1).then(|| non_ts[0])
             })
+            .or_else(|| {
+                // A *cross-series* aggregate (`sum by (job) (…)`) emits the group
+                // labels (Utf8) alongside the single numeric value column (e.g.
+                // `[job:Utf8, sum:Float64]`), so the sole-non-ts fallback above is
+                // ambiguous. The PromQL sample value of such a vector is that one
+                // numeric column — the labels are keys, not values. Pick it when
+                // it is the unique non-timestamp numeric column, so an outer
+                // ranking (`topk(k, sum by (job) (…))`) resolves its sort key.
+                let numeric: Vec<ColumnId> = (0..schema.columns.len())
+                    .filter(|&i| Some(i) != schema.time_index)
+                    .filter(|&i| {
+                        matches!(schema.columns[i].dtype, DataType::Float64 | DataType::Int64)
+                    })
+                    .collect();
+                (numeric.len() == 1).then(|| numeric[0])
+            })
             .ok_or_else(|| ResolveError::NoSampleValue {
                 available: schema.columns.iter().map(|c| c.name.clone()).collect(),
             }),
@@ -218,6 +234,31 @@ mod tests {
         assert_eq!(resolve_column_ref(&ColumnRef::SampleValue, &s), Ok(1));
     }
 
+    #[test]
+    fn sample_value_resolves_to_sole_numeric_after_cross_series_aggregate() {
+        // A cross-series aggregate output `[job:Utf8, sum:Float64]` has no `value`
+        // column and two non-ts columns (ambiguous), but exactly one numeric
+        // column — the sample value an outer `topk` ranks by.
+        let s = Schema::new(vec![
+            Column::new("job", DataType::Utf8, true),
+            Column::new("sum", DataType::Float64, false),
+        ]);
+        assert_eq!(resolve_column_ref(&ColumnRef::SampleValue, &s), Ok(1));
+    }
+
+    #[test]
+    fn sample_value_ambiguous_when_two_numeric_columns() {
+        // Two numeric non-ts columns → genuinely ambiguous → NoSampleValue.
+        let s = Schema::new(vec![
+            Column::new("a", DataType::Float64, false),
+            Column::new("b", DataType::Int64, false),
+        ]);
+        assert!(matches!(
+            resolve_column_ref(&ColumnRef::SampleValue, &s),
+            Err(ResolveError::NoSampleValue { .. })
+        ));
+    }
+
     #[test]
     fn resolve_unknown_name_errors() {
         let s = infer_source_schema("m");

crates/e2e/tests/nested.rs

@@ -145,6 +145,36 @@ fn q25_div_over_complex_subtrees() {
     );
 }
 
+// #27 — outer cross-series reduction over a nested per-group reduction over a
+//   per-series rate: `max(sum by (job) (rate(m[5m])))`. Three stacked levels —
+//   the arbitrary function nesting the old two-level template could not express.
+//   The inner `sum by (job)` resolves job at col 2 against rate's
+//   label-preserving output schema; the outer `max` has no grouping.
+#[test]
+fn q27_max_over_sum_by_job_over_rate() {
+    let scan = QueryExpr::Scan {
+        source: Source::TimeSeries {
+            metric: "http_requests_total".into(),
+        },
+        predicates: vec![],
+        schema: metric_schema(&["job"]),
+    };
+    let inner_rate = agg(
+        vec![],
+        AggIntent::Rate,
+        QueryExpr::TimeRange {
+            range: Duration::from_secs(300),
+            child: Box::new(scan),
+        },
+    );
+    let sum_by_job = agg(vec![2], AggIntent::Sum { col: None }, inner_rate);
+    let expected = agg(vec![], AggIntent::Max { col: None }, sum_by_job);
+    assert_eq!(
+        lower("max(sum by (job) (rate(http_requests_total[5m])))"),
+        expected
+    );
+}
+
 // #24 — sum by job over rate over a filtered scan
 //   same schema [ts, value, job, status]; rate is label-preserving,
 //   so outer sum by job still finds job at col 2

crates/lower/src/promql.rs

@@ -16,6 +16,8 @@
 //! | `quantile_over_time(φ, m{f}[w])` | `Aggregate{[Quantile(φ)], Window{w, Filter(Source)}}` |
 //! | `histogram_quantile(φ, <expr>)` | `Aggregate{[Quantile(φ)]}` over the fully-lowered `<expr>` (preserves any `sum by (le)`/`rate`) |
 //! | `OUTER_op(inner_func(m[w]))` (e.g. `sum(rate(m[w]))`) | `Aggregate{[OUTER_op]}` over `Aggregate{[inner_func]}` — two levels |
+//! | `OUTER_op(<any expr>)` (e.g. `max(sum by (job) (rate(m[w])))`, `sum(rate(a[w]) + rate(b[w]))`) | `Aggregate{[OUTER_op]}` over the fully-lowered `<any expr>` — arbitrary function nesting (issue #27) |
+//! | `topk(k, <non-count expr>)` / `bottomk(k, <any expr>)` | `Sort{value} → Limit{k}` over the fully-lowered argument |
 //! | `avg/min/max/sum_over_time(m[w])` | `Aggregate{[Avg/Min/Max/Sum], Window{w}}` |
 //! | `stddev/stdvar_over_time(m[w])` | `Aggregate{[StdDev/Variance], Window{w}}` |
 //! | `count_over_time(m[w])` | `Aggregate{[Count], Window{w}}` |
@@ -190,10 +192,38 @@ fn walk(expr: &Expr) -> Result<L2> {
 
 fn walk_aggregate(agg: &AggregateExpr) -> Result<L2> {
     let keys = resolve_group(agg)?;
-    let inner = lower_inner(&agg.expr)?;
+    let outer = outer_kind(agg)?;
+
+    // Fast path — the argument is a bare selector or a single range-vector
+    // function (`rate`/`increase`/`*_over_time`). `lower_inner` lowers it via the
+    // flat selector/call template, which also recognises the heavy-hitter
+    // `topk(k, count_over_time(...))` shape. This is the common two-level case
+    // (`sum by (job) (rate(m[5m]))`).
+    if let Ok(inner) = lower_inner(&agg.expr) {
+        return build(inner, keys, outer);
+    }
+
+    // General nesting — the argument is itself a composite expression: another
+    // aggregate (`max(sum by (job) (rate(m[5m])))`), a binary op
+    // (`sum(rate(a[5m]) + rate(b[5m]))`), a sub-query, or a function lowered
+    // elsewhere (`sum(histogram_quantile(0.9, …))`). Lower it recursively with
+    // the same `walk` used at the top level, then wrap it in the outer
+    // aggregation. This is the path that lifts the old two-level limit to
+    // arbitrary function nesting (issue #27). If the inner expression is itself
+    // unsupported (e.g. unary negation), `walk` surfaces that error, so a
+    // genuinely unsupported query is still cleanly rejected rather than
+    // mislowered.
+    let child = walk(&agg.expr)?;
+    build_over_subtree(outer, keys, child)
+}
+
+/// Map an `AggregateExpr`'s operator (`sum`/`avg`/`topk`/…) to the [`Outer`]
+/// shape, independent of what the argument is — so both the flat fast path and
+/// the general recursive path share one operator-dispatch.
+fn outer_kind(agg: &AggregateExpr) -> Result<Outer> {
     let op = agg.op.id();
 
-    let outer = if op == token::T_TOPK {
+    Ok(if op == token::T_TOPK {
         Outer::TopK {
             k: count_param(agg)?,
             descending: true,
@@ -230,9 +260,41 @@ fn walk_aggregate(agg: &AggregateExpr) -> Result<L2> {
         return Err(LoweringError::UnsupportedAggregateOp(format!(
             "aggregate token {op}"
         )));
-    };
+    })
+}
 
-    build(inner, keys, outer)
+/// Wrap an already-lowered L2 subtree in the outer aggregation. This is the
+/// general-nesting counterpart to [`build`]: where `build` assembles the
+/// two-level shape from a flat [`Inner`], this composes the outer operator over
+/// an arbitrary child (`max(sum by (job) (…))`, `sum(a + b)`, …).
+///
+/// A heavy-hitter `TopK` is only recognised on the flat `count_over_time` shape
+/// (handled in `build`); over a general subtree, `topk`/`bottomk` is a generic
+/// order-by-value + limit — the same `Sort{partition_by} → Limit` pair `build`
+/// emits for any non-heavy-hitter ranking.
+fn build_over_subtree(outer: Outer, keys: Vec<ColumnRef>, child: L2) -> Result<L2> {
+    Ok(match outer {
+        // `walk_aggregate` always passes a real aggregator; `None` can't occur.
+        Outer::None => child,
+        Outer::Plain(intent) => outer_aggregate(keys, outer_func(&intent), child),
+        Outer::Count => outer_aggregate(keys, AggFunc::CountDistinct, child),
+        Outer::TopK { k, descending } => {
+            let sorted = L2::Sort {
+                keys: vec![L2SortKey {
+                    expr: L2Expr::Column(ColumnRef::SampleValue),
+                    ascending: !descending,
+                    nulls_first: false,
+                }],
+                partition_by: keys,
+                input: Box::new(child),
+            };
+            L2::Limit {
+                n: k,
+                offset: 0,
+                input: Box::new(sorted),
+            }
+        }
+    })
 }
 
 /// `histogram_quantile(φ, <expr>)` lowers `<expr>` in full — preserving any

crates/lower/tests/promql_conformance.rs

@@ -577,11 +577,62 @@ fn bottomk_is_generic_sort_limit() {
 }
 
 #[test]
-fn topk_over_aggregate_arg_is_rejected__GAP() {
+fn topk_over_nested_aggregate_is_generic_sort_limit() {
     // SEMANTICS (PromQL): `topk(3, sum by(x)(rate(...)))` is extremely common.
-    // Our aggregate-argument lowering only accepts selectors/calls, not a
-    // nested aggregate, so this is rejected today.
-    let _ = rejected("topk(3, sum by(instance) (rate(node_cpu_seconds_total[5m])))");
+    // The argument is a nested aggregate (not the heavy-hitter `count` shape),
+    // so it ranks by sample value → generic `Sort{value desc} → Limit{k}` over
+    // the fully-lowered inner aggregate (issue #27: arbitrary function nesting).
+    let qe = ok("topk(3, sum by(instance) (rate(node_cpu_seconds_total[5m])))");
+    let QueryExpr::Limit { n, child, .. } = &qe else {
+        panic!("expected outer Limit, got {qe:?}");
+    };
+    assert_eq!(*n, 3);
+    let QueryExpr::Sort { keys, child, .. } = child.as_ref() else {
+        panic!("expected Sort under Limit, got {child:?}");
+    };
+    assert!(!keys[0].ascending, "topk ranks descending by value");
+    // The inner `sum by (instance)` survives as a cross-series Aggregate over the
+    // per-series rate — the nesting the old two-level template could not express.
+    assert!(
+        has(child, |i| matches!(i, AggIntent::Sum { .. }))
+            && has(child, |i| matches!(i, AggIntent::Rate)),
+        "inner sum-over-rate preserved, got {:?}",
+        intents(child)
+    );
+}
+
+#[test]
+fn outer_aggregate_over_nested_aggregate_nests() {
+    // `max(sum by (job) (rate(m[5m])))` — an outer cross-series reduction over a
+    // nested per-group reduction over a per-series rate: three stacked levels the
+    // flat two-level template rejected. Each level survives into L3 (issue #27).
+    let qe = ok("max(sum by (job) (rate(http_requests_total[5m])))");
+    let QueryExpr::Aggregate { aggs, child, .. } = &qe else {
+        panic!("expected outer Aggregate, got {qe:?}");
+    };
+    assert!(matches!(aggs.as_slice(), [AggIntent::Max { .. }]));
+    let QueryExpr::Aggregate { by, aggs, .. } = child.as_ref() else {
+        panic!("expected inner `sum by (job)` Aggregate, got {child:?}");
+    };
+    assert_eq!(by, &vec![2], "job grouping survives on the inner aggregate");
+    assert!(matches!(aggs.as_slice(), [AggIntent::Sum { .. }]));
+    assert!(has(&qe, |i| matches!(i, AggIntent::Rate)), "rate preserved");
+}
+
+#[test]
+fn aggregate_over_binary_op_nests() {
+    // `sum(rate(a[5m]) + rate(b[5m]))` — an aggregate whose argument is a binary
+    // op over two range vectors. The old template only accepted a single inner
+    // selector/call; now the binary op lowers and the outer sum wraps it.
+    let qe = ok("sum(rate(a[5m]) + rate(b[5m]))");
+    let QueryExpr::Aggregate { aggs, child, .. } = &qe else {
+        panic!("expected outer Aggregate, got {qe:?}");
+    };
+    assert!(matches!(aggs.as_slice(), [AggIntent::Sum { .. }]));
+    assert!(
+        matches!(child.as_ref(), QueryExpr::BinaryOp { .. }),
+        "argument lowers as a BinaryOp, got {child:?}"
+    );
 }
 
 // ─────────────────────────────────────────────────────────────────────────────