mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
feat(promql): expose selector scan requirements and input bindings (#9870)
I hereby agree to the terms of the [GreptimeDB CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md). ## Refer to a related PR or issue link (optional) First of three; vector–vector lowering (#9871) follows on this input contract, and the multi-scan execution wiring (#9896) follows that. Base is the fused wiring branch, so the three form one reviewable stack. ## What's changed and what's your intention? Expose selector-specific scan requirements and executor input bindings from the native PromQL compiler without expanding its supported expression set. - Add `compile_selectors`, whose callback receives metric, matchers, resolved offset/anchor, actual evaluation endpoints, and instant lookback or range-window duration, and returns the selected raw schema. - Return each requirement and schema with its ordering and executor `NodePath`, obtained through the existing tree walker. - Provide checked bounding scan intervals. Window and positive-lookback lower edges are excluded; zero-lookback instant selection includes the timestamp. The interval may contain gaps not used by individual evaluation windows. - An anchored selector takes **both** edges from its anchor, not just the upper one. Taking the lower edge from the grid would scan from the statement's start, which is the wrong side of the data for an anchor before the statement and skips the selected sample for one after it. Covered by `an_anchored_instant_selector_scans_from_its_anchor`, which pins all three cases: anchored inside the statement, anchored before it, and a zero lookback. - Preserve the existing `compile` API and selector/range/math behavior, including the base's scalar–vector broadcasting (`broadcast` / `scalar_matching`). A vector–vector operation stays refused here, naming the missing per-input identity. Tests cover the callback requirement, schema identity in the actual input node, nested input paths, callback errors, offset/anchor bounds, anchored instant bounds, an unaligned query end, lookback boundaries, and overflow. Validation on `d984231b6cb`: - `cargo nextest run -p promql -p query`: 1800 passed, 2 skipped. - `cargo nextest run -p query --lib promql::native_compile`: 93 passed. - `cargo clippy -p query --all-targets --all-features -- -D warnings`: passed. No persisted or wire-format changes. ## PR Checklist - [x] I have written the necessary rustdoc comments. - [x] I have added the necessary unit tests and integration tests. - [ ] This PR requires documentation updates. - [x] API changes are backward compatible. - [x] Schema or data changes are backward compatible. - [ ] This PR needs to be backported to release branches, and I have added the `backport-<target>` labels (e.g. `backport-v1.3` targets `release/v1.3`). Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
@@ -151,6 +151,7 @@ fn rate_tree_of(schema: Arc<SeriesSchema>) -> NativeTree {
|
||||
input: Box::new(NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output: raw,
|
||||
|
||||
@@ -58,6 +58,7 @@ fn raw_input() -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
@@ -210,6 +211,7 @@ fn an_input_reports_whether_it_is_raw_and_what_ordering_it_guarantees() {
|
||||
let evaluated_input = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: contiguous.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Contiguous,
|
||||
}),
|
||||
output: contiguous,
|
||||
@@ -227,6 +229,7 @@ fn an_input_reports_whether_it_is_raw_and_what_ordering_it_guarantees() {
|
||||
let raw_input = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: interleaved.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Interleaved,
|
||||
}),
|
||||
output: interleaved,
|
||||
|
||||
@@ -110,6 +110,7 @@ fn raw_input_with_ordering(ordering: SeriesOrdering) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering,
|
||||
}),
|
||||
output,
|
||||
@@ -122,6 +123,7 @@ fn raw_input_of(count: usize) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
@@ -139,6 +141,7 @@ fn evaluated_input_with_ordering(grid: Arc<EvalGrid>, ordering: SeriesOrdering)
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering,
|
||||
}),
|
||||
output,
|
||||
@@ -4694,6 +4697,7 @@ fn evaluated_input_of(count: usize, grid: Arc<EvalGrid>) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
|
||||
@@ -129,6 +129,12 @@ pub struct AggregateStateSchema {
|
||||
/// rejecting an interleaved input; no other consumer reads it yet.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BoundInput {
|
||||
/// Child index in the DataFusion plan. The compiler initially assigns selector-occurrence
|
||||
/// indices; the caller maps them to child indices when constructing the plan.
|
||||
///
|
||||
/// Multiple logical inputs may reference the same child after plan collapse. This is not a
|
||||
/// runtime `InputId` or an identity for an input leaf.
|
||||
pub source_index: usize,
|
||||
/// What the input produces: raw samples from a scan, or state from an exchange boundary.
|
||||
pub output: NativeType,
|
||||
/// The ordering the input guarantees.
|
||||
|
||||
@@ -62,6 +62,7 @@ fn input() -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
@@ -75,6 +76,7 @@ fn input_on(grid: Arc<EvalGrid>) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Contiguous,
|
||||
}),
|
||||
output,
|
||||
@@ -90,6 +92,7 @@ fn input_with_fields(count: usize) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
@@ -102,6 +105,7 @@ fn interleaved_input() -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Interleaved,
|
||||
}),
|
||||
output,
|
||||
@@ -1190,6 +1194,7 @@ fn an_input_must_report_the_binding_it_was_bound_to() {
|
||||
let node = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output: raw(),
|
||||
@@ -1205,6 +1210,7 @@ fn an_input_can_bind_aggregate_state() {
|
||||
let bound = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: state.output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Contiguous,
|
||||
}),
|
||||
output: state.output.clone(),
|
||||
@@ -1218,6 +1224,7 @@ fn an_input_can_bind_aggregate_state() {
|
||||
let other = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: state.output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::Contiguous,
|
||||
}),
|
||||
output: partial_aggregate(window_over(input(), None)).output,
|
||||
|
||||
@@ -110,6 +110,7 @@ fn source(schema: &Arc<SeriesSchema>) -> NativeNode {
|
||||
NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: output.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output,
|
||||
|
||||
@@ -718,6 +718,7 @@ mod tests {
|
||||
let input = NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw.clone(),
|
||||
source_index: 0,
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output: raw,
|
||||
|
||||
@@ -128,6 +128,7 @@
|
||||
//! chunks of different series that interleave, so an `Interleaved` ordering is refused once, at the
|
||||
//! entry, and [`NativeNode::validate`] refuses the same ordering for a tree built some other way.
|
||||
|
||||
use std::cell::RefCell;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -146,7 +147,7 @@ use promql::native::plan::{
|
||||
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};
|
||||
use promql_parser::label::MatchOp as ParserMatchOp;
|
||||
use promql_parser::label::{MatchOp as ParserMatchOp, Matchers};
|
||||
use promql_parser::parser::token::TokenType;
|
||||
use promql_parser::parser::value::ValueType;
|
||||
use promql_parser::parser::{
|
||||
@@ -157,6 +158,68 @@ use promql_parser::parser::{
|
||||
|
||||
use crate::promql::label_values::signed_millis_since_epoch;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum SelectorTiming {
|
||||
Instant { lookback_ms: i64 },
|
||||
Window { range_ms: i64 },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SelectorRequirement {
|
||||
pub metric: Option<String>,
|
||||
pub matchers: Matchers,
|
||||
pub offset_ms: i64,
|
||||
pub at_ms: Option<i64>,
|
||||
pub grid_start_ms: i64,
|
||||
pub grid_end_ms: i64,
|
||||
pub timing: SelectorTiming,
|
||||
}
|
||||
|
||||
impl SelectorRequirement {
|
||||
/// Returns `(start_ms, end_ms, include_start)` for the scan hull covering this selector.
|
||||
///
|
||||
/// The upper endpoint is inclusive. The hull can include timestamps unused by a sparse
|
||||
/// evaluation grid; `grid_end_ms` is its actual last point, not the statement's end. The lower
|
||||
/// endpoint is inclusive only for a zero-lookback instant selector.
|
||||
pub fn raw_scan_bounds(&self) -> Result<(i64, i64, bool)> {
|
||||
let end = self
|
||||
.at_ms
|
||||
.unwrap_or(self.grid_end_ms)
|
||||
.checked_sub(self.offset_ms)
|
||||
.ok_or_else(|| DataFusionError::Plan("selector scan end overflows".into()))?;
|
||||
let lower = self
|
||||
.at_ms
|
||||
.unwrap_or(self.grid_start_ms)
|
||||
.checked_sub(self.offset_ms)
|
||||
.ok_or_else(|| DataFusionError::Plan("selector scan start overflows".into()))?;
|
||||
let (duration, inclusive) = match self.timing {
|
||||
SelectorTiming::Instant { lookback_ms } => (lookback_ms, lookback_ms == 0),
|
||||
SelectorTiming::Window { range_ms } => (range_ms, false),
|
||||
};
|
||||
let start = lower
|
||||
.checked_sub(duration)
|
||||
.ok_or_else(|| DataFusionError::Plan("selector scan start overflows".into()))?;
|
||||
Ok((start, end, inclusive))
|
||||
}
|
||||
}
|
||||
|
||||
/// One compiled selector occurrence and the scan layout bound for it.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct CompiledInput {
|
||||
/// Occurrence number assigned in compiler traversal order, starting at zero. The caller maps this
|
||||
/// occurrence to the appropriate scan; different occurrences may map to a shared scan child.
|
||||
pub source_index: usize,
|
||||
pub requirement: SelectorRequirement,
|
||||
pub schema: Arc<SeriesSchema>,
|
||||
pub ordering: SeriesOrdering,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct CompiledTree {
|
||||
pub tree: NativeTree,
|
||||
pub inputs: Arc<[CompiledInput]>,
|
||||
}
|
||||
|
||||
/// The shapes this compiler accepts, quoted in every rejection.
|
||||
///
|
||||
/// It is one sentence of grammar rather than a list of examples, because it is what a refused query
|
||||
@@ -208,18 +271,32 @@ const SUPPORTED: &str = "a selector — `x`, optionally with an `offset` — a r
|
||||
/// [`NativeNode::validate`] refuses the same ordering in its `Instant` and `Window` branches, as a
|
||||
/// backstop for a tree built some other way.
|
||||
///
|
||||
/// The tree carries no metric name and no matchers, so `rate(x{job="a"}[5m])` and `rate(y[5m])`
|
||||
/// compile to identical trees, and an `EXPLAIN` built on this display cannot tell them apart. That
|
||||
/// is fine while a tree has one selector; a tree with two inputs will need each `Input` to name its
|
||||
/// selector.
|
||||
/// This compatibility entry accepts exactly one selector. Call [`compile_selectors`] to retain
|
||||
/// each selector's requirement and binding when compiling more than one input.
|
||||
pub fn compile(
|
||||
stmt: &EvalStmt,
|
||||
input: Arc<SeriesSchema>,
|
||||
ordering: SeriesOrdering,
|
||||
) -> Result<NativeTree> {
|
||||
// Every shape this compiler accepts reads one series to its last chunk before the next, so raw
|
||||
// samples whose chunks of different series interleave cannot be fed to any of them. The refusal
|
||||
// is here rather than in a shape's branch so that a shape added later inherits it.
|
||||
let compiled = compile_selectors(stmt, &|_| Ok(Arc::clone(&input)), ordering)?;
|
||||
if compiled.inputs.len() != 1 {
|
||||
return Err(not_implemented(&format!(
|
||||
"`compile` with {} selector inputs; use `compile_selectors` to bind each input",
|
||||
compiled.inputs.len()
|
||||
)));
|
||||
}
|
||||
Ok(compiled.tree)
|
||||
}
|
||||
|
||||
/// Compiles a statement and returns a binding record for each selector occurrence.
|
||||
///
|
||||
/// `bindings` is called once per selector with its requirements. The resulting source indices are
|
||||
/// occurrence numbers in compiler traversal order; the caller maps them to scan children.
|
||||
pub fn compile_selectors(
|
||||
stmt: &EvalStmt,
|
||||
bindings: &dyn Fn(&SelectorRequirement) -> Result<Arc<SeriesSchema>>,
|
||||
ordering: SeriesOrdering,
|
||||
) -> Result<CompiledTree> {
|
||||
if ordering == SeriesOrdering::Interleaved {
|
||||
return Err(not_implemented(
|
||||
"raw samples whose series interleave, which this compiler cannot consume: a native \
|
||||
@@ -227,11 +304,6 @@ pub fn compile(
|
||||
series may not interleave",
|
||||
));
|
||||
}
|
||||
// A tree reads a scan and reports evaluated samples, so an expression whose value is not a
|
||||
// vector has nothing to compile into: `1 + 2` is a scalar, and the scalar node is bound as a
|
||||
// constant next to a vector that gives it something to broadcast over rather than on its own.
|
||||
// The check is on the statement's type rather than on the compiled node, so it does not depend on
|
||||
// which kinds of expression happen to compile to a scalar today.
|
||||
if stmt.expr.value_type() != ValueType::Vector {
|
||||
return Err(not_implemented(&format!(
|
||||
"the expression `{}`, which is a {} where a tree reports evaluated samples",
|
||||
@@ -239,38 +311,67 @@ pub fn compile(
|
||||
stmt.expr.value_type()
|
||||
)));
|
||||
}
|
||||
let inputs = RefCell::new(Vec::new());
|
||||
let context = Context {
|
||||
stmt,
|
||||
input,
|
||||
bindings,
|
||||
ordering,
|
||||
// The grid is built once for the whole compilation rather than per node: every node of a
|
||||
// tree without a subquery lands on this one, and a subquery replaces it for its own subtree.
|
||||
inputs: &inputs,
|
||||
grid: statement_grid(stmt)?,
|
||||
};
|
||||
// The whole tree is validated once, here: a node compiled as another node's input is checked as
|
||||
// part of the tree rather than on its own, so the same node cannot be legal in one place and
|
||||
// illegal in another.
|
||||
NativeTree::try_new(compile_expr(&context, &stmt.expr)?)
|
||||
let tree = NativeTree::try_new(compile_expr(&context, &stmt.expr)?)?;
|
||||
let inputs = inputs.into_inner();
|
||||
validate_inputs(tree.root(), &inputs)?;
|
||||
Ok(CompiledTree {
|
||||
tree,
|
||||
inputs: inputs.into(),
|
||||
})
|
||||
}
|
||||
|
||||
/// What every node of one compilation derives from the statement rather than from the expression.
|
||||
fn validate_inputs(root: &NativeNode, inputs: &[CompiledInput]) -> Result<()> {
|
||||
fn check(node: &NativeNode, inputs: &[CompiledInput], next: &mut usize) -> bool {
|
||||
match &node.expr {
|
||||
NativeExpr::Input(bound) => {
|
||||
let Some(input) = inputs.get(*next) else {
|
||||
return false;
|
||||
};
|
||||
*next += 1;
|
||||
input.source_index == *next - 1
|
||||
&& bound.source_index == input.source_index
|
||||
&& bound.ordering == input.ordering
|
||||
&& bound
|
||||
.output
|
||||
.samples()
|
||||
.is_some_and(|schema| Arc::ptr_eq(schema, &input.schema))
|
||||
}
|
||||
NativeExpr::Scalar(_) => true,
|
||||
NativeExpr::Instant { input, .. }
|
||||
| NativeExpr::Window { input, .. }
|
||||
| NativeExpr::Map { input, .. }
|
||||
| NativeExpr::Vector { input, .. }
|
||||
| NativeExpr::Aggregate { input, .. } => check(input, inputs, next),
|
||||
NativeExpr::Binary { left, right, .. } => {
|
||||
check(left, inputs, next) && check(right, inputs, next)
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut next = 0;
|
||||
if check(root, inputs, &mut next) && next == inputs.len() {
|
||||
return Ok(());
|
||||
}
|
||||
Err(DataFusionError::Internal(
|
||||
"a compiled tree's recorded inputs do not match its input nodes".into(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Shared statement, per-selector binding callback, occurrence records, ordering and current grid.
|
||||
///
|
||||
/// The three are the same for the whole tree: one statement decides the grid, the lookback and the
|
||||
/// anchors, one caller binds the raw input, and one scan guarantees the ordering. Passing them as
|
||||
/// one value is what lets [`compile_expr`] recurse without restating what it was handed.
|
||||
/// A subquery substitutes its own grid while reusing the statement-level bindings and input list.
|
||||
struct Context<'a> {
|
||||
/// The statement being compiled.
|
||||
stmt: &'a EvalStmt,
|
||||
/// The raw series the caller bound; shared by every input node rather than copied per node.
|
||||
input: Arc<SeriesSchema>,
|
||||
/// The ordering the caller's scan guarantees.
|
||||
bindings: &'a dyn Fn(&SelectorRequirement) -> Result<Arc<SeriesSchema>>,
|
||||
ordering: SeriesOrdering,
|
||||
/// The grid the expression being compiled is evaluated on.
|
||||
///
|
||||
/// It is the statement's grid for everything except the inside of a subquery, which opens a grid
|
||||
/// of its own: `rate(x[5m:1m])` evaluates `x` at one-minute steps over the five-minute lookback,
|
||||
/// not at the statement's steps. Carrying it here is what lets [`compile_expr`] recurse into a
|
||||
/// subquery without every node builder having to ask which grid it is on.
|
||||
inputs: &'a RefCell<Vec<CompiledInput>>,
|
||||
grid: Arc<EvalGrid>,
|
||||
}
|
||||
|
||||
@@ -364,9 +465,7 @@ fn binary_operation(context: &Context, binary: &BinaryExpr) -> Result<NativeNode
|
||||
if binary.lhs.value_type() == ValueType::Vector && binary.rhs.value_type() == ValueType::Vector
|
||||
{
|
||||
return Err(not_implemented(&format!(
|
||||
"the binary operation `{binary}` between two vectors, which needs each input to name \
|
||||
the selector it reads: an input binding carries its ordering and what it produces, but \
|
||||
not which selector it was bound to"
|
||||
"the binary operation `{binary}` between two vectors, whose matching and lowering are not supported yet"
|
||||
)));
|
||||
}
|
||||
// What is left has a constant on one side, and a modifier is refused rather than dropped. This
|
||||
@@ -920,22 +1019,31 @@ fn window_selection(
|
||||
// never pays for a grid nothing would read.
|
||||
let output_grid = Arc::clone(&context.grid);
|
||||
|
||||
// A range function reports one value per input field, so the layout is the input's.
|
||||
let (grid_start_ms, grid_end_ms) = grid_endpoints(&output_grid)?;
|
||||
let (source_index, schema) = bind_selector(
|
||||
context,
|
||||
SelectorRequirement {
|
||||
metric: vs.name.clone(),
|
||||
matchers: vs.matchers.clone(),
|
||||
offset_ms,
|
||||
at_ms: anchor_ms,
|
||||
grid_start_ms,
|
||||
grid_end_ms,
|
||||
timing: SelectorTiming::Window { range_ms },
|
||||
},
|
||||
)?;
|
||||
let output = NativeType::Samples(Arc::new(SeriesSchema {
|
||||
fields: Arc::clone(&context.input.fields),
|
||||
fields: Arc::clone(&schema.fields),
|
||||
eval_grid: Some(Arc::clone(&output_grid)),
|
||||
}));
|
||||
|
||||
// One allocation shared by the binding and the node's output: `NativeTree::try_new` compares
|
||||
// them by identity, because a compiler that meant one schema has no reason to allocate it
|
||||
// twice.
|
||||
let raw = NativeType::Samples(Arc::clone(&context.input));
|
||||
let raw = NativeType::Samples(Arc::clone(&schema));
|
||||
|
||||
Ok(NativeNode {
|
||||
expr: NativeExpr::Window {
|
||||
input: Box::new(NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw.clone(),
|
||||
source_index,
|
||||
ordering: context.ordering,
|
||||
}),
|
||||
output: raw,
|
||||
@@ -1033,21 +1141,26 @@ fn subquery_window(
|
||||
in its anchored interval or the statement's range"
|
||||
))
|
||||
})?;
|
||||
// Keep subquery grids deferred: selector binding needs only their spec endpoints, not a
|
||||
// materialized timestamp vector proportional to the range.
|
||||
let child_grid = Arc::new(EvalGrid::try_deferred(child_spec).map_err(grid_error)?);
|
||||
let inner = compile_expr(
|
||||
&Context {
|
||||
stmt: context.stmt,
|
||||
input: Arc::clone(&context.input),
|
||||
bindings: context.bindings,
|
||||
inputs: context.inputs,
|
||||
ordering: context.ordering,
|
||||
grid: Arc::clone(&child_grid),
|
||||
},
|
||||
expr,
|
||||
)?;
|
||||
|
||||
// A range function reports one value per input field, and the inner expression reports one per
|
||||
// field of the same raw input, so the layout is the input's.
|
||||
let inner_schema = inner
|
||||
.output
|
||||
.samples()
|
||||
.ok_or_else(|| DataFusionError::Internal("subquery body did not produce samples".into()))?;
|
||||
let output = NativeType::Samples(Arc::new(SeriesSchema {
|
||||
fields: Arc::clone(&context.input.fields),
|
||||
fields: Arc::clone(&inner_schema.fields),
|
||||
eval_grid: Some(Arc::clone(&context.grid)),
|
||||
}));
|
||||
|
||||
@@ -1085,21 +1198,31 @@ fn instant_selection(context: &Context, selector: &VectorSelector) -> Result<Nat
|
||||
let lookback_ms = lookback_millis(context.stmt.lookback_delta.as_millis())?;
|
||||
let offset_ms = selector_offset_ms(&selector.offset)?;
|
||||
let output_grid = Arc::clone(&context.grid);
|
||||
|
||||
// An instant selection reports one value per input field, so the layout is the input's.
|
||||
let (grid_start_ms, grid_end_ms) = grid_endpoints(&output_grid)?;
|
||||
let (source_index, schema) = bind_selector(
|
||||
context,
|
||||
SelectorRequirement {
|
||||
metric: selector.name.clone(),
|
||||
matchers: selector.matchers.clone(),
|
||||
offset_ms,
|
||||
at_ms: None,
|
||||
grid_start_ms,
|
||||
grid_end_ms,
|
||||
timing: SelectorTiming::Instant { lookback_ms },
|
||||
},
|
||||
)?;
|
||||
let output = NativeType::Samples(Arc::new(SeriesSchema {
|
||||
fields: Arc::clone(&context.input.fields),
|
||||
fields: Arc::clone(&schema.fields),
|
||||
eval_grid: Some(Arc::clone(&output_grid)),
|
||||
}));
|
||||
|
||||
// One allocation shared by the binding and the node's output, as in the window shape.
|
||||
let raw = NativeType::Samples(Arc::clone(&context.input));
|
||||
let raw = NativeType::Samples(Arc::clone(&schema));
|
||||
|
||||
Ok(NativeNode {
|
||||
expr: NativeExpr::Instant {
|
||||
input: Box::new(NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
output: raw.clone(),
|
||||
source_index,
|
||||
ordering: context.ordering,
|
||||
}),
|
||||
output: raw,
|
||||
@@ -1114,6 +1237,29 @@ fn instant_selection(context: &Context, selector: &VectorSelector) -> Result<Nat
|
||||
})
|
||||
}
|
||||
|
||||
fn bind_selector(
|
||||
context: &Context,
|
||||
requirement: SelectorRequirement,
|
||||
) -> Result<(usize, Arc<SeriesSchema>)> {
|
||||
let schema = (context.bindings)(&requirement)?;
|
||||
let source_index = context.inputs.borrow().len();
|
||||
context.inputs.borrow_mut().push(CompiledInput {
|
||||
source_index,
|
||||
requirement,
|
||||
schema: Arc::clone(&schema),
|
||||
ordering: context.ordering,
|
||||
});
|
||||
Ok((source_index, schema))
|
||||
}
|
||||
|
||||
fn grid_endpoints(grid: &EvalGrid) -> Result<(i64, i64)> {
|
||||
let spec = grid.spec();
|
||||
let end = i128::from(spec.start_ms) + i128::from(spec.len - 1) * i128::from(spec.step_ms);
|
||||
let end = i64::try_from(end)
|
||||
.map_err(|_| DataFusionError::Internal("evaluation grid endpoint overflows".into()))?;
|
||||
Ok((spec.start_ms, end))
|
||||
}
|
||||
|
||||
/// The grid the statement is evaluated on: `(end - start) / step + 1` points, the last of which
|
||||
/// does not pass `end`.
|
||||
///
|
||||
|
||||
@@ -27,8 +27,8 @@ use promql::evaluated::EvaluatedChunk;
|
||||
use promql::map::MathOp;
|
||||
use promql::native::execute::{InputId, NativeExecutionState, NativeExecutor};
|
||||
use promql::native::plan::{
|
||||
AggregateKind, AggregateStage, BoundAggregate, BoundInstant, BoundMap, BoundWindow, Grouping,
|
||||
MapOp, NativeExpr, NativeNode, NativeTree, NativeType,
|
||||
AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundWindow,
|
||||
Grouping, MapOp, NativeExpr, NativeNode, NativeTree, NativeType,
|
||||
};
|
||||
use promql::range_eval::RangeFunction;
|
||||
use promql::vector_match::{BinaryMatching, BinaryOp};
|
||||
@@ -37,7 +37,9 @@ use promql_parser::parser::{
|
||||
VectorMatchFillValues,
|
||||
};
|
||||
|
||||
use super::compile;
|
||||
use super::{
|
||||
CompiledInput, SelectorRequirement, SelectorTiming, compile, compile_selectors, validate_inputs,
|
||||
};
|
||||
|
||||
/// The raw series schema a caller binds from the table: one field and no evaluation grid.
|
||||
fn raw_schema() -> Arc<SeriesSchema> {
|
||||
@@ -1621,9 +1623,7 @@ fn a_constant_operand_binds_the_value_it_folds_to() {
|
||||
}
|
||||
}
|
||||
|
||||
/// An operation between two vectors is refused because the tree cannot say which selector each
|
||||
/// `Input` reads: `a + b` and `a + a` would compile to the same tree. The refusal names the missing
|
||||
/// identity rather than the query it found, so whoever adds it knows what to add.
|
||||
/// Binary operations between vectors remain unsupported in the native lowering stage.
|
||||
#[test]
|
||||
fn an_operation_between_two_vectors_is_refused_for_the_missing_selector_identity() {
|
||||
for query in [
|
||||
@@ -4327,3 +4327,296 @@ fn an_invalid_label_replace_expression_fails_the_statement() {
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, DataFusionError::Plan(_)), "{error}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_requirements_capture_bounds_and_source_occurrence() {
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt("rate(x[10s] offset 2s @ 50)"),
|
||||
&|requirement| {
|
||||
assert_eq!(requirement.metric.as_deref(), Some("x"));
|
||||
assert_eq!(requirement.raw_scan_bounds()?, (38_000, 48_000, false));
|
||||
Ok(raw_schema())
|
||||
},
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(compiled.inputs.len(), 1);
|
||||
assert_eq!(compiled.inputs[0].source_index, 0);
|
||||
assert_eq!(
|
||||
compiled.inputs[0].requirement.timing,
|
||||
SelectorTiming::Window { range_ms: 10_000 }
|
||||
);
|
||||
let NativeExpr::Window { input, .. } = &compiled.tree.root().expr else {
|
||||
panic!("not window")
|
||||
};
|
||||
let NativeExpr::Input(bound) = &input.expr else {
|
||||
panic!("not input")
|
||||
};
|
||||
assert_eq!(bound.source_index, 0);
|
||||
}
|
||||
|
||||
fn selector_requirement(
|
||||
grid_start_ms: i64,
|
||||
grid_end_ms: i64,
|
||||
offset_ms: i64,
|
||||
at_ms: Option<i64>,
|
||||
timing: SelectorTiming,
|
||||
) -> SelectorRequirement {
|
||||
SelectorRequirement {
|
||||
metric: None,
|
||||
matchers: promql_parser::label::Matchers::empty(),
|
||||
offset_ms,
|
||||
at_ms,
|
||||
grid_start_ms,
|
||||
grid_end_ms,
|
||||
timing,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_scan_bounds_cover_instant_ranges_anchors_and_overflow() {
|
||||
let cases = [
|
||||
("x", (-5_000, 95_000, false)), // 5s lookback, end before non-aligned statement end
|
||||
("rate(x[10s])", (-10_000, 95_000, false)),
|
||||
("rate(x[10s] offset 2s)", (-12_000, 93_000, false)),
|
||||
("rate(x[10s] offset -2s)", (-8_000, 97_000, false)),
|
||||
("rate(x[10s] @ start())", (-10_000, 0, false)),
|
||||
("rate(x[10s] @ end())", (87_000, 97_000, false)),
|
||||
("rate(x[10s] @ end() offset 2s)", (85_000, 95_000, false)),
|
||||
];
|
||||
for (query, expected) in cases {
|
||||
let mut stmt = eval_stmt_over(query, Duration::from_secs(97));
|
||||
if query == "x" {
|
||||
stmt.lookback_delta = Duration::from_secs(5);
|
||||
}
|
||||
let tree = compile_selectors(&stmt, &|_| Ok(raw_schema()), SeriesOrdering::KeyThenTime)
|
||||
.unwrap_or_else(|e| panic!("{query}: {e}"));
|
||||
assert_eq!(
|
||||
tree.inputs[0].requirement.raw_scan_bounds().unwrap(),
|
||||
expected,
|
||||
"{query}"
|
||||
);
|
||||
}
|
||||
|
||||
let mut zero = eval_stmt("x");
|
||||
zero.lookback_delta = Duration::ZERO;
|
||||
let tree =
|
||||
compile_selectors(&zero, &|_| Ok(raw_schema()), SeriesOrdering::KeyThenTime).unwrap();
|
||||
assert_eq!(
|
||||
tree.inputs[0].requirement.raw_scan_bounds().unwrap(),
|
||||
(0, 100_000, true)
|
||||
);
|
||||
|
||||
for (req, edge) in [
|
||||
(
|
||||
selector_requirement(
|
||||
0,
|
||||
0,
|
||||
-1,
|
||||
Some(i64::MAX),
|
||||
SelectorTiming::Instant { lookback_ms: 0 },
|
||||
),
|
||||
"end",
|
||||
),
|
||||
(
|
||||
selector_requirement(
|
||||
i64::MIN,
|
||||
0,
|
||||
1,
|
||||
None,
|
||||
SelectorTiming::Instant { lookback_ms: 0 },
|
||||
),
|
||||
"start",
|
||||
),
|
||||
(
|
||||
selector_requirement(i64::MIN, 0, 0, None, SelectorTiming::Window { range_ms: 1 }),
|
||||
"duration",
|
||||
),
|
||||
] {
|
||||
assert!(
|
||||
matches!(req.raw_scan_bounds(), Err(DataFusionError::Plan(_))),
|
||||
"{edge}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subquery_selector_bounds_use_inner_grid_and_selector_offset() {
|
||||
for (query, expected_grid, expected_scan) in [
|
||||
(
|
||||
"max_over_time(rate(x[1m] offset 10s)[5m:30s])",
|
||||
(-270_000, 90_000),
|
||||
(-340_000, 80_000, false),
|
||||
),
|
||||
(
|
||||
"max_over_time(rate(x[1m])[5m:30s] @ 1234 offset 1m)",
|
||||
(900_000, 1_170_000),
|
||||
(840_000, 1_170_000, false),
|
||||
),
|
||||
] {
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt(query),
|
||||
&|_| Ok(raw_schema()),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap_or_else(|e| panic!("{query}: {e}"));
|
||||
let req = &compiled.inputs[0].requirement;
|
||||
assert_eq!(
|
||||
(req.grid_start_ms, req.grid_end_ms),
|
||||
expected_grid,
|
||||
"{query}"
|
||||
);
|
||||
assert_eq!(req.raw_scan_bounds().unwrap(), expected_scan, "{query}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_binding_callback_error_is_returned_once() {
|
||||
let calls = std::cell::Cell::new(0);
|
||||
let error = compile_selectors(
|
||||
&eval_stmt("rate(x[1m])"),
|
||||
&|_| {
|
||||
calls.set(calls.get() + 1);
|
||||
Err(DataFusionError::Plan("bind failed".into()))
|
||||
},
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert_eq!(calls.get(), 1);
|
||||
assert!(matches!(error, DataFusionError::Plan(_)));
|
||||
assert!(error.to_string().contains("bind failed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_inputs_rejects_misbound_occurrences_and_accepts_scalar_wrappers() {
|
||||
for query in ["x * 2", "2 * x"] {
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt(query),
|
||||
&|_| Ok(raw_schema()),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
validate_inputs(compiled.tree.root(), &compiled.inputs).unwrap();
|
||||
}
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt("x * 2"),
|
||||
&|_| Ok(raw_schema()),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
let original = compiled.inputs[0].clone();
|
||||
|
||||
for records in [Vec::new(), vec![original.clone(), original.clone()]] {
|
||||
assert!(matches!(
|
||||
validate_inputs(compiled.tree.root(), &records),
|
||||
Err(DataFusionError::Internal(_))
|
||||
));
|
||||
}
|
||||
let mut wrong_source_record = original.clone();
|
||||
wrong_source_record.source_index = 1;
|
||||
assert!(validate_inputs(compiled.tree.root(), &[wrong_source_record]).is_err());
|
||||
|
||||
let scalar_output = compiled.tree.root().output.clone();
|
||||
fn set_input_source(node: &mut NativeNode, source: usize) {
|
||||
match &mut node.expr {
|
||||
NativeExpr::Input(bound) => bound.source_index = source,
|
||||
NativeExpr::Instant { input, .. }
|
||||
| NativeExpr::Window { input, .. }
|
||||
| NativeExpr::Map { input, .. }
|
||||
| NativeExpr::Vector { input, .. }
|
||||
| NativeExpr::Aggregate { input, .. } => set_input_source(input, source),
|
||||
NativeExpr::Binary { left, right, .. } => {
|
||||
set_input_source(left, source);
|
||||
set_input_source(right, source);
|
||||
}
|
||||
NativeExpr::Scalar(_) => {}
|
||||
}
|
||||
}
|
||||
let mut wrong_bound_source = compiled.tree.root().clone();
|
||||
set_input_source(&mut wrong_bound_source, 1);
|
||||
assert!(validate_inputs(&wrong_bound_source, std::slice::from_ref(&original)).is_err());
|
||||
|
||||
let wrong_schema = Arc::new(SeriesSchema {
|
||||
fields: Arc::clone(&original.schema.fields),
|
||||
eval_grid: None,
|
||||
});
|
||||
let mut wrong_record = original.clone();
|
||||
wrong_record.schema = wrong_schema;
|
||||
assert!(validate_inputs(compiled.tree.root(), &[wrong_record]).is_err());
|
||||
let mut wrong_ordering = original;
|
||||
wrong_ordering.ordering = SeriesOrdering::Contiguous;
|
||||
assert!(validate_inputs(compiled.tree.root(), &[wrong_ordering]).is_err());
|
||||
|
||||
// A scalar is not an input leaf and must not consume an input record.
|
||||
let scalar = NativeNode {
|
||||
expr: NativeExpr::Scalar(promql::native::plan::BoundScalar { value: 2.0 }),
|
||||
output: scalar_output,
|
||||
};
|
||||
assert!(validate_inputs(&scalar, &[]).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_inputs_accepts_distinct_occurrences_sharing_a_schema() {
|
||||
let schema = raw_schema();
|
||||
let input = |source_index| NativeNode {
|
||||
expr: NativeExpr::Input(BoundInput {
|
||||
source_index,
|
||||
output: NativeType::Samples(Arc::clone(&schema)),
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
}),
|
||||
output: NativeType::Samples(Arc::clone(&schema)),
|
||||
};
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt("x"),
|
||||
&|_| Ok(Arc::clone(&schema)),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
let scalar_tree = compile_selectors(
|
||||
&eval_stmt("x + 1"),
|
||||
&|_| Ok(Arc::clone(&schema)),
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
let NativeExpr::Binary { matching, .. } = &scalar_tree.tree.root().expr else {
|
||||
panic!("not binary")
|
||||
};
|
||||
let tree = NativeNode {
|
||||
expr: NativeExpr::Binary {
|
||||
left: Box::new(input(0)),
|
||||
right: Box::new(input(1)),
|
||||
matching: matching.clone(),
|
||||
},
|
||||
output: NativeType::Samples(Arc::clone(&schema)),
|
||||
};
|
||||
let requirements = [0, 1].map(|source_index| CompiledInput {
|
||||
source_index,
|
||||
requirement: compiled.inputs[0].requirement.clone(),
|
||||
schema: Arc::clone(&schema),
|
||||
ordering: SeriesOrdering::KeyThenTime,
|
||||
});
|
||||
validate_inputs(&tree, &requirements).unwrap();
|
||||
let mut wrong = requirements;
|
||||
wrong[1].source_index = 0;
|
||||
assert!(validate_inputs(&tree, &wrong).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_binding_callback_receives_matchers_once() {
|
||||
let calls = std::cell::Cell::new(0);
|
||||
let compiled = compile_selectors(
|
||||
&eval_stmt("rate(x{job=\"api\"}[1m])"),
|
||||
&|requirement| {
|
||||
calls.set(calls.get() + 1);
|
||||
assert_eq!(requirement.metric.as_deref(), Some("x"));
|
||||
assert_eq!(requirement.matchers.matchers.len(), 1);
|
||||
assert_eq!(requirement.matchers.matchers[0].name, "job");
|
||||
assert_eq!(requirement.matchers.matchers[0].value, "api");
|
||||
Ok(raw_schema())
|
||||
},
|
||||
SeriesOrdering::KeyThenTime,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(calls.get(), 1);
|
||||
assert_eq!(compiled.inputs.len(), 1);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user