From a722ee8e662d0601c97a497566b18af91fffdbbe Mon Sep 17 00:00:00 2001 From: zz_y Date: Fri, 19 Jun 2026 11:13:25 -0600 Subject: [PATCH] feat(lower): support arbitrary PromQL function nesting in aggregates (#27) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PromQL front end lowered an aggregate argument through a flat two-level template (`Inner`/`Outer`): an outer aggregator over exactly one inner selector or range-vector function. That rejected the very common case of an aggregate over a *composite* expression — a nested 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, …))`). The L3 intent algebra, the L2→L3 converter, and the Binder are already fully recursive — the only ceiling was this front-end template. `walk_aggregate` now tries the flat fast path first (preserving the heavy-hitter `topk(k, count_over_time(...))` recognition) and, when the argument is a composite expression, recurses via the same `walk` used at the top level and wraps the result in the outer aggregation. Genuinely unsupported inner expressions (e.g. unary negation) still surface their error, so nothing is silently mislowered. `ColumnRef::SampleValue` now also resolves to the sole *numeric* non-timestamp column, so an outer ranking over a cross-series aggregate (`topk(k, sum by (job) (…))`, whose value column is renamed `sum`) finds its sort key. Tests: flips the `topk_over_aggregate_arg_is_rejected__GAP` conformance gap to a passing test; adds nested-aggregate, aggregate-over-binary-op, and outer-over- nested-aggregate conformance tests, an exact-tree e2e test, and SampleValue resolution unit tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/intent_algebra/column_resolution.rs | 41 +++++++++++ crates/e2e/tests/nested.rs | 30 ++++++++ crates/lower/src/promql.rs | 70 +++++++++++++++++-- crates/lower/tests/promql_conformance.rs | 59 ++++++++++++++-- 4 files changed, 192 insertions(+), 8 deletions(-) diff --git a/crates/core/src/intent_algebra/column_resolution.rs b/crates/core/src/intent_algebra/column_resolution.rs index 2b5177a3..904c0cc5 100644 --- a/crates/core/src/intent_algebra/column_resolution.rs +++ b/crates/core/src/intent_algebra/column_resolution.rs @@ -76,6 +76,22 @@ pub fn resolve_column_ref(col: &ColumnRef, schema: &Schema) -> Result = (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"); diff --git a/crates/e2e/tests/nested.rs b/crates/e2e/tests/nested.rs index 81ef1d86..6ae02c30 100644 --- a/crates/e2e/tests/nested.rs +++ b/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 diff --git a/crates/lower/src/promql.rs b/crates/lower/src/promql.rs index cd7d5408..68fcf08a 100644 --- a/crates/lower/src/promql.rs +++ b/crates/lower/src/promql.rs @@ -16,6 +16,8 @@ //! | `quantile_over_time(φ, m{f}[w])` | `Aggregate{[Quantile(φ)], Window{w, Filter(Source)}}` | //! | `histogram_quantile(φ, )` | `Aggregate{[Quantile(φ)]}` over the fully-lowered `` (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()` (e.g. `max(sum by (job) (rate(m[w])))`, `sum(rate(a[w]) + rate(b[w]))`) | `Aggregate{[OUTER_op]}` over the fully-lowered `` — arbitrary function nesting (issue #27) | +//! | `topk(k, )` / `bottomk(k, )` | `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 { fn walk_aggregate(agg: &AggregateExpr) -> Result { 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 { 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 { 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, child: L2) -> Result { + 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(φ, )` lowers `` in full — preserving any diff --git a/crates/lower/tests/promql_conformance.rs b/crates/lower/tests/promql_conformance.rs index f84d71f6..87bc5840 100644 --- a/crates/lower/tests/promql_conformance.rs +++ b/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:?}" + ); } // ─────────────────────────────────────────────────────────────────────────────