From 63a7d814a055e4f7782b312587aba242e23ec477 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Thu, 22 Oct 2026 04:27:19 +0800 Subject: [PATCH 1/3] feat(promql): run topk, bottomk, limitk and limit_ratio as a native select node Signed-off-by: Dennis Zhuang --- src/promql/src/native/display.rs | 32 +- src/promql/src/native/display/tests.rs | 48 ++- src/promql/src/native/execute.rs | 486 ++++++++++++++++++++----- src/promql/src/native/execute/tests.rs | 186 +++++++++- src/promql/src/native/metrics.rs | 3 +- src/promql/src/native/plan.rs | 35 ++ src/promql/src/native/plan/tests.rs | 17 +- src/promql/src/select.rs | 10 +- src/promql/src/select/tests.rs | 25 ++ 9 files changed, 731 insertions(+), 111 deletions(-) diff --git a/src/promql/src/native/display.rs b/src/promql/src/native/display.rs index 86d40946cc1..35ccb27542f 100644 --- a/src/promql/src/native/display.rs +++ b/src/promql/src/native/display.rs @@ -28,8 +28,8 @@ use datatypes::series::time::EvalGrid; use crate::map::MathOp; use crate::native::metrics::{self, NodePath}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, - NativeType, + AggregateKind, AggregateStage, BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, + NativeTree, NativeType, }; use crate::range_eval::RangeFunction; use crate::vector_match::{BinaryOp, Cardinality, FillValues, MatchingLabels}; @@ -102,18 +102,21 @@ fn describe(node: &NativeNode) -> String { let ops: Vec = op.ops.iter().map(math_op).collect(); ops.join(", ") } + NativeExpr::Select { op, .. } => match op { + BoundSelect::TopK { k, top, group } => format!( + "{} {k} {}", + if *top { "topk" } else { "bottomk" }, + grouping(group) + ), + BoundSelect::LimitK { k, group } => format!("limitk {k} {}", grouping(group)), + BoundSelect::LimitRatio { ratio } => format!("limit_ratio {ratio}"), + }, NativeExpr::Aggregate { op, .. } => format!( - "{} {} [{}] {}", + "{} {} {}", match op.kind { AggregateKind::Sum => "sum", }, - match &op.group { - Grouping::By(_) => "by", - Grouping::Without(_) => "without", - }, - label_list(match &op.group { - Grouping::By(labels) | Grouping::Without(labels) => labels, - }), + grouping(&op.group), match op.stage { AggregateStage::Complete => "complete", AggregateStage::Partial => "partial", @@ -252,6 +255,14 @@ fn binary_op(op: BinaryOp) -> &'static str { } } +/// A grouping as `by [labels]` or `without [labels]`. +fn grouping(group: &Grouping) -> String { + match group { + Grouping::By(labels) => format!("by [{}]", label_list(labels)), + Grouping::Without(labels) => format!("without [{}]", label_list(labels)), + } +} + /// What kind of node this is, as its line names it. fn kind(expr: &NativeExpr) -> &'static str { match expr { @@ -261,6 +272,7 @@ fn kind(expr: &NativeExpr) -> &'static str { NativeExpr::Window { .. } => "Window", NativeExpr::Map { .. } => "Map", NativeExpr::Aggregate { .. } => "Aggregate", + NativeExpr::Select { .. } => "Select", NativeExpr::Binary { .. } => "Binary", } } diff --git a/src/promql/src/native/display/tests.rs b/src/promql/src/native/display/tests.rs index 316f6a39a4e..46e9098cf5f 100644 --- a/src/promql/src/native/display/tests.rs +++ b/src/promql/src/native/display/tests.rs @@ -20,7 +20,7 @@ use datatypes::series::time::{EvalGrid, GridSpec}; use super::super::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::map::MathOp; use crate::range_eval::{RangeFunction, RangeWindow}; @@ -394,3 +394,49 @@ fn a_without_grouping_does_not_read_as_by() { ]) ); } + +/// A selection names its operator, constant and grouping. +#[test] +fn a_selection_names_its_binding() { + let out = grid(0, 60_000, 4); + let window = window_over( + raw_input(), + RangeFunction::Rate, + 300_000, + 0, + None, + Arc::clone(&out), + ); + let lines = [ + ( + BoundSelect::TopK { + k: 3.0, + top: false, + group: Grouping::Without(Arc::new(["job".into()])), + }, + " root Select bottomk 3 without [job]", + ), + ( + BoundSelect::LimitK { + k: 2.0, + group: Grouping::By(Arc::new([])), + }, + " root Select limitk 2 by []", + ), + ( + BoundSelect::LimitRatio { ratio: -0.5 }, + " root Select limit_ratio -0.5", + ), + ]; + for (op, line) in lines { + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Select { + input: Box::new(window.clone()), + op, + }, + output: evaluated(Arc::clone(&out)), + }) + .unwrap(); + assert_eq!(tree.to_string().lines().nth(1), Some(line)); + } +} diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 7684ee2aa2b..3e379c1bdb6 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -40,7 +40,9 @@ //! operands and runs the matching kernel once the input has ended, because the operation pairs two //! whole vectors point by point and nothing is decided for a point until both sides are complete. //! The chunks of one logical series are assembled into the single series the kernel matches before it -//! runs, so a series that arrived in many chunks is one candidate, not one per chunk. +//! runs, so a series that arrived in many chunks is one candidate, not one per chunk. A selection +//! (`topk`, `limitk`, `limit_ratio`) reads its one input the same way, because which samples it +//! keeps at a point depends on every series present there. //! //! A node kind this runtime does not execute yet is rejected when the execution is built, naming //! the kind and its position. An unsupported tree therefore never starts, rather than answering @@ -96,8 +98,11 @@ use crate::instant_eval::select_instants; use crate::map::{MathOp, drop_metric_name, map_math}; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::{NodeMetrics, NodePath}; -use crate::native::plan::{BoundInstant, NativeExpr, NativeNode, NativeTree}; +use crate::native::plan::{ + BoundInstant, BoundSelect, Grouping, NativeExpr, NativeNode, NativeTree, +}; use crate::range_eval::RangeEvaluator; +use crate::select; use crate::vector_match::{ BinaryMatching, binary_vectors_with_reservation, validate_binary_operand, }; @@ -379,6 +384,8 @@ enum NodeTag { Map, /// Matches two evaluated vectors and applies one operation to every matched pair of series. Binary, + /// Keeps some samples of its whole input vector at each grid point. + Select, } impl NodeTag { @@ -390,6 +397,7 @@ impl NodeTag { Self::Window => "window", Self::Map => "map", Self::Binary => "binary", + Self::Select => "selection", } } } @@ -1243,19 +1251,20 @@ impl InstantRuntime { } } -/// One operand of a binary node, chunk by chunk, until it ends. +/// An input a node reads whole before it computes -- one operand of a binary node, or the input of a +/// selection -- chunk by chunk, until it ends. /// /// What is kept is the series each chunk carries, not a concatenation of them: a chunk is retained /// as it arrives and nothing is copied on the way in, because the chunks of a series can share an -/// allocation and the figure the operand charges for is the data it holds rather than a running copy +/// allocation and the figure the input charges for is the data it holds rather than a running copy /// of it. /// -/// The kernel matches whole series by label set, so the chunks are assembled into one series per -/// logical series once the input has ended: a series that arrived in a thousand chunks is otherwise -/// a thousand candidates on this side for every candidate of the other, and its own rows would be -/// paired a chunk at a time. +/// The kernels read whole series, so the chunks are assembled into one series per logical series +/// once the input has ended: a series that arrived in a thousand chunks is otherwise a thousand +/// candidates for a binary match or a ranking, and its own rows would be read a chunk at a time. +/// #[derive(Default)] -struct BinaryOperand { +struct RetainedInput { /// The series whose last chunk did not close it, with the position after its last row. /// /// `None` when no series is open: nothing has arrived yet, or the last series closed. A label @@ -1296,7 +1305,7 @@ struct BinaryOperand { held: usize, } -impl BinaryOperand { +impl RetainedInput { /// The bytes `chunk` would add to what this operand holds, read without retaining it. /// /// It is what [`Self::enqueue`] adds for the same chunk: what the chunk holds on its own, the @@ -1319,6 +1328,62 @@ impl BinaryOperand { .saturating_add(end) } + /// Retains one chunk after checking that it continues what this input handed over. + /// + /// An input hands over one series at a time, each chunk of it after the rows the series already + /// reached; a chunk that starts another series while one is open, or goes back inside the open + /// one, is refused before it is retained. + fn retain(&mut self, chunk: EvaluatedChunk, reader: Reader<'_>) -> Result<()> { + // An empty chunk closes a series without holding a row, so it says nothing about where that + // series ended: only a chunk that carries rows can be read as going backwards. + let (first, last) = ( + chunk.series.time.first_position(), + chunk.series.time.last_position(), + ); + let metric = chunk.series.metric.clone(); + // The position a chunk of the open series may start at: the row after the last one that + // series handed over. Nothing is open after a series closed, so the next one starts over. + let end = self.open.as_ref().map_or(0, |(_, end)| *end); + if let Some((open, _)) = &self.open + && *open != metric + { + let (owner, input) = (reader.owner(), reader.input()); + return Err(DataFusionError::Execution(format!( + "{owner} reads whole series, and {input} handed over {metric:?} before {open:?} \ + ended" + ))); + } + if let Some(first) = first + && first < end + { + let (owner, input) = (reader.owner(), reader.input()); + return Err(DataFusionError::Execution(format!( + "{owner} was handed chunks of {metric:?} by {input} that do not continue each \ + other: one starts at position {first}, after the series already reached {end}" + ))); + } + self.open = (!chunk.is_last).then(|| (metric, last.map_or(end, |last| last + 1))); + self.enqueue(chunk); + Ok(()) + } + + /// Refuses an input that ended inside a series. + /// + /// An input hands over one series at a time, so the only series that can be open is the last + /// one it handed over: a chunk that did not say `is_last` leaves it open, and nothing came after + /// it to close it. + fn ensure_closed(&self, reader: Reader<'_>) -> Result<()> { + match &self.open { + Some((metric, _)) => Err(DataFusionError::Execution(format!( + "the input of {} ended before the last chunk of a series: {} handed over \ + {metric:?} without the chunk that closes it", + reader.owner(), + reader.input() + ))), + None => Ok(()), + } + } + /// Retains one chunk, folding the estimate it is charged with into what the operand holds. fn enqueue(&mut self, chunk: EvaluatedChunk) { let closing = chunk.is_last; @@ -1349,7 +1414,7 @@ impl BinaryOperand { /// The chunk ranges of the logical series this operand closed, in arrival order. /// /// Read after the input ended: a series that is still open has no end, and - /// [`BinaryRuntime::closed`] refuses such an input before anything is assembled. + /// [`Self::ensure_closed`] refuses such an input before anything is assembled. fn boundaries(&self) -> impl Iterator> + '_ { self.ends.iter().scan(0, |start, &end| { let range = *start..end; @@ -1426,24 +1491,25 @@ impl BinaryOperand { /// the kernel matches and reads, so it must be on the schema's grid and hold as many rows as its /// axis places, which is what [`Series::try_new`] checks. What it allocates is /// [`Self::assembly_bytes`], which a caller reserves before this runs. - fn assemble(&self, schema: &SeriesSchema, path: &NodePath, side: usize) -> Result> { + fn assemble(&self, schema: &SeriesSchema, reader: Reader<'_>) -> Result> { let grid = schema.eval_grid.clone().ok_or_else(|| { + let (owner, input) = (reader.owner(), reader.input()); DataFusionError::Execution(format!( - "the binary node at {path} assembles its {} operand on an evaluation grid, and the \ - operand's schema names none", - side_name(side) + "{owner} assembles {input} on an evaluation grid, and its schema names none" )) })?; let mut assembled = Vec::with_capacity(self.ends.len()); for boundary in self.boundaries() { let chunks = &self.series[boundary]; - let metric = chunks.first().map(|chunk| chunk.metric.clone()).ok_or_else(|| { - DataFusionError::Execution(format!( - "the binary node at {path} closed a logical series of its {} operand without a \ - chunk to assemble", - side_name(side) - )) - })?; + let metric = chunks + .first() + .map(|chunk| chunk.metric.clone()) + .ok_or_else(|| { + let (owner, input) = (reader.owner(), reader.input()); + DataFusionError::Execution(format!( + "{owner} closed a logical series of {input} without a chunk to assemble" + )) + })?; let runs = chunks .iter() .map(|chunk| chunk.time.runs().count()) @@ -1459,10 +1525,10 @@ impl BinaryOperand { let mut segments = Vec::with_capacity(chunks.len()); for chunk in chunks { let field = chunk.fields.get(index).ok_or_else(|| { + let (owner, input) = (reader.owner(), reader.input()); DataFusionError::Execution(format!( - "the binary node at {path} cannot assemble its {} operand: a chunk of \ - it carries {} fields, not {}", - side_name(side), + "{owner} cannot assemble {input}: a chunk of it carries {} fields, not \ + {}", chunk.fields.len(), schema.fields.len() )) @@ -1507,7 +1573,7 @@ struct BinaryRuntime { /// operand and `1` the right one, the order the tree declares them in. schemas: [Arc; 2], /// The retained chunks of each operand, in the same order. - operands: [BinaryOperand; 2], + operands: [RetainedInput; 2], /// Whether the kernel has run. It runs at most once: the operands it read are gone afterwards. computed: bool, } @@ -1520,13 +1586,38 @@ fn side_name(index: usize) -> &'static str { } } +/// The node a [`RetainedInput`] belongs to, which its messages name; formatted only when one is. +#[derive(Clone, Copy)] +enum Reader<'a> { + /// The operand at `side` of the binary node at `path`. + Operand { path: &'a NodePath, side: usize }, + /// The input of the selection at the path. + Select(&'a NodePath), +} + +impl Reader<'_> { + fn owner(self) -> String { + match self { + Self::Operand { path, .. } => format!("the binary node at {path}"), + Self::Select(path) => format!("the selection at {path}"), + } + } + + fn input(self) -> String { + match self { + Self::Operand { side, .. } => format!("its {} operand", side_name(side)), + Self::Select(_) => "its input".to_string(), + } + } +} + impl BinaryRuntime { /// The run-time state of the binary operation `matching` over its two operands. fn new(matching: Arc, schemas: [Arc; 2]) -> Self { Self { matching, schemas, - operands: [BinaryOperand::default(), BinaryOperand::default()], + operands: [RetainedInput::default(), RetainedInput::default()], computed: false, } } @@ -1552,8 +1643,8 @@ impl BinaryRuntime { let [left, right] = &self.operands; let [left_schema, right_schema] = &self.schemas; Ok([ - left.assemble(left_schema, path, 0)?, - right.assemble(right_schema, path, 1)?, + left.assemble(left_schema, Reader::Operand { path, side: 0 })?, + right.assemble(right_schema, Reader::Operand { path, side: 1 })?, ]) } @@ -1571,20 +1662,11 @@ impl BinaryRuntime { /// Reports an operand that ended inside a series. /// - /// A side hands over one series at a time, so the only series that can be open is the last one - /// it handed over: a chunk that did not say `is_last` leaves it open, and nothing came after it - /// to close it. The kernel is not run over such an operand: it would answer the pairs of a - /// series whose remaining rows were never read. + /// The kernel is not run over such an operand: it would answer the pairs of a series whose + /// remaining rows were never read. fn closed(&self, path: &NodePath) -> Result<()> { for (side, operand) in self.operands.iter().enumerate() { - let Some((metric, _)) = &operand.open else { - continue; - }; - return Err(DataFusionError::Execution(format!( - "the input of the binary node at {path} ended before the last chunk of a series: \ - its {} operand handed over {metric:?} without the chunk that closes it", - side_name(side) - ))); + operand.ensure_closed(Reader::Operand { path, side })?; } Ok(()) } @@ -1607,41 +1689,10 @@ impl BinaryRuntime { )) }, )?; - // An empty chunk closes a series without holding a row, so it says nothing about where that - // series ended: only a chunk that carries rows can be read as going backwards. - let (first, last) = ( - chunk.series.time.first_position(), - chunk.series.time.last_position(), - ); - let metric = chunk.series.metric.clone(); let Some(operand) = self.operands.get_mut(side) else { return Err(no_binary(path)); }; - // The position a chunk of the open series may start at: the row after the last one that - // series handed over. Nothing is open after a series closed, so the next one starts over. - let end = operand.open.as_ref().map_or(0, |(_, end)| *end); - if let Some((open, _)) = &operand.open - && *open != metric - { - return Err(DataFusionError::Execution(format!( - "the binary node at {path} combines whole series, and its {} operand handed over \ - {metric:?} before {open:?} ended", - side_name(side) - ))); - } - if let Some(first) = first - && first < end - { - return Err(DataFusionError::Execution(format!( - "the chunks of {metric:?} handed to the {} operand of the binary node at {path} do \ - not continue each other: one starts at position {first}, after the series already \ - reached {end}", - side_name(side) - ))); - } - operand.open = (!chunk.is_last).then(|| (metric, last.map_or(end, |last| last + 1))); - operand.enqueue(chunk); - Ok(()) + operand.retain(chunk, Reader::Operand { path, side }) } /// Runs the kernel over both operands and returns the series it produced. @@ -1725,6 +1776,143 @@ impl BinaryRuntime { } } +/// What a selection node holds: its whole input until it ends. +/// +/// Which samples `topk` or `limitk` keeps at a point depends on every series present there, so, +/// like a binary node, this node retains every chunk its input hands over, assembles the chunks of +/// each logical series, and runs its kernel once, after the input has ended. The results are queued +/// in the node's own output queue and handed over from there. +struct SelectRuntime { + op: BoundSelect, + /// The layout of the input, which is also what the node produces. + schema: Arc, + input: RetainedInput, + /// Whether the kernel has run. It runs at most once: the input it read is gone afterwards. + computed: bool, +} + +impl SelectRuntime { + fn new(op: BoundSelect, schema: Arc) -> Self { + Self { + op, + schema, + input: RetainedInput::default(), + computed: false, + } + } + + /// Runs the kernel over the whole input and returns the series it kept, or `None` when it + /// already ran. + /// + /// The reservation covers the retained chunks, the assembled series and a rough estimate of the kernel's workspace and of the series it rebuilds + /// ([`selection_workspace_bytes`]). The kernels do not report their workspace as they allocate, + /// so the estimate is reserved once, before the kernel runs. + fn compute(&mut self, path: &NodePath, buffer: &mut NodeBuffer) -> Result>> { + if self.computed { + return Ok(None); + } + self.input.ensure_closed(Reader::Select(path))?; + let retained = self.input.held; + let assembly = self + .input + .assembly_bytes(&self.schema) + .ok_or_else(|| select_overflow(path))?; + buffer.grow(assembly)?; + let assembled = match self.input.assemble(&self.schema, Reader::Select(path)) { + Ok(assembled) => assembled, + Err(error) => { + buffer.resize(retained)?; + 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)); + } + buffer.grow(selection_workspace_bytes( + &assembled, + self.schema.fields.len(), + ))?; + let metrics = Arc::clone(&buffer.metrics); + let output = { + let _timer = metrics.timer(); + select_series(&self.op, &self.schema, &assembled)? + }; + drop(assembled); + self.input.release(); + self.computed = true; + Ok(Some(output)) + } +} + +/// The error for a selection whose reservation does not fit an address space. +fn select_overflow(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!( + "the input of the selection at {path} does not fit one address space: the bytes it would \ + reserve overflow" + )) +} + +/// Runs the kernel `op` binds over the whole input. +fn select_series( + op: &BoundSelect, + schema: &SeriesSchema, + series: &[Series], +) -> Result> { + let grouping = |group: &Grouping| match group { + Grouping::By(labels) => select::Grouping::By(labels.to_vec()), + Grouping::Without(labels) => select::Grouping::Without(labels.to_vec()), + }; + match op { + BoundSelect::TopK { k, top, group } => { + select::topk(*k, *top, &grouping(group), schema, series) + } + BoundSelect::LimitK { k, group } => select::limitk(*k, &grouping(group), schema, series), + BoundSelect::LimitRatio { ratio } => select::limit_ratio(*ratio, schema, series), + } +} + +/// A rough estimate of what a selection kernel allocates over `series`, beside the input it reads. +/// +/// Per row, the `(point, row)` a kept sample is recorded as, the candidate it is ranked as and the +/// value each field rebuilds it into; per series, its place in the label order and its group, its +/// label pairs in the group key, its cursor and the series and field descriptors of its result; and +/// the histograms `limitk` and `limit_ratio` copy into their results, at the widest histogram of +/// each field. Like [`merge_workspace_bytes`] it is not an allocation bound: vector slack and the +/// hash table's own layout are not measured. +fn selection_workspace_bytes(series: &[Series], fields: usize) -> usize { + let row = std::mem::size_of::<(usize, usize)>() + + std::mem::size_of::<(usize, usize, f64)>() + + fields * std::mem::size_of::(); + let per_series = 2 * std::mem::size_of::() + + std::mem::size_of::>() + + std::mem::size_of::>() + + std::mem::size_of::() + + fields * std::mem::size_of::(); + let label = std::mem::size_of::<(Arc, Arc)>(); + series.iter().fold(0usize, |total, item| { + let rows = item.time.len(); + let labels = item.metric.tags().len() + 1; + let histograms = item + .fields + .iter() + .filter_map(histogram_row_bytes) + .fold(0usize, |bytes, widest| { + bytes.saturating_add(widest.saturating_mul(rows)) + }); + total + .saturating_add(rows.saturating_mul(row)) + .saturating_add(per_series) + .saturating_add(labels.saturating_mul(label)) + .saturating_add(histograms) + }) +} + /// The buffered chunks as the one source axis they make up: their timestamps in arrival order. /// /// The axis type validates the result as strictly increasing, which is also what refuses a series @@ -1873,7 +2061,8 @@ struct NodeRuntime { /// The chunks this node holds for its consumer, each with the estimate it was queued with. /// /// An input node's queue is filled by the caller, a window's by the chunks it computed while - /// flushing its collected input, and a binary node's by the results its kernel produced; no node + /// flushing its collected input, and a binary or selection node's by the results its kernel + /// produced; no node /// uses it for both. A chunk is measured once, as it enters, and the figure it was queued with /// is given back as it leaves, so neither the queue nor the handover walks it. queued: VecDeque<(NativeChunk, usize)>, @@ -1887,6 +2076,8 @@ struct NodeRuntime { map: Option, /// What a binary node holds. binary: Option, + /// What a selection node holds. + select: Option, /// The layout the caller's chunks of an input node are checked against. /// /// A tree with a binary node reads two inputs, and each of them is bound to its own layout, so @@ -1936,10 +2127,10 @@ impl NodeRuntime { /// Hands over the next queued chunk, if there is one, giving back what it was queued with. /// /// It is the one way a chunk leaves this node's queue: an input node hands over what the caller - /// queued, and a binary node hands over what its kernel produced. `values` is how that node - /// counts the values it hands over, because the two do not count the same thing: an input's - /// figure has always been its positions, while a binary result holds one value per field and a - /// field may have none at a position it decided. + /// queued, and a binary or selection node hands over what its kernel produced. `values` is how + /// that node counts the values it hands over, because the two do not count the same thing: an + /// input's figure has always been its positions, while a kernel's result holds one value per + /// field and a field may have none at a position it decided. fn pop_queued(&mut self, values: fn(&NativeChunk) -> usize) -> Result> { let Some((chunk, bytes)) = self.queued.pop_front() else { return Ok(None); @@ -2135,6 +2326,56 @@ impl NodeRuntime { let Some(output) = output else { return Ok(()); }; + self.queue_results(output) + } + + /// Retains one chunk of a selection's input, reserved for before it is retained, as + /// [`Self::retain_binary`] retains an operand's. + fn retain_select(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { + let NativeChunk::Evaluated(chunk) = chunk else { + return Err(DataFusionError::Execution(format!( + "the selection at {path} consumes evaluated chunks, not raw samples" + ))); + }; + let select = self.select.as_mut().ok_or_else(|| no_select(path))?; + // What a bound tree says the input produces is not proof of what a chunk carries, so the + // chunk is checked against the layout before it is retained. + validate_binary_operand(select.schema.as_ref(), std::slice::from_ref(&chunk.series)) + .map_err(|error| { + DataFusionError::Execution(format!( + "the selection at {path} rejected its input: {error}" + )) + })?; + let pending = select.input.pending(&chunk); + let positions = chunk.series.time.len(); + self.buffer.grow(pending)?; + let metrics = Arc::clone(&self.buffer.metrics); + { + let _timer = metrics.timer(); + let select = self.select.as_mut().ok_or_else(|| no_select(path))?; + select.input.retain(chunk, Reader::Select(path))?; + let retained = select.input.held; + self.buffer.resize(retained)?; + } + metrics.record_input(1, positions); + Ok(()) + } + + /// Runs this node's selection kernel and queues the series it kept. + fn compute_select(&mut self, path: &NodePath) -> Result<()> { + let output = { + let select = self.select.as_mut().ok_or_else(|| no_select(path))?; + select.compute(path, &mut self.buffer)? + }; + let Some(output) = output else { + return Ok(()); + }; + self.queue_results(output) + } + + /// Queues the series a kernel that reads its whole input produced, and settles the reservation + /// on the queue. + fn queue_results(&mut self, output: Vec) -> Result<()> { // The results are queued the way a node that produces a chunk per input queues what it // hands on: the queue is the one source of chunk order and of the bytes queued chunks hold. let mut closed = 0usize; @@ -2142,19 +2383,19 @@ impl NodeRuntime { // The kernel decided every point of a result series at once, so the one chunk it is // handed over as closes it. let chunk = NativeChunk::Evaluated(EvaluatedChunk::try_new(series, true)?); - // The result is already admitted: the kernel kept it inside the workspace its callback - // reserved. What the queue keeps is the estimate of the chunk that holds it. + // The result is already admitted, within the workspace reserved for the kernel. What the + // queue keeps is the estimate of the chunk that holds it. let bytes = chunk.bytes(); self.queued.push_back((chunk, bytes)); closed += 1; } - // This node's own accounting of closed series is the results it closed, one per matched - // series: an operand chunk that closed its input series is a closure of the node below, and + // This node's own accounting of closed series is the results it closed, one per result + // series: an input chunk that closed its series is a closure of the node below, and // counting it here as well would report the same work twice. self.buffer.metrics.record_completed_series(closed); - // The operands the reservation covered are gone and the queue holds what the kernel - // returned, so the reservation settles on the queue: what a queued result is given back - // with as it leaves is exactly the figure it is held for here. + // The input the reservation covered is gone and the queue holds what the kernel returned, + // so the reservation settles on the queue: what a queued result is given back with as it + // leaves is exactly the figure it is held for here. let queued = self .queued .iter() @@ -2354,6 +2595,7 @@ impl NodeRuntime { self.window = None; self.instant = None; self.binary = None; + self.select = None; self.source_schema = None; self.emitted = None; self.merge = None; @@ -2442,6 +2684,7 @@ impl NativeExecutionState { // types than the kernel takes, so the shape is checked where the operands' layouts // are read. NativeExpr::Binary { .. } => NodeTag::Binary, + NativeExpr::Select { .. } => NodeTag::Select, // The kinds below are legal in a bound tree, and `validate` stays the authority on // that; this runtime simply does not execute them yet. Rejected before a single // chunk is read: a tree the runtime cannot execute must not start, because a @@ -2457,6 +2700,7 @@ impl NativeExecutionState { window: None, map: None, binary: None, + select: None, source_schema: None, emitted: None, merge: None, @@ -2542,6 +2786,21 @@ impl NativeExecutionState { // produced at points that do not collide; the kernel groups and merges them // itself, which is why this node keeps no ledger of what it handed out. } + NativeExpr::Select { op, .. } => { + let schema = node.output.samples().cloned().ok_or_else(|| { + DataFusionError::Execution(format!( + "the selection at {path} does not produce samples" + )) + })?; + // `topk`, `bottomk` and `limitk` rank one value per sample, so their kernels + // take one field; `limit_ratio` keeps whole series and takes any number. + if !matches!(op, BoundSelect::LimitRatio { .. }) && schema.fields.len() != 1 { + return Err(not_implemented("multi-field Select", &path)); + } + runtime.select = Some(SelectRuntime::new(op.clone(), schema)); + // A selection keeps the labels of the series it keeps, so it keeps no ledger of + // what it handed out. + } NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { // Rejected above, where the pass has the position to name. } @@ -2556,7 +2815,8 @@ impl NativeExecutionState { } NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } - | NativeExpr::Map { input, .. } => { + | NativeExpr::Map { input, .. } + | NativeExpr::Select { input, .. } => { pending.push((input.as_ref(), path.input())); true } @@ -2623,6 +2883,13 @@ impl NativeExecutionState { // them until the input ends -- so its flag is false whatever theirs are. holds_shared[index] = false; } + NodeTag::Select => { + // 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; + } NodeTag::Window => { let (drops_metric_name, input_schema) = match expr { NativeExpr::Window { input, op } => ( @@ -2798,12 +3065,12 @@ impl NativeExecutionState { /// Reports a stream that ended inside a series, without closing any node. /// - /// Closing is destructive -- a window's evaluator takes itself apart -- and the binary nodes - /// still have to compute after this check, so what is read here is only the state that follows - /// what each node was handed: an input node reports the series it handed over without the chunk - /// that ends it, and a binary node reports an operand that is still open on either side. A - /// truncation is therefore refused before a binary kernel allocates its workspace, rather than - /// after one has already been computed on half a vector. + /// Closing is destructive -- a window's evaluator takes itself apart -- and the binary and + /// selection nodes still have to compute after this check, so what is read here is only the + /// state that follows what each node was handed: an input node reports the series it handed + /// over without the chunk that ends it, and a binary or selection node reports an input that is + /// still open. A truncation is therefore refused before a kernel allocates its workspace, rather + /// than after one has already been computed on half a vector. fn unclosed(&self) -> Result<()> { // A source series that never ended is refused before anything is merged: the close a merge // synthesises would make it look like one that did. @@ -2814,6 +3081,9 @@ impl NativeExecutionState { if let Some(binary) = &runtime.binary { binary.closed(path)?; } + if let Some(select) = &runtime.select { + select.input.ensure_closed(Reader::Select(path))?; + } } Ok(()) } @@ -2934,7 +3204,8 @@ impl NativeExecutor { /// is materialised on their edge -- except below a window that collects its input, which holds /// that input until [`Self::finish`]. A binary node is the other exception: it pairs two whole /// vectors point by point, so it retains every chunk each operand hands over and reports `None` - /// until [`Self::finish`] ends the input, which is when it computes its result. + /// until [`Self::finish`] ends the input, which is when it computes its result. A selection does + /// the same with its one input. /// /// A chunk's `is_last` says the series whose data it carries ended. It does not say its label /// set is finished: two source series whose labels were rewritten into one set are two pieces @@ -3116,6 +3387,15 @@ impl NativeExecutor { NodeTag::Map | NodeTag::Window | NodeTag::Instant => self.state.nodes[index] .input .is_none_or(|child| self.settled(child)), + NodeTag::Select => { + runtime + .select + .as_ref() + .is_some_and(|select| select.computed) + && self.state.nodes[index] + .input + .is_none_or(|child| self.settled(child)) + } NodeTag::Binary => { runtime .binary @@ -3283,6 +3563,22 @@ impl NativeExecutor { // does, so its values are counted rather than read as its positions. runtime.pop_queued(evaluated_values) } + NodeTag::Select => { + let child = input.ok_or_else(|| missing_input(tag))?; + // Which samples are kept at a point is not decided until the whole input has + // arrived, so the input is read to the end before anything is handed over. + while let Some(chunk) = self.pull(child)? { + let path = &self.state.paths[index]; + self.state.runtimes[index].retain_select(path, chunk)?; + } + if !self.inputs_finished || !self.settled(child) { + return Ok(None); + } + let path = &self.state.paths[index]; + let runtime = &mut self.state.runtimes[index]; + runtime.compute_select(path)?; + runtime.pop_queued(evaluated_values) + } } } @@ -3576,6 +3872,10 @@ fn no_binary(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no binary operation")) } +fn no_select(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!("the node at {path} holds no selection")) +} + /// The error for a reservation the platform cannot express. /// /// The binary node reserves its retained operands and the kernel's live workspace together, and diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index f31cf86be64..a9d762ad816 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -45,7 +45,7 @@ use crate::map::MathOp; use crate::native::metrics::{NodeMetric, NodeMetrics, NodePath}; use crate::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::vector_match::{ @@ -482,6 +482,7 @@ fn source_runtime(pool: &Arc) -> NodeRuntime { window: None, map: None, binary: None, + select: None, source_schema: None, emitted: None, merge: None, @@ -501,6 +502,7 @@ fn binary_runtime(pool: &Arc, schema: &Arc) -> Nod Arc::new(binary_matching(BinaryOp::Add)), [Arc::clone(schema), Arc::clone(schema)], )), + select: None, source_schema: None, emitted: None, merge: None, @@ -6125,7 +6127,8 @@ fn an_operand_chunk_that_rewinds_is_refused_with_or_without_values() { let message = error.to_string(); assert!( message.contains("do not continue each other") - && message.contains("the left operand of the binary node at root") + && message.contains("the binary node at root") + && message.contains("by its left operand") && message.contains("after the series already reached 2"), "{message}" ); @@ -6417,3 +6420,182 @@ fn an_all_missing_match_group_fits_its_pool() { assert_eq!(pool.reserved(), 0); assert_eq!(execution.state().reserved_bytes(), 0); } + +fn select(inner: NativeNode, op: BoundSelect) -> NativeNode { + let output = inner.output.clone(); + NativeNode { + expr: NativeExpr::Select { + input: Box::new(inner), + op, + }, + output, + } +} + +/// The positions a series decided, with their values. +type Decided = Vec<(usize, Option)>; + +/// The values a chunk stream decided, by the `job` label of the series that holds them. +fn decided_by_job(chunks: &[EvaluatedChunk]) -> Vec<(String, Decided)> { + chunks + .iter() + .map(|chunk| { + ( + chunk + .series + .metric + .get("job") + .unwrap_or_default() + .to_string(), + decided(std::slice::from_ref(chunk)), + ) + }) + .collect() +} + +/// `topk` ranks the whole vector at each point, so a series that arrives in several chunks is one +/// candidate, and the reservation is given back once the kept series are handed over. +#[test] +fn topk_ranks_whole_series_at_each_point() { + let grid = grid(0, 60_000, 3); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(select( + evaluated_input(Arc::clone(&grid)), + BoundSelect::TopK { + k: 1.0, + top: true, + group: Grouping::By(Arc::from([])), + }, + )) + .unwrap(); + let pool = pool(); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + for chunk in [ + evaluated_chunk_of( + &schema, + named_metric("m", "a"), + 0, + vec![Some(1.0), Some(5.0)], + false, + ), + evaluated_chunk_of(&schema, named_metric("m", "a"), 2, vec![Some(7.0)], true), + evaluated_chunk_of( + &schema, + named_metric("m", "b"), + 0, + vec![Some(3.0), Some(4.0), Some(6.0)], + true, + ), + ] { + execution.push_evaluated(chunk).unwrap(); + } + let chunks = run(&mut execution); + assert_eq!( + decided_by_job(&chunks), + vec![ + ("a".to_string(), vec![(1, Some(5.0)), (2, Some(7.0))]), + ("b".to_string(), vec![(0, Some(3.0))]), + ] + ); + assert!(chunks.iter().all(|chunk| chunk.is_last)); + assert_eq!(chunks[0].series.metric.name(), Some("m")); + 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. +#[test] +fn limitk_takes_series_in_label_order() { + let grid = grid(0, 60_000, 2); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(select( + map(evaluated_input(Arc::clone(&grid)), vec![MathOp::Abs]), + BoundSelect::LimitK { + k: 1.0, + group: Grouping::By(Arc::from([])), + }, + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for chunk in [ + evaluated_chunk_of( + &schema, + named_metric("b", "y"), + 0, + 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), + ] { + 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))]), + ] + ); +} + +/// A binary node above a selection reads what the selection kept, once the selection has read its +/// whole input: `topk` keeps `{job="b"}` at point 0 and `{job="a"}` at point 1, so only point 1 +/// matches the right side's `{job="a"}`. +#[test] +fn a_binary_node_reads_what_a_selection_kept() { + let grid = grid(0, 60_000, 2); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(binary_of( + select( + evaluated_input(Arc::clone(&grid)), + BoundSelect::TopK { + k: 1.0, + top: true, + group: Grouping::By(Arc::from([])), + }, + ), + evaluated_input(Arc::clone(&grid)), + binary_matching(BinaryOp::Add), + )) + .unwrap(); + let (left, right) = (NodePath::root().left().input(), NodePath::root().right()); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for (path, chunk) in [ + ( + &left, + evaluated_chunk_of( + &schema, + named_metric("m", "a"), + 0, + vec![Some(1.0), Some(5.0)], + true, + ), + ), + ( + &right, + evaluated_chunk_of( + &schema, + named_metric("m", "a"), + 0, + vec![Some(10.0), Some(10.0)], + true, + ), + ), + ( + &left, + evaluated_chunk_of( + &schema, + named_metric("m", "b"), + 0, + vec![Some(3.0), Some(4.0)], + true, + ), + ), + ] { + execution.push_evaluated_at(path, chunk).unwrap(); + } + assert_eq!(decided(&run(&mut execution)), vec![(1, Some(15.0))]); +} diff --git a/src/promql/src/native/metrics.rs b/src/promql/src/native/metrics.rs index 6e36153fe2e..6e1a2635d6c 100644 --- a/src/promql/src/native/metrics.rs +++ b/src/promql/src/native/metrics.rs @@ -100,7 +100,8 @@ pub fn walk(root: &NativeNode, visit: &mut dyn FnMut(&NativeNode, &NodePath)) { NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } - | NativeExpr::Aggregate { input, .. } => go(input, &path.input(), visit), + | NativeExpr::Aggregate { input, .. } + | NativeExpr::Select { input, .. } => go(input, &path.input(), visit), NativeExpr::Binary { left, right, .. } => { go(left, &path.left(), visit); go(right, &path.right(), visit); diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index f5715b21f72..42bd6a0f892 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -222,6 +222,22 @@ pub struct BoundAggregate { pub stage: AggregateStage, } +/// An operator that keeps some of the input's samples at each grid point, with their series' labels +/// and values: which samples it keeps depends on the whole vector at that point, not on one series. +/// +/// Its constants are bound here, as a window's are: a `k` or a ratio only known at run time has no +/// field. +#[derive(Clone, Debug, PartialEq)] +pub enum BoundSelect { + /// `topk` (`top`) or `bottomk`: the `k` largest or smallest values of each group. + TopK { k: f64, top: bool, group: Grouping }, + /// `limitk`: the first `k` series of each group that hold a sample, in input order. + LimitK { k: f64, group: Grouping }, + /// `limit_ratio`: the series whose label hash falls in the ratio; it keeps whole series, so it + /// has no grouping. + LimitRatio { ratio: f64 }, +} + /// What a node computes. Every variant holds a bound operation, never a name to resolve. #[derive(Clone, Debug)] pub enum NativeExpr { @@ -249,6 +265,11 @@ pub enum NativeExpr { input: Box, op: BoundAggregate, }, + /// A selection of samples out of the whole input vector. + Select { + input: Box, + op: BoundSelect, + }, /// A binary operation between two inputs, with its matching rules. /// /// Scalar operands are broadcast, so `matching` only applies to vector–vector operations. @@ -390,6 +411,20 @@ impl NativeNode { ensure_same_fields(&input.output, &self.output, "a map")?; ensure_grid(&self.output, input_grid, "a map")?; } + NativeExpr::Select { input, .. } => { + input.validate()?; + // A selection hands on samples of its input unchanged, so it consumes and produces + // the same evaluated vector. + let Some((_, input_grid)) = evaluated_vector(&input.output) else { + return Err(invalid(&format!( + "a selection consumes an evaluated vector, not {}", + input.output.describe() + ))); + }; + ensure_samples(&self.output, "a selection")?; + ensure_same_fields(&input.output, &self.output, "a selection")?; + ensure_grid(&self.output, input_grid, "a selection")?; + } NativeExpr::Aggregate { input, op } => { input.validate()?; // A stage's contract is a pair of types, not only what it produces: `Partial` diff --git a/src/promql/src/native/plan/tests.rs b/src/promql/src/native/plan/tests.rs index ba01eb53cdd..f4f24a99dc7 100644 --- a/src/promql/src/native/plan/tests.rs +++ b/src/promql/src/native/plan/tests.rs @@ -20,7 +20,7 @@ use datatypes::series::time::{EvalGrid, GridSpec}; use super::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::map::MathOp; use crate::range_eval::{RangeFunction, RangeWindow}; @@ -1218,3 +1218,18 @@ fn an_input_can_bind_aggregate_state() { let error = other.validate().unwrap_err(); assert!(error.to_string().contains("not the binding"), "{error}"); } + +/// A selection ranks the values at each grid point, so raw samples, which have no shared points, +/// are not something it can consume. +#[test] +fn a_selection_cannot_consume_raw_samples() { + let node = NativeNode { + expr: NativeExpr::Select { + input: Box::new(input()), + op: BoundSelect::LimitRatio { ratio: 0.5 }, + }, + output: evaluated(Arc::new([]), grid(0, 60_000, 4)), + }; + let error = node.validate().unwrap_err(); + assert!(error.to_string().contains("an evaluated vector"), "{error}"); +} diff --git a/src/promql/src/select.rs b/src/promql/src/select.rs index c65a7955c36..6645862e097 100644 --- a/src/promql/src/select.rs +++ b/src/promql/src/select.rs @@ -59,7 +59,8 @@ pub enum SortOrder { /// `topk(k, ...)` (`top`) or `bottomk(k, ...)` with a constant `k`. /// /// At every point each group keeps the `k` largest (or smallest) values, NaN ranking last for -/// both; the selected samples keep their series' labels and values. Prometheus breaks ties by +/// both; the selected samples keep their series' labels and values. Histogram samples are not +/// ranked, as Prometheus ignores them here. Prometheus breaks ties by /// the order its storage returns series in, which is not a property of the series; here ties /// go to the series whose label set sorts first. pub fn topk( @@ -102,7 +103,10 @@ pub fn topk( let Some(row) = series[member].time.row_of(point) else { continue; }; - if let Some(value) = cursor_value(&mut cursors[slot], row)? { + if let SampleRef::Float(value) = cursors[slot] + .at(row) + .expect("the axis and the fields are aligned") + { candidates.push((member, row, value)); } } @@ -538,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. -fn compare_labels(a: &MetricName, b: &MetricName) -> Ordering { +pub(crate) fn compare_labels(a: &MetricName, b: &MetricName) -> Ordering { labels_with_name(a).cmp(&labels_with_name(b)) } diff --git a/src/promql/src/select/tests.rs b/src/promql/src/select/tests.rs index 3bac6739214..4433020e2c2 100644 --- a/src/promql/src/select/tests.rs +++ b/src/promql/src/select/tests.rs @@ -269,6 +269,31 @@ fn limit_operators_keep_histogram_samples() { assert_eq!(kinds(&kept[0], 0), [(1, "histogram")]); } +/// `topk` ranks floats only: a histogram sample is skipped, as Prometheus ignores it, and the +/// float below it is kept in its place. +#[test] +fn topk_skips_histogram_samples() { + let histogram = + OwnedHistogram::parse("0/0/0/3ff0000000000000/3ff0000000000000/0/1_1/3ff0000000000000///"); + let schema = schema(1, 1); + let mut builder = GridSeriesBuilder::try_new(Arc::new(schema.clone())).unwrap(); + builder + .push_samples( + 0, + &[Some(GridSample::Histogram(Box::new( + FloatHistogram::try_from_view(&histogram.view()).unwrap(), + )))], + ) + .unwrap(); + let series = vec![ + builder.finish(metric("__name__=a")).unwrap().unwrap(), + build(&schema, metric("__name__=b"), &vec![vec![Some(1.0)]]).unwrap(), + ]; + let kept = topk(1.0, true, &Grouping::By(Vec::new()), &schema, &series).unwrap(); + assert_eq!(kept.len(), 1); + assert_eq!(kept[0].metric.name(), Some("b")); +} + /// A selection hands back the sample it selected, start timestamp included. /// /// `limitk` and `limit_ratio` choose samples rather than compute them, so a histogram that arrives From bec5f76637828eeacd9b410be7bc59035472361d Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Thu, 22 Oct 2026 13:39:12 +0800 Subject: [PATCH 2/3] fix(promql): keep the input order for limitk and merge label-set pieces instead of sorting Signed-off-by: Dennis Zhuang --- src/promql/src/native/execute.rs | 48 ++++++++++++++++++++------ src/promql/src/native/execute/tests.rs | 17 ++++----- src/promql/src/select.rs | 2 +- 3 files changed, 45 insertions(+), 22 deletions(-) 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)) } From 90478ec4931c032d31c97bcdcd508c1a269f8069 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Thu, 22 Oct 2026 13:46:11 +0800 Subject: [PATCH 3/3] refactor(promql): name the whole-vector node Vector Signed-off-by: Dennis Zhuang --- src/promql/src/native/display.rs | 12 +- src/promql/src/native/display/tests.rs | 20 ++-- src/promql/src/native/execute.rs | 153 ++++++++++++------------- src/promql/src/native/execute/tests.rs | 26 ++--- src/promql/src/native/metrics.rs | 2 +- src/promql/src/native/plan.rs | 20 ++-- src/promql/src/native/plan/tests.rs | 10 +- 7 files changed, 120 insertions(+), 123 deletions(-) diff --git a/src/promql/src/native/display.rs b/src/promql/src/native/display.rs index 35ccb27542f..bb3c20181e7 100644 --- a/src/promql/src/native/display.rs +++ b/src/promql/src/native/display.rs @@ -28,7 +28,7 @@ use datatypes::series::time::EvalGrid; use crate::map::MathOp; use crate::native::metrics::{self, NodePath}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, + AggregateKind, AggregateStage, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::range_eval::RangeFunction; @@ -102,14 +102,14 @@ fn describe(node: &NativeNode) -> String { let ops: Vec = op.ops.iter().map(math_op).collect(); ops.join(", ") } - NativeExpr::Select { op, .. } => match op { - BoundSelect::TopK { k, top, group } => format!( + NativeExpr::Vector { op, .. } => match op { + BoundVectorOp::TopK { k, top, group } => format!( "{} {k} {}", if *top { "topk" } else { "bottomk" }, grouping(group) ), - BoundSelect::LimitK { k, group } => format!("limitk {k} {}", grouping(group)), - BoundSelect::LimitRatio { ratio } => format!("limit_ratio {ratio}"), + BoundVectorOp::LimitK { k, group } => format!("limitk {k} {}", grouping(group)), + BoundVectorOp::LimitRatio { ratio } => format!("limit_ratio {ratio}"), }, NativeExpr::Aggregate { op, .. } => format!( "{} {} {}", @@ -272,7 +272,7 @@ fn kind(expr: &NativeExpr) -> &'static str { NativeExpr::Window { .. } => "Window", NativeExpr::Map { .. } => "Map", NativeExpr::Aggregate { .. } => "Aggregate", - NativeExpr::Select { .. } => "Select", + NativeExpr::Vector { .. } => "Vector", NativeExpr::Binary { .. } => "Binary", } } diff --git a/src/promql/src/native/display/tests.rs b/src/promql/src/native/display/tests.rs index 46e9098cf5f..61953e7c36a 100644 --- a/src/promql/src/native/display/tests.rs +++ b/src/promql/src/native/display/tests.rs @@ -20,7 +20,7 @@ use datatypes::series::time::{EvalGrid, GridSpec}; use super::super::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::map::MathOp; use crate::range_eval::{RangeFunction, RangeWindow}; @@ -395,9 +395,9 @@ fn a_without_grouping_does_not_read_as_by() { ); } -/// A selection names its operator, constant and grouping. +/// A vector operation names its operator, constant and grouping. #[test] -fn a_selection_names_its_binding() { +fn a_vector_operation_names_its_binding() { let out = grid(0, 60_000, 4); let window = window_over( raw_input(), @@ -409,28 +409,28 @@ fn a_selection_names_its_binding() { ); let lines = [ ( - BoundSelect::TopK { + BoundVectorOp::TopK { k: 3.0, top: false, group: Grouping::Without(Arc::new(["job".into()])), }, - " root Select bottomk 3 without [job]", + " root Vector bottomk 3 without [job]", ), ( - BoundSelect::LimitK { + BoundVectorOp::LimitK { k: 2.0, group: Grouping::By(Arc::new([])), }, - " root Select limitk 2 by []", + " root Vector limitk 2 by []", ), ( - BoundSelect::LimitRatio { ratio: -0.5 }, - " root Select limit_ratio -0.5", + BoundVectorOp::LimitRatio { ratio: -0.5 }, + " root Vector limit_ratio -0.5", ), ]; for (op, line) in lines { let tree = NativeTree::try_new(NativeNode { - expr: NativeExpr::Select { + expr: NativeExpr::Vector { input: Box::new(window.clone()), op, }, diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index fb5a5b05c21..f847e2d0f1a 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -40,7 +40,7 @@ //! operands and runs the matching kernel once the input has ended, because the operation pairs two //! whole vectors point by point and nothing is decided for a point until both sides are complete. //! The chunks of one logical series are assembled into the single series the kernel matches before it -//! runs, so a series that arrived in many chunks is one candidate, not one per chunk. A selection +//! runs, so a series that arrived in many chunks is one candidate, not one per chunk. A vector operation //! (`topk`, `limitk`, `limit_ratio`) reads its one input the same way, because which samples it //! keeps at a point depends on every series present there. //! @@ -99,7 +99,7 @@ use crate::map::{MathOp, drop_metric_name, map_math}; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::{NodeMetrics, NodePath}; use crate::native::plan::{ - BoundInstant, BoundSelect, Grouping, NativeExpr, NativeNode, NativeTree, + BoundInstant, BoundVectorOp, Grouping, NativeExpr, NativeNode, NativeTree, }; use crate::range_eval::RangeEvaluator; use crate::select; @@ -385,7 +385,7 @@ enum NodeTag { /// Matches two evaluated vectors and applies one operation to every matched pair of series. Binary, /// Keeps some samples of its whole input vector at each grid point. - Select, + Vector, } impl NodeTag { @@ -397,7 +397,7 @@ impl NodeTag { Self::Window => "window", Self::Map => "map", Self::Binary => "binary", - Self::Select => "selection", + Self::Vector => "vector operation", } } } @@ -1252,7 +1252,7 @@ impl InstantRuntime { } /// An input a node reads whole before it computes -- one operand of a binary node, or the input of a -/// selection -- chunk by chunk, until it ends. +/// vector operation -- chunk by chunk, until it ends. /// /// What is kept is the series each chunk carries, not a concatenation of them: a chunk is retained /// as it arrives and nothing is copied on the way in, because the chunks of a series can share an @@ -1591,22 +1591,22 @@ fn side_name(index: usize) -> &'static str { enum Reader<'a> { /// The operand at `side` of the binary node at `path`. Operand { path: &'a NodePath, side: usize }, - /// The input of the selection at the path. - Select(&'a NodePath), + /// The input of the vector operation at the path. + Vector(&'a NodePath), } impl Reader<'_> { fn owner(self) -> String { match self { Self::Operand { path, .. } => format!("the binary node at {path}"), - Self::Select(path) => format!("the selection at {path}"), + Self::Vector(path) => format!("the vector operation at {path}"), } } fn input(self) -> String { match self { Self::Operand { side, .. } => format!("its {} operand", side_name(side)), - Self::Select(_) => "its input".to_string(), + Self::Vector(_) => "its input".to_string(), } } } @@ -1776,14 +1776,14 @@ impl BinaryRuntime { } } -/// What a selection node holds: its whole input until it ends. +/// What a vector node holds: its whole input until it ends. /// /// Which samples `topk` or `limitk` keeps at a point depends on every series present there, so, /// like a binary node, this node retains every chunk its input hands over, assembles the chunks of /// each logical series, and runs its kernel once, after the input has ended. The results are queued /// in the node's own output queue and handed over from there. -struct SelectRuntime { - op: BoundSelect, +struct VectorRuntime { + op: BoundVectorOp, /// The layout of the input, which is also what the node produces. schema: Arc, input: RetainedInput, @@ -1794,8 +1794,8 @@ struct SelectRuntime { computed: bool, } -impl SelectRuntime { - fn new(op: BoundSelect, schema: Arc) -> Self { +impl VectorRuntime { + fn new(op: BoundVectorOp, schema: Arc) -> Self { Self { op, schema, @@ -1809,20 +1809,20 @@ impl SelectRuntime { /// already ran. /// /// The reservation covers the retained chunks, the assembled series and a rough estimate of the kernel's workspace and of the series it rebuilds - /// ([`selection_workspace_bytes`]). The kernels do not report their workspace as they allocate, + /// ([`vector_workspace_bytes`]). The kernels do not report their workspace as they allocate, /// so the estimate is reserved once, before the kernel runs. fn compute(&mut self, path: &NodePath, buffer: &mut NodeBuffer) -> Result>> { if self.computed { return Ok(None); } - self.input.ensure_closed(Reader::Select(path))?; + self.input.ensure_closed(Reader::Vector(path))?; let retained = self.input.held; let assembly = self .input .assembly_bytes(&self.schema) - .ok_or_else(|| select_overflow(path))?; + .ok_or_else(|| vector_overflow(path))?; buffer.grow(assembly)?; - let assembled = match self.input.assemble(&self.schema, Reader::Select(path)) { + let assembled = match self.input.assemble(&self.schema, Reader::Vector(path)) { Ok(assembled) => assembled, Err(error) => { buffer.resize(retained)?; @@ -1845,14 +1845,11 @@ impl SelectRuntime { } else { assembled }; - buffer.grow(selection_workspace_bytes( - &assembled, - self.schema.fields.len(), - ))?; + buffer.grow(vector_workspace_bytes(&assembled, self.schema.fields.len()))?; let metrics = Arc::clone(&buffer.metrics); let output = { let _timer = metrics.timer(); - select_series(&self.op, &self.schema, &assembled)? + vector_series(&self.op, &self.schema, &assembled)? }; drop(assembled); self.input.release(); @@ -1861,17 +1858,17 @@ impl SelectRuntime { } } -/// The error for a selection whose reservation does not fit an address space. -fn select_overflow(path: &NodePath) -> DataFusionError { +/// The error for a vector operation whose reservation does not fit an address space. +fn vector_overflow(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!( - "the input of the selection at {path} does not fit one address space: the bytes it would \ + "the input of the vector operation at {path} does not fit one address space: the bytes it would \ reserve overflow" )) } /// Runs the kernel `op` binds over the whole input. -fn select_series( - op: &BoundSelect, +fn vector_series( + op: &BoundVectorOp, schema: &SeriesSchema, series: &[Series], ) -> Result> { @@ -1880,15 +1877,15 @@ fn select_series( Grouping::Without(labels) => select::Grouping::Without(labels.to_vec()), }; match op { - BoundSelect::TopK { k, top, group } => { + BoundVectorOp::TopK { k, top, group } => { select::topk(*k, *top, &grouping(group), schema, series) } - BoundSelect::LimitK { k, group } => select::limitk(*k, &grouping(group), schema, series), - BoundSelect::LimitRatio { ratio } => select::limit_ratio(*ratio, schema, series), + BoundVectorOp::LimitK { k, group } => select::limitk(*k, &grouping(group), schema, series), + BoundVectorOp::LimitRatio { ratio } => select::limit_ratio(*ratio, schema, series), } } -/// A rough estimate of what a selection kernel allocates over `series`, beside the input it reads. +/// A rough estimate of what a vector kernel allocates over `series`, beside the input it reads. /// /// Per row, the `(point, row)` a kept sample is recorded as, the candidate it is ranked as and the /// value each field rebuilds it into; per series, its place in the label order and its group, its @@ -1896,7 +1893,7 @@ fn select_series( /// the histograms `limitk` and `limit_ratio` copy into their results, at the widest histogram of /// each field. Like [`merge_workspace_bytes`] it is not an allocation bound: vector slack and the /// hash table's own layout are not measured. -fn selection_workspace_bytes(series: &[Series], fields: usize) -> usize { +fn vector_workspace_bytes(series: &[Series], fields: usize) -> usize { let row = std::mem::size_of::<(usize, usize)>() + std::mem::size_of::<(usize, usize, f64)>() + fields * std::mem::size_of::(); @@ -2072,7 +2069,7 @@ struct NodeRuntime { /// The chunks this node holds for its consumer, each with the estimate it was queued with. /// /// An input node's queue is filled by the caller, a window's by the chunks it computed while - /// flushing its collected input, and a binary or selection node's by the results its kernel + /// flushing its collected input, and a binary or vector node's by the results its kernel /// produced; no node /// uses it for both. A chunk is measured once, as it enters, and the figure it was queued with /// is given back as it leaves, so neither the queue nor the handover walks it. @@ -2087,8 +2084,8 @@ struct NodeRuntime { map: Option, /// What a binary node holds. binary: Option, - /// What a selection node holds. - select: Option, + /// What a vector node holds. + vector: Option, /// The layout the caller's chunks of an input node are checked against. /// /// A tree with a binary node reads two inputs, and each of them is bound to its own layout, so @@ -2138,7 +2135,7 @@ impl NodeRuntime { /// Hands over the next queued chunk, if there is one, giving back what it was queued with. /// /// It is the one way a chunk leaves this node's queue: an input node hands over what the caller - /// queued, and a binary or selection node hands over what its kernel produced. `values` is how + /// queued, and a binary or vector node hands over what its kernel produced. `values` is how /// that node counts the values it hands over, because the two do not count the same thing: an /// input's figure has always been its positions, while a kernel's result holds one value per /// field and a field may have none at a position it decided. @@ -2340,43 +2337,43 @@ impl NodeRuntime { self.queue_results(output) } - /// Retains one chunk of a selection's input, reserved for before it is retained, as + /// Retains one chunk of a vector node's input, reserved for before it is retained, as /// [`Self::retain_binary`] retains an operand's. - fn retain_select(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { + fn retain_vector(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { let NativeChunk::Evaluated(chunk) = chunk else { return Err(DataFusionError::Execution(format!( - "the selection at {path} consumes evaluated chunks, not raw samples" + "the vector operation at {path} consumes evaluated chunks, not raw samples" ))); }; - let select = self.select.as_mut().ok_or_else(|| no_select(path))?; + let vector = self.vector.as_mut().ok_or_else(|| no_vector(path))?; // What a bound tree says the input produces is not proof of what a chunk carries, so the // chunk is checked against the layout before it is retained. - validate_binary_operand(select.schema.as_ref(), std::slice::from_ref(&chunk.series)) + validate_binary_operand(vector.schema.as_ref(), std::slice::from_ref(&chunk.series)) .map_err(|error| { DataFusionError::Execution(format!( - "the selection at {path} rejected its input: {error}" + "the vector operation at {path} rejected its input: {error}" )) })?; - let pending = select.input.pending(&chunk); + let pending = vector.input.pending(&chunk); let positions = chunk.series.time.len(); self.buffer.grow(pending)?; let metrics = Arc::clone(&self.buffer.metrics); { let _timer = metrics.timer(); - let select = self.select.as_mut().ok_or_else(|| no_select(path))?; - select.input.retain(chunk, Reader::Select(path))?; - let retained = select.input.held; + let vector = self.vector.as_mut().ok_or_else(|| no_vector(path))?; + vector.input.retain(chunk, Reader::Vector(path))?; + let retained = vector.input.held; self.buffer.resize(retained)?; } metrics.record_input(1, positions); Ok(()) } - /// Runs this node's selection kernel and queues the series it kept. - fn compute_select(&mut self, path: &NodePath) -> Result<()> { + /// Runs this node's vector kernel and queues the series it kept. + fn compute_vector(&mut self, path: &NodePath) -> Result<()> { let output = { - let select = self.select.as_mut().ok_or_else(|| no_select(path))?; - select.compute(path, &mut self.buffer)? + let vector = self.vector.as_mut().ok_or_else(|| no_vector(path))?; + vector.compute(path, &mut self.buffer)? }; let Some(output) = output else { return Ok(()); @@ -2606,7 +2603,7 @@ impl NodeRuntime { self.window = None; self.instant = None; self.binary = None; - self.select = None; + self.vector = None; self.source_schema = None; self.emitted = None; self.merge = None; @@ -2695,7 +2692,7 @@ impl NativeExecutionState { // types than the kernel takes, so the shape is checked where the operands' layouts // are read. NativeExpr::Binary { .. } => NodeTag::Binary, - NativeExpr::Select { .. } => NodeTag::Select, + NativeExpr::Vector { .. } => NodeTag::Vector, // The kinds below are legal in a bound tree, and `validate` stays the authority on // that; this runtime simply does not execute them yet. Rejected before a single // chunk is read: a tree the runtime cannot execute must not start, because a @@ -2711,7 +2708,7 @@ impl NativeExecutionState { window: None, map: None, binary: None, - select: None, + vector: None, source_schema: None, emitted: None, merge: None, @@ -2797,18 +2794,18 @@ impl NativeExecutionState { // produced at points that do not collide; the kernel groups and merges them // itself, which is why this node keeps no ledger of what it handed out. } - NativeExpr::Select { op, .. } => { + NativeExpr::Vector { op, .. } => { let schema = node.output.samples().cloned().ok_or_else(|| { DataFusionError::Execution(format!( - "the selection at {path} does not produce samples" + "the vector operation at {path} does not produce samples" )) })?; // `topk`, `bottomk` and `limitk` rank one value per sample, so their kernels // take one field; `limit_ratio` keeps whole series and takes any number. - if !matches!(op, BoundSelect::LimitRatio { .. }) && schema.fields.len() != 1 { - return Err(not_implemented("multi-field Select", &path)); + if !matches!(op, BoundVectorOp::LimitRatio { .. }) && schema.fields.len() != 1 { + return Err(not_implemented("multi-field Vector", &path)); } - runtime.select = Some(SelectRuntime::new(op.clone(), schema)); + runtime.vector = Some(VectorRuntime::new(op.clone(), schema)); // A selection keeps the labels of the series it keeps, so it keeps no ledger of // what it handed out. } @@ -2827,7 +2824,7 @@ impl NativeExecutionState { NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } - | NativeExpr::Select { input, .. } => { + | NativeExpr::Vector { input, .. } => { pending.push((input.as_ref(), path.input())); true } @@ -2894,7 +2891,7 @@ impl NativeExecutionState { // them until the input ends -- so its flag is false whatever theirs are. holds_shared[index] = false; } - NodeTag::Select => { + NodeTag::Vector => { // 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 @@ -2902,14 +2899,14 @@ impl NativeExecutionState { // order, so it merges the pieces first and hands over whole series. let limitk = matches!( expr, - NativeExpr::Select { - op: BoundSelect::LimitK { .. }, + NativeExpr::Vector { + op: BoundVectorOp::LimitK { .. }, .. } ); if limitk && input_holds_shared { - if let Some(select) = runtimes[index].select.as_mut() { - select.merge_input = true; + if let Some(vector) = runtimes[index].vector.as_mut() { + vector.merge_input = true; } holds_shared[index] = false; } else { @@ -3092,9 +3089,9 @@ impl NativeExecutionState { /// Reports a stream that ended inside a series, without closing any node. /// /// Closing is destructive -- a window's evaluator takes itself apart -- and the binary and - /// selection nodes still have to compute after this check, so what is read here is only the + /// vector nodes still have to compute after this check, so what is read here is only the /// state that follows what each node was handed: an input node reports the series it handed - /// over without the chunk that ends it, and a binary or selection node reports an input that is + /// over without the chunk that ends it, and a binary or vector node reports an input that is /// still open. A truncation is therefore refused before a kernel allocates its workspace, rather /// than after one has already been computed on half a vector. fn unclosed(&self) -> Result<()> { @@ -3107,8 +3104,8 @@ impl NativeExecutionState { if let Some(binary) = &runtime.binary { binary.closed(path)?; } - if let Some(select) = &runtime.select { - select.input.ensure_closed(Reader::Select(path))?; + if let Some(vector) = &runtime.vector { + vector.input.ensure_closed(Reader::Vector(path))?; } } Ok(()) @@ -3230,7 +3227,7 @@ impl NativeExecutor { /// is materialised on their edge -- except below a window that collects its input, which holds /// that input until [`Self::finish`]. A binary node is the other exception: it pairs two whole /// vectors point by point, so it retains every chunk each operand hands over and reports `None` - /// until [`Self::finish`] ends the input, which is when it computes its result. A selection does + /// until [`Self::finish`] ends the input, which is when it computes its result. A vector node does /// the same with its one input. /// /// A chunk's `is_last` says the series whose data it carries ended. It does not say its label @@ -3413,11 +3410,11 @@ impl NativeExecutor { NodeTag::Map | NodeTag::Window | NodeTag::Instant => self.state.nodes[index] .input .is_none_or(|child| self.settled(child)), - NodeTag::Select => { + NodeTag::Vector => { runtime - .select + .vector .as_ref() - .is_some_and(|select| select.computed) + .is_some_and(|vector| vector.computed) && self.state.nodes[index] .input .is_none_or(|child| self.settled(child)) @@ -3589,20 +3586,20 @@ impl NativeExecutor { // does, so its values are counted rather than read as its positions. runtime.pop_queued(evaluated_values) } - NodeTag::Select => { + NodeTag::Vector => { let child = input.ok_or_else(|| missing_input(tag))?; // Which samples are kept at a point is not decided until the whole input has // arrived, so the input is read to the end before anything is handed over. while let Some(chunk) = self.pull(child)? { let path = &self.state.paths[index]; - self.state.runtimes[index].retain_select(path, chunk)?; + self.state.runtimes[index].retain_vector(path, chunk)?; } if !self.inputs_finished || !self.settled(child) { return Ok(None); } let path = &self.state.paths[index]; let runtime = &mut self.state.runtimes[index]; - runtime.compute_select(path)?; + runtime.compute_vector(path)?; runtime.pop_queued(evaluated_values) } } @@ -3898,8 +3895,8 @@ fn no_binary(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no binary operation")) } -fn no_select(path: &NodePath) -> DataFusionError { - DataFusionError::Execution(format!("the node at {path} holds no selection")) +fn no_vector(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!("the node at {path} holds no vector operation")) } /// The error for a reservation the platform cannot express. diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index 35821903b07..1af5da32f20 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -45,7 +45,7 @@ use crate::map::MathOp; use crate::native::metrics::{NodeMetric, NodeMetrics, NodePath}; use crate::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::vector_match::{ @@ -482,7 +482,7 @@ fn source_runtime(pool: &Arc) -> NodeRuntime { window: None, map: None, binary: None, - select: None, + vector: None, source_schema: None, emitted: None, merge: None, @@ -502,7 +502,7 @@ fn binary_runtime(pool: &Arc, schema: &Arc) -> Nod Arc::new(binary_matching(BinaryOp::Add)), [Arc::clone(schema), Arc::clone(schema)], )), - select: None, + vector: None, source_schema: None, emitted: None, merge: None, @@ -6421,10 +6421,10 @@ fn an_all_missing_match_group_fits_its_pool() { assert_eq!(execution.state().reserved_bytes(), 0); } -fn select(inner: NativeNode, op: BoundSelect) -> NativeNode { +fn vector_op(inner: NativeNode, op: BoundVectorOp) -> NativeNode { let output = inner.output.clone(); NativeNode { - expr: NativeExpr::Select { + expr: NativeExpr::Vector { input: Box::new(inner), op, }, @@ -6459,9 +6459,9 @@ fn decided_by_job(chunks: &[EvaluatedChunk]) -> Vec<(String, Decided)> { fn topk_ranks_whole_series_at_each_point() { let grid = grid(0, 60_000, 3); let schema = evaluated_schema(Arc::clone(&grid)); - let tree = NativeTree::try_new(select( + let tree = NativeTree::try_new(vector_op( evaluated_input(Arc::clone(&grid)), - BoundSelect::TopK { + BoundVectorOp::TopK { k: 1.0, top: true, group: Grouping::By(Arc::from([])), @@ -6509,9 +6509,9 @@ fn topk_ranks_whole_series_at_each_point() { 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( + let tree = NativeTree::try_new(vector_op( map(evaluated_input(Arc::clone(&grid)), vec![MathOp::Abs]), - BoundSelect::LimitK { + BoundVectorOp::LimitK { k: 1.0, group: Grouping::By(Arc::from([])), }, @@ -6538,17 +6538,17 @@ fn limitk_takes_series_in_input_order() { ); } -/// A binary node above a selection reads what the selection kept, once the selection has read its +/// A binary node above a vector node reads what `topk` kept, once the vector node has read its /// whole input: `topk` keeps `{job="b"}` at point 0 and `{job="a"}` at point 1, so only point 1 /// matches the right side's `{job="a"}`. #[test] -fn a_binary_node_reads_what_a_selection_kept() { +fn a_binary_node_reads_what_a_vector_node_kept() { let grid = grid(0, 60_000, 2); let schema = evaluated_schema(Arc::clone(&grid)); let tree = NativeTree::try_new(binary_of( - select( + vector_op( evaluated_input(Arc::clone(&grid)), - BoundSelect::TopK { + BoundVectorOp::TopK { k: 1.0, top: true, group: Grouping::By(Arc::from([])), diff --git a/src/promql/src/native/metrics.rs b/src/promql/src/native/metrics.rs index 6e1a2635d6c..c894c9607b8 100644 --- a/src/promql/src/native/metrics.rs +++ b/src/promql/src/native/metrics.rs @@ -101,7 +101,7 @@ pub fn walk(root: &NativeNode, visit: &mut dyn FnMut(&NativeNode, &NodePath)) { | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } | NativeExpr::Aggregate { input, .. } - | NativeExpr::Select { input, .. } => go(input, &path.input(), visit), + | NativeExpr::Vector { input, .. } => go(input, &path.input(), visit), NativeExpr::Binary { left, right, .. } => { go(left, &path.left(), visit); go(right, &path.right(), visit); diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 42bd6a0f892..8cc1cdfd7be 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -228,7 +228,7 @@ pub struct BoundAggregate { /// Its constants are bound here, as a window's are: a `k` or a ratio only known at run time has no /// field. #[derive(Clone, Debug, PartialEq)] -pub enum BoundSelect { +pub enum BoundVectorOp { /// `topk` (`top`) or `bottomk`: the `k` largest or smallest values of each group. TopK { k: f64, top: bool, group: Grouping }, /// `limitk`: the first `k` series of each group that hold a sample, in input order. @@ -265,10 +265,10 @@ pub enum NativeExpr { input: Box, op: BoundAggregate, }, - /// A selection of samples out of the whole input vector. - Select { + /// An operation over the whole input vector at each grid point. + Vector { input: Box, - op: BoundSelect, + op: BoundVectorOp, }, /// A binary operation between two inputs, with its matching rules. /// @@ -411,19 +411,19 @@ impl NativeNode { ensure_same_fields(&input.output, &self.output, "a map")?; ensure_grid(&self.output, input_grid, "a map")?; } - NativeExpr::Select { input, .. } => { + NativeExpr::Vector { input, .. } => { input.validate()?; - // A selection hands on samples of its input unchanged, so it consumes and produces + // A vector operation keeps the input's fields and grid, so it consumes and produces // the same evaluated vector. let Some((_, input_grid)) = evaluated_vector(&input.output) else { return Err(invalid(&format!( - "a selection consumes an evaluated vector, not {}", + "a vector operation consumes an evaluated vector, not {}", input.output.describe() ))); }; - ensure_samples(&self.output, "a selection")?; - ensure_same_fields(&input.output, &self.output, "a selection")?; - ensure_grid(&self.output, input_grid, "a selection")?; + ensure_samples(&self.output, "a vector operation")?; + ensure_same_fields(&input.output, &self.output, "a vector operation")?; + ensure_grid(&self.output, input_grid, "a vector operation")?; } NativeExpr::Aggregate { input, op } => { input.validate()?; diff --git a/src/promql/src/native/plan/tests.rs b/src/promql/src/native/plan/tests.rs index f4f24a99dc7..f4d14b58b9c 100644 --- a/src/promql/src/native/plan/tests.rs +++ b/src/promql/src/native/plan/tests.rs @@ -20,7 +20,7 @@ use datatypes::series::time::{EvalGrid, GridSpec}; use super::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundSelect, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::map::MathOp; use crate::range_eval::{RangeFunction, RangeWindow}; @@ -1219,14 +1219,14 @@ fn an_input_can_bind_aggregate_state() { assert!(error.to_string().contains("not the binding"), "{error}"); } -/// A selection ranks the values at each grid point, so raw samples, which have no shared points, +/// A vector operation reads the values at each grid point, so raw samples, which have no shared points, /// are not something it can consume. #[test] -fn a_selection_cannot_consume_raw_samples() { +fn a_vector_operation_cannot_consume_raw_samples() { let node = NativeNode { - expr: NativeExpr::Select { + expr: NativeExpr::Vector { input: Box::new(input()), - op: BoundSelect::LimitRatio { ratio: 0.5 }, + op: BoundVectorOp::LimitRatio { ratio: 0.5 }, }, output: evaluated(Arc::new([]), grid(0, 60_000, 4)), };