From 3385cbba496fc3cb061843fd25cdf49a13e82be4 Mon Sep 17 00:00:00 2001 From: zz_y Date: Sat, 4 Jul 2026 09:03:33 -0600 Subject: [PATCH] feat(ir): shared L3 canonicalization pass for heavy-hitter topk (#34) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Semantically equivalent SQL and PromQL queries produced structurally different L3, and the heavy-hitter gate was duplicated across both front ends with divergent recognition logic. This adds a single post-lowering canonicalization pass both languages run through. - `asap_l2::canonicalize` — a bottom-up L3 rewrite, run at the end of the shared `convert_root` so both front ends get it for free. It promotes a count-ranked `Limit{Sort{[Project] Aggregate([Count])}}` into the canonical heavy-hitter shape: an outer `Aggregate([TopK{k}])` (grouped by the sort's partition) over the *explicit* inner `Aggregate([Count])` — the shape PromQL's `topk(k, count_over_time(...))` already produced. The match is positional (the DESC key must land on the Count's output column), so it is oblivious to whether the count was aliased. - Remove the SQL front-end heavy-hitter gate (`heavy_hitter_topk` / `lower_as_topk` and their helpers). SQL now emits a plain Sort+Limit and lets canonicalize recognise the count-ranked shape. This fixes the alias blind spot (#20 — `ORDER BY cnt DESC` now promotes like `ORDER BY COUNT(*)`) and closes the SQL/PromQL structural gap (#25 — SQL's implicit-count `TopK{Scan}` is gone; both languages now emit `TopK` over an explicit `Count`). The outer `TopK.by` is the ranking partition ([] for a global `ORDER BY … LIMIT k`), with the grouping on the inner `Count`. Tests: - `asap-l2` unit tests for the pass: promotion (with/without a passthrough projection), idempotency, and the negative cases (ascending, OFFSET, ranking a group key, non-Count aggregate). - `crates/lower/tests/cross_language.rs` — the executable spec: SQL S2 and PromQL P1 reach the same canonical shape; aliased ≡ inline (#20); non-count / OFFSET stay generic in both languages. - Updated the SQL `count_ranked_topk_is_heavy_hitter` expectation to the canonical two-level form. Co-Authored-By: Claude Opus 4.8 --- crates/frontend-sql/src/sql/mod.rs | 115 +-------- crates/frontend-sql/tests/sql_lowering.rs | 43 +++- crates/l2/src/canonicalize.rs | 279 ++++++++++++++++++++++ crates/l2/src/lib.rs | 2 + crates/l2/src/lower.rs | 5 +- crates/lower/tests/cross_language.rs | 128 ++++++++++ 6 files changed, 460 insertions(+), 112 deletions(-) create mode 100644 crates/l2/src/canonicalize.rs create mode 100644 crates/lower/tests/cross_language.rs diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index feda353c..5d548866 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -378,17 +378,12 @@ impl<'a> SqlLowerer<'a> { } fn lower_sort(&self, sort: &logical_expr::Sort) -> Result { - // Heavy-hitter TopK only when ranking DESC by a single COUNT aggregate - // (the frequency sketch the `TopK` intent represents). Any other - // ranking — by a SUM/AVG/… output or a group column — keeps the real - // Aggregate under a generic Sort+Limit so its aggregate isn't discarded. - if let Some(k) = sort.fetch { - if let Some(agg) = find_aggregate(strip_projections_and_aliases(&sort.input)) { - if heavy_hitter_topk(sort, agg) { - return self.lower_as_topk(agg, k as u64); - } - } - } + // A count-ranked `ORDER BY … LIMIT k` is the frequency heavy-hitter the + // `TopK` intent represents, but that promotion now happens in the shared + // L3 `canonicalize` pass (issue #34) — the same one both front ends run — + // so SQL emits a plain `Sort` (+ `Limit`) here and lets canonicalization + // recognise the count-ranked shape positionally. This removes the gate's + // alias blind spot (#20). let keys = sort .expr .iter() @@ -410,19 +405,8 @@ impl<'a> SqlLowerer<'a> { } fn lower_limit(&self, limit: &logical_expr::Limit) -> Result { - // Heavy-hitter TopK only for a count-ranked Limit-over-Sort-over-Aggregate - // with no OFFSET (see `lower_sort`). Otherwise fall through to Limit+Sort. - if let Some(k) = eval_fetch(&limit.fetch) { - if eval_fetch(&limit.skip).unwrap_or(0) == 0 { - if let LogicalPlan::Sort(sort) = strip_aliases(&limit.input) { - if let Some(agg) = find_aggregate(strip_projections_and_aliases(&sort.input)) { - if heavy_hitter_topk(sort, agg) { - return self.lower_as_topk(agg, k as u64); - } - } - } - } - } + // Count-ranked `LIMIT k` over a `Sort` is promoted to the heavy-hitter + // `TopK` by the shared L3 `canonicalize` pass (issue #34), not here. Ok(L2::Limit { n: eval_fetch(&limit.fetch).unwrap_or(usize::MAX) as u64, offset: eval_fetch(&limit.skip).unwrap_or(0) as u64, @@ -430,18 +414,6 @@ impl<'a> SqlLowerer<'a> { }) } - fn lower_as_topk(&self, agg: &logical_expr::Aggregate, k: u64) -> Result { - let by = agg - .group_expr - .iter() - .map(expr_to_group_ref) - .collect::, _>>()?; - Ok(L2::TopK { - k, - by, - input: Box::new(self.lower_plan(&agg.input)?), - }) - } } // ── Aggregate / group-key helpers ─────────────────────────────────────────────── @@ -575,53 +547,6 @@ fn extract_percentile_q(args: &[Expr]) -> Result { } } -/// True iff `sort` ranks **descending by a single `COUNT` aggregate** of `agg` -/// — the only shape the heavy-hitter (frequency) `TopK` sketch is correct for. -/// -/// Requires `agg` to have exactly one aggregate (a plain, non-`DISTINCT` -/// `COUNT`) and the sole sort key to reference *that* output column (not a -/// group key, a `SUM`/`AVG`/… output, or a multi-aggregate select). Anything -/// else stays a generic `Sort` + `Limit` over the real `Aggregate`, mirroring -/// the PromQL gate (`topk` is heavy-hitter only over `count_over_time`). -fn heavy_hitter_topk(sort: &logical_expr::Sort, agg: &logical_expr::Aggregate) -> bool { - let [key] = sort.expr.as_slice() else { - return false; - }; - if key.asc { - return false; - } - if agg.aggr_expr.len() != 1 || !is_count_aggregate(&agg.aggr_expr[0]) { - return false; - } - // The DESC key must rank by the count's output column, not a group key. The - // aggregate schema is `[group fields …, aggregate fields …]`, so the single - // count output sits at index `group_expr.len()`. - let count_name = agg - .schema - .fields() - .get(agg.group_expr.len()) - .map(|f| f.name().clone()); - column_name(&key.expr) == count_name -} - -/// The referenced column name of a bare/aliased column expression, else `None`. -fn column_name(expr: &Expr) -> Option { - match expr { - Expr::Column(c) => Some(c.name.clone()), - Expr::Alias(a) => column_name(&a.expr), - _ => None, - } -} - -/// Whether `expr` is a plain (non-`DISTINCT`) `COUNT` aggregate. -fn is_count_aggregate(expr: &Expr) -> bool { - match expr { - Expr::Alias(a) => is_count_aggregate(&a.expr), - Expr::AggregateFunction(f) => f.func.name().eq_ignore_ascii_case("count") && !f.distinct, - _ => false, - } -} - // ── LogicalPlan navigation helpers ────────────────────────────────────────────── fn eval_fetch(expr_opt: &Option>) -> Option { @@ -633,30 +558,6 @@ fn eval_fetch(expr_opt: &Option>) -> Option { }) } -fn strip_aliases(plan: &LogicalPlan) -> &LogicalPlan { - match plan { - LogicalPlan::SubqueryAlias(a) => strip_aliases(&a.input), - _ => plan, - } -} - -/// Strip Projection + SubqueryAlias for TopK pattern-matching only. -fn strip_projections_and_aliases(plan: &LogicalPlan) -> &LogicalPlan { - match plan { - LogicalPlan::SubqueryAlias(a) => strip_projections_and_aliases(&a.input), - LogicalPlan::Projection(p) => strip_projections_and_aliases(&p.input), - _ => plan, - } -} - -fn find_aggregate(plan: &LogicalPlan) -> Option<&logical_expr::Aggregate> { - match plan { - LogicalPlan::Aggregate(agg) => Some(agg), - LogicalPlan::Projection(p) => find_aggregate(&p.input), - LogicalPlan::SubqueryAlias(a) => find_aggregate(&a.input), - _ => None, - } -} /// Map a DataFusion window-function definition to the L3 [`WindowFuncKind`]. /// `NthValue` is returned with `None`; `lower_window` fills in `n` from args. diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index f5fea816..15db83cc 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -175,19 +175,54 @@ async fn single_agg_group_by_keeps_key_in_output_schema() { #[tokio::test] async fn count_ranked_topk_is_heavy_hitter() { // `ORDER BY COUNT(*) DESC LIMIT k` over a single COUNT aggregate is the one - // case the heavy-hitter (frequency) sketch is correct for. (The key must - // reference the count output directly; an alias would safely fall back to a - // generic Sort+Limit.) + // case the heavy-hitter (frequency) sketch is correct for. The shared L3 + // `canonicalize` pass (issue #34) promotes it to the canonical two-level + // form: an outer global `TopK` (by: []) over the explicit inner `Count` + // grouped by `service`. let qe = lower( "SELECT service, COUNT(*) FROM metrics GROUP BY service ORDER BY COUNT(*) DESC LIMIT 10", ) .await; let (by, aggs) = find_aggregate(&qe).expect("expected an Aggregate"); - assert_eq!(by, &vec![1], "GROUP BY service → col 1"); + assert!( + by.is_empty(), + "outer TopK is a global ranking (by: []), got {by:?}" + ); assert!( matches!(aggs.as_slice(), [AggIntent::TopK { k: 10, .. }]), "count-ranked topk → heavy-hitter TopK, got {aggs:?}" ); + // The inner child is the explicit Count, grouped by service (col 1). + let QueryExpr::Aggregate { child, .. } = &qe else { + panic!("expected outer Aggregate, got {qe:?}"); + }; + let (inner_by, inner_aggs) = find_aggregate(child).expect("expected inner Count aggregate"); + assert_eq!(inner_by, &vec![1], "inner Count grouped by service → col 1"); + assert!( + matches!(inner_aggs.as_slice(), [AggIntent::Count { .. }]), + "inner aggregate is the explicit Count, got {inner_aggs:?}" + ); +} + +#[tokio::test] +async fn count_ranked_topk_via_alias_is_also_heavy_hitter() { + // Regression for #20: aliasing `COUNT(*)` in the ORDER BY used to defeat the + // SQL front-end gate. The positional `canonicalize` pass now promotes it too, + // so the aliased and inline forms produce identical L3. + let inline = lower( + "SELECT service, COUNT(*) FROM metrics GROUP BY service ORDER BY COUNT(*) DESC LIMIT 10", + ) + .await; + let aliased = lower( + "SELECT service, COUNT(*) AS cnt FROM metrics GROUP BY service ORDER BY cnt DESC LIMIT 10", + ) + .await; + assert_eq!(inline, aliased, "aliased count-ranked topk must match the inline form"); + let (_, aggs) = find_aggregate(&aliased).expect("expected an Aggregate"); + assert!( + matches!(aggs.as_slice(), [AggIntent::TopK { k: 10, .. }]), + "aliased count-ranked topk → heavy-hitter TopK, got {aggs:?}" + ); } #[tokio::test] diff --git a/crates/l2/src/canonicalize.rs b/crates/l2/src/canonicalize.rs new file mode 100644 index 00000000..2ba428c8 --- /dev/null +++ b/crates/l2/src/canonicalize.rs @@ -0,0 +1,279 @@ +//! Shared post-lowering canonicalization of the L3 [`QueryExpr`]. +//! +//! Both language front ends funnel through [`convert_root`](crate::convert_root), +//! which runs this pass over the converted tree. Its job is to erase +//! *structural* differences between semantically identical queries so an L4 rule +//! matching on the intent algebra sees one canonical spelling regardless of +//! source language (issue #34). +//! +//! ## Heavy-hitter promotion +//! +//! A count-ranked "order by the count, take the top k" is the frequency +//! heavy-hitter the [`AggIntent::TopK`] intent represents. Front ends emit it as +//! an ordinary `Limit { Sort { … Aggregate([Count]) } }` (SQL, and PromQL's +//! generic `topk` path); this pass promotes that shape to the canonical +//! +//! ```text +//! Aggregate { by: , aggs: [TopK{k}], +//! child: Aggregate { aggs: [Count], … } } +//! ``` +//! +//! — an outer `TopK` over the *explicit* inner `Count` (the shape PromQL's +//! `topk(k, count_over_time(…))` already produces). Because the match is +//! **positional** (it checks that the DESC sort key lands on the count's output +//! column) it is oblivious to whether the count was aliased in the source, which +//! is exactly the SQL alias gap (#20) the old front-end gate missed. + +use asap_ir::intent_algebra::agg_intent::AggIntent; +use asap_ir::intent_algebra::query_expr::{GroupKeys, QueryExpr, SortKey}; +use asap_ir::intent_algebra::L3Expr; + +/// Rewrite `expr` into its canonical form (bottom-up). Idempotent: a tree that +/// is already canonical is returned unchanged. +pub fn canonicalize(mut expr: QueryExpr) -> QueryExpr { + canon(&mut expr); + expr +} + +fn canon(expr: &mut QueryExpr) { + // Bottom-up: canonicalize every child before matching at this node, so an + // inner heavy-hitter is promoted before an enclosing rewrite inspects it. + for child in children_mut(expr) { + canon(child); + } + if let Some(promoted) = try_promote_heavy_hitter(expr) { + *expr = promoted; + } +} + +/// Mutable references to the direct `QueryExpr` children of a node. +fn children_mut(expr: &mut QueryExpr) -> Vec<&mut QueryExpr> { + use QueryExpr::*; + match expr { + Scan { .. } | Ref { .. } | Scalar(_) | EvalTime => vec![], + VectorFromScalar(c) | ScalarFromVector(c) => vec![c.as_mut()], + Relabel { child, .. } + | Filter { child, .. } + | Project { child, .. } + | Aggregate { child, .. } + | Window { child, .. } + | Distinct { child, .. } + | Subquery { child, .. } + | TimeRange { child, .. } + | WindowFunc { child, .. } + | Sort { child, .. } + | Limit { child, .. } => vec![child.as_mut()], + LetBinding { expr, child, .. } => vec![expr.as_mut(), child.as_mut()], + Merge { children } => children.iter_mut().collect(), + Join { left, right, .. } | SetOp { left, right, .. } => { + vec![left.as_mut(), right.as_mut()] + } + BinaryOp { lhs, rhs, .. } => vec![lhs.as_mut(), rhs.as_mut()], + } +} + +/// Recognise a count-ranked `Limit { Sort { [Project] Aggregate([Count]) } }` +/// and rewrite it to the canonical heavy-hitter `Aggregate([TopK])` over the +/// explicit inner `Count`. Returns `None` when the shape does not match. +fn try_promote_heavy_hitter(expr: &QueryExpr) -> Option { + // Limit k, no offset (an OFFSET means "not the top k"). + let QueryExpr::Limit { + n: k, + offset: 0, + child, + } = expr + else { + return None; + }; + // A single DESC ordering key. + let QueryExpr::Sort { + keys, + partition_by, + child: sort_child, + } = child.as_ref() + else { + return None; + }; + let [SortKey { + expr: L3Expr::Column(sort_col), + ascending: false, + .. + }] = keys.as_slice() + else { + return None; + }; + + // The ordered relation is an `Aggregate`, optionally behind a passthrough + // projection (a bare-column SELECT list). Map the sort key through the + // projection to the aggregate's own output column. + let (agg_expr, ranked_col) = match sort_child.as_ref() { + QueryExpr::Project { cols, child, .. } => { + let L3Expr::Column(underlying) = &cols.get(*sort_col)?.expr else { + return None; + }; + (child.as_ref(), *underlying) + } + other => (other, *sort_col), + }; + + // Exactly one `Count` aggregate, and the DESC key must rank by *its* output + // column — the count sits at index `by.len()` (after the group keys). + let QueryExpr::Aggregate { by, aggs, .. } = agg_expr else { + return None; + }; + let [AggIntent::Count { accuracy }] = aggs.as_slice() else { + return None; + }; + if ranked_col != by.len() { + return None; + } + + // Outer heavy-hitter `TopK`, grouped by the ranking's partition (empty for a + // global `ORDER BY … LIMIT k`; the `by` labels for a partitioned `topk by`), + // over the *unchanged* inner `Count` aggregate. + Some(QueryExpr::Aggregate { + by: GroupKeys(partition_by.to_vec()), + aggs: vec![AggIntent::TopK { + k: *k, + accuracy: accuracy.clone(), + }], + output_names: Vec::new(), + having: None, + child: Box::new(agg_expr.clone()), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_ir::intent_algebra::query_expr::{ProjectItem, Source}; + use asap_ir::intent_algebra::schema::{Column, DataType, Schema}; + use asap_ir::types::AccuracyTarget; + + fn scan() -> QueryExpr { + QueryExpr::Scan { + source: Source::TimeSeries { + metric: "m".into(), + }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Column::new("ts", DataType::Timestamp, false), + Column::new("service", DataType::Utf8, false), + Column::new("value", DataType::Float64, false), + ], + 0, + vec![], + ), + } + } + + /// `Aggregate{ by: [1], [Count] }` over the scan — output cols `[service, count]`. + fn count_by_service() -> QueryExpr { + QueryExpr::Aggregate { + by: GroupKeys(vec![1]), + aggs: vec![AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }], + output_names: vec![], + having: None, + child: Box::new(scan()), + } + } + + fn desc(col: usize) -> Vec { + vec![SortKey { + expr: L3Expr::Column(col), + ascending: false, + nulls_first: false, + }] + } + + fn limit(n: usize, offset: usize, child: QueryExpr) -> QueryExpr { + QueryExpr::Limit { + n, + offset, + child: Box::new(child), + } + } + + fn sort(keys: Vec, child: QueryExpr) -> QueryExpr { + QueryExpr::Sort { + keys, + partition_by: GroupKeys(vec![]), + child: Box::new(child), + } + } + + fn is_topk_over_count(qe: &QueryExpr) -> bool { + matches!(qe, + QueryExpr::Aggregate { aggs, child, .. } + if matches!(aggs.as_slice(), [AggIntent::TopK { k: 5, .. }]) + && matches!(child.as_ref(), QueryExpr::Aggregate { aggs, .. } + if matches!(aggs.as_slice(), [AggIntent::Count { .. }]))) + } + + #[test] + fn promotes_count_ranked_limit_sort() { + // Limit 5 { Sort DESC by count-col (1) { Aggregate[Count] by [1] } }. + let q = limit(5, 0, sort(desc(1), count_by_service())); + assert!(is_topk_over_count(&canonicalize(q))); + } + + #[test] + fn promotes_through_a_passthrough_projection() { + // …with a `SELECT service, count` projection between the Sort and the Agg. + let proj = QueryExpr::Project { + cols: vec![ + ProjectItem { alias: None, expr: L3Expr::Column(0) }, + ProjectItem { alias: Some("c".into()), expr: L3Expr::Column(1) }, + ], + qualifier: None, + child: Box::new(count_by_service()), + }; + let q = limit(5, 0, sort(desc(1), proj)); + assert!(is_topk_over_count(&canonicalize(q))); + } + + #[test] + fn is_idempotent() { + let q = limit(5, 0, sort(desc(1), count_by_service())); + let once = canonicalize(q); + let twice = canonicalize(once.clone()); + assert_eq!(once, twice, "canonicalize must be idempotent"); + } + + #[test] + fn does_not_promote_ascending_sort() { + let asc = vec![SortKey { expr: L3Expr::Column(1), ascending: true, nulls_first: false }]; + let q = limit(5, 0, sort(asc, count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q))); + } + + #[test] + fn does_not_promote_with_offset() { + let q = limit(5, 2, sort(desc(1), count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q))); + } + + #[test] + fn does_not_promote_ranking_by_a_group_key() { + // DESC by col 0 (the `service` group key), not the count → not a + // frequency heavy-hitter. + let q = limit(5, 0, sort(desc(0), count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q))); + } + + #[test] + fn does_not_promote_non_count_aggregate() { + let sum = QueryExpr::Aggregate { + by: GroupKeys(vec![1]), + aggs: vec![AggIntent::Sum { col: None }], + output_names: vec![], + having: None, + child: Box::new(scan()), + }; + let q = limit(5, 0, sort(desc(1), sum)); + assert!(!is_topk_over_count(&canonicalize(q))); + } +} diff --git a/crates/l2/src/lib.rs b/crates/l2/src/lib.rs index c1360b15..5102519b 100644 --- a/crates/l2/src/lib.rs +++ b/crates/l2/src/lib.rs @@ -11,11 +11,13 @@ //! pull this lowering machinery. pub mod binder; +pub mod canonicalize; pub mod column_resolution; pub mod lower; pub mod relational; pub use binder::{Binder, SchemaCatalog, UsageDerivedCatalog}; +pub use canonicalize::canonicalize; pub use column_resolution::{ infer_schema_for_root, infer_source_schema, output_schema_for_aggregate, resolve_column_ref, resolve_column_refs, resolve_expr, ResolveError, diff --git a/crates/l2/src/lower.rs b/crates/l2/src/lower.rs index 78f5d128..6a5e72fc 100644 --- a/crates/l2/src/lower.rs +++ b/crates/l2/src/lower.rs @@ -63,7 +63,10 @@ pub fn convert_root( accuracy: &AccuracyTarget, ) -> Result { let fallback = Binder::new().bind(legacy); - convert(legacy, &fallback, accuracy) + let l3 = convert(legacy, &fallback, accuracy)?; + // Both language front ends end here, so this is the one place to normalize + // structural differences between equivalent queries (issue #34). + Ok(crate::canonicalize::canonicalize(l3)) } /// Convert a Layer-2 tree to canonical L3. diff --git a/crates/lower/tests/cross_language.rs b/crates/lower/tests/cross_language.rs new file mode 100644 index 00000000..71b6dba9 --- /dev/null +++ b/crates/lower/tests/cross_language.rs @@ -0,0 +1,128 @@ +//! Cross-language L3 equivalence (issue #34). +//! +//! Semantically equivalent SQL and PromQL queries must lower to the **same +//! canonical intent algebra**, so an L4 rule matching on `AggIntent` sees one +//! spelling regardless of source language. These tests are the executable spec +//! for the shared [`canonicalize`](asap_l2::canonicalize) pass: they pin the +//! canonical heavy-hitter shape and assert both front ends reach it. +//! +//! A literal `lower_sql(S) == lower_promql(P)` cannot hold — the two count +//! *different* things (SQL rows vs. time-series samples over a window) and read +//! from different sources, so their leaves differ by design (#25). What must +//! match is the **shape above the leaf**: an outer `Aggregate([TopK{k}])` over an +//! explicit inner `Aggregate([Count])`. + +use asap_ir::intent_algebra::schema::{Column, DataType, Schema}; +use asap_ir::intent_algebra::{AggIntent, GroupKeys, QueryExpr}; +use asap_ir::types::AccuracyTarget; +use asap_lower::{lower_promql, lower_sql, SqlCatalog}; + +fn col(name: &str, dtype: DataType) -> Column { + Column::new(name, dtype, false) +} + +fn catalog() -> SqlCatalog { + SqlCatalog::new().with_table( + "metrics", + Schema::with_time_index( + vec![ + col("ts", DataType::Timestamp), + col("service", DataType::Utf8), + col("region", DataType::Utf8), + col("latency", DataType::Float64), + col("bytes", DataType::Int64), + ], + 0, + vec![], + ), + ) +} + +async fn sql(q: &str) -> QueryExpr { + lower_sql(q, &catalog(), AccuracyTarget::Exact) + .await + .unwrap_or_else(|e| panic!("SQL {q:?} failed to lower: {e:?}")) +} + +fn promql(q: &str) -> QueryExpr { + lower_promql(q, AccuracyTarget::Exact) + .unwrap_or_else(|e| panic!("PromQL {q:?} failed to lower: {e:?}")) +} + +/// The canonical heavy-hitter shape: an outer `Aggregate([TopK{k}])` (grouped by +/// `by`) over an inner `Aggregate([Count])`. Returns `(k, outer_by)`. +fn heavy_hitter(qe: &QueryExpr) -> Option<(usize, GroupKeys)> { + let QueryExpr::Aggregate { + by, aggs, child, .. + } = qe + else { + return None; + }; + let [AggIntent::TopK { k, .. }] = aggs.as_slice() else { + return None; + }; + // The child must be the explicit inner Count (not a raw Scan) — this is the + // structural unification #25 asked for. + let QueryExpr::Aggregate { aggs: inner, .. } = child.as_ref() else { + return None; + }; + matches!(inner.as_slice(), [AggIntent::Count { .. }]).then(|| (*k, by.clone())) +} + +#[tokio::test] +async fn sql_and_promql_heavy_hitter_share_the_canonical_shape() { + // S2 (SQL, global count-topk) and P1 (PromQL, global count-topk) both express + // "top-5 by count". They must reach the same canonical shape: outer global + // TopK{5} (by: []) over an explicit inner Count. + let s2 = sql( + "SELECT service, COUNT(*) FROM metrics GROUP BY service ORDER BY COUNT(*) DESC LIMIT 5", + ) + .await; + let p1 = promql("topk(5, count_over_time(http_requests_total[5m]))"); + + let (sk, sby) = heavy_hitter(&s2).expect("SQL S2 is a canonical heavy-hitter"); + let (pk, pby) = heavy_hitter(&p1).expect("PromQL P1 is a canonical heavy-hitter"); + assert_eq!(sk, 5); + assert_eq!(pk, 5); + assert!(sby.is_empty(), "S2 is a global topk"); + assert!(pby.is_empty(), "P1 is a global topk"); +} + +#[tokio::test] +async fn sql_aliased_and_inline_count_topk_are_identical() { + // #20: aliasing the COUNT in the ORDER BY must not change the L3. + let inline = + sql("SELECT service, COUNT(*) FROM metrics GROUP BY service ORDER BY COUNT(*) DESC LIMIT 5") + .await; + let aliased = sql( + "SELECT service, COUNT(*) AS c FROM metrics GROUP BY service ORDER BY c DESC LIMIT 5", + ) + .await; + assert_eq!(inline, aliased); +} + +#[tokio::test] +async fn non_count_ranked_topk_stays_generic_in_both_languages() { + // Ranking by a non-count (SUM / a raw value) is NOT a frequency heavy-hitter + // in either language — it must stay a generic Sort+Limit, never a TopK. + let s = sql( + "SELECT service, SUM(bytes) FROM metrics GROUP BY service ORDER BY SUM(bytes) DESC LIMIT 5", + ) + .await; + let p = promql("topk(5, http_requests_total)"); + assert!(heavy_hitter(&s).is_none(), "SUM-ranked is not a heavy-hitter: {s:?}"); + assert!( + !matches!(&s, QueryExpr::Aggregate { aggs, .. } if matches!(aggs.as_slice(), [AggIntent::TopK { .. }])), + ); + assert!(heavy_hitter(&p).is_none(), "value-ranked topk is not a count heavy-hitter: {p:?}"); +} + +#[tokio::test] +async fn offset_defeats_heavy_hitter_promotion() { + // `LIMIT k OFFSET n` is not "the top k" — it must stay a Sort+Limit. + let s = sql( + "SELECT service, COUNT(*) FROM metrics GROUP BY service ORDER BY COUNT(*) DESC LIMIT 5 OFFSET 3", + ) + .await; + assert!(heavy_hitter(&s).is_none(), "OFFSET must not promote to TopK: {s:?}"); +}