rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(promql): compile and run set operators natively

Browse Source 
Bind `and`, `or` and `unless` (with `on`/`ignoring`) to a new
`NativeExpr::Set` node instead of refusing them. The node runs the
existing MATCH-SET kernel through the binary runtime, which retains both
operands until the input ends, so a set operation reads every series of
both sides in one execution.

The set kernel now reports its workspace like the binary kernel does:
kept rows, the filter side's match groups, the sort and the merge are
charged to the node's reservation.

The compiler refuses a fill value on a set operator, which Prometheus
rejects and the parser accepts, and sides with different field counts.
Same-table set operations over two scans take the existing two-scan
native route.

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
2026-10-26 16:23:05 +08:00
parent 1cab5f1624
commit 1be0df046c
16 changed files with 1010 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/promql/src/extension_plan/aggregate_finalize.rs
+12 -1
View File
@@ -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<Vec<FieldSpec>> {
| 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))
}
}
src/promql/src/extension_plan/aggregate_finalize/tests.rs
+1 -1
View File
@@ -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);
}
src/promql/src/extension_plan/fused_series.rs
+1 -1
View File
@@ -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);
}
src/promql/src/extension_plan/fused_series/tests.rs
+1 -1
View File
@@ -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);
}
src/promql/src/native/display.rs
+20 -6
View File
@@ -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")?;
}
src/promql/src/native/display/tests.rs
+47 -2
View File
@@ -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<EvalGrid> {
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
src/promql/src/native/execute.rs
+69 -24
View File
@@ -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<BinaryMatching>,
/// 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<SeriesSchema>; 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<BinaryMatching>),
/// `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<BinaryMatching>, schemas: [Arc<SeriesSchema>; 2]) -> Self {
/// The run-time state of `kernel` over its two operands.
fn new(kernel: PairKernel, schemas: [Arc<SeriesSchema>; 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
src/promql/src/native/execute/tests.rs
+175 -6
View File
@@ -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::<Vec<_>>()
.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`.
// ---------------------------------------------------------------------------
src/promql/src/native/plan.rs
+58 -8
View File
@@ -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<str>, 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<NativeNode>,
matching: BinaryMatching,
},
/// A set operation between two vectors, which keeps samples of either side as they are.
Set {
left: Box<NativeNode>,
right: Box<NativeNode>,
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(())
}
src/promql/src/vector_match.rs
+166 -14
View File
@@ -1143,6 +1143,45 @@ pub fn set_vectors(
lhs: &[Series],
rhs_schema: &SeriesSchema,
rhs: &[Series],
) -> Result<Vec<Series>> {
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<Vec<Series>> {
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<Vec<Series>> {
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::<usize>())
.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<Results> {
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::<RowCursor<'_>>()))?;
}
// 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::<String>().saturating_add(message.capacity()))?;
}
errors.push((duplicate.point, 0, message));
}
Ok(results)
}
src/promql/src/vector_match/tests.rs
+39 -1
View File
@@ -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<usize> = (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}"
);
}
src/query/src/promql/native_compile.rs
+86 -27
View File
@@ -35,6 +35,7 @@
//! <expression> <operator> <constant> e.g. x * 100, x > bool 0
//! <constant> <operator> <expression> e.g. 1 / x, 0 < x
//! -<expression> e.g. -x
//! <expression> <set operator> <expression> e.g. x and on(job) y, x unless y
//! <aggregator> [(by|without)(<labels>)] (<expression>)
//! e.g. avg by(job)(rate(x[5m]) * 100)
//! <aggregator> (<expression>) [(by|without)(<labels>)]
@@ -55,7 +56,8 @@
//! offset 1m`
//! ```
//!
//! where `<aggregator>` is `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` or
//! where `<set operator>` is `and`, `or` or `unless`, optionally with `on(...)` or `ignoring(...)`,
//! `<aggregator>` is `sum`, `avg`, `count`, `min`, `max`, `group`, `stddev`, `stdvar` or
//! `quantile`, `<selector aggregator>` 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 @@
//! `<expression>` → `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(<constant>, <expression>)` — 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<NativeNode> {
/// 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<NativeNode> {
// 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<NativeNode
))
}
/// Compiles `and`, `or` or `unless` between two vectors, matched by the labels the query wrote.
///
/// The parser already refuses a scalar operand, `bool` and `group_left` / `group_right` on a set
/// operator. It does not refuse a fill value, which Prometheus does ("filling in missing series not
/// allowed for set operators"), so that is refused here along with any cardinality other than the
/// many-to-many the parser assigns.
fn set_operation(context: &Context, binary: &BinaryExpr, op: SetOp) -> Result<NativeNode> {
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<BinaryMatching> {
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<BinaryMatching>
})
}
/// The set operator a PromQL token denotes, or `None` when it is not one.
fn set_op(op: &TokenType) -> Option<SetOp> {
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<Arc<str>> {
labels.iter().map(|name| Arc::from(name.as_str())).collect()
}
src/query/src/promql/native_compile/tests.rs
+92 -30
View File
@@ -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<f64> {
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);
}
src/query/src/promql/planner/native.rs
+15 -10
View File
@@ -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);
}
src/query/src/promql/planner/test.rs
+1 -1
View File
@@ -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);
}
tests-integration/src/tests/promql_test.rs
+227
View File
@@ -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<Row>) -> 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::<TimestampMillisecondArray>()
.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::<StringArray>().unwrap();
let kinds = kinds.as_any().downcast_ref::<StringArray>().unwrap();
let values = batch
.column_by_name("val")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.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<Item = usize>, kind: &str, value: f64| {
hosts
.map(|index| (5_000, host(index), kind.to_string(), value))
.collect::<Vec<_>>()
};
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::<CooperativeExec>()
.map_or(children[0], |cooperative| cooperative.input());
let left = left.downcast_ref::<SeriesScanExec>().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::<std::collections::HashSet<_>>();
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() {