diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 17232a8fb56..a339bb67820 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -151,6 +151,7 @@ fn rate_tree_of(schema: Arc) -> NativeTree { input: Box::new(NativeNode { expr: NativeExpr::Input(BoundInput { output: raw.clone(), + source_index: 0, ordering: SeriesOrdering::KeyThenTime, }), output: raw, diff --git a/src/promql/src/native/display/tests.rs b/src/promql/src/native/display/tests.rs index 89351e184b9..6342510f967 100644 --- a/src/promql/src/native/display/tests.rs +++ b/src/promql/src/native/display/tests.rs @@ -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, diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index 1f0b4e49315..04309090262 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -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, 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) -> NativeNode { NativeNode { expr: NativeExpr::Input(BoundInput { output: output.clone(), + source_index: 0, ordering: SeriesOrdering::KeyThenTime, }), output, diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 7409d4b0b56..69142174fee 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -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. diff --git a/src/promql/src/native/plan/tests.rs b/src/promql/src/native/plan/tests.rs index 36ffb24bf7b..c921b44a49d 100644 --- a/src/promql/src/native/plan/tests.rs +++ b/src/promql/src/native/plan/tests.rs @@ -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) -> 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, diff --git a/src/promql/tests/native_binary.rs b/src/promql/tests/native_binary.rs index 561be21f964..cfa6cf22a7a 100644 --- a/src/promql/tests/native_binary.rs +++ b/src/promql/tests/native_binary.rs @@ -110,6 +110,7 @@ fn source(schema: &Arc) -> NativeNode { NativeNode { expr: NativeExpr::Input(BoundInput { output: output.clone(), + source_index: 0, ordering: SeriesOrdering::KeyThenTime, }), output, diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 4221b009390..34e7c25ce56 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -718,6 +718,7 @@ mod tests { let input = NativeNode { expr: NativeExpr::Input(BoundInput { output: raw.clone(), + source_index: 0, ordering: SeriesOrdering::KeyThenTime, }), output: raw, diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index ef5abdea27b..5df9f1afad6 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -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, + pub matchers: Matchers, + pub offset_ms: i64, + pub at_ms: Option, + 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, + 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, ordering: SeriesOrdering, ) -> Result { - // 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>, + ordering: SeriesOrdering, +) -> Result { 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, - /// The ordering the caller's scan guarantees. + bindings: &'a dyn Fn(&SelectorRequirement) -> Result>, 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>, grid: Arc, } @@ -364,9 +465,7 @@ fn binary_operation(context: &Context, binary: &BinaryExpr) -> Result Result Result Result<(usize, Arc)> { + 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`. /// diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 1a64367fb4a..ec085371907 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -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 { @@ -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, + 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); +}