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}" + ); +}