diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 3e379c1bdb6..fb5a5b05c21 100644 --- a/src/promql/src/native/execute.rs +++ b/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, 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 { diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index a9d762ad816..35821903b07 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/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))])] ); } diff --git a/src/promql/src/select.rs b/src/promql/src/select.rs index 6645862e097..c4f22fd57c4 100644 --- a/src/promql/src/select.rs +++ b/src/promql/src/select.rs @@ -542,7 +542,7 @@ fn labels_with_name(metric: &MetricName) -> Vec<(Arc, Arc)> { } /// 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)) }