fix: exact top-k answers keep the chosen series' labels

crates/integration-tests/tests/planner_layering_example1.rs

@@ -411,7 +411,8 @@ fn stage2_physical_asap_optimization_exact_topk_is_sort_then_limit() {
 }
 
 /// Exact Q2 sorts each job's series by the ranked value, descending, and a
-/// `Project` over the limit restores the logical `[job, topk_10]` columns.
+/// `Project` over the limit restores the logical `[job, topk_10]` columns,
+/// then keeps each chosen series' identity (its full label set).
 #[test]
 fn stage2_physical_asap_optimization_exact_topk_returns_job_and_topk_10() {
     let run = pipeline();
@@ -420,7 +421,12 @@ fn stage2_physical_asap_optimization_exact_topk_returns_job_and_topk_10() {
         panic!("{}: Q2 is an operator root", p.id)
     };
     let names: Vec<_> = root.schema.fields.iter().map(|f| f.name.as_str()).collect();
-    assert_eq!(names, ["job", "topk_10"], "{}", p.id);
+    assert_eq!(
+        names,
+        ["job", "topk_10", "$promql_series_identity"],
+        "{}",
+        p.id
+    );
     let NonASAPOp::Project { child: limit, .. } = root.expect_non_asap() else {
         panic!("{}: Q2's root is a Project", p.id)
     };

crates/physical-asap-optimization/src/lib.rs

@@ -73,7 +73,9 @@
 //! per-group limit, the shape the PromQL frontend uses for generic `topk`.
 //! A `Project` over the limited rows restores the logical `[keys.., topk_k]`
 //! columns, so an exact root has the same column names as a summary
-//! candidate's and as the frontend declares, and no input column leaks.
+//! candidate's and as the frontend declares. Over PromQL rows it also keeps
+//! the series identity, after them: `topk` answers whole input series, so
+//! each answer carries the chosen series' full label set.
 //! `topk_k` carries the ranked value's type and nullability (the frontend
 //! declares a non-null `Utf8` placeholder); casting the value to `Utf8` would
 //! not make it the same answer.
@@ -765,6 +767,19 @@ fn implement_topk(
         alias: Some(ranked.name.clone()),
         expr: ScalarExpr::Column(value),
     });
+    // PromQL `topk` answers whole input series: rows that carry the series
+    // identity keep it, after the logical columns, so each answer has the
+    // chosen series' full label set and readers of `[keys.., topk_k]` still
+    // find those columns where they were.
+    if let Ok(identity) = resolve_column_ref(
+        &ColumnRef::Named(asap_types::ir::schema::PROMQL_SERIES_IDENTITY.into()),
+        &child.schema,
+    ) {
+        cols.push(ProjectItem {
+            alias: Some(asap_types::ir::schema::PROMQL_SERIES_IDENTITY.into()),
+            expr: ScalarExpr::Column(identity),
+        });
+    }
     Ok(Rc::new(
         OperatorNode::new(Operator::NonASAP(NonASAPOp::Project {
             cols,

crates/planner/tests/runtime_install.rs

@@ -1032,3 +1032,121 @@ async fn panes_of_an_hour_share_one_definition_with_distinct_times() {
         );
     }
 }
+
+/// PromQL `topk` returns whole input series: `topk by (region) (2,
+/// sum_over_time(data[1m]))` answers, per region, the two series with the
+/// largest sums, each with its own full label set (region, service and
+/// instance; `sum_over_time` drops the metric name), not only `region`.
+/// Series `e1`..`e3` are in `east` with sums 6, 60 and 30; `w1`, `w2` in
+/// `west` with sums 9 and 3.
+#[tokio::test]
+async fn topk_answers_carry_the_chosen_series_labels() {
+    let workload: PlanningWorkload = serde_json::from_str(
+        r#"{"query_workload": {"language": "promql", "repeating_queries": [{
+              "query": "topk by (region) (2, sum_over_time(data[1m]))",
+              "demand": {"fixed_interval": 60000},
+              "requirements": {},
+              "predictability": {"predictable": {"known_at": null}}}]},
+            "data_workload": {"arrival": {"ingesting": {"rows_per_second": 0.5}},
+              "sources": {"data": {"series_count": 5, "sample_interval": 10000}}}}"#,
+    )
+    .unwrap();
+    let plan_json = plan_json_for(workload, &DeploymentCapabilities::default()).await;
+    let samples = topk_samples();
+    let run = run_data_plane(
+        &plan_json,
+        &samples,
+        EVALUATION_MS,
+        &mut SummaryStore::new(),
+    );
+    assert_full_series_answers(&run, 4);
+    let labels = |instance: &str, region: &str, service: &str| {
+        BTreeMap::from([
+            ("instance".to_string(), instance.to_string()),
+            ("region".to_string(), region.to_string()),
+            ("service".to_string(), service.to_string()),
+        ])
+    };
+    let mut answered = run.answers_per_query[0].clone();
+    answered.sort_by(|a, b| a.0.cmp(&b.0));
+    let mut expected = vec![
+        (labels("e2", "east", "web"), 60.0),
+        (labels("e3", "east", "db"), 30.0),
+        (labels("w1", "west", "api"), 9.0),
+        (labels("w2", "west", "web"), 3.0),
+    ];
+    expected.sort_by(|a, b| a.0.cmp(&b.0));
+    assert_eq!(answered.len(), expected.len(), "{answered:?}");
+    for ((labels, value), (want_labels, want)) in answered.iter().zip(&expected) {
+        assert_eq!(labels, want_labels);
+        assert!((value - want).abs() < 1e-9, "{labels:?}: {value}");
+    }
+    assert_eq!(run.answers_per_query, run.reference_answers_per_query);
+}
+
+/// Samples of `data` for `topk_answers_carry_the_chosen_series_labels`:
+/// three per series inside the last minute, each a third of its sum.
+fn topk_samples() -> Vec<Sample> {
+    let mut samples = Vec::new();
+    for (instance, region, service, sum) in [
+        ("e1", "east", "api", 6.0),
+        ("e2", "east", "web", 60.0),
+        ("e3", "east", "db", 30.0),
+        ("w1", "west", "api", 9.0),
+        ("w2", "west", "web", 3.0),
+    ] {
+        // In time order, as a counter's samples arrive.
+        for step in (1..=3).rev() {
+            let mut sample = sample("data", instance, EVALUATION_MS - step * 15_000, sum / 3.0);
+            sample.labels.insert("region".into(), region.into());
+            sample.labels.insert("service".into(), service.into());
+            samples.push(sample);
+        }
+    }
+    samples
+}
+
+/// The first query's `expected` answers, each a whole input series (it
+/// carries its `instance`), and the same as the query-time reference's.
+fn assert_full_series_answers(run: &DataPlaneRun, expected: usize) {
+    let answered = &run.answers_per_query[0];
+    assert_eq!(answered.len(), expected, "{answered:?}");
+    for (labels, _) in answered {
+        assert!(labels.contains_key("instance"), "{labels:?}");
+        assert!(labels.contains_key("service"), "{labels:?}");
+    }
+    assert_eq!(run.answers_per_query, run.reference_answers_per_query);
+}
+
+/// Every top-k and bottom-k form, grouped or not, over sums, counts and
+/// rates, answers whole input series: `topk(2, ...)` the two of five,
+/// `by (region)` per region, `bottomk` likewise.
+#[tokio::test]
+async fn every_topk_form_answers_whole_series() {
+    for (query, expected) in [
+        ("topk(2, sum_over_time(data[1m]))", 2),
+        ("topk by (region) (1, count_over_time(data[1m]))", 2),
+        ("topk by (region) (1, rate(data[1m]))", 2),
+        ("bottomk by (region) (1, sum_over_time(data[1m]))", 2),
+        ("bottomk(3, sum_over_time(data[1m]))", 3),
+    ] {
+        let workload: PlanningWorkload = serde_json::from_str(&format!(
+            r#"{{"query_workload": {{"language": "promql", "repeating_queries": [{{
+                  "query": {query:?},
+                  "demand": {{"fixed_interval": 60000}},
+                  "requirements": {{}},
+                  "predictability": {{"predictable": {{"known_at": null}}}}}}]}},
+                "data_workload": {{"arrival": {{"ingesting": {{"rows_per_second": 0.5}}}},
+                  "sources": {{"data": {{"series_count": 5, "sample_interval": 10000}}}}}}}}"#
+        ))
+        .unwrap();
+        let plan_json = plan_json_for(workload, &DeploymentCapabilities::default()).await;
+        let run = run_data_plane(
+            &plan_json,
+            &topk_samples(),
+            EVALUATION_MS,
+            &mut SummaryStore::new(),
+        );
+        assert_full_series_answers(&run, expected);
+    }
+}