From 6a42124ac591250a2e961f139986588e2b2c7350 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:23:49 +0000 Subject: [PATCH] fix(physical): exact top-k answers keep the chosen series' identity PromQL topk and bottomk answer whole input series. Stage 2's exact top-k implementation projected the limited rows to the logical [keys.., topk_k] columns and dropped the series identity, so `topk by (G) (k, ...)` (and its count and rate forms) answered with the G labels only. The Project now keeps $promql_series_identity after the logical columns when the ranked rows carry it: each answer has the chosen series' full label set, and readers of [keys.., topk_k] find those columns where they were. Tests: topk by (region) (2, sum_over_time(data[1m])) over five series answers the four chosen series with their full labels and hand-computed sums, equal to the query-time reference; every topk/bottomk form, grouped or not, over sums, counts and rates answers whole series. The example-1 fixture and its exact top-k shape test now include the column. --- .../tests/planner_layering_example1.rs | 10 +- crates/physical-asap-optimization/src/lib.rs | 17 +- crates/planner/tests/runtime_install.rs | 118 ++++++++ .../examples/planner-layering-example1.json | 266 ++++++++++++++++++ 4 files changed, 408 insertions(+), 3 deletions(-) diff --git a/crates/integration-tests/tests/planner_layering_example1.rs b/crates/integration-tests/tests/planner_layering_example1.rs index 3db57eb64..54b27f858 100644 --- a/crates/integration-tests/tests/planner_layering_example1.rs +++ b/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) }; diff --git a/crates/physical-asap-optimization/src/lib.rs b/crates/physical-asap-optimization/src/lib.rs index ab32d3b65..988dca581 100644 --- a/crates/physical-asap-optimization/src/lib.rs +++ b/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, diff --git a/crates/planner/tests/runtime_install.rs b/crates/planner/tests/runtime_install.rs index c7bcc230f..ac22574eb 100644 --- a/crates/planner/tests/runtime_install.rs +++ b/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 { + 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); + } +} diff --git a/tools/dag-viewer/examples/planner-layering-example1.json b/tools/dag-viewer/examples/planner-layering-example1.json index 3241349d9..e7e36edf9 100644 --- a/tools/dag-viewer/examples/planner-layering-example1.json +++ b/tools/dag-viewer/examples/planner-layering-example1.json @@ -24000,6 +24000,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -24025,6 +24033,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -25589,6 +25603,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -25614,6 +25636,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -27041,6 +27069,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -27066,6 +27102,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -28744,6 +28786,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -28769,6 +28819,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -30449,6 +30505,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -30474,6 +30538,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -32026,6 +32096,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -32051,6 +32129,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -33731,6 +33815,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -33756,6 +33848,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -35434,6 +35532,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -35459,6 +35565,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -37139,6 +37251,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -37164,6 +37284,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -38716,6 +38842,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -38741,6 +38875,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -40421,6 +40561,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -40446,6 +40594,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -42112,6 +42266,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -42137,6 +42299,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -43805,6 +43973,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -43830,6 +44006,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -45498,6 +45680,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -45523,6 +45713,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -47193,6 +47389,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -47218,6 +47422,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -48758,6 +48968,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -48783,6 +49001,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -50451,6 +50675,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -50476,6 +50708,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -52144,6 +52382,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -52169,6 +52415,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null @@ -53713,6 +53965,14 @@ "name": "topk_10", "nullable": false, "table": null + }, + { + "dtype": { + "Plain": "utf8" + }, + "name": "$promql_series_identity", + "nullable": false, + "table": null } ], "time_index": null, @@ -53738,6 +53998,12 @@ "expr": { "Column": 1 } + }, + { + "alias": "$promql_series_identity", + "expr": { + "Column": 3 + } } ], "qualifier": null