From d1003ac8feaf7f2d4242d923b110b3beb37ee23e Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 11:46:22 +0800 Subject: [PATCH] feat(promql): run a bound tree over raw series in one fused node `FusedSeriesExec` evaluates a compiled `NativeTree` directly over the raw series its child produces, instead of materialising a row batch through the per-operator chain. It binds the time column, the fields and the tags once, checks the input against that binding when the plan is built, and converts the chunks the tree settles back into the statement's output columns. The node's partition count is the child's, one tree per partition. The child must be a raw-series scan that assigns every series whole to one partition: `PerSeries` distribution is what makes a per-partition evaluation equivalent to evaluating the whole series. A zero-partition input is refused; anything else is the child's business. A label set can recur across partitions once `rate` and the map kernels drop `__name__`, so the occupancy map is plan state rather than per-stream. Two series reaching one label set is an error only when their occupied grid points overlap, which is Prometheus' range-query rule; disjoint points are legal and are merged into one occupancy record. The reservation for that map is attached to the first execution's pool and released when the plan is dropped, because a later partition still has to collide with what an earlier one emitted. `occupied_runs` reads a chunk's absolute positions rather than assuming one contiguous run: a sparse evaluated chunk may hold several, and treating the first as the whole range would count the holes as rows. A poll does a bounded amount of work and wakes itself when the budget runs out, so a long run of chunks that settle nothing cannot monopolise the executor thread. Signed-off-by: discord9 --- src/promql/src/extension_plan.rs | 2 + src/promql/src/extension_plan/fused_series.rs | 1352 +++++++++++++++ .../src/extension_plan/fused_series/tests.rs | 1543 +++++++++++++++++ 3 files changed, 2897 insertions(+) create mode 100644 src/promql/src/extension_plan/fused_series.rs create mode 100644 src/promql/src/extension_plan/fused_series/tests.rs diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 750ee41b5e8..baf57a275c8 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -14,6 +14,7 @@ mod absent; mod empty_metric; +mod fused_series; mod histogram_fold; mod instant_manipulate; mod normalize; @@ -37,6 +38,7 @@ use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; +pub use fused_series::{FusedSeries, FusedSeriesBinding, FusedSeriesExec, FusedSeriesStream}; pub use histogram_fold::{ HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, }; diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs new file mode 100644 index 00000000000..7ab19c1d163 --- /dev/null +++ b/src/promql/src/extension_plan/fused_series.rs @@ -0,0 +1,1352 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! One node that runs a compiled native tree over raw series. +//! +//! [`FusedSeries`] is the logical half: it holds the PromQL text the tree was compiled from, the +//! tree itself and the binding its input was compiled against. [`FusedSeriesExec`] is the physical +//! half: it owns one execution of that tree per partition and converts its input batches into raw +//! series chunks and its evaluated chunks back into batches. +//! +//! The two halves show different things in `EXPLAIN`, which is the point of having both. The +//! logical node prints the query a reader wrote ([`FusedSeries`]'s `fmt_for_explain`), and the +//! physical node prints the tree it will run ([`FusedSeriesExec`]'s `fmt_as`, which embeds +//! [`NativeTree`]'s own display rather than a second renderer). A reader comparing the two plans +//! can therefore see whether a query was compiled the way they expected without decoding node +//! names. +//! +//! Scope of this layer: the input is a stream of ordinary record batches carrying one timestamp +//! column, the series' field columns and the series' tag columns, with each complete label set +//! contiguous and its timestamps ascending. The output is one batch per evaluated chunk, with the +//! statement's grid timestamps in the declared time column. Everything the tree cannot do is +//! rejected when the execution plan is built, not when the first batch arrives: a tree the native +//! runtime does not execute, a binding that does not match the tree, or an input or output schema +//! the conversion cannot classify. + +use std::cmp::Ordering; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use common_query::series_frame::{InputMapping, record_batches_to_raw_series}; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::buffer::ScalarBuffer; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::DFSchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation, UnboundedMemoryPool}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use datatypes::series::MetricName; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::{ + DenseSamplesRef, FieldSpec, SampleRef, Samples, Series, SeriesSchema, +}; +use datatypes::series::time::EvalGrid; +use futures::Stream; +use store_api::region_engine::SendableRawSeriesStream; + +use crate::evaluated::EvaluatedChunk; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::metrics::{self, NodePath}; +use crate::native::plan::{NativeExpr, NativeTree}; + +#[cfg(test)] +mod tests; + +/// The label a metric name is stored under, which is not a tag column of the input. +const METRIC_NAME: &str = "__name__"; + +/// How many input steps one poll of [`FusedSeriesStream`] runs before it wakes itself. +/// +/// A step polls the input once and hands what it settled to the converter. Settling is not +/// guaranteed: raw samples before the first evaluation point settle nothing, so an input of many +/// such chunks would otherwise be read to its end inside one poll, holding the task's thread for +/// an unbounded time and delaying both other tasks on it and the cancellation the consumer asked +/// for. The budget makes that a poll that yields instead, and the waker makes sure it is polled +/// again, so nothing is lost by stopping early. +const POLL_BUDGET: usize = 64; + +/// What the tree's input was compiled against, and how its output columns are named. +/// +/// The binding travels with the tree: a tree is legal only against the schema it was bound to, and +/// a reader of the output needs to know which grid a chunk's absolute positions are on. Keeping it +/// in one value is what lets both halves of the node be built from the same description. +/// +/// What the node checks is that the input's columns are exactly the ones the binding accounts for: +/// the time column, the fields and the tags. What it cannot check is whether the caller meant a +/// *different* column of the same type. A binding whose [`metric_name`](Self::metric_name) is not +/// the metric the scan selects, or whose [`time_index`](Self::time_index) names another millisecond +/// timestamp column of the input, is indistinguishable from a correct one from here: the node sees +/// one timestamp column and one constant name, so picking them is the caller's obligation. +#[derive(Clone, Debug)] +pub struct FusedSeriesBinding { + /// The raw series schema the tree's input node was bound to, without an evaluation grid. + pub series_schema: Arc, + /// The statement's grid, which is the grid every evaluated chunk is addressed on. + pub grid: Arc, + /// The metric name every input series is bound to, when the input reads one named metric. + /// + /// The name is a binding rather than a column because a raw series chunk carries one metric + /// name for all of its samples; a name that varied inside a series would not be one series. + pub metric_name: Option>, + /// The input column holding the timestamps. + pub time_index: String, + /// The logical tag names the output carries, read from each series' labels. + /// + /// [`METRIC_NAME`] is special: it is resolved to the bound metric name instead of a column, + /// which is how a reader sees that a range function dropped it. + pub tags: Vec, +} + +impl FusedSeriesBinding { + /// Whether `other` binds the same fields under the same names. + fn same_fields(&self, other: &SeriesSchema) -> bool { + self.series_schema.fields.len() == other.fields.len() + && self + .series_schema + .fields + .iter() + .zip(other.fields.iter()) + .all(|(mine, theirs)| mine.name == theirs.name) + } +} + +impl PartialEq for FusedSeriesBinding { + fn eq(&self, other: &Self) -> bool { + // The schemas are compared by the field names they declare, and the grids by their spec: + // neither type has structural equality, and neither identity nor allocation is meant here. + self.metric_name == other.metric_name + && self.time_index == other.time_index + && self.tags == other.tags + && self.grid.spec() == other.grid.spec() + && self.same_fields(&other.series_schema) + } +} + +impl Eq for FusedSeriesBinding {} + +impl Hash for FusedSeriesBinding { + fn hash(&self, state: &mut H) { + self.metric_name.hash(state); + self.time_index.hash(state); + self.tags.hash(state); + let spec = self.grid.spec(); + (spec.start_ms, spec.step_ms, spec.len).hash(state); + for field in self.series_schema.fields.iter() { + field.name.hash(state); + } + } +} + +impl PartialOrd for FusedSeriesBinding { + fn partial_cmp(&self, other: &Self) -> Option { + // A total order over the same fields equality reads, so equal bindings compare equal. + match self + .metric_name + .as_deref() + .partial_cmp(&other.metric_name.as_deref()) + { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.time_index.partial_cmp(&other.time_index) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tags.partial_cmp(&other.tags) { + Some(Ordering::Equal) => {} + ord => return ord, + } + let (mine, theirs) = (self.grid.spec(), other.grid.spec()); + (mine.start_ms, mine.step_ms, mine.len).partial_cmp(&( + theirs.start_ms, + theirs.step_ms, + theirs.len, + )) + } +} + +/// The logical plan of one node that runs a compiled native tree. +/// +/// The node adds nothing to the tree: it records the query the tree was compiled from, so +/// `EXPLAIN` shows the reader their own query rather than the shape it compiled to, and it carries +/// the binding the tree's input was compiled against so the execution plan can be built from it. +#[derive(Debug)] +pub struct FusedSeries { + /// The PromQL text this tree was compiled from. + promql: Arc, + /// The compiled tree to run. + tree: NativeTree, + /// What the tree's input was compiled against. + binding: FusedSeriesBinding, + /// The plan producing the input series. + input: LogicalPlan, + /// The columns this node produces. + output_schema: DFSchemaRef, +} + +impl FusedSeries { + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeries" + } + + /// Creates the node, checking that `tree` is legal against `binding`. + pub fn try_new( + promql: Arc, + tree: NativeTree, + binding: FusedSeriesBinding, + input: LogicalPlan, + output_schema: DFSchemaRef, + ) -> DataFusionResult { + validate_binding(&tree, &binding)?; + Ok(Self { + promql, + tree, + binding, + input, + output_schema, + }) + } + + /// The PromQL text this node was compiled from. + pub fn promql(&self) -> &str { + &self.promql + } + + /// The compiled tree this node runs. + pub fn tree(&self) -> &NativeTree { + &self.tree + } + + /// What the tree's input was compiled against. + pub fn binding(&self) -> &FusedSeriesBinding { + &self.binding + } + + /// The physical node that runs `exec_input` as this node's input. + /// + /// Building it rejects a tree the runtime cannot execute and a schema the conversion cannot + /// classify, so a plan that is built is a plan that runs. + pub fn to_execution_plan( + &self, + exec_input: Arc, + ) -> DataFusionResult> { + let output_schema: SchemaRef = self.output_schema.inner().clone(); + Ok(Arc::new(FusedSeriesExec::try_new( + Arc::clone(&self.promql), + Arc::new(self.tree.clone()), + Arc::new(self.binding.clone()), + exec_input, + output_schema, + )?)) + } +} + +impl UserDefinedLogicalNodeCore for FusedSeries { + fn name(&self) -> &str { + Self::name() + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + // The tree's operands are bound inside it; there is no expression for an optimizer to + // rewrite separately from the tree that owns it. + Vec::new() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeries: {}", self.promql) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DataFusionResult { + if inputs.len() != 1 { + return Err(DataFusionError::Internal( + "FusedSeries should have at exact one input".to_string(), + )); + } + + let input = inputs.pop().unwrap(); + Ok(Self { + promql: Arc::clone(&self.promql), + tree: self.tree.clone(), + binding: self.binding.clone(), + input, + output_schema: Arc::clone(&self.output_schema), + }) + } +} + +impl PartialEq for FusedSeries { + fn eq(&self, other: &Self) -> bool { + // `NativeTree` carries no structural equality, so the rendered tree stands in for it. The + // renderer names every node with its bound operation and grids, which is what makes two + // nodes the same computation; comparing the text alone would call two different trees + // equal whenever their PromQL happens to read the same. + self.promql == other.promql + && self.tree.to_string() == other.tree.to_string() + && self.binding == other.binding + && self.input == other.input + && self.output_schema == other.output_schema + } +} + +impl Eq for FusedSeries {} + +impl Hash for FusedSeries { + fn hash(&self, state: &mut H) { + self.promql.hash(state); + self.binding.hash(state); + self.input.hash(state); + self.output_schema.hash(state); + } +} + +impl PartialOrd for FusedSeries { + fn partial_cmp(&self, other: &Self) -> Option { + match self.promql.partial_cmp(&other.promql) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.binding.partial_cmp(&other.binding) { + Some(Ordering::Equal) => {} + ord => return ord, + } + self.input.partial_cmp(&other.input) + } +} + +/// Checks that `tree` reads the schema `binding` declares, on the grid it declares. +/// +/// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule +/// [`NativeTree`]'s own validation and the kernels apply, and a chunk's timestamps already come +/// from its absolute grid positions, so a separately allocated grid with the same spec is the same +/// grid to everything downstream. +/// +/// Returns the position of the tree's input node, which is what a caller names when it hands a raw +/// chunk over: a position is how the runtime addresses a node, so nothing on the node has to carry +/// an identity that the tree itself does not. +fn validate_binding(tree: &NativeTree, binding: &FusedSeriesBinding) -> DataFusionResult { + if binding.series_schema.eval_grid.is_some() { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema, which cannot have an evaluation grid" + .to_string(), + )); + } + if binding.tags.iter().any(|tag| tag == METRIC_NAME) && binding.metric_name.is_none() { + return Err(DataFusionError::Plan(format!( + "FusedSeries declares a {METRIC_NAME:?} tag column but binds no metric name" + ))); + } + + // The tree's own input node says which schema it reads, so the binding is checked against the + // tree rather than trusted: a mismatch would evaluate one schema and label the result as + // another. + let mut sources = 0; + let mut bound = false; + let mut input_path = None; + metrics::walk(tree.root(), &mut |node, path| { + let NativeExpr::Input(input) = &node.expr else { + return; + }; + sources += 1; + input_path.get_or_insert_with(|| path.clone()); + if let Some(schema) = input.output.samples() { + // Raw, not evaluated: this node feeds the tree from a scan, so a tree whose input was + // compiled against an already-evaluated schema would read the wrong thing. + bound = binding.same_fields(schema) && schema.eval_grid.is_none(); + } + }); + if sources != 1 { + return Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree with {sources} inputs, and it needs exactly one" + ))); + } + if !bound { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema the tree's input was not compiled against" + .to_string(), + )); + } + let input_path = input_path.ok_or_else(|| { + DataFusionError::Plan("FusedSeries runs a tree with no input node".to_string()) + })?; + + match tree.root().output.samples() { + Some(schema) + if schema + .eval_grid + .as_ref() + .is_some_and(|grid| grid.spec() == binding.grid.spec()) => + { + Ok(input_path) + } + _ => Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on the statement's grid" + .to_string(), + )), + } +} + +/// Where one output column of [`FusedSeriesExec`] comes from. +#[derive(Clone, Copy, Debug)] +enum OutputColumn { + /// The evaluation grid timestamps of the current chunk. + Time, + /// Series field by index in the bound series schema. + Field(usize), + /// Declared logical tag by index in the caller-supplied tag list. + Tag(usize), +} + +/// Classifies every output column, rejecting anything the conversion cannot produce. +fn classify_output( + output_schema: &SchemaRef, + series_schema: &SeriesSchema, + tags: &[String], + time_column: &str, +) -> DataFusionResult> { + for (index, tag) in tags.iter().enumerate() { + if tags[..index].contains(tag) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {tag:?} in FusedSeriesExec output schema" + ))); + } + if tag == time_column { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec cannot be the time column" + ))); + } + if series_schema + .fields + .iter() + .any(|field| field.name.as_ref() == tag) + { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec clashes with a series field" + ))); + } + } + + let mut time_seen = false; + let mut field_seen = vec![false; series_schema.fields.len()]; + let mut tag_seen = vec![false; tags.len()]; + let mut classification = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { + let name = field.name(); + if name == time_column { + if std::mem::replace(&mut time_seen, true) { + return Err(DataFusionError::Plan(format!( + "duplicate time column {name:?} in FusedSeriesExec output schema" + ))); + } + match field.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, None) => {} + other => { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must be Timestamp(Millisecond, \ + None), got {other:?}" + ))); + } + } + if field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must not be nullable" + ))); + } + classification.push(OutputColumn::Time); + continue; + } + if let Some(index) = series_schema + .fields + .iter() + .position(|field| field.name.as_ref() == name) + { + if std::mem::replace(&mut field_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate field column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be Float64, got {:?}", + field.data_type() + ))); + } + // A grid point inside the statement's interval can have no value in this field. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be nullable, because a grid \ + point can be missing" + ))); + } + classification.push(OutputColumn::Field(index)); + continue; + } + if let Some(index) = tags.iter().position(|tag| tag == name) { + if std::mem::replace(&mut tag_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Utf8 { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be Utf8, got {:?}", + field.data_type() + ))); + } + // `rate` and the map kernels drop `__name__`, so a tag can be absent from a row. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be nullable, because range \ + functions and map kernels drop the metric name" + ))); + } + classification.push(OutputColumn::Tag(index)); + continue; + } + return Err(DataFusionError::Plan(format!( + "unknown column {name:?} in FusedSeriesExec output schema" + ))); + } + + if !time_seen { + return Err(DataFusionError::Plan(format!( + "missing time column {time_column:?} in FusedSeriesExec output schema" + ))); + } + for (index, seen) in field_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing field column {:?} in FusedSeriesExec output schema", + series_schema.fields[index].name + ))); + } + } + for (index, seen) in tag_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing tag column {:?} in FusedSeriesExec output schema", + tags[index] + ))); + } + } + Ok(classification) +} + +/// Maps the input columns this node reads onto the raw-series roles it needs. +/// +/// The input has to be exactly the columns the binding accounts for: the time column, the fields +/// and the tags. A column outside them is part of some series' identity that the raw chunks would +/// not carry, so two rows a Prometheus query keeps apart could be answered as one series with the +/// wrong labels; such an input is refused rather than read past. Whether a *different* column of +/// the same type was meant is not something this function can see — that is the caller commitment +/// [`FusedSeriesBinding`] documents. +fn input_mapping( + binding: &FusedSeriesBinding, + schema: &SchemaRef, +) -> DataFusionResult { + // A name that appears more than once says less than it looks: every one of those columns + // passes the coverage check below, because the check reads names, while the mapping reads + // `index_of`, which answers with one index. The other column would then be part of a series' + // identity that the raw chunks never carry -- rows a Prometheus query keeps apart answered as + // one series with the wrong labels, and no error. Which of the two the caller meant is not + // visible from here, so the schema is refused rather than one of them chosen. + let mut named = HashSet::with_capacity(schema.fields().len()); + for field in schema.fields() { + if !named.insert(field.name()) { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec input has more than one column named {:?}: it cannot say which \ + one its binding maps, and the other would take no part in the series it holds", + field.name() + ))); + } + } + if schema.column_with_name(METRIC_NAME).is_some() { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec cannot read a {METRIC_NAME:?} column: a raw series' metric name is \ + bound on the node, not varied inside one series" + ))); + } + let index_of = |name: &str| -> DataFusionResult { + schema.index_of(name).map_err(|_| { + DataFusionError::Plan(format!("FusedSeriesExec input has no column {name:?}")) + }) + }; + + // Every column of the input has to be a role the binding reads. The binding cannot express + // "ignore this column": a column it does not read is one the raw series loses. + for field in schema.fields() { + let name = field.name(); + let read = name == binding.time_index.as_str() + || binding.tags.iter().any(|tag| tag == name) + || binding + .series_schema + .fields + .iter() + .any(|spec| spec.name.as_ref() == name); + if !read { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec input column {name:?} is neither the time column {:?} nor one of \ + the fields and tags its binding maps, so it would take no part in the series' \ + identity", + binding.time_index + ))); + } + } + + let tags = binding + .tags + .iter() + .filter(|tag| tag.as_str() != METRIC_NAME) + .map(|tag| Ok((Arc::from(tag.as_str()), index_of(tag)?))) + .collect::, usize)>>>()?; + let fields = binding + .series_schema + .fields + .iter() + .map(|spec: &FieldSpec| { + let index = index_of(&spec.name)?; + // A native-histogram column would reach the output converter as a histogram sample, + // which only fails there, once execution has started. This node answers in floats, so + // the column it reads has to be one. + if schema.field(index).data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec reads field {:?} as Float64, and the input column is {}", + spec.name, + schema.field(index).data_type() + ))); + } + Ok((spec.clone(), index)) + }) + .collect::>>()?; + let mapping = InputMapping { + time_index: index_of(&binding.time_index)?, + metric_name: binding.metric_name.clone(), + tags, + fields, + }; + // The shared converter is the authority on which column types it accepts, so it is asked + // rather than copied: handing it an empty stream of this schema makes it validate the mapping + // with no batch to convert. + validate_mapping_types(&mapping, schema)?; + Ok(mapping) +} + +/// Rejects a mapping the shared converter would refuse, naming the column at fault. +/// +/// The types accepted here are the converter's own, not a copy of its rules: [`converter_accepts`] +/// runs the converter's validation, so a time or tag type this node accepts is a type `execute` +/// converts, and one it refuses can only be refused there too. +fn validate_mapping_types(mapping: &InputMapping, schema: &SchemaRef) -> DataFusionResult<()> { + let Err(error) = converter_accepts(mapping, schema) else { + return Ok(()); + }; + // The converter reports the mapping as a whole, so the column at fault is found by asking it + // again with the time column alone and then with one tag at a time. The field columns are + // already checked above, and they are the only other columns this mapping maps. + let mut without_tags = mapping.clone(); + let tags = std::mem::take(&mut without_tags.tags); + if converter_accepts(&without_tags, schema).is_err() { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec cannot read {:?} as its time column: {error}", + schema.field(mapping.time_index).name() + ))); + } + for tag in tags { + let mut one = without_tags.clone(); + one.tags = vec![tag.clone()]; + if converter_accepts(&one, schema).is_err() { + return Err(DataFusionError::Plan(format!( + "FusedSeriesExec cannot read tag {:?} from its input column: {error}", + tag.0 + ))); + } + } + // Something the converter refuses that no single role accounts for: its own error says it. + Err(error) +} + +/// Whether the shared converter validates `mapping` against `schema`. +fn converter_accepts(mapping: &InputMapping, schema: &SchemaRef) -> DataFusionResult<()> { + record_batches_to_raw_series( + Box::pin(EmptyRecordBatchStream::new(Arc::clone(schema))), + mapping.clone(), + ) + .map(|_| ()) +} + +/// The physical node that runs a compiled native tree over its input's raw series. +pub struct FusedSeriesExec { + /// The PromQL text the tree was compiled from, which is what the plan line shows. + promql: Arc, + /// The tree this node runs, shared by every partition's execution. + tree: Arc, + /// What the tree's input was compiled against, and what the output columns are called. + binding: Arc, + /// The position of the tree's input node, where every raw chunk a partition reads enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The plan producing the raw series. + input: Arc, + /// The columns this node produces. + output_schema: SchemaRef, + /// The grid positions every label set in this plan's output already occupies. + /// + /// Plan-level rather than per-partition: see [`SharedLabelSets`]. + labels: SharedLabelSets, + /// Per-node metrics, registered by each execution of the tree. + metric: ExecutionPlanMetricsSet, + /// This node's partitioning, emission and boundedness. + properties: Arc, +} + +impl FusedSeriesExec { + /// Creates the node, rejecting whatever this node cannot execute or convert. + pub fn try_new( + promql: Arc, + tree: Arc, + binding: Arc, + input: Arc, + output_schema: SchemaRef, + ) -> DataFusionResult { + let input_path = validate_binding(&tree, &binding)?; + // Building a state is what proves the runtime executes this tree: a node kind it does not + // run yet is refused here, so a plan that is built is a plan that starts. The throwaway + // state is dropped at once, returning its (unbounded) reservations. + let _ = NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + )?; + let classification = classify_output( + &output_schema, + &binding.series_schema, + &binding.tags, + &binding.time_index, + )?; + // The mapping is checked against the declared input schema, so a column this node reads + // but the input does not produce is reported when the plan is built. + input_mapping(&binding, &input.schema())?; + // One evaluator runs per input partition, so the node's partition count is the child's. + // The child must therefore be a raw-series scan that assigns every series whole to one + // partition: `PerSeries` distribution is what makes a per-partition evaluation equivalent + // to evaluating the whole series, because no series is split across partitions and no + // series appears in two of them. + // + // One window holds the state of the series it is reading, and a raw series arrives as one + // contiguous run. The scan guarantees that per partition, not across them: a scan that put + // half of a series in one partition would leave each tree answering half of it. + let input_properties = input.properties(); + let partition_count = input_properties.partitioning.partition_count(); + if partition_count == 0 { + return Err(DataFusionError::Plan( + "FusedSeriesExec input has no partitions".into(), + )); + } + + // Emitted one batch per evaluated chunk at a time, and as many partitions as the input + // hands over: the node runs one evaluator per partition and answers each as it settles. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + input_properties.emission_type, + input_properties.boundedness, + )); + // Shared by every partition's stream: a label set colliding is a property of the whole + // output vector, not of one partition. Creating it needs no pool; each stream reserves the + // bytes it adds in its own reservation. + let labels: SharedLabelSets = Arc::new(Mutex::new(HashMap::new())); + Ok(Self { + promql, + tree, + binding, + input_path, + classification: classification.into(), + input, + output_schema, + labels, + metric: ExecutionPlanMetricsSet::new(), + properties, + }) + } + + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeriesExec" + } +} + +impl std::fmt::Debug for FusedSeriesExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("FusedSeriesExec").finish() + } +} + +impl DisplayAs for FusedSeriesExec { + /// The query, then the tree it compiles to: the tree is rendered by [`NativeTree`]'s own + /// `Display`, so `EXPLAIN` and the node's metrics name the same nodes. + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree) + } +} + +impl ExecutionPlan for FusedSeriesExec { + fn name(&self) -> &str { + Self::name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + // The tree's operands are bound inside it; this node owns no expression of its own. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + let [child] = <[Arc; 1]>::try_from(children).map_err(|children| { + DataFusionError::Plan(format!( + "FusedSeriesExec requires exactly one child, got {}", + children.len() + )) + })?; + // Rebuilt rather than patched: the new child's schema is what the input mapping is checked + // against, so a child that changed shape is refused instead of read with stale indices. + Ok(Arc::new(Self::try_new( + Arc::clone(&self.promql), + Arc::clone(&self.tree), + Arc::clone(&self.binding), + child, + self.output_schema.clone(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let input = self.input.execute(partition, Arc::clone(&context))?; + let mapping = input_mapping(&self.binding, &self.input.schema())?; + // Batches become raw series chunks through the shared frame converter, which is what groups + // rows into one chunk per label set and closes each series exactly once. + let input = record_batches_to_raw_series(input, mapping)?; + let state = NativeExecutionState::try_new( + &self.tree, + Arc::clone(&context.runtime_env().memory_pool), + &self.metric, + partition, + )?; + // What this stream adds to the shared label-set map is its own to account for; the map + // itself outlives it, so the entries are not given back with the reservation. + let reservation = + MemoryConsumer::new("FusedSeriesStream").register(&context.runtime_env().memory_pool); + Ok(Box::pin(FusedSeriesStream { + input: Some(input), + executor: Some(NativeExecutor::new(state)), + input_path: self.input_path.clone(), + classification: Arc::clone(&self.classification), + binding: Arc::clone(&self.binding), + output_schema: self.output_schema.clone(), + labels: Arc::clone(&self.labels), + reservation, + ready: VecDeque::new(), + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metric.clone_inner()) + } +} + +/// Converts one evaluated chunk into one output batch. +/// +/// The time column is the absolute grid positions the chunk carries, missing samples are nulls and +/// ordinary NaN stays a value. +fn series_to_record_batch( + classification: &[OutputColumn], + tags: &[String], + output_schema: &SchemaRef, + chunk: &EvaluatedChunk, +) -> DataFusionResult { + let series = &chunk.series; + // A chunk that forms one run is read from its range, which lends the grid's buffers and so + // allocates nothing. A chunk whose rows form several runs has no single range to read + // ([`EvaluatedChunk::range`] answers `None` for it), and the positions between its runs are + // holes rather than rows: the axis is walked for its carried rows instead, which is the one + // traversal that visits exactly those. + let timestamps = match chunk.range() { + Some(range) => range.timestamps(), + None => ScalarBuffer::from(series.time.timestamps().into_owned()), + }; + let rows = timestamps.len(); + let mut columns: Vec = Vec::with_capacity(classification.len()); + for (position, column) in classification.iter().enumerate() { + match column { + OutputColumn::Time => columns.push(Arc::new(TimestampMillisecondArray::new( + timestamps.clone(), + None, + ))), + OutputColumn::Field(index) => { + let samples = series.fields.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "evaluated series has no field at index {index}" + )) + })?; + if samples.len() != rows { + return Err(DataFusionError::Execution(format!( + "evaluated series field {index} has {} rows, time axis has {rows}", + samples.len() + ))); + } + columns.push(field_column(samples)?); + } + OutputColumn::Tag(index) => { + // `__name__` resolves to the metric name, every other tag through the label set. + let name = tags.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "output column {position} declares unavailable tag {index}" + )) + })?; + let value = series.metric.get(name); + columns.push(Arc::new(StringArray::from(vec![value; rows]))); + } + } + } + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) +} + +/// Converts one evaluated series field into one `Float64` column. +/// +/// Dense float samples are reused as they are: the array already carries exactly the right nulls +/// and keeps NaN a value, so copying it through `Samples::iter()` would only cost an allocation +/// and a re-encode. Chunked samples go through the iterator, and histograms have no float column. +fn field_column(samples: &Samples) -> DataFusionResult { + match samples.dense() { + Some(DenseSamplesRef::Float(values)) => Ok(Arc::new(values.clone())), + Some(DenseSamplesRef::Histogram(_)) | Some(DenseSamplesRef::Mixed(_)) => { + Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )) + } + None => { + let values = samples + .iter() + .map(|sample| match sample { + SampleRef::Float(value) => Ok(Some(value)), + SampleRef::Missing => Ok(None), + SampleRef::Histogram(_) => Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )), + }) + .collect::>>()?; + Ok(Arc::new(Float64Array::from(values))) + } + } +} + +/// A canonical, collision-free encoding of a final label set. +/// +/// [`MetricName`] implements `PartialEq + Eq` but not `Hash`, and it is a dependency-owned type, +/// so the duplicate-label-set state keys on these bytes instead of on the label set itself. Every +/// component carries a length prefix: label names and values are arbitrary UTF-8, so a separator +/// byte could occur inside a component and let two different label sets encode to the same bytes. +/// [`MetricName`] keeps its tags sorted by name, so one label set has exactly one encoding. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct LabelsKey(Vec); + +impl LabelsKey { + /// Encodes the metric name, marked for presence, and then every label, name before value. + fn new(metric: &MetricName) -> Self { + let mut bytes = Vec::new(); + match metric.name() { + // The marker keeps an absent metric name distinct from an empty one. + Some(name) => { + bytes.push(1); + push_component(&mut bytes, name); + } + None => bytes.push(0), + } + for tag in metric.tags() { + push_component(&mut bytes, &tag.name); + push_component(&mut bytes, &tag.value); + } + Self(bytes) + } + + /// Number of bytes this key occupies, for the memory reservation. + fn len(&self) -> usize { + self.0.len() + } +} + +/// Appends one length-prefixed UTF-8 component. +fn push_component(bytes: &mut Vec, component: &str) { + bytes.extend_from_slice(&component.len().to_le_bytes()); + bytes.extend_from_slice(component.as_bytes()); +} + +/// Renders a label set the way Prometheus prints one, for the duplicate-label-set error. +fn render_labels(metric: &MetricName) -> String { + let mut parts = Vec::new(); + if let Some(name) = metric.name() { + parts.push(format!("__name__={name:?}")); + } + parts.extend( + metric + .tags() + .iter() + .map(|tag| format!("{}={:?}", tag.name, tag.value)), + ); + format!("{{{}}}", parts.join(", ")) +} + +/// The grid positions each final label set already occupies, as merged, sorted runs. +/// +/// Shared by every partition's stream, because a label-set collision is a property of the whole +/// output vector and not of one partition. `PerSeries` distribution keeps each series whole inside +/// one partition, but two series that carry the same labels after the range function and its maps +/// drop a label can still land in different partitions. Checking only within a partition would +/// miss exactly those collisions. +/// +/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed +/// for the partitions still running, so the entries are only released with the plan itself. +type SharedLabelSets = Arc>>>; + +/// The grid positions the rows of one chunk hold a value at, as sorted, disjoint runs. +/// +/// Each run is `(start, end_exclusive)` over absolute grid indices. A row whose every field is +/// missing is not occupied: Prometheus has no point there, so it cannot collide with another +/// series. +fn occupied_runs(series: &Series) -> DataFusionResult> { + // Rows are addressed on the statement's grid, so a row sits at the absolute position two + // series can collide on. An explicit axis has no such position, and a chunk never carries + // one: this is the refusal `EvaluatedChunk` already makes when the chunk is built. + if series.time.grid().is_none() { + return Err(DataFusionError::Execution( + "FusedSeriesExec requires evaluated series on a time grid".to_string(), + )); + } + let rows = series.time.len(); + // A series with one dense, null-free float field holds a value on every row it carries, so a + // chunk that forms one run occupies exactly the range it reports and no row has to be + // inspected. This is the shape the fused evaluator produces for a range function over + // ordinary float samples. + if rows > 0 + && series.fields.len() == 1 + && let Some(DenseSamplesRef::Float(values)) = + series.fields.iter().next().and_then(Samples::dense) + && values.len() == rows + && values.null_count() == 0 + && let Some(range) = series.time.as_grid_range() + { + return Ok(vec![(range.start(), range.start() + rows)]); + } + // `positions` walks a chunk that forms one run and one whose rows form several the same way, + // in position order, so the holes between two runs are never read as rows. + let mut runs: Vec<(usize, usize)> = Vec::new(); + for (position, row) in series.time.positions() { + if !series.fields.iter().any(|samples| has_value(samples, row)) { + continue; + } + match runs.last_mut() { + Some(run) if run.1 == position => run.1 += 1, + _ => runs.push((position, position + 1)), + } + } + Ok(runs) +} + +/// Whether one field holds a value, rather than a missing position, at `index`. +fn has_value(samples: &Samples, index: usize) -> bool { + match samples.dense() { + Some(DenseSamplesRef::Float(values)) => values.is_valid(index), + Some(DenseSamplesRef::Histogram(values)) => values.indices().is_valid(index), + Some(DenseSamplesRef::Mixed(values)) => { + values.floats().is_valid(index) || values.histograms().indices().is_valid(index) + } + None => !matches!(samples.sample(index), SampleRef::Missing), + } +} + +/// Whether two sorted, disjoint run lists share at least one grid index. +fn runs_overlap(stored: &[(usize, usize)], new: &[(usize, usize)]) -> bool { + let (mut left, mut right) = (0, 0); + while left < stored.len() && right < new.len() { + let (stored_start, stored_end) = stored[left]; + let (new_start, new_end) = new[right]; + if new_start < stored_end && stored_start < new_end { + return true; + } + if stored_end <= new_start { + left += 1; + } else { + right += 1; + } + } + false +} + +/// Adds sorted, disjoint `new` runs to the sorted, disjoint `stored` runs, coalescing neighbours. +fn merge_runs(stored: &mut Vec<(usize, usize)>, new: &[(usize, usize)]) { + let mut merged: Vec<(usize, usize)> = Vec::with_capacity(stored.len() + new.len()); + let (mut left, mut right) = (0, 0); + while left < stored.len() || right < new.len() { + let run = match (stored.get(left), new.get(right)) { + (Some(stored), Some(new)) if stored.0 <= new.0 => { + left += 1; + *stored + } + (Some(_), Some(new)) => { + right += 1; + *new + } + (Some(stored), None) => { + left += 1; + *stored + } + (None, Some(new)) => { + right += 1; + *new + } + (None, None) => break, + }; + match merged.last_mut() { + Some(last) if run.0 <= last.1 => last.1 = last.1.max(run.1), + _ => merged.push(run), + } + } + *stored = merged; +} + +/// The batches one execution of [`FusedSeriesExec`] yields. +pub struct FusedSeriesStream { + /// The input batches, read as raw series chunks, or `None` once the stream ended in an error. + /// + /// The converter behind it holds its current batch, the chunk it is building and its own + /// source, so a terminated stream that kept it alive would hold upstream resources until the + /// consumer drops the stream. + input: Option, + /// The execution of this node's tree, or `None` once it has been closed. + executor: Option, + /// The position of the tree's input node, which is where every chunk read enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The tag column names, indexed by `OutputColumn::Tag`. + binding: Arc, + /// The schema of every batch this stream yields. + output_schema: SchemaRef, + /// The label-set occupancy shared with the other partitions' streams. + labels: SharedLabelSets, + /// Accounts for what this stream adds to `labels`; freed when the stream ends. + reservation: MemoryReservation, + /// The batches the executor settled since the last poll of the input. + ready: VecDeque, + /// Whether the input ended and everything it settled has been handed out. + done: bool, +} + +impl FusedSeriesStream { + /// Hands one raw chunk over and queues what it settled. + fn push(&mut self, chunk: RawSeriesChunk) -> DataFusionResult<()> { + let Some(executor) = self.executor.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec stream pushed a chunk after its execution closed".to_string(), + )); + }; + // The chunk enters at the tree's input node by position: a position is how the runtime + // addresses a node, and a tree that gains a second input names them the same way. + executor.push_raw_at(&self.input_path, chunk)?; + let mut settled = Vec::new(); + while let Some(chunk) = executor.next_chunk()? { + settled.push(chunk); + } + self.convert(settled) + } + + /// Ends the input and queues what closing settled. + fn close(&mut self) -> DataFusionResult<()> { + let Some(mut executor) = self.executor.take() else { + return Ok(()); + }; + self.convert(executor.finish()?) + } + + /// Ends the stream in an error, releasing the execution and the input at once. + /// + /// An ordinary drop of the stream releases all of this as well, so this is not about the last + /// release: it is about a consumer that keeps the terminated stream, which would otherwise + /// hold the evaluator, the converter's held batch and chunk, and the plan upstream of it. + fn fail(&mut self) { + self.done = true; + self.executor = None; + self.input = None; + // The label-set map is shared, so a stream stopping early must not clear entries another + // partition is still checking against: its own reservation is what this stream owns, and + // that is what it gives back. + self.reservation.free(); + // A batch queued before the error was settled by a run that has now stopped. + self.ready.clear(); + } + + /// Claims the occupied grid positions of one chunk of a final series for its label set. + /// + /// A collision is an error only when both series hold a sample at the same evaluation + /// timestamp, which is Prometheus' range-query rule. Series with disjoint points are legal: + /// Prometheus merges them into one output series, while this node emits them as separate + /// batches and defers that merge to matrix assembly, so every sample survives. + fn claim_label_set(&mut self, series: &Series) -> DataFusionResult<()> { + let runs = occupied_runs(series)?; + let key = LabelsKey::new(&series.metric); + let mut seen = self.labels.lock().unwrap_or_else(|e| e.into_inner()); + match seen.entry(key) { + Entry::Occupied(mut entry) => { + let stored = entry.get_mut(); + if runs_overlap(stored, &runs) { + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset: {}", + render_labels(&series.metric) + ))); + } + merge_runs(stored, &runs); + } + Entry::Vacant(entry) => { + // Only the label-set state is this node's own cross-poll state that grows with + // the output, so it is the only thing reserved here. + let bytes = entry.key().len() + runs.len() * 2 * size_of::(); + self.reservation.try_grow(bytes)?; + entry.insert(runs); + } + } + Ok(()) + } + + /// Converts the chunks the executor settled into batches, in the order it produced them. + fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { + for chunk in &chunks { + // The label set is claimed before its batch is queued, so a collision surfaces at the + // chunk that reached it and nothing this plan would refuse is handed out after it. + self.claim_label_set(&chunk.series)?; + let batch = series_to_record_batch( + &self.classification, + &self.binding.tags, + &self.output_schema, + chunk, + )?; + self.ready.push_back(batch); + } + Ok(()) + } +} + +impl RecordBatchStream for FusedSeriesStream { + fn schema(&self) -> SchemaRef { + self.output_schema.clone() + } +} + +impl Stream for FusedSeriesStream { + type Item = DataFusionResult; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll>> { + let this = &mut *self; + // A batch that is already settled is handed over before anything is read, so a bounded + // poll never delays one that is ready behind more input. + for _ in 0..POLL_BUDGET { + if let Some(batch) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(batch))); + } + if this.done { + return Poll::Ready(None); + } + // An error is the only thing that clears the input, and it ends the stream with it. + let Some(input) = this.input.as_mut() else { + return Poll::Ready(None); + }; + let polled = input.as_mut().poll_next(cx); + match polled { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + // The input ended, so every series it opened has to be closed: what closing + // settles is handed out before the stream ends. + this.done = true; + if let Err(error) = this.close() { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + // An input error stops the execution: a stopped evaluator must not be reused. + this.fail(); + return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); + } + Poll::Ready(Some(Ok(chunk))) => { + if let Err(error) = this.push(chunk) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + } + // The budget ran out on input that settled nothing. Everything read so far is kept, so the + // next poll continues where this one stopped; waking first is what asks for that poll. + cx.waker().wake_by_ref(); + Poll::Pending + } +} diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs new file mode 100644 index 00000000000..c4335420059 --- /dev/null +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -0,0 +1,1543 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! What this node has to get right: the two `EXPLAIN` lines it shows, the batches it runs end to +//! end, and what it refuses to run. +//! +//! The trees here are built node by node with the same bindings the PromQL compiler produces for +//! `rate(x[5m])`: the compiler itself lives in the `query` crate, which depends on this one, so a +//! test in this crate cannot call it. +//! +//! Kernel numerics are not the subject: `rate` is checked against a counter that advances one +//! second per second, so the value it has to produce is exactly 1 (or 2) at every grid point and +//! any position or finalization mistake shows up as a wrong, missing or duplicated row. + +use std::collections::VecDeque; +use std::fmt::{self, Display, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::ToDFSchema; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::UnboundedMemoryPool; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, collect, displayable, +}; +use datafusion::prelude::SessionContext; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, Series, SeriesFields, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; +use datatypes::series::{Label, MetricName, SeriesOrdering}; +use futures::{Stream, StreamExt, TryStreamExt}; + +use super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, series_to_record_batch}; +use crate::evaluated::EvaluatedChunk; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundWindow, Grouping, NativeExpr, + NativeNode, NativeTree, NativeType, +}; +use crate::range_eval::{RangeFunction, RangeWindow}; + +/// The query every tree here is compiled from. +const PROMQL: &str = "rate(x[5m])"; +/// The range the statement's window selects, which is the `5m` of [`PROMQL`]. +const RANGE_MS: i64 = 300_000; +/// The statement's grid: `rate(x[5m])` evaluated every minute from five minutes in. +const GRID_START: i64 = 300_000; +const GRID_STEP: i64 = 60_000; +const GRID_LEN: u64 = 5; +/// One raw sample per minute, from 0 to 9 minutes: enough to fill every window on the grid. +const SAMPLE_STEP: i64 = 60_000; +const SAMPLES: i64 = 10; +/// How many one-sample batches the quiet series below arrives in. +/// +/// Two bounds hold at once: the series reaches the end of the last grid point's window (540s, so +/// indices `0..=540`), which is what lets it rate at 1 over the whole grid, and it arrives in more +/// batches than one poll's budget, which is what makes a poll that ignored the budget visible. +const QUIET_SAMPLES: i64 = 541; +/// One second between those samples, so the series rates at 1 once it is read to its end. +/// +/// The step is a second and not a minute on purpose: at a minute the series would arrive in ten +/// batches, fewer than [`POLL_BUDGET`], and the budget the test above exists to check would never +/// be reached. +const QUIET_STEP: i64 = 1_000; + +const TIME: &str = "time"; +const VALUE: &str = "value"; +const HOST: &str = "host"; +const ZONE: &str = "zone"; +const NAME: &str = "__name__"; +const METRIC: &str = "http_requests_total"; + +fn grid() -> Arc { + Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ) +} + +/// The raw schema the tree's input is bound to: one float field called [`VALUE`]. +fn raw_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: VALUE.into() }]), + eval_grid: None, + }) +} + +/// The tree `rate(x[5m])` compiles to: raw samples folded by one window onto the grid. +fn rate_tree() -> NativeTree { + let grid = grid(); + // One schema shared by the window's binding and its output, which is what `NativeTree`'s + // validation compares by identity. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + // The input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: grid, + }, + }, + output: evaluated, + }) + .unwrap() +} + +/// The `rate(x[5m])` tree with a `sum by (job)` over it, which the runtime does not execute. +fn aggregate_tree() -> NativeTree { + let rate = rate_tree(); + let window = rate.root().clone(); + let grid = window + .output + .samples() + .expect("a window produces samples") + .eval_grid + .clone() + .expect("a window's output is evaluated on a grid"); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(window), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: Grouping::By(Arc::new(["job".into()])), + stage: AggregateStage::Complete, + }, + }, + output: NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(grid), + })), + }) + .expect("the bound tree is legal") +} + +/// What the tree's input was compiled against, read off the tree itself. +fn binding(tree: &NativeTree) -> FusedSeriesBinding { + FusedSeriesBinding { + series_schema: raw_schema(), + grid: tree + .root() + .output + .samples() + .expect("this tree produces samples") + .eval_grid + .clone() + .expect("this tree's output is evaluated on a grid"), + metric_name: Some(Arc::from(METRIC)), + time_index: TIME.to_string(), + tags: vec![HOST.to_string(), NAME.to_string()], + } +} + +/// The batches the input plan produces: the time, the field and the tag column. +fn input_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + ])) +} + +/// The columns this node produces: the field, the tags and the grid timestamps. +fn output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +fn input_plan() -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: input_schema().to_dfschema_ref().unwrap(), + }) +} + +/// A physical input producing `batches` in one partition, with `schema` as its schema. +fn memory_exec_with(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) +} + +/// A physical input producing `batches` in one partition, on the input schema of these tests. +fn memory_exec(batches: Vec) -> Arc { + memory_exec_with(input_schema(), batches) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own. +fn memory_exec_partitioned(partitions: Vec>) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&partitions, input_schema(), None) + .expect("every partition holds batches of the input schema"), + ))) +} + +/// A physical input whose one partition's stream `make` builds. +fn staged( + schema: SchemaRef, + make: impl Fn() -> SendableRecordBatchStream + Send + Sync + 'static, +) -> Arc { + Arc::new(StagedPlan { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + make: Box::new(make), + }) +} + +/// A physical input, so a test can stage what a memory source cannot: a pending poll, an error +/// after some rows, and a stream whose drop it can see. +struct StagedPlan { + make: Box SendableRecordBatchStream + Send + Sync>, + properties: Arc, +} + +impl std::fmt::Debug for StagedPlan { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("StagedPlan").finish() + } +} + +impl DisplayAs for StagedPlan { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { + Ok(()) + } +} + +impl ExecutionPlan for StagedPlan { + fn name(&self) -> &str { + "StagedPlan" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + // A scan answers the partitions it declares and no others, which is what makes an index + // past this plan's count an error the node above has to pass on. + if partition >= self.properties.partitioning.partition_count() { + return Err(DataFusionError::Execution(format!( + "StagedPlan has no partition {partition}" + ))); + } + Ok((self.make)()) + } +} + +/// The stream a [`StagedPlan`] hands over: `batches`, then an optional error, with a pending poll +/// and a drop a test can observe. +struct StagedStream { + schema: SchemaRef, + batches: VecDeque, + handed: usize, + /// Hand this many batches, then answer [`Poll::Pending`] once. + pending_after: Option, + /// The error to report once the batches are exhausted. + error: Option<&'static str>, + /// Set when this stream is dropped, so a test can see the input being released. + dropped: Option>, + /// Counts the batches this stream has handed over, so a test can see how far one poll read. + handed_out: Option>, +} + +impl StagedStream { + fn new(schema: SchemaRef, batches: Vec) -> Self { + Self { + schema, + batches: batches.into(), + handed: 0, + pending_after: None, + error: None, + dropped: None, + handed_out: None, + } + } + + fn pending_after(mut self, batches: usize) -> Self { + self.pending_after = Some(batches); + self + } + + fn error(mut self, message: &'static str) -> Self { + self.error = Some(message); + self + } + + fn dropped(mut self, flag: &Arc) -> Self { + self.dropped = Some(Arc::clone(flag)); + self + } + + fn counted(mut self, counter: &Arc) -> Self { + self.handed_out = Some(Arc::clone(counter)); + self + } +} + +impl Stream for StagedStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if this.pending_after == Some(this.handed) { + this.pending_after = None; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + match this.batches.pop_front() { + Some(batch) => { + this.handed += 1; + if let Some(counter) = &this.handed_out { + counter.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(Some(Ok(batch))) + } + None => match this.error.take() { + Some(message) => { + Poll::Ready(Some(Err(DataFusionError::Execution(message.to_string())))) + } + None => Poll::Ready(None), + }, + } + } +} + +impl RecordBatchStream for StagedStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Drop for StagedStream { + fn drop(&mut self) { + if let Some(dropped) = &self.dropped { + dropped.store(true, Ordering::SeqCst); + } + } +} + +fn fused(tree: NativeTree) -> FusedSeries { + let binding = binding(&tree); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// One batch of one series' `(timestamp, value)` rows. +fn series_batch(host: &str, rows: &[(i64, f64)]) -> RecordBatch { + let times = rows.iter().map(|(time, _)| *time).collect::>(); + let values = rows.iter().map(|(_, value)| *value).collect::>(); + let hosts = rows.iter().map(|_| host).collect::>(); + RecordBatch::try_new( + input_schema(), + vec![ + Arc::new(TimestampMillisecondArray::from(times)), + Arc::new(Float64Array::from(values)), + Arc::new(StringArray::from(hosts)), + ], + ) + .unwrap() +} + +/// One series' samples: one per minute from 0, `scale` seconds per second. +fn series_samples(scale: f64) -> Vec<(i64, f64)> { + (0..SAMPLES) + .map(|index| { + let time = index * SAMPLE_STEP; + (time, scale * time as f64 / 1000.0) + }) + .collect() +} + +/// Two series, each split over two batches. A series' batches stay adjacent, which is the input +/// contract: a label set's rows are contiguous across the whole stream. +fn two_series_batches() -> Vec { + let first = series_samples(1.0); + let second = series_samples(2.0); + vec![ + series_batch("web-1", &first[..3]), + series_batch("web-1", &first[3..]), + series_batch("web-2", &second[..3]), + series_batch("web-2", &second[3..]), + ] +} + +/// One series split over three batches, so its chunks have to be stitched. +fn one_series_batches() -> Vec { + let rows = series_samples(1.0); + vec![ + series_batch("web-1", &rows[..4]), + series_batch("web-1", &rows[4..7]), + series_batch("web-1", &rows[7..]), + ] +} + +/// The rows of `batches`: timestamp, value, host and metric name. +type Row = (i64, Option, Option, Option); + +fn rows(batches: &[RecordBatch]) -> Vec { + let mut rows = Vec::new(); + for batch in batches { + let time = batch + .column_by_name(TIME) + .expect("every batch has a time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let value = batch + .column_by_name(VALUE) + .expect("every batch has the field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + let host = batch + .column_by_name(HOST) + .expect("every batch has the tag column") + .as_any() + .downcast_ref::() + .expect("the tag column is Utf8"); + let name = batch + .column_by_name(NAME) + .expect("every batch has the metric name column") + .as_any() + .downcast_ref::() + .expect("the metric name column is Utf8"); + for row in 0..batch.num_rows() { + rows.push(( + time.value(row), + value.is_valid(row).then(|| value.value(row)), + host.is_valid(row).then(|| host.value(row).to_string()), + name.is_valid(row).then(|| name.value(row).to_string()), + )); + } + } + rows +} + +/// The input schema of the collision below: the time, the field, `host` and `zone`. +fn zone_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + ])) +} + +/// The two rows of the collision: one series, once with an empty `zone` and once with no `zone`. +fn zone_batches() -> Vec { + let schema = zone_schema(); + let samples = series_samples(1.0); + let times: Vec<_> = samples.iter().map(|(time, _)| *time).collect(); + let values: Vec<_> = samples.iter().map(|(_, value)| *value).collect(); + let hosts: Vec<_> = times.iter().map(|_| "web-1").collect(); + // The two rows differ only in the zone column: empty against null. + let zones: [ArrayRef; 2] = [ + Arc::new(StringArray::from(vec![""; times.len()])), + Arc::new(StringArray::from(vec![None::<&str>; times.len()])), + ]; + zones + .into_iter() + .map(|zone| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(TimestampMillisecondArray::from(times.clone())), + Arc::new(Float64Array::from(values.clone())), + Arc::new(StringArray::from(hosts.clone())), + zone, + ], + ) + .expect("one batch of the two rows") + }) + .collect() +} + +/// `rate(x[5m])` over an input that carries `zone` as well as `host`, with both in the output. +fn zone_node() -> FusedSeries { + let schema = zone_schema(); + let tree = rate_tree(); + let binding = FusedSeriesBinding { + tags: vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ..binding(&tree) + }; + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::clone(&schema).to_dfschema_ref().unwrap(), + }), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. +fn zone_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// The canonical label set these tests collide on, as the runtime renders it in its refusal. +const CLOSED_LABEL_SET: &str = + "the node at root already closed the label set [(\"host\", \"web-1\")]"; + +/// Two rows that are already one Prometheus identity, carried under two different raw keys, are +/// refused by the runtime that would have to answer both. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one identity to Prometheus: an empty label +/// value is no label, so the two rows are the same series and `rate` creates no collision here. +/// What the two *raw* keys differ in is the empty `zone` against the null one, which the scan +/// encodes differently and the converter therefore opens two input series for. `rate` drops the +/// metric name and the empty value, so the second raw series arrives on the canonical label set the +/// runtime has already closed; the refusal names that label set rather than answering it, because +/// answering would report one series twice. +#[tokio::test] +async fn differing_raw_keys_that_close_on_one_label_set_are_refused() { + let exec = zone_node() + .to_execution_plan(memory_exec_with(zone_schema(), zone_batches())) + .expect("the runtime executes a rate tree"); + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second raw key cannot be answered under the first one's label set"); + assert!( + error.to_string().contains(CLOSED_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// A runtime error after buffered rows is reported, and the terminated stream releases its input. +/// +/// The converter holds the batch it is reading, the chunk it is building and its own source, so a +/// stream that kept its input after the error would hold all three until the consumer drops it. The +/// collision above is the error: it happens on the second series, after the first one's rows have +/// been handed out. +#[tokio::test] +async fn a_runtime_error_after_rows_releases_the_input() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = zone_schema(); + let batches = zone_batches(); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).dropped(&dropped)) + }); + let mut stream = zone_node() + .to_execution_plan(source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!( + "the collision was answered instead of refused: {:?}", + rows(&answered) + ), + } + }; + assert!( + error.to_string().contains(CLOSED_LABEL_SET), + "refused for another reason: {error}" + ); + // The first series' rows were handed out before the error, and the input is gone although the + // terminated stream is still alive here. + assert!( + !answered.is_empty(), + "the first series is answered before the error" + ); + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// Two series with *different* names collapse once `rate` drops the name. +/// +/// This is the collision the canonical label set exists for, at the layer that sees it: +/// `http_requests_total{host="web-1"}` and `http_errors_total{host="web-1"}` reach the window as two +/// identities and leave it as `{host="web-1"}`. [`FusedSeries`] binds one constant metric name for +/// every series it reads, so this collision cannot be staged through it; the runtime is driven +/// directly, with the names on the raw chunks. +#[test] +fn two_differently_named_series_collide_once_rate_drops_the_name() { + let tree = rate_tree(); + let mut execution = NativeExecutor::new( + NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + ) + .expect("the runtime executes a rate tree"), + ); + let samples = series_samples(1.0); + execution + .push_raw(raw_chunk(METRIC, &samples)) + .expect("the first series is a legal raw chunk"); + execution + .next_chunk() + .expect("the first series is answered") + .expect("rate decides every grid point of the first series"); + execution + .push_raw(raw_chunk("http_errors_total", &samples)) + .expect("the second series is a legal raw chunk"); + let error = execution + .next_chunk() + .expect_err("the second series lands on the label set the first one closed"); + assert!( + error.to_string().contains(CLOSED_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// One raw chunk of the series `name{host="web-1"}`: one sample per minute, closing its series. +fn raw_chunk(name: &str, samples: &[(i64, f64)]) -> RawSeriesChunk { + let times = samples.iter().map(|(time, _)| *time).collect::>(); + let values = samples.iter().map(|(_, value)| *value).collect::>(); + RawSeriesChunk::try_new( + &raw_schema(), + Buffer::from_slice_ref(name.as_bytes()), + Some( + MetricName::try_new( + Some(Arc::from(name)), + vec![Label { + name: Arc::from(HOST), + value: Arc::from("web-1"), + }], + ) + .expect("the labels of this test are legal"), + ), + RawTimestamps::try_new(ScalarBuffer::from(times)).expect("ascending timestamps"), + SeriesFields::one(Samples::from_float(Float64Array::from(values))), + true, + ) + .expect("one series' samples are a legal raw chunk") +} + +/// The grid timestamps of the statement, in order. +fn grid_timestamps() -> Vec { + (0..GRID_LEN) + .map(|index| GRID_START + index as i64 * GRID_STEP) + .collect() +} + +/// A waker a test can ask whether it was woken. +/// +/// A poll that stops on its own budget has to wake itself, and the default waker a `poll!` uses +/// swallows that, so the tests that check it poll with this one. +#[derive(Default)] +struct CountingWake(AtomicBool); + +impl futures::task::ArcWake for CountingWake { + fn wake_by_ref(woken: &Arc) { + woken.0.store(true, Ordering::SeqCst); + } +} + +impl CountingWake { + fn woken(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + +/// One series as many chunks that settle nothing: one sample per batch, all before the grid starts. +/// +/// The counter advances one second per second, so `rate` answers 1 at every grid point once the +/// series is read to its end; until a sample reaches the grid's first point, a push decides +/// nothing and settles no chunk. +fn quiet_series_batches() -> Vec { + (0..QUIET_SAMPLES) + .map(|index| series_batch("web-1", &[(index * QUIET_STEP, index as f64)])) + .collect() +} + +/// Asserts that `rows` are `rates`' series, each over the whole grid in order. +fn assert_series_rates(rows: &[Row], rates: &[(f64, &str)]) { + let expected = rates + .iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, *rate, *host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } +} + +fn count(metrics: &MetricsSet, name: &str) -> usize { + match metrics.sum_by_name(name) { + Some(MetricValue::Count { count, .. }) => count.value(), + other => panic!("no count named {name}: {other:?}"), + } +} + +/// Prints `node` the way the logical `EXPLAIN` prints it. +struct Explain<'a>(&'a FusedSeries); + +impl Display for Explain<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.0.fmt_for_explain(f) + } +} + +/// The logical node prints the query a reader wrote, not the tree it compiled to. +#[test] +fn fmt_for_explain_prints_the_promql() { + assert_eq!( + "FusedSeries: rate(x[5m])", + Explain(&fused(rate_tree())).to_string() + ); +} + +/// The physical node prints the tree it was built from, rendered by the tree's own display. +#[test] +fn fmt_as_prints_the_native_tree() { + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec(Vec::new())) + .expect("the runtime executes a rate tree"); + let text = displayable(exec.as_ref()).indent(false).to_string(); + + assert!( + text.starts_with("FusedSeriesExec: rate(x[5m])"), + "the node line names the query first: {text}" + ); + // A distinctive line of the compiled tree: the window's bound range, offset, anchor and grid. + assert!( + text.contains("root Window rate [300000ms] offset 0 anchor - on grid 300000/60000x5"), + "the node line carries the tree it runs: {text}" + ); + assert!( + text.contains("input Input raw samples"), + "the node line carries the tree's source: {text}" + ); +} + +/// `rate(x[5m])` runs end to end over two raw series: values, tags and the dropped metric name. +#[tokio::test] +async fn runs_rate_over_two_series() { + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec(two_series_batches())) + .expect("the runtime executes a rate tree"); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // A counter advancing one second per second rates at 1, and at 2 for the series that advances + // two: the value at each grid point says the right window was read for the right series. + let expected = [(1.0, "web-1"), (2.0, "web-2")] + .into_iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, rate, host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } + // `rate` drops `__name__`, so the tag column that reports it is empty for every row. + assert!( + rows.iter().all(|(_, _, _, name)| name.is_none()), + "rate drops the metric name: {rows:?}" + ); +} + +/// A series that spans several input chunks keeps its absolute grid positions and closes once. +#[tokio::test] +async fn a_series_spanning_chunks_keeps_absolute_positions_and_closes_once() { + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec(one_series_batches())) + .expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // Every grid position exactly once, in order: a chunk's range is never read as a row number. + let times = rows.iter().map(|row| row.0).collect::>(); + assert_eq!(grid_timestamps(), times); + assert!( + rows.iter().all(|(_, value, ..)| value.is_some()), + "every grid point of this series has a value: {rows:?}" + ); + + // The source closed the series on exactly one chunk, and the window saw exactly that one + // closing: one `is_last`, not one per input batch. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(1, count(&metrics, "native.root.input.completed_series")); + assert_eq!(1, count(&metrics, "native.root.completed_series")); +} + +/// A tree the runtime does not execute is refused when the execution plan is built. +#[tokio::test] +async fn a_tree_the_runtime_cannot_run_is_rejected_when_the_exec_is_built() { + let node = fused(aggregate_tree()); + let error = node + .to_execution_plan(memory_exec(Vec::new())) + .expect_err("the runtime does not execute aggregate nodes"); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "the runtime refuses the tree it cannot run: {error}" + ); + assert!(error.to_string().contains("Aggregate"), "{error}"); +} + +/// A tree evaluated on another grid than the node declares is refused before it runs. +#[test] +fn a_tree_bound_to_another_grid_is_rejected() { + let tree = rate_tree(); + let other_grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, GRID_STEP, GRID_LEN).unwrap()).unwrap()); + let binding = FusedSeriesBinding { + grid: other_grid, + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect_err("the tree is not evaluated on the declared grid"); + assert!(error.to_string().contains("grid"), "{error}"); +} + +/// A grid allocated separately from the tree's, with the same spec, is the tree's grid. +/// +/// The comparison is by spec, which is the rule [`NativeTree`]'s own validation and the kernels +/// apply, and a chunk's timestamps come from its absolute grid positions rather than from the +/// allocation, so nothing downstream needs the `Arc` to be shared. This pairs with +/// [`a_tree_bound_to_another_grid_is_rejected`]: a different spec is still refused. +#[test] +fn a_separately_allocated_grid_with_an_equal_spec_is_accepted() { + let tree = rate_tree(); + let equal = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ); + assert!( + !Arc::ptr_eq(&equal, &binding(&tree).grid), + "the two grids are separate allocations" + ); + let binding = FusedSeriesBinding { + grid: equal, + ..binding(&tree) + }; + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("an equal grid spec is the same grid"); +} + +/// An input that ends without a series still runs the tree, and yields no rows. +#[tokio::test] +async fn an_empty_input_runs_the_tree_and_yields_no_rows() { + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec(Vec::new())) + .expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("a tree with no series in it still closes"); + assert!(rows(&batches).is_empty(), "{:?}", rows(&batches)); + // The run closed on no series, so nothing was completed and nothing was decided. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(0, count(&metrics, "native.root.completed_series")); + assert_eq!(0, count(&metrics, "native.root.input.completed_series")); +} + +/// An upstream error after some rows is reported, and the terminated stream releases that input. +#[tokio::test] +async fn an_upstream_error_after_rows_is_reported() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = input_schema(); + let batches = one_series_batches(); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + .error("the source failed") + .dropped(&dropped), + ) + } + }); + let mut stream = fused(rate_tree()) + .to_execution_plan(source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The error is reported rather than swallowed by the rows the input already handed over. + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!( + "the input's error was swallowed after {:?}", + rows(&answered) + ), + } + }; + assert!( + error.to_string().contains("the source failed"), + "reported as another error: {error}" + ); + // Still alive here: the release is the stream's, not the drop of a temporary. + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// A pending poll loses nothing when the next-item future is dropped and the stream polled again. +#[tokio::test] +async fn a_pending_boundary_loses_no_rows() { + let expected = rows( + &collect( + fused(rate_tree()) + .to_execution_plan(memory_exec(one_series_batches())) + .expect("the runtime executes a rate tree"), + SessionContext::default().task_ctx(), + ) + .await + .expect("the node runs"), + ); + assert!(!expected.is_empty(), "the fixture decides some grid points"); + + let schema = input_schema(); + let batches = one_series_batches(); + let source = staged(schema.clone(), move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + // One batch is in the converter's hands when the stream reports pending. + .pending_after(1), + ) + }); + let mut stream = fused(rate_tree()) + .to_execution_plan(source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The future is dropped at the end of this block, as a consumer that gives up waiting does. + { + let next = stream.next(); + futures::pin_mut!(next); + assert!( + futures::poll!(next.as_mut()).is_pending(), + "the input's pending poll reaches the consumer" + ); + } + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the pending poll left it"); + assert_eq!(expected, rows(&batches)); +} + +/// A tag the input has no column for is refused when the plan is built, naming the column. +#[test] +fn a_missing_tag_column_is_rejected_when_the_plan_is_built() { + let tree = rate_tree(); + let binding = FusedSeriesBinding { + tags: vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the logical node does not see the input's columns") + .to_execution_plan(memory_exec(Vec::new())) + .expect_err("the input has no zone column"); + assert!( + error.to_string().contains("no column \"zone\""), + "refused for another reason: {error}" + ); +} + +/// A tag column whose type the converter does not read as a label is refused when the plan is +/// built, naming the column. +#[test] +fn a_wrong_typed_tag_column_is_rejected_when_the_plan_is_built() { + // `zone` is numeric here, which no label of a raw series can be read from. + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Float64, true), + ])); + let tree = rate_tree(); + let binding = FusedSeriesBinding { + tags: vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::clone(&schema).to_dfschema_ref().unwrap(), + }), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the logical node does not see the input's column types") + .to_execution_plan(memory_exec_with(schema, Vec::new())) + .expect_err("a numeric column cannot carry a label"); + assert!( + error.to_string().contains("tag \"zone\""), + "refused without naming the column: {error}" + ); +} + +/// An input column the binding does not map is refused when the plan is built, naming the column. +/// +/// `zone` is part of these series' identity to the scan but not to the binding, so reading the +/// input would drop it: rows a Prometheus query keeps apart would be answered as one series with +/// the wrong labels. +#[test] +fn an_input_column_the_binding_does_not_map_is_rejected() { + let tree = rate_tree(); + let binding = binding(&tree); + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: zone_schema().to_dfschema_ref().unwrap(), + }), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the logical node does not see the input's columns") + .to_execution_plan(memory_exec_with(zone_schema(), Vec::new())) + .expect_err("the binding does not account for the zone column"); + assert!( + error.to_string().contains("column \"zone\""), + "refused without naming the column: {error}" + ); +} + +/// An input with two columns of one name is refused when the plan is built. +/// +/// The coverage check reads names, so both `host` columns pass it as a role the binding maps, while +/// the mapping itself reads `index_of`, which answers with one index: the other column would be +/// part of these rows' identity that a raw series never carries, and rows a Prometheus query keeps +/// apart would be answered as one series with the wrong labels -- and no error. Nothing about the +/// two columns says which one was meant, so the schema is refused rather than one of them chosen. +#[test] +fn an_input_with_two_columns_of_one_name_is_rejected() { + let duplicated = Arc::new(Schema::new(vec![ + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + // The second `host` column, which the name-based coverage check accepts just as well. + Field::new(HOST, DataType::Utf8, true), + ])); + let error = fused(rate_tree()) + .to_execution_plan(memory_exec_with(duplicated, Vec::new())) + .expect_err("the binding maps one of the two host columns"); + assert!( + error + .to_string() + .contains("more than one column named \"host\""), + "refused without naming the duplicate: {error}" + ); +} + +/// One poll reads at most [`POLL_BUDGET`] chunks of input, and stopping there loses nothing. +/// +/// `rate(x[5m])` decides nothing until a sample reaches the grid's first point, so a series can +/// arrive as many chunks that settle no batch at all: a poll that kept reading them would hold the +/// task's thread until the series ended, which is what stops other tasks on that thread and delays +/// the cancellation a consumer asked for. The first series here is [`QUIET_SAMPLES`] such chunks; +/// the second one covers the grid, so a poll that resumes properly still answers both. +#[tokio::test] +async fn one_poll_reads_at_most_the_budget_and_loses_nothing() { + let mut batches = quiet_series_batches(); + batches.push(series_batch("web-2", &series_samples(2.0))); + let total = batches.len(); + assert!( + total > POLL_BUDGET, + "the fixture has to outlast one poll: {total} batches" + ); + + let pulled = Arc::new(AtomicUsize::new(0)); + let woken = Arc::new(CountingWake::default()); + let schema = input_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = fused(rate_tree()) + .to_execution_plan(source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The first poll: the quiet series settles nothing, so the budget is what ends this poll. + let waker = futures::task::waker(Arc::clone(&woken)); + let mut context = Context::from_waker(&waker); + assert!( + stream.as_mut().poll_next(&mut context).is_pending(), + "chunks that settle nothing cannot answer a poll" + ); + let read = pulled.load(Ordering::SeqCst); + assert!(read > 0, "the poll did read its input"); + // One step polls the converter once, and the converter reads at most two batches per poll: the + // key it is holding and the first row of the next, which is what tells it the first key ended. + // The bound is therefore two batches per step, and the budget is what it is measured in. + assert!( + read <= 2 * POLL_BUDGET + 1, + "one poll read {read} of {total} batches, past its budget of {POLL_BUDGET} steps" + ); + assert!( + read < total, + "one poll read the whole input ({read} batches): the budget did not stop it" + ); + assert!( + woken.woken(), + "a poll that stopped on its budget has to wake itself" + ); + + // Nothing the budget stopped is lost: the stream resumes where the poll left it. + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the bounded poll left it"); + assert_series_rates(&rows(&batches), &[(1.0, "web-1"), (2.0, "web-2")]); +} + +/// A batch that is ready is handed over before any batch behind it is read. +/// +/// Each series below is one batch, and a series closes on the first batch of the next one, so the +/// first answer needs two batches and nothing more. A poll that read on would answer the same batch +/// only after the whole input, which is exactly what the budget above exists to prevent. +#[tokio::test] +async fn a_ready_batch_is_returned_before_later_input_is_read() { + let batches = (0..20) + .map(|index| series_batch(&format!("web-{index}"), &series_samples(1.0))) + .collect::>(); + let total = batches.len(); + let pulled = Arc::new(AtomicUsize::new(0)); + let schema = input_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = fused(rate_tree()) + .to_execution_plan(source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let first = stream + .next() + .await + .expect("the first series is answered") + .expect("the first series is answered without an error"); + let answered = rows(&[first]); + assert_series_rates(&answered, &[(1.0, "web-0")]); + // Two batches: the first series' own, and the one whose first key closes it. + assert_eq!( + 2, + pulled.load(Ordering::SeqCst), + "the answer waited for input behind the series it reports" + ); + assert!(total > 2, "the fixture has to outlast the first answer"); +} + +/// A chunk whose rows form several runs is converted row by row from the rows it carries. +/// +/// A sparse result is one chunk holding two runs with a hole between them +/// ([`EvaluatedChunk::range`] answers `None` for it). Reading the hole as rows would invent grid +/// points this series holds no value at, and reading positions as compact rows would shift every +/// row after the hole onto the wrong timestamp: the batch has to hold exactly the carried rows, at +/// their absolute grid times, with nulls nowhere else. +#[test] +fn a_chunk_of_several_runs_converts_its_own_rows_only() { + let grid = grid(); + let axis = TimeAxis::from_grid_spans( + Arc::clone(&grid), + // Grid positions 0..2 and 4..5: position 2 and 3 hold no row at all. + vec![GridSpan::new(0, 2), GridSpan::new(4, 1)], + ) + .expect("the spans are increasing and inside the grid"); + let chunk = EvaluatedChunk::try_new( + Series { + metric: MetricName::try_new(Some(Arc::from(METRIC)), Vec::new()) + .expect("one metric name is a legal identity"), + time: axis, + fields: SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from( + // One value per carried row, so the second run's value sits at grid position 4. + vec![7.0, 8.0, 11.0], + ))]) + .expect("one field is a legal field list"), + }, + true, + ) + .expect("a chunk on a grid"); + + // The columns of `output_schema`, in its order: the field, the two tags, then the time. + let batch = series_to_record_batch( + &[ + OutputColumn::Field(0), + OutputColumn::Tag(0), + OutputColumn::Tag(1), + OutputColumn::Time, + ], + &[HOST.to_string(), NAME.to_string()], + &output_schema(), + &chunk, + ) + .expect("the chunk converts"); + + let stamps = batch + .column_by_name(TIME) + .expect("the batch has its time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let values = batch + .column_by_name(VALUE) + .expect("the batch has its field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + assert_eq!(3, batch.num_rows(), "only the carried rows are rows"); + let grid_stamps = grid.timestamps(); + assert_eq!( + vec![grid_stamps[0], grid_stamps[1], grid_stamps[4]], + stamps.values()[..3].to_vec(), + "the carried rows keep their absolute grid positions" + ); + assert_eq!(vec![7.0, 8.0, 11.0], values.values()[..3].to_vec()); +} + +/// Every message in `error`'s chain, so an error a layer above wrapped is still read. +fn error_chain(error: &DataFusionError) -> String { + let mut message = error.to_string(); + let mut source = std::error::Error::source(error); + while let Some(cause) = source { + message.push_str(": "); + message.push_str(&cause.to_string()); + source = cause.source(); + } + message +} + +/// Two partitions, one series each, and both series leave the window on the same label set. +/// +/// `rate` drops `__name__`, so `http_requests_total{host="web-1"}` reads as `{host="web-1"}` in +/// the output. `PerSeries` distribution keeps each of the two whole inside its own partition, so a +/// check that lived inside one execution would see one series each and report nothing: only state +/// shared by the partitions sees that the output vector holds the label set twice. +#[tokio::test] +async fn a_duplicate_label_set_across_partitions_is_refused() { + let series = vec![series_batch("web-1", &series_samples(1.0))]; + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec_partitioned(vec![series.clone(), series])) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second partition lands on the label set the first one claimed"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// The same label set on points that do not overlap is legal, and every point survives. +/// +/// A rate needs two samples in its window, so the first series only has a value at the first grid +/// point and the second only at the last: the two share a final label set but no evaluation +/// timestamp. Prometheus merges such series into one; this node emits both batches and leaves the +/// merge to matrix assembly, so the check has to allow what it would otherwise refuse. +#[tokio::test] +async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { + let first = vec![series_batch( + "web-1", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + )]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![series_batch( + "web-1", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + )]; + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("disjoint grid points are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + let mut rows = rows(&batches); + rows.sort_by_key(|row| row.0); + assert_eq!( + vec![GRID_START, last], + rows.iter().map(|row| row.0).collect::>(), + "both series must survive with the one point each holds" + ); + for ((_, value, host, _), rate) in rows.iter().zip([1.0, 1.0]) { + assert_eq!(Some("web-1"), host.as_deref()); + let value = value.expect("the point each series holds has a value"); + assert!((rate - value).abs() < 1e-9, "{value} is not {rate}"); + } +} + +/// Each partition runs its own series and answers it, with no label set claimed twice. +#[tokio::test] +async fn every_partition_evaluates_its_own_series() { + let partitions = [1.0, 2.0, 3.0] + .into_iter() + .enumerate() + .map(|(index, scale)| { + vec![series_batch( + &format!("web-{index}"), + &series_samples(scale), + )] + }) + .collect::>(); + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec_partitioned(partitions)) + .expect("the runtime executes a rate tree"); + assert_eq!( + 3, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("each partition runs its own tree"); + assert_eq!(3, batches.len(), "one batch per partition"); + let rows = rows(&batches); + assert_eq!( + 3 * GRID_LEN as usize, + rows.len(), + "every partition reports its own series over the whole grid" + ); + // Batches arrive in whatever order the partitions settle, so the series are read apart. + for (index, scale) in [1.0, 2.0, 3.0].into_iter().enumerate() { + let host = format!("web-{index}"); + let mut series = rows + .iter() + .filter(|row| row.2.as_deref() == Some(host.as_str())) + .cloned() + .collect::>(); + series.sort_by_key(|row| row.0); + assert_eq!(GRID_LEN as usize, series.len(), "{host}: {series:?}"); + assert_series_rates(&series, &[(scale, host.as_str())]); + } +} + +/// A partition index the input does not have is refused rather than answered. +/// +/// The node runs one evaluator per input partition and adds none of its own, so it forwards the +/// index to its input: the count it declares is the input's, and an index past that is the input's +/// refusal, not a series read from another partition. +#[tokio::test] +async fn a_partition_past_the_inputs_count_is_refused() { + let schema = input_schema(); + let source = staged(schema.clone(), move || { + Box::pin(StagedStream::new(schema.clone(), Vec::new())) + }); + let exec = fused(rate_tree()) + .to_execution_plan(source) + .expect("the runtime executes a rate tree"); + assert_eq!( + 1, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + exec.execute(0, SessionContext::default().task_ctx()) + .expect("the one partition the input has"); + let error = match exec.execute(1, SessionContext::default().task_ctx()) { + Ok(_) => panic!("a partition the input does not have is answered"), + Err(error) => error, + }; + assert!( + error.to_string().contains("partition"), + "refused for another reason: {error}" + ); +}