fix(promql): keep the input order for limitk and merge label-set pieces instead of sorting

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

src/promql/src/native/execute.rs

@@ -1787,6 +1787,9 @@ struct SelectRuntime {
     /// The layout of the input, which is also what the node produces.
     schema: Arc<SeriesSchema>,
     input: RetainedInput,
+    /// Whether the input's pieces of one label set are merged before the kernel reads them, which
+    /// `limitk` needs when its input can hold them.
+    merge_input: bool,
     /// Whether the kernel has run. It runs at most once: the input it read is gone afterwards.
     computed: bool,
 }
@@ -1797,6 +1800,7 @@ impl SelectRuntime {
             op,
             schema,
             input: RetainedInput::default(),
+            merge_input: false,
             computed: false,
         }
     }
@@ -1825,15 +1829,22 @@ impl SelectRuntime {
                 return Err(error);
             }
         };
-        let mut assembled = assembled;
-        if matches!(self.op, BoundSelect::LimitK { .. }) {
-            // `limitk` keeps the first series in input order. Prometheus reads them in label order
-            // once a query spans its head and a block, while a scan here need not order distinct
-            // series at all, so label order is what makes the choice the same on every run. The
-            // sort is stable, so the pieces of one label set stay in the order they arrived and,
-            // holding disjoint points, are read as the one series they make up.
-            assembled.sort_by(|a, b| select::compare_labels(&a.metric, &b.metric));
-        }
+        // `limitk` keeps the first series in the order its input hands them over, as Prometheus
+        // takes the order of the vector it is given -- a sort below it decides that order. Pieces
+        // of one label set are one series there, placed where the label set first arrived, so
+        // they are merged first; ranked apart, a later piece would be counted after the series
+        // that arrived between them.
+        let assembled = if self.merge_input {
+            let rows = assembled.iter().map(|series| series.time.len()).sum();
+            buffer.grow(merge_workspace_bytes(
+                rows,
+                assembled.len(),
+                self.schema.fields.len(),
+            ))?;
+            merge_same_labelset(&self.schema, assembled)?
+        } else {
+            assembled
+        };
         buffer.grow(selection_workspace_bytes(
             &assembled,
             self.schema.fields.len(),
@@ -2887,8 +2898,23 @@ impl NativeExecutionState {
                         // A selection keeps samples of the series it read, so pieces of one label
                         // set stay pieces: they hold disjoint points, so ranking them apart picks
                         // what ranking the merged series would, and what is above treats them as
-                        // it treats a map's.
-                        holds_shared[index] = input_holds_shared;
+                        // it treats a map's. `limitk` is the exception: it counts series in input
+                        // order, so it merges the pieces first and hands over whole series.
+                        let limitk = matches!(
+                            expr,
+                            NativeExpr::Select {
+                                op: BoundSelect::LimitK { .. },
+                                ..
+                            }
+                        );
+                        if limitk && input_holds_shared {
+                            if let Some(select) = runtimes[index].select.as_mut() {
+                                select.merge_input = true;
+                            }
+                            holds_shared[index] = false;
+                        } else {
+                            holds_shared[index] = input_holds_shared;
+                        }
                     }
                     NodeTag::Window => {
                         let (drops_metric_name, input_schema) = match expr {

src/promql/src/native/execute/tests.rs

@@ -6502,11 +6502,11 @@ fn topk_ranks_whole_series_at_each_point() {
     assert_eq!(pool.reserved(), 0);
 }
 
-/// `limitk` takes series in label order, not in the order they arrived: `{job="x"}` is kept at
-/// both points although `{job="y"}` arrived first. The two series a map folds into `{job="x"}`
-/// stay pieces, each kept where it holds the point.
+/// `limitk` takes series in the order its input hands them over, not in label order, and the two
+/// series a map folds into `{job="z"}` are one series placed where it first arrived: it is kept at
+/// both points although `{job="y"}` sorts first and arrived before the piece that fills point 0.
 #[test]
-fn limitk_takes_series_in_label_order() {
+fn limitk_takes_series_in_input_order() {
     let grid = grid(0, 60_000, 2);
     let schema = evaluated_schema(Arc::clone(&grid));
     let tree = NativeTree::try_new(select(
@@ -6519,6 +6519,7 @@ fn limitk_takes_series_in_label_order() {
     .unwrap();
     let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new());
     for chunk in [
+        evaluated_chunk_of(&schema, named_metric("a", "z"), 1, vec![Some(1.0)], true),
         evaluated_chunk_of(
             &schema,
             named_metric("b", "y"),
@@ -6526,18 +6527,14 @@ fn limitk_takes_series_in_label_order() {
             vec![Some(2.0), Some(2.0)],
             true,
         ),
-        evaluated_chunk_of(&schema, named_metric("a", "x"), 1, vec![Some(1.0)], true),
-        evaluated_chunk_of(&schema, named_metric("c", "x"), 0, vec![Some(3.0)], true),
+        evaluated_chunk_of(&schema, named_metric("c", "z"), 0, vec![Some(3.0)], true),
     ] {
         execution.push_evaluated(chunk).unwrap();
     }
     let chunks = run(&mut execution);
     assert_eq!(
         decided_by_job(&chunks),
-        vec![
-            ("x".to_string(), vec![(1, Some(1.0))]),
-            ("x".to_string(), vec![(0, Some(3.0))]),
-        ]
+        vec![("z".to_string(), vec![(0, Some(3.0)), (1, Some(1.0))])]
     );
 }
 

src/promql/src/select.rs

@@ -542,7 +542,7 @@ fn labels_with_name(metric: &MetricName) -> Vec<(Arc<str>, Arc<str>)> {
 }
 
 /// Prometheus `labels.Compare`: label pairs in name order, then the shorter set first.
-pub(crate) fn compare_labels(a: &MetricName, b: &MetricName) -> Ordering {
+fn compare_labels(a: &MetricName, b: &MetricName) -> Ordering {
     labels_with_name(a).cmp(&labels_with_name(b))
 }