diff --git a/src/promql/src/extension_plan/aggregate_finalize.rs b/src/promql/src/extension_plan/aggregate_finalize.rs index 382ee3e705e..8dcb7394aa1 100644 --- a/src/promql/src/extension_plan/aggregate_finalize.rs +++ b/src/promql/src/extension_plan/aggregate_finalize.rs @@ -466,6 +466,17 @@ fn split(node: &NativeNode, fields: &[FieldSpec], time_index: &str) -> DfResult< }; Ok(Some(inner)) } + // A set operation keeps samples by whether their match group is present on the other side, + // which the merged values of one side cannot answer. + NativeExpr::Set { left, right, .. } => { + match ( + split(left, fields, time_index)?, + split(right, fields, time_index)?, + ) { + (None, None) => Ok(None), + _ => Err(declined("an aggregate under a set operation")), + } + } } } @@ -491,7 +502,7 @@ fn prefix_fields(tree: &NativeTree) -> DfResult> { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => raw_fields(input), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { raw_fields(left).or_else(|| raw_fields(right)) } } diff --git a/src/promql/src/extension_plan/aggregate_finalize/tests.rs b/src/promql/src/extension_plan/aggregate_finalize/tests.rs index 09d156a42f1..3554e6e1572 100644 --- a/src/promql/src/extension_plan/aggregate_finalize/tests.rs +++ b/src/promql/src/extension_plan/aggregate_finalize/tests.rs @@ -545,7 +545,7 @@ fn the_prefix_and_the_suffix_each_name_their_own_source_zero() { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => inputs(input, found), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { inputs(left, found); inputs(right, found); } diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index aebc42ef3e4..60ed584aa8f 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -618,7 +618,7 @@ fn tree_inputs(tree: &NativeTree) -> Vec<&crate::native::plan::BoundInput> { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => collect(input, inputs), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { collect(left, inputs); collect(right, inputs); } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 8ebc9b035b3..417487230c3 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -360,7 +360,7 @@ fn declaring_orderings(tree: &NativeTree, orderings: &[SeriesOrdering]) -> Nativ | NativeExpr::Vector { input, .. } => { **input = rewrite(input, orderings, seen); } - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { **left = rewrite(left, orderings, seen); **right = rewrite(right, orderings, seen); } diff --git a/src/promql/src/native/display.rs b/src/promql/src/native/display.rs index 42e0a69bfd5..d51f89e2410 100644 --- a/src/promql/src/native/display.rs +++ b/src/promql/src/native/display.rs @@ -32,7 +32,7 @@ use crate::native::plan::{ }; use crate::range_eval::RangeFunction; use crate::select::SortOrder; -use crate::vector_match::{BinaryOp, Cardinality, FillValues, MatchingLabels}; +use crate::vector_match::{BinaryOp, Cardinality, FillValues, MatchingLabels, SetOp}; /// The header line of every tree. const HEADER: &str = "NativeTree"; @@ -156,10 +156,7 @@ fn describe(node: &NativeNode) -> String { "{}{} {} {}{}", binary_op(matching.op), if matching.return_bool { " bool" } else { "" }, - match &matching.labels { - MatchingLabels::On(labels) => format!("on [{}]", label_list(labels)), - MatchingLabels::Ignoring(labels) => format!("ignoring [{}]", label_list(labels)), - }, + matching_labels(&matching.labels), match &matching.cardinality { Cardinality::OneToOne => "one-to-one".to_string(), Cardinality::ManyToOne(labels) => format!("many-to-one [{}]", label_list(labels)), @@ -167,6 +164,22 @@ fn describe(node: &NativeNode) -> String { }, fill(matching.fill) ), + NativeExpr::Set { op, .. } => format!( + "{} {}", + match op.op { + SetOp::And => "and", + SetOp::Or => "or", + SetOp::Unless => "unless", + }, + matching_labels(&op.labels) + ), + } +} + +fn matching_labels(labels: &MatchingLabels) -> String { + match labels { + MatchingLabels::On(labels) => format!("on [{}]", label_list(labels)), + MatchingLabels::Ignoring(labels) => format!("ignoring [{}]", label_list(labels)), } } @@ -349,6 +362,7 @@ fn kind(expr: &NativeExpr) -> &'static str { NativeExpr::Aggregate { .. } => "Aggregate", NativeExpr::Vector { .. } => "Vector", NativeExpr::Binary { .. } => "Binary", + NativeExpr::Set { .. } => "Set", } } @@ -373,7 +387,7 @@ fn display_node( | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => display_node(f, input, depth + 1, "input")?, - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { display_node(f, left, depth + 1, "left")?; display_node(f, right, depth + 1, "right")?; } diff --git a/src/promql/src/native/display/tests.rs b/src/promql/src/native/display/tests.rs index c7a8eb97be2..bcb4e44babd 100644 --- a/src/promql/src/native/display/tests.rs +++ b/src/promql/src/native/display/tests.rs @@ -20,13 +20,16 @@ use datatypes::series::{Label, SeriesLabels, SeriesOrdering}; use super::super::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + BoundSetOp, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, + NativeType, }; use crate::map::MathOp; use crate::native::plan::MapOp; use crate::range_eval::{RangeFunction, RangeModifier, RangeWindow}; use crate::select::SortOrder; -use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels}; +use crate::vector_match::{ + BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, SetOp, +}; fn grid(start: i64, step: i64, len: u64) -> Arc { Arc::new(EvalGrid::try_new(GridSpec::try_new(start, step, len).unwrap()).unwrap()) @@ -371,6 +374,48 @@ fn an_aggregate_and_a_binary_name_their_binding() { ); } +/// A set operation names its operator and its matching labels: the logical plan compares fused +/// nodes by this rendering, so `and`, `or` and `unless`, and `on` and `ignoring`, must not read +/// alike. +#[test] +fn a_set_operation_names_its_operator_and_labels() { + let out = grid(0, 60_000, 4); + let operand = || { + window_over( + raw_input(), + RangeFunction::Rate, + 300_000, + 0, + None, + Arc::clone(&out), + ) + }; + let root_line = |op: SetOp, labels: MatchingLabels| { + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Set { + left: Box::new(operand()), + right: Box::new(operand()), + op: BoundSetOp { op, labels }, + }, + output: evaluated(Arc::clone(&out)), + }) + .unwrap(); + tree.to_string().lines().nth(1).unwrap().to_string() + }; + assert_eq!( + root_line(SetOp::And, MatchingLabels::On(vec!["job".into()])), + " root Set and on [job]" + ); + assert_eq!( + root_line(SetOp::Or, MatchingLabels::Ignoring(vec!["job".into()])), + " root Set or ignoring [job]" + ); + assert_eq!( + root_line(SetOp::Unless, MatchingLabels::Ignoring(Vec::new())), + " root Set unless ignoring []" + ); +} + /// A `without` grouping reads as `without [...]`, not as `by [...]`. /// /// The two are different queries over the same label list: `by(job)` keeps `job` and drops the diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 6151430203a..ba0e8f19aec 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -52,8 +52,8 @@ use crate::native::aggregate::AggregateRuntime; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::NodeMetrics; use crate::native::plan::{ - BoundInstant, BoundScalar, BoundVectorOp, Grouping, MapOp, NativeExpr, NativeNode, NativeTree, - NativeType, + BoundInstant, BoundScalar, BoundSetOp, BoundVectorOp, Grouping, MapOp, NativeExpr, NativeNode, + NativeTree, NativeType, }; #[cfg(test)] pub(crate) use crate::native::runs::absorb_rebuilds; @@ -62,7 +62,7 @@ use crate::range_eval::RangeEvaluator; use crate::select::{self, SortOrder}; use crate::vector_match::{ BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, - binary_vectors_with_reservation, validate_binary_operand, + binary_vectors_with_reservation, set_vectors_with_reservation, validate_binary_operand, }; #[cfg(test)] @@ -1388,7 +1388,7 @@ impl RetainedInput { } } -/// What a binary node holds: both operands until they end. +/// What a binary or set node holds: both operands until they end. /// /// The operation pairs two whole vectors point by point, so this node cannot answer from the chunks /// in hand the way a map does: it retains every chunk each operand hands over, assembles the chunks @@ -1396,8 +1396,8 @@ impl RetainedInput { /// input has ended and both operands closed. The chunks it produces are queued in the node's own /// output queue and handed over from there. struct BinaryRuntime { - /// The operation and its modifiers, bound by the compiler. - matching: Arc, + /// The kernel and the modifiers the compiler bound it with. + kernel: PairKernel, /// The layout each operand is bound to, what a chunk of it is checked against: `0` is the left /// operand and `1` the right one, the order the tree declares them in. schemas: [Arc; 2], @@ -1407,6 +1407,14 @@ struct BinaryRuntime { computed: bool, } +/// The kernel a [`BinaryRuntime`] runs over its two assembled operands. +enum PairKernel { + /// An arithmetic or comparison operation between matched series. + Binary(Arc), + /// `and`, `or` or `unless`, which keeps samples of either side as they are. + Set(BoundSetOp), +} + /// What side `index` is called in an error. fn side_name(index: usize) -> &'static str { match index { @@ -1441,10 +1449,10 @@ impl Reader { } impl BinaryRuntime { - /// The run-time state of the binary operation `matching` over its two operands. - fn new(matching: Arc, schemas: [Arc; 2]) -> Self { + /// The run-time state of `kernel` over its two operands. + fn new(kernel: PairKernel, schemas: [Arc; 2]) -> Self { Self { - matching, + kernel, schemas, operands: [RetainedInput::default(), RetainedInput::default()], computed: false, @@ -1571,20 +1579,32 @@ impl BinaryRuntime { // The timer covers this node's own work only: it is dropped before the driver polls // the input again. let _timer = metrics.timer(); - binary_vectors_with_reservation( - &self.matching, - left_schema, - left, - right_schema, - right, - &mut |workspace| { - let bytes = retained - .checked_add(assembly) - .and_then(|held| held.checked_add(workspace)) - .ok_or_else(reservation_overflow)?; - buffer.resize(bytes) - }, - ) + let mut reserve = |workspace: usize| { + let bytes = retained + .checked_add(assembly) + .and_then(|held| held.checked_add(workspace)) + .ok_or_else(reservation_overflow)?; + buffer.resize(bytes) + }; + match &self.kernel { + PairKernel::Binary(matching) => binary_vectors_with_reservation( + matching, + left_schema, + left, + right_schema, + right, + &mut reserve, + ), + PairKernel::Set(set) => set_vectors_with_reservation( + set.op, + &set.labels, + left_schema, + left, + right_schema, + right, + &mut reserve, + ), + } }; // The operands and the assembled series have been handed to the kernel and the reservation // covers the workspace it allocated past them, so a failure is the driver's to release: it @@ -2300,7 +2320,7 @@ impl RuntimeNode { { ( Some(BinaryRuntime::new( - Arc::new(matching.clone()), + PairKernel::Binary(Arc::new(matching.clone())), [Arc::clone(left_schema), Arc::clone(right_schema)], )), None, @@ -2348,6 +2368,31 @@ impl RuntimeNode { } } } + NativeExpr::Set { left, right, op } => { + let (Some(left_schema), Some(right_schema)) = + (left.output.samples(), right.output.samples()) + else { + return Err(DataFusionError::Execution( + "a set operation consumes two vectors".to_string(), + )); + }; + let set = BinaryRuntime::new( + PairKernel::Set(op.clone()), + [Arc::clone(left_schema), Arc::clone(right_schema)], + ); + let (left_node, _) = Self::build(left, pool, inputs)?; + let (right_node, _) = Self::build(right, pool, inputs)?; + // The kernel files one series per label set, so its output shares none. + ( + RuntimeOp::Binary { + left: Box::new(left_node), + right: Box::new(right_node), + binary: Some(set), + state: NodeState::new("native.set", Arc::clone(&metrics), pool), + }, + false, + ) + } NativeExpr::Aggregate { input, op } => { let schema = node .output diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index 3ccb23d2875..b33554da7e3 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -34,8 +34,8 @@ use datatypes::series::{Label, SeriesLabels, SeriesOrdering}; use super::{ BinaryRuntime, Emitted, InputId, InputSlot, NativeChunk, NativeExecutionState, NativeExecutor, - NodeState, RuntimeNode, RuntimeOp, absorb_rebuilds, evaluated_bytes, grid_end_ms, map_chunk, - operand_chunk_bytes, raw_bytes, vector_workspace_bytes, + NodeState, PairKernel, RuntimeNode, RuntimeOp, absorb_rebuilds, evaluated_bytes, grid_end_ms, + map_chunk, operand_chunk_bytes, raw_bytes, vector_workspace_bytes, }; use crate::evaluated::EvaluatedChunk; use crate::grid_series::merge_same_labelset; @@ -46,13 +46,13 @@ use crate::map_label::LabelRewrite; use crate::native::metrics::{NodeMetric, NodeMetrics}; use crate::native::plan::{ AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, BoundInstant, - BoundMap, BoundScalar, BoundVectorOp, BoundWindow, Grouping, MapOp, NativeExpr, NativeNode, - NativeTree, NativeType, + BoundMap, BoundScalar, BoundSetOp, BoundVectorOp, BoundWindow, Grouping, MapOp, NativeExpr, + NativeNode, NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::select::SortOrder; use crate::vector_match::{ - BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, reset_visited_pairs, + BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, SetOp, reset_visited_pairs, visited_pairs, }; @@ -527,7 +527,7 @@ fn binary_runtime( ( NodeState::new("native.binary", Arc::new(NodeMetrics::default()), pool), BinaryRuntime::new( - Arc::new(binary_matching(BinaryOp::Add)), + PairKernel::Binary(Arc::new(binary_matching(BinaryOp::Add))), [Arc::clone(schema), Arc::clone(schema)], ), ) @@ -7189,6 +7189,175 @@ fn an_all_missing_match_group_fits_its_pool() { assert_eq!(execution.state().reserved_bytes(), 0); } +// --------------------------------------------------------------------------- +// Set operations, driven through `NativeExecutor`. +// --------------------------------------------------------------------------- + +/// A set node over `left` and `right`, bound the way the compiler binds one: the left side's fields +/// on the grid both operands share. +fn set_of(left: NativeNode, right: NativeNode, op: SetOp, labels: MatchingLabels) -> NativeNode { + let output = left.output.clone(); + NativeNode { + expr: NativeExpr::Set { + left: Box::new(left), + right: Box::new(right), + op: BoundSetOp { op, labels }, + }, + output, + } +} + +/// Every series a chunk stream holds, by its rendered label set with the metric name first. +fn decided_by_label_set(chunks: &[EvaluatedChunk]) -> Vec<(String, Decided)> { + let mut series: Vec<(String, Decided)> = chunks + .iter() + .map(|chunk| { + let metric = &chunk.series.metric; + let labels = metric + .name() + .into_iter() + .map(|name| format!("__name__={name}")) + .chain( + metric + .tags() + .iter() + .map(|tag| format!("{}={}", tag.name, tag.value)), + ) + .collect::>() + .join(","); + (labels, decided(std::slice::from_ref(chunk))) + }) + .collect(); + series.sort_by(|a, b| a.0.cmp(&b.0)); + series +} + +/// Each operator matches per grid point, by the labels it was bound with, and hands samples on with +/// the labels of the side they came from. +/// +/// `l{job="a"}` holds every point and `l{job="b"}` only the last two; `r{job="a"}` holds the middle +/// two and `r{job="b"}` every point. The pods differ on every pair, so the default matching pairs +/// nothing, while `on(job)` and `ignoring(pod)` pair by job alone. +#[test] +fn set_operators_match_by_their_labels_at_each_point() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let left = [ + (pod_metric("l", "a", "x"), 0, vec![1.0, 2.0, 3.0, 4.0]), + (pod_metric("l", "b", "y"), 2, vec![7.0, 8.0]), + ]; + let right = [ + (pod_metric("r", "a", "z"), 1, vec![10.0, 20.0]), + (pod_metric("r", "b", "w"), 0, vec![30.0, 40.0, 50.0, 60.0]), + ]; + let on_job = || MatchingLabels::On(vec!["job".into()]); + let la = "__name__=l,job=a,pod=x".to_string(); + let lb = "__name__=l,job=b,pod=y".to_string(); + let rb = "__name__=r,job=b,pod=w".to_string(); + for (op, labels, expected) in [ + ( + SetOp::And, + on_job(), + vec![ + (la.clone(), vec![(1, Some(2.0)), (2, Some(3.0))]), + (lb.clone(), vec![(2, Some(7.0)), (3, Some(8.0))]), + ], + ), + (SetOp::And, MatchingLabels::Ignoring(Vec::new()), vec![]), + ( + SetOp::Unless, + MatchingLabels::Ignoring(vec!["pod".into()]), + vec![(la.clone(), vec![(0, Some(1.0)), (3, Some(4.0))])], + ), + // `r{job="a"}` finds the left side's job at each of its points and is dropped; `r{job="b"}` + // fills the two points where the left side has no `job="b"`, under its own labels. + ( + SetOp::Or, + on_job(), + vec![ + ( + la.clone(), + vec![ + (0, Some(1.0)), + (1, Some(2.0)), + (2, Some(3.0)), + (3, Some(4.0)), + ], + ), + (lb.clone(), vec![(2, Some(7.0)), (3, Some(8.0))]), + (rb.clone(), vec![(0, Some(30.0)), (1, Some(40.0))]), + ], + ), + ] { + let tree = NativeTree::try_new(set_of( + evaluated_input_with_source_index(Arc::clone(&grid), 0), + evaluated_input_with_source_index(Arc::clone(&grid), 1), + op, + labels.clone(), + )) + .unwrap(); + let mut execution = executor(&tree, pool()); + for (source, side) in [(0, &left), (1, &right)] { + let id = execution.state().input_ids(source).next().unwrap(); + for (metric, start, values) in side.iter() { + execution + .push_evaluated( + id, + evaluated_chunk_of( + &schema, + metric.clone(), + *start, + values.iter().copied().map(Some).collect(), + true, + ), + ) + .unwrap(); + } + } + assert_eq!( + decided_by_label_set(&run(&mut execution)), + expected, + "{op:?} {labels:?}" + ); + assert_eq!(execution.state().reserved_bytes(), 0, "{op:?} {labels:?}"); + } +} + +/// The rows a set operation keeps are reserved from the pool like a binary operation's workspace: +/// a pool that holds both operands, their assembly and the kept output, but not the rows the kernel +/// files and merges on the way, refuses the operation and gets back what the node held. +#[test] +fn an_oversized_set_workspace_is_refused_by_the_pool() { + let points = 10_000; + let grid = grid(0, 60_000, points); + let schema = evaluated_schema(Arc::clone(&grid)); + let operand = || evaluated_chunk(&schema, 0, vec![1.0; points as usize], true); + let (retained, assembly) = binary_charge(&schema, &operand(), &operand()); + // Room for a copy of the one-field output on top of what the node holds before the kernel runs, + // which is far less than the rows `and` files for every point and the merge reads them into. + let pool = bounded_pool(retained + assembly + 2 * points as usize * 9); + let tree = NativeTree::try_new(set_of( + evaluated_input_with_source_index(Arc::clone(&grid), 0), + evaluated_input_with_source_index(Arc::clone(&grid), 1), + SetOp::And, + MatchingLabels::Ignoring(Vec::new()), + )) + .unwrap(); + let mut execution = executor(&tree, Arc::clone(&pool)); + for source in [0, 1] { + let id = execution.state().input_ids(source).next().unwrap(); + execution.push_evaluated(id, operand()).unwrap(); + } + let error = execution.finish().unwrap_err(); + assert!( + matches!(error, DataFusionError::ResourcesExhausted(_)), + "the set workspace was not refused by the pool: {error}" + ); + assert_eq!(pool.reserved(), 0); + assert_eq!(execution.state().reserved_bytes(), 0); + assert!(!root_holds(&execution, holds_binary)); +} + // --------------------------------------------------------------------------- // Scalar broadcasts, driven through `NativeExecutor`. // --------------------------------------------------------------------------- diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 37ea484b7ea..df72e816711 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -23,10 +23,9 @@ //! operation states about its input, which is what lets a new combination of supported operations //! be an ordinary tree instead of another hand-written shape. //! -//! What the tree cannot express yet is deliberate for this stage: [`NativeExpr::Binary`] holds an -//! arithmetic or comparison operation, so `and`, `or` and `unless` have no node, and -//! [`NativeExpr::Scalar`] holds a constant, so `time()` and `scalar(v)` have no node either. The -//! legality below is therefore about the operations that exist, not about PromQL's whole surface. +//! What the tree cannot express yet is deliberate for this stage: [`NativeExpr::Scalar`] holds a +//! constant, so `time()` and `scalar(v)` have no node. The legality below is therefore about the +//! operations that exist, not about PromQL's whole surface. use std::sync::Arc; @@ -40,7 +39,7 @@ use crate::map::MathOp; use crate::map_label::LabelRewrite; use crate::range_eval::{RangeFunction, RangeWindow}; use crate::select::SortOrder; -use crate::vector_match::BinaryMatching; +use crate::vector_match::{BinaryMatching, MatchingLabels, SetOp}; /// What one node produces. /// @@ -308,6 +307,18 @@ pub enum BoundVectorOp { CountValues { label: Arc, group: Grouping }, } +/// A set operator and the labels its two vectors are matched by. +/// +/// It has no cardinality, `bool` or fill: PromQL matches a set operator many-to-many and refuses +/// the other modifiers on it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BoundSetOp { + /// `and`, `or` or `unless`. + pub op: SetOp, + /// The labels a sample's match group is read from on either side. + pub labels: MatchingLabels, +} + /// What a node computes. Every variant holds a bound operation, never a name to resolve. #[derive(Clone, Debug)] pub enum NativeExpr { @@ -350,6 +361,12 @@ pub enum NativeExpr { right: Box, matching: BinaryMatching, }, + /// A set operation between two vectors, which keeps samples of either side as they are. + Set { + left: Box, + right: Box, + op: BoundSetOp, + }, } /// One node of a bound tree. @@ -363,8 +380,8 @@ pub struct NativeNode { /// Whether `node` reads the whole input vector at a point rather than one series at a time. /// -/// Three kinds do: an `Aggregate` folds every series of its group, a `Binary` pairs two vectors' -/// series by label, and a `Vector` operation (`sort`, `absent`, `topk`, `limitk`, +/// Four kinds do: an `Aggregate` folds every series of its group, a `Binary` and a `Set` pair two +/// vectors' series by label, and a `Vector` operation (`sort`, `absent`, `topk`, `limitk`, /// `histogram_quantile`, ...) reorders or reads the whole vector at a grid point. Everything else - /// an input, a scalar, a selection, a window, a map, and a binary operation with a constant on one /// side, whose broadcast reaches each series the other side answers - answers one series from one @@ -390,7 +407,7 @@ pub fn reads_across_series(node: &NativeNode) -> bool { // Two vectors are matched by label, which pairs series of the whole vector. (_, _) => true, }, - NativeExpr::Aggregate { .. } | NativeExpr::Vector { .. } => true, + NativeExpr::Aggregate { .. } | NativeExpr::Vector { .. } | NativeExpr::Set { .. } => true, } } @@ -695,6 +712,39 @@ impl NativeNode { } } } + NativeExpr::Set { left, right, .. } => { + left.validate()?; + right.validate()?; + // A set operation keeps samples of whole vectors on the grid both share: a scalar + // has no labels to match, and PromQL refuses it on either side. + let Some((left_schema, left_grid)) = evaluated_vector(&left.output) else { + return Err(invalid(&format!( + "a set operation consumes evaluated vectors, not {} on its left", + left.output.describe() + ))); + }; + let Some((right_schema, right_grid)) = evaluated_vector(&right.output) else { + return Err(invalid(&format!( + "a set operation consumes evaluated vectors, not {} on its right", + right.output.describe() + ))); + }; + if left_grid.spec() != right_grid.spec() { + return Err(invalid("a set operation aligns two different grids")); + } + // `or` hands the right side's samples on under the left side's fields, so the + // kernel pairs every field by position and refuses sides of different counts. + if left_schema.fields.len() != right_schema.fields.len() { + return Err(invalid(&format!( + "a set operation pairs fields by position, and its sides carry {} and {}", + left_schema.fields.len(), + right_schema.fields.len() + ))); + } + ensure_samples(&self.output, "a set operation")?; + ensure_same_fields(&left.output, &self.output, "a set operation")?; + ensure_grid(&self.output, left_grid, "a set operation")?; + } } Ok(()) } diff --git a/src/promql/src/vector_match.rs b/src/promql/src/vector_match.rs index a26d9d069ee..66a534cd12b 100644 --- a/src/promql/src/vector_match.rs +++ b/src/promql/src/vector_match.rs @@ -1143,6 +1143,45 @@ pub fn set_vectors( lhs: &[Series], rhs_schema: &SeriesSchema, rhs: &[Series], +) -> Result> { + set_vectors_inner(op, labels, lhs_schema, lhs, rhs_schema, rhs, None) +} + +/// [`set_vectors`], handing the live workspace it holds to `reserve` as it allocates. +/// +/// The figures follow [`binary_vectors_with_reservation`]: what the results keep is held as their +/// containers grow, each field's walk and sort is a phase that replaces the one before it, and the +/// merge is the last phase. A total `reserve` refuses ends the operation with that error. +pub(crate) fn set_vectors_with_reservation( + op: SetOp, + labels: &MatchingLabels, + lhs_schema: &SeriesSchema, + lhs: &[Series], + rhs_schema: &SeriesSchema, + rhs: &[Series], + reserve: &mut dyn FnMut(usize) -> Result<()>, +) -> Result> { + let mut budget = Budget::new(reserve); + set_vectors_inner( + op, + labels, + lhs_schema, + lhs, + rhs_schema, + rhs, + Some(&mut budget), + ) +} + +/// The one set kernel, whether or not a caller accounts for its workspace. +fn set_vectors_inner( + op: SetOp, + labels: &MatchingLabels, + lhs_schema: &SeriesSchema, + lhs: &[Series], + rhs_schema: &SeriesSchema, + rhs: &[Series], + mut budget: Option<&mut Budget<'_>>, ) -> Result> { let grid = same_grid(lhs_schema, rhs_schema)?; if lhs_schema.fields.len() != rhs_schema.fields.len() { @@ -1150,55 +1189,140 @@ pub fn set_vectors( "set operators between vectors with different field counts".into(), )); } + let width = lhs_schema.fields.len(); + if let Some(budget) = budget.as_deref_mut() { + budget.hold(output_bytes(width))?; + } let output = Arc::new(SeriesSchema { fields: lhs_schema.fields.clone(), eval_grid: Some(grid.clone()), }); let mut errors = Vec::new(); - let mut fields = Vec::with_capacity(output.fields.len()); - for field in 0..output.fields.len() { + let mut fields = Vec::with_capacity(width); + for field in 0..width { let lhs = Operand::try_new(lhs_schema, lhs, field)?; let rhs = Operand::try_new(rhs_schema, rhs, field)?; - let mut results = set_field(op, labels, &lhs, &rhs)?; - if let Some(duplicate) = results.sort_rows() { - let timestamp = (grid.spec().start_ms as i128 - + duplicate.point as i128 * grid.spec().step_ms as i128) - as i64; - errors.push((duplicate.point, 0, same_labelset(&duplicate.key, timestamp))); + let results = set_field( + op, + labels, + &grid, + &lhs, + &rhs, + &mut errors, + budget.as_deref_mut(), + )?; + // The field's group table and sort are gone with its frame; what it filed stays live. + if let Some(budget) = budget.as_deref_mut() { + budget.phase(0)?; } fields.push(results); } earliest_error(errors)?; - // A set operator accounts for no workspace, so its merge is charged nothing either. - merge_fields(fields, output, None) + merge_fields(fields, output, budget) } +/// The figures one set field's phases are charged from, read once from its operands. +struct SetFieldMetadata { + /// The filter side's match groups. + tables: usize, + /// The widest key one row builds: the match key a kept series is looked up by, or the key a row + /// is filed under. + scratch: usize, + /// The one message the field can report, a fixed frame around one rendering. + error: usize, +} + +fn set_field_metadata( + labels: &MatchingLabels, + kept: &Operand<'_>, + filter: &Operand<'_>, +) -> SetFieldMetadata { + let mut tables = 0usize; + for series in filter.series { + tables = tables + .saturating_add(match_key_bytes(&series.metric, labels)) + .saturating_add(2 * std::mem::size_of::()) + .saturating_add(GROUP_SLOT_BYTES); + } + // `or` files the filter side's rows too, so both sides' label sets bound a row's keys and the + // text of the message. + let (mut scratch, mut text) = (0usize, 0usize); + for series in kept.series.iter().chain(filter.series) { + scratch = scratch + .max(match_key_bytes(&series.metric, labels)) + .max(result_key_bytes(&series.metric)); + text = text.max(label_text_bytes(&series.metric)); + } + SetFieldMetadata { + tables, + scratch, + error: ERROR_TEXT_BYTES.saturating_add(text), + } +} + +/// Filters one field of the kept side by the presence of its match group on the filter side, and +/// appends the first point where two kept rows share a label set as `(point, rank, message)`. fn set_field( op: SetOp, labels: &MatchingLabels, + grid: &EvalGrid, lhs: &Operand<'_>, rhs: &Operand<'_>, + errors: &mut Vec<(usize, u8, String)>, + mut budget: Option<&mut Budget<'_>>, ) -> Result { let (kept, filter, keep_matched) = match op { SetOp::And => (lhs, rhs, true), SetOp::Unless => (lhs, rhs, false), SetOp::Or => (rhs, lhs, false), }; + // The walk is one phase, entered before the group table exists. A kept series reads its + // candidates through one cursor each, so the widest group's cursors join it once it does. + let metadata = match budget.as_deref_mut() { + Some(budget) => { + let metadata = set_field_metadata(labels, kept, filter); + budget.phase(metadata.tables.saturating_add(metadata.scratch))?; + Some(metadata) + } + None => None, + }; let filter_groups = group(filter, labels); - // A set operator hands on the label sets of its operands and no caller accounts for its - // workspace, so it files its rows without charging them. + if let Some(budget) = budget.as_deref_mut() { + let widest = filter_groups.values().map(Vec::len).max().unwrap_or(0); + budget.grow(widest.saturating_mul(std::mem::size_of::>()))?; + } + // A set operator hands on its operands' label sets unchanged, so a row is filed under the label + // set of the series it came from. + let charged = budget.is_some(); + let label_set = |metric: &SeriesLabels| { + if charged { + result_set_bytes(metric, metric, &[]) + } else { + 0 + } + }; let mut results = Results::default(); if op == SetOp::Or { for index in 0..lhs.series.len() { + let metric = &lhs.series[index].metric; + let label_set = label_set(metric); let mut cursor = lhs.cursor(index); for (point, row) in lhs.positions(index) { if let Some(value) = owned(cursor.at(row).expect("a row never goes back"))? { - results.insert(&lhs.series[index].metric, 0, point, Some(value), 0, None)?; + results.insert( + metric, + 0, + point, + Some(value), + label_set, + budget.as_deref_mut(), + )?; } } } } for (index, series) in kept.series.iter().enumerate() { + let label_set = label_set(&series.metric); let candidates = filter_groups.get(&group_key(&series.metric, labels)); let candidate_indices: &[usize] = candidates.map_or(&[], Vec::as_slice); // One forward cursor per candidate: `kept`'s positions only increase, so the row a @@ -1223,10 +1347,38 @@ fn set_field( } } if matched == keep_matched { - results.insert(&series.metric, 0, point, owned(sample)?, 0, None)?; + results.insert( + &series.metric, + 0, + point, + owned(sample)?, + label_set, + budget.as_deref_mut(), + )?; } } } + // The group table is gone before the sort phase, which is charged without it: the message the + // sort may append and one slot per row it sorts, entered before it allocates either. + drop(filter_groups); + if let (Some(budget), Some(metadata)) = (budget.as_deref_mut(), &metadata) { + budget.phase( + metadata + .error + .saturating_add(metadata.scratch) + .saturating_add(results.rows.saturating_mul(SORT_ROW_BYTES)), + )?; + } + if let Some(duplicate) = results.sort_rows() { + let timestamp = (grid.spec().start_ms as i128 + + duplicate.point as i128 * grid.spec().step_ms as i128) as i64; + let message = same_labelset(&duplicate.key, timestamp); + // The message stays live until the earliest one is reported. + if let Some(budget) = budget { + budget.hold(std::mem::size_of::().saturating_add(message.capacity()))?; + } + errors.push((duplicate.point, 0, message)); + } Ok(results) } diff --git a/src/promql/src/vector_match/tests.rs b/src/promql/src/vector_match/tests.rs index 02cb8549c86..fecc9578e05 100644 --- a/src/promql/src/vector_match/tests.rs +++ b/src/promql/src/vector_match/tests.rs @@ -28,7 +28,7 @@ use rand::{Rng, SeedableRng}; use super::{ BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, RESULT_ROW_BYTES, SetOp, - binary_vectors, binary_vectors_with_reservation, set_vectors, + binary_vectors, binary_vectors_with_reservation, set_vectors, set_vectors_with_reservation, }; use crate::grid_series::reference::{ Dense, OwnedHistogram, Points, assert_histogram_close, assert_result_eq, golden_cases, @@ -3404,3 +3404,41 @@ fn a_set_operation_keeps_the_payload_of_the_sample_it_keeps() { "`or` did not hand back each side's own sample: {or:?}" ); } + +/// A set operation reserves the rows it keeps where they are filed, the way a binary operation +/// does: `and` keeps every left point here and `unless` none, and the two operations differ in +/// nothing else, so what the results hold once the field's tables are released differs by at least +/// one filed row per kept point. +#[test] +fn a_set_operation_holds_the_rows_it_keeps() { + const POINTS: usize = 64; + let grid = budget_grid(POINTS); + let schema = Arc::new(schema(1, &grid)); + let all: Vec = (0..POINTS).collect(); + let lhs = vec![budget_operand(&schema, "a", &all)]; + let rhs = vec![budget_operand(&schema, "b", &all)]; + let retained = |op: SetOp| { + let mut seen = Vec::new(); + let output = set_vectors_with_reservation( + op, + &MatchingLabels::Ignoring(Vec::new()), + &schema, + &lhs, + &schema, + &rhs, + &mut |total| { + seen.push(total); + Ok(()) + }, + ) + .unwrap(); + (output.len(), retained_total(&seen)) + }; + let (kept, and) = retained(SetOp::And); + let (dropped, unless) = retained(SetOp::Unless); + assert_eq!((kept, dropped), (1, 0)); + assert!( + and >= unless + POINTS * FILED_ROW, + "`and` holds {and} bytes and `unless` {unless}" + ); +} diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index eb6928cbafa..e417b923654 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -35,6 +35,7 @@ //! e.g. x * 100, x > bool 0 //! e.g. 1 / x, 0 < x //! - e.g. -x +//! e.g. x and on(job) y, x unless y //! [(by|without)()] () //! e.g. avg by(job)(rate(x[5m]) * 100) //! () [(by|without)()] @@ -55,7 +56,8 @@ //! offset 1m` //! ``` //! -//! where `` is `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` or +//! where `` is `and`, `or` or `unless`, optionally with `on(...)` or `ignoring(...)`, +//! `` is `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` or //! `quantile`, `` is `topk`, `bottomk`, `limitk` or `limit_ratio`, the //! function names are [`range_signature`]'s and [`math_signature`]'s, and which side //! of the matrix a range function's constants go on is [`RangeSignature`]'s: before it for @@ -78,7 +80,8 @@ //! `` → `Map` (a pointwise math, native histogram or label function over what its //! argument evaluates to) and, over //! any of those, `Binary` (an arithmetic or comparison operation between vectors or with a -//! constant on one side), `Scalar` (that constant), `Aggregate` (one of the nine aggregators that +//! constant on one side), `Scalar` (that constant), `Set` (`and`, `or` or `unless` between two +//! vectors), `Aggregate` (one of the nine aggregators that //! fold a group's values into one value per grid point, optionally grouped) and `Vector` (an //! operation over the whole vector at each point: a sort, `absent`, a classic histogram function, a //! selecting aggregator or `count_values`). Everything else is a @@ -141,12 +144,13 @@ use promql::map_label::LabelRewrite; use promql::map_time::DateFunction; use promql::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundVectorOp, BoundWindow, Grouping, MapOp, NativeExpr, NativeNode, NativeTree, NativeType, + BoundSetOp, BoundVectorOp, BoundWindow, Grouping, MapOp, NativeExpr, NativeNode, NativeTree, + NativeType, }; use promql::range_eval::{RangeFunction, RangeWindow, subquery_grid}; use promql::select::{LabelMatcher, MatchOp, SortOrder, absent_labels}; use promql::vector_match::{ - BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, binary_output_schema, + BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, SetOp, binary_output_schema, }; use promql_parser::label::{MatchOp as ParserMatchOp, Matchers}; use promql_parser::parser::token::TokenType; @@ -239,9 +243,10 @@ const SUPPORTED: &str = "a selector — `x`, optionally with an `offset` — a r function over any of them \ (`label_replace(x, \"a\", \"$1\", \"b\", \"(.*)\")`), a binary \ operation between one of them and a constant (`2 / (x + 1)`, \ - `x > bool 0`), a negation of one of them (`-x`), or one of the aggregators \ - that folds a group's values into one value per grid point — `sum`, `avg`, \ - `count`, `min`, `max`, `group`, `stddev`, `stdvar` and \ + `x > bool 0`), a negation of one of them (`-x`), a set operation between \ + two of them (`x and on(job) y`, `x or y`, `x unless y`), or one of the \ + aggregators that folds a group's values into one value per grid point — \ + `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` and \ `quantile(, )` — over any of them, optionally \ grouped by `by` or `without` as in `avg by(job)(rate(x[5m]) * 100)`; \ `topk`, `bottomk`, `limitk`, `limit_ratio` or `count_values` over any of \ @@ -368,7 +373,7 @@ fn validate_inputs(root: &NativeNode, inputs: &[CompiledInput]) -> Result<()> { | NativeExpr::Timestamp { input } | NativeExpr::Vector { input, .. } | NativeExpr::Aggregate { input, .. } => check(input, inputs, next), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { check(left, inputs, next) && check(right, inputs, next) } } @@ -456,17 +461,8 @@ fn negate(context: &Context, operand: &PromExpr) -> Result { /// A scalar carries no labels, so broadcasting it over a vector is unambiguous. Vector operands /// retain their own input bindings and matching, including a fill on either side. fn binary_operation(context: &Context, binary: &BinaryExpr) -> Result { - // A set operator combines two vectors' series rather than their values, and the tree has no node - // for it at all. Its name is reported before the operands are looked at, so a set operation - // between two vectors is refused as the set operation it is rather than as a binary operation - // the compiler has no node for. - if binary.op.is_set_operator() { - return Err(not_implemented(&format!( - "the set operator `{}` in `{binary}`, which has no native node: it combines two \ - vectors' series, and every binary node here is an arithmetic or comparison operation \ - over values", - binary.op - ))); + if let Some(op) = set_op(&binary.op) { + return set_operation(context, binary, op); } let Some(op) = binary_op(&binary.op) else { return Err(not_implemented(&format!( @@ -533,6 +529,54 @@ fn binary_operation(context: &Context, binary: &BinaryExpr) -> Result Result { + let modifier = binary.modifier.as_ref(); + if modifier.is_some_and(|modifier| { + modifier.card != VectorMatchCardinality::ManyToMany + || modifier.fill_values != VectorMatchFillValues::default() + }) { + return Err(not_implemented(&format!( + "the modifiers in `{binary}`: a set operator matches many-to-many and takes no fill \ + value" + ))); + } + let labels = matching_labels(modifier); + let left = compile_expr(context, &binary.lhs)?; + let right = compile_expr(context, &binary.rhs)?; + let (Some(lhs), Some(rhs)) = (left.output.samples(), right.output.samples()) else { + return Err(DataFusionError::Internal(format!( + "the set operation `{binary}` compiled an operand to something other than samples: the \ + parser only accepts a set operator between two vectors" + ))); + }; + // `or` reports the right side's samples under the left side's fields, which the kernel can only + // do field by field. + if lhs.fields.len() != rhs.fields.len() { + return Err(not_implemented(&format!( + "the set operation `{binary}` between vectors of {} and {} fields", + lhs.fields.len(), + rhs.fields.len() + ))); + } + // A set operation hands on the samples it keeps unchanged, so it reports the left side's fields + // on the grid both sides share. + let output = left.output.clone(); + Ok(NativeNode { + expr: NativeExpr::Set { + left: Box::new(left), + right: Box::new(right), + op: BoundSetOp { op, labels }, + }, + output, + }) +} + /// Wraps two compiled operands in the binary node that applies `op` to them. /// /// The output is the vector side's own binding, because a scalar carries no fields: a broadcast keeps @@ -586,18 +630,12 @@ fn scalar_matching(op: BinaryOp, return_bool: bool) -> BinaryMatching { /// cardinality is one-to-one. Fill fields retain the parser's left/right values, which name the /// sides as written: the kernel swaps them together with the operands under `group_right`. /// -/// `ManyToMany` is the cardinality of a set operator, which [`binary_operation`] refuses before +/// `ManyToMany` is the cardinality of a set operator, which [`binary_operation`] compiles before /// asking for a matching, so reaching it means the two disagree; it is reported as a shape with no /// node rather than panicking on a path a query cannot take today. fn vector_matching(binary: &BinaryExpr, op: BinaryOp) -> Result { let modifier = binary.modifier.as_ref(); - let labels = match modifier.and_then(|modifier| modifier.matching.as_ref()) { - None => MatchingLabels::Ignoring(Vec::new()), - Some(LabelModifier::Include(labels)) => MatchingLabels::On(label_names(&labels.labels)), - Some(LabelModifier::Exclude(labels)) => { - MatchingLabels::Ignoring(label_names(&labels.labels)) - } - }; + let labels = matching_labels(modifier); let cardinality = match modifier.map(|modifier| &modifier.card) { None | Some(VectorMatchCardinality::OneToOne) => Cardinality::OneToOne, Some(VectorMatchCardinality::ManyToOne(labels)) => { @@ -627,6 +665,27 @@ fn vector_matching(binary: &BinaryExpr, op: BinaryOp) -> Result }) } +/// The set operator a PromQL token denotes, or `None` when it is not one. +fn set_op(op: &TokenType) -> Option { + match op.id() { + token::T_LAND => Some(SetOp::And), + token::T_LOR => Some(SetOp::Or), + token::T_LUNLESS => Some(SetOp::Unless), + _ => None, + } +} + +/// The labels `on(...)` or `ignoring(...)` names; with neither, every label but `__name__`. +fn matching_labels(modifier: Option<&BinModifier>) -> MatchingLabels { + match modifier.and_then(|modifier| modifier.matching.as_ref()) { + None => MatchingLabels::Ignoring(Vec::new()), + Some(LabelModifier::Include(labels)) => MatchingLabels::On(label_names(&labels.labels)), + Some(LabelModifier::Exclude(labels)) => { + MatchingLabels::Ignoring(label_names(&labels.labels)) + } + } +} + fn label_names(labels: &[String]) -> Vec> { labels.iter().map(|name| Arc::from(name.as_str())).collect() } diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 1ae05ff55ec..dd81670bbd3 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -27,11 +27,11 @@ use promql::evaluated::EvaluatedChunk; use promql::map::MathOp; use promql::native::execute::{InputId, NativeExecutionState, NativeExecutor}; use promql::native::plan::{ - AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundWindow, - Grouping, MapOp, NativeExpr, NativeNode, NativeTree, NativeType, + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundSetOp, + BoundWindow, Grouping, MapOp, NativeExpr, NativeNode, NativeTree, NativeType, }; use promql::range_eval::RangeFunction; -use promql::vector_match::{BinaryMatching, BinaryOp, Cardinality, MatchingLabels}; +use promql::vector_match::{BinaryMatching, BinaryOp, Cardinality, MatchingLabels, SetOp}; use promql_parser::parser::{ AtModifier, BinModifier, EvalStmt, Expr as PromExpr, NumberLiteral, Offset, VectorMatchFillValues, @@ -232,6 +232,14 @@ fn binary_parts(node: &NativeNode) -> (&NativeNode, &NativeNode, &BinaryMatching (left, right, matching) } +/// The two operands and the binding of a `Set` node. +fn set_parts(node: &NativeNode) -> (&NativeNode, &NativeNode, &BoundSetOp) { + let NativeExpr::Set { left, right, op } = &node.expr else { + panic!("the node is not a set operation: {:?}", node.expr); + }; + (left, right, op) +} + /// The constant a `Scalar` node binds, if it is one. fn constant(node: &NativeNode) -> Option { match &node.expr { @@ -1826,39 +1834,93 @@ fn vector_binary_matching_keeps_ignoring_and_group_left_separate() { ); } +/// A set operator binds as a set node over its two operands, with the labels the query wrote, and +/// nests like any other expression: on either side of it and above it. #[test] -fn vector_binary_set_operators_remain_rejected_explicitly() { - let error = compile_selectors( - &eval_stmt("a and b"), +fn set_operators_bind_their_operator_and_labels() { + for (query, op, on, labels) in [ + ("a and b", SetOp::And, false, vec![]), + ("a or on(job) b", SetOp::Or, true, vec!["job"]), + ( + "a unless ignoring(pod) b", + SetOp::Unless, + false, + vec!["pod"], + ), + ] { + let compiled = compile_selectors( + &eval_stmt(query), + &|_| Ok(raw_schema()), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`{query}` was refused: {error}")); + assert_eq!(compiled.inputs.len(), 2, "{query}"); + let (left, right, bound) = set_parts(compiled.tree.root()); + assert_eq!(bound.op, op, "{query}"); + assert_eq!(matching_labels(&bound.labels), (on, labels), "{query}"); + assert!(matches!(left.expr, NativeExpr::Instant { .. }), "{query}"); + assert!(matches!(right.expr, NativeExpr::Instant { .. }), "{query}"); + } + + let compiled = compile_selectors( + &eval_stmt("abs(rate(x[5m]) and on(job) sum by (job) (up)) * 2"), &|_| Ok(raw_schema()), SeriesOrdering::KeyThenTime, ) - .expect_err("set operators have no binary value node"); - assert!(error.to_string().contains("set operator `and`"), "{error}"); + .unwrap(); + let (vector, _, _) = binary_parts(compiled.tree.root()); + let NativeExpr::Map { input, .. } = &vector.expr else { + panic!( + "the multiplication does not read the `abs`: {:?}", + vector.expr + ); + }; + let (left, right, _) = set_parts(input); + assert!( + matches!(left.expr, NativeExpr::Window { .. }), + "{:?}", + left.expr + ); + assert!( + matches!(right.expr, NativeExpr::Aggregate { .. }), + "{:?}", + right.expr + ); } -/// A set operator has no node at all, so it is refused as the set operator it is rather than as an -/// arithmetic operation the compiler could not bind: `and`, `or` and `unless` combine two vectors' -/// series, which is not what a binary value node does. +/// What a set operator cannot carry is refused rather than dropped: Prometheus refuses a fill value +/// on a set operator, which this parser accepts, and `or` cannot hand on a right side whose fields +/// do not pair with the left side's. #[test] -fn a_set_operator_is_refused_as_a_set_operator() { - for (query, name) in [ - ("a and b", "and"), - ("a or b", "or"), - ("a unless b", "unless"), - ] { - let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) - .expect_err(&format!("`{query}` compiled, and it should not have")); - assert!( - matches!(error, DataFusionError::NotImplemented(_)), - "`{query}` did not report NotImplemented: {error}" - ); - let message = error.to_string(); - assert!( - message.contains(&format!("set operator `{name}`")), - "`{query}` did not name the operator it found: {message}" - ); - } +fn a_set_operation_refuses_a_fill_value_and_unpaired_fields() { + let error = compile_selectors( + &with_binary_modifier("a and b", |modifier| { + modifier.fill_values = VectorMatchFillValues::new(0.0, 0.0); + }), + &|_| Ok(raw_schema()), + SeriesOrdering::KeyThenTime, + ) + .expect_err("a fill value on `and` compiled"); + assert!( + matches!(&error, DataFusionError::NotImplemented(message) if message.contains("no fill")), + "{error}" + ); + + let error = compile_selectors( + &eval_stmt("a or b"), + &|requirement| { + Ok(match requirement.metric.as_deref() { + Some("a") => raw_schema_with_fields(2), + _ => raw_schema(), + }) + }, + SeriesOrdering::KeyThenTime, + ) + .expect_err("`or` over two and one fields compiled"); + assert!( + matches!(&error, DataFusionError::NotImplemented(message) if message.contains("2 and 1")), + "{error}" + ); } /// A statement with no vector has nothing to read, so it is refused: a tree reads a scan and reports @@ -4778,7 +4840,7 @@ fn validate_inputs_rejects_misbound_occurrences_and_accepts_scalar_wrappers() { | NativeExpr::Timestamp { input } | NativeExpr::Vector { input, .. } | NativeExpr::Aggregate { input, .. } => set_input_source(input, source), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { set_input_source(left, source); set_input_source(right, source); } diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 05083212150..0330ceca389 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -644,15 +644,19 @@ fn scan_time_filter(start: i64, end: i64, inclusive: bool, time_column: &str) -> lower.and(column.lt_eq(upper)) } -/// Whether `root` is a supported one-to-one match between two independent scan occurrences. +/// Whether `root` is a supported one-to-one match or set operation between two independent scan +/// occurrences. pub(crate) fn independent_scan_pair(root: &NativeNode) -> bool { - let NativeExpr::Binary { - left, - right, - matching, - } = &root.expr - else { - return false; + let (left, right, matching) = match &root.expr { + NativeExpr::Binary { + left, + right, + matching, + } => (left, right, matching), + NativeExpr::Set { left, right, .. } => { + return pointwise_operand(left) && pointwise_operand(right); + } + _ => return false, }; let supported_operator = matches!( @@ -741,7 +745,8 @@ pub(crate) fn pointwise_operand(node: &NativeNode) -> bool { NativeExpr::Input(_) | NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } - | NativeExpr::Vector { .. } => false, + | NativeExpr::Vector { .. } + | NativeExpr::Set { .. } => false, } } @@ -928,7 +933,7 @@ mod tests { | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } | NativeExpr::Timestamp { input } => collect(input, sources), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { collect(left, sources); collect(right, sources); } diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 24ed505a645..c089331f7c7 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2513,7 +2513,7 @@ async fn native_fused_binary_keeps_independent_scan_occurrences() { | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } | NativeExpr::Timestamp { input } => sources(input, out), - NativeExpr::Binary { left, right, .. } => { + NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { sources(left, out); sources(right, out); } diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 7fb49064a3b..4365580d978 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -2649,6 +2649,233 @@ async fn native_promql_binary_join_uses_two_independent_scans() { assert_eq!(native_rows, legacy_rows); } +/// Set operators run over both whole scans in one execution (#9896 gathers every partition of each +/// source). A right side of one series sits in one scan partition while the left side spans +/// several, so a per-partition answer would differ for `and on()` and `unless on()`. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_set_operators_read_both_scans_whole() { + const HOSTS: usize = 4_608; + const POINTS: i64 = 6; + const BATCH_ROWS: usize = 512; + + fn insert_request(table: &str, schema: &[ColumnSchema], rows: Vec) -> RowInsertRequests { + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table.to_string(), + rows: Some(Rows { + schema: schema.to_vec(), + rows, + }), + }], + } + } + + /// `(ts, host, kind, value)` of every row, sorted, since a set result keeps `kind`. + async fn rows(output: Output) -> Vec<(i64, String, String, f64)> { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("TQL EVAL returns record batches"), + }; + let mut rows = Vec::new(); + for batch in batches.iter() { + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let text = |name: &str| { + cast( + batch.column_by_name(name).unwrap().as_ref(), + &DataType::Utf8, + ) + .unwrap() + }; + let (hosts, kinds) = (text("host"), text("kind")); + let hosts = hosts.as_any().downcast_ref::().unwrap(); + let kinds = kinds.as_any().downcast_ref::().unwrap(); + let values = batch + .column_by_name("val") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + timestamps.value(row), + hosts.value(row).to_string(), + kinds.value(row).to_string(), + values.value(row), + )); + } + } + rows.sort_by(|left, right| (left.0, &left.1, &left.2).cmp(&(right.0, &right.1, &right.2))); + rows + } + + common_telemetry::init_default_ut_logging(); + let config = mito2::config::MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new("native_promql_set_operators_e2e") + .with_mito_config(config), + ) + .await; + let instance = standalone.fe_instance().clone(); + let table = "native_set_pair"; + execute_all( + &instance, + &format!( + "CREATE TABLE {table} (host STRING, kind STRING, ts TIMESTAMP TIME INDEX, val DOUBLE, PRIMARY KEY(host, kind))" + ), + QueryContext::arc(), + ) + .await; + let schema = vec![ + ColumnSchema { + column_name: "host".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "kind".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "val".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ]; + let mut batch = Vec::with_capacity(BATCH_ROWS); + for host in 0..HOSTS { + for (kind, intercept) in [("a", 100.0), ("b", 1_000.0)] { + for point in 0..POINTS { + batch.push(Row { + values: vec![ + ValueData::StringValue(format!("h{host:04}")).into(), + ValueData::StringValue(kind.to_string()).into(), + ValueData::TimestampMillisecondValue(point * 1_000).into(), + ValueData::F64Value(intercept + point as f64).into(), + ], + }); + if batch.len() == BATCH_ROWS { + let rows = std::mem::replace(&mut batch, Vec::with_capacity(BATCH_ROWS)); + instance + .handle_row_inserts( + insert_request(table, &schema, rows), + QueryContext::arc(), + false, + false, + ) + .await + .unwrap(); + } + } + } + } + if !batch.is_empty() { + instance + .handle_row_inserts( + insert_request(table, &schema, batch), + QueryContext::arc(), + false, + false, + ) + .await + .unwrap(); + } + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{table}')"), + QueryContext::arc(), + ) + .await; + + let host = |host: usize| format!("h{host:04}"); + let at_five = |hosts: &mut dyn Iterator, kind: &str, value: f64| { + hosts + .map(|index| (5_000, host(index), kind.to_string(), value)) + .collect::>() + }; + let mut union = at_five(&mut (0..10), "a", 105.0); + union.extend(at_five(&mut (0..100), "b", 1_005.0)); + union.sort_by(|left, right| (left.0, &left.1, &left.2).cmp(&(right.0, &right.1, &right.2))); + // `spread` marks the queries whose left scan has to span several partitions: those are the + // ones a per-partition answer would get wrong. + for (query, spread, expected) in [ + ( + format!("{table}{{kind=\"a\"}} and on() {table}{{kind=\"b\",host=\"h0000\"}}"), + true, + at_five(&mut (0..HOSTS), "a", 105.0), + ), + ( + format!("{table}{{kind=\"a\"}} unless on() {table}{{kind=\"b\",host=\"h0000\"}}"), + true, + Vec::new(), + ), + ( + format!("{table}{{kind=\"a\"}} unless on(host) {table}{{kind=\"b\",host=~\"h00..\"}}"), + true, + at_five(&mut (100..HOSTS), "a", 105.0), + ), + ( + format!( + "{table}{{kind=\"a\",host=~\"h000.\"}} or {table}{{kind=\"b\",host=~\"h00..\"}}" + ), + false, + union, + ), + ] { + let sql = format!("TQL EVAL (5, 5, '1s') {query}"); + observer.clear(); + let native = rows(native_tql_output(&instance, &sql, true).await).await; + let plan = observer.take(); + assert_eq!(plan.name(), "FusedSeriesExec", "{query}"); + assert_eq!(plan.output_partitioning().partition_count(), 1, "{query}"); + let children = plan.children(); + assert_eq!(children.len(), 2, "{query}"); + let left = children[0] + .downcast_ref::() + .map_or(children[0], |cooperative| cooperative.input()); + let left = left.downcast_ref::().unwrap(); + assert_eq!(left.properties().partitioning.partition_count(), 4); + let nonempty = left + .metrics() + .unwrap() + .iter() + .filter(|metric| { + matches!( + metric.value(), + datafusion::physical_plan::metrics::MetricValue::OutputRows(count) + if count.value() > 0 + ) + }) + .filter_map(|metric| metric.partition()) + .collect::>(); + assert!( + !spread || nonempty.len() > 1, + "the left scan of `{query}` must spread over partitions: {nonempty:?}" + ); + assert_eq!(native, expected, "native: {query}"); + let legacy = rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(native, legacy, "legacy: {query}"); + } +} + /// Repeated selectors retain correct binary results while using independent physical scans. #[tokio::test(flavor = "multi_thread")] async fn native_promql_repeated_selector_binary_uses_independent_scans() {