diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 9f7ddbc77dd..3bdcc8a74c0 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -41,7 +41,9 @@ 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 fused_series::{ + FusedOutputFormat, 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 index e34821b98aa..f35a5d4b01b 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -44,7 +44,9 @@ use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; -use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; +use common_query::series_frame::{ + ensure_frame_matches_label_source, pack_evaluated_series, unpack_raw_series, +}; use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::ScalarBuffer; use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; @@ -64,6 +66,7 @@ use datafusion::physical_plan::{ Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use datatypes::series::SeriesOrdering; +use datatypes::series::frame::{FrameLabelMode, SeriesFrameBuilder, SeriesFrameCursor}; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesSchema}; use datatypes::series::time::EvalGrid; @@ -164,6 +167,48 @@ impl PartialOrd for FusedSeriesBinding { } } +/// The output representation fixed when a fused plan is built. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum FusedOutputFormat { + /// The existing tabular query result. + PlainRows, + /// A self-describing evaluated SeriesFrame. + SeriesFrame, +} + +impl FusedOutputFormat { + const fn as_str(self) -> Option<&'static str> { + match self { + Self::PlainRows => None, + Self::SeriesFrame => Some("series_frame"), + } + } +} + +fn evaluated_output(tree: &NativeTree) -> DataFusionResult> { + match &tree.root().output { + NativeType::Samples(schema) if schema.eval_grid.is_some() => Ok(Arc::clone(schema)), + NativeType::Samples(_) => Err(DataFusionError::Plan( + "FusedSeries frame output requires evaluated samples with a grid".into(), + )), + NativeType::AggregateState(_) => Err(DataFusionError::Plan( + "FusedSeries frame output does not support an aggregate-state root".into(), + )), + NativeType::Scalar(_) => Err(DataFusionError::Plan( + "FusedSeries frame output does not support a scalar root".into(), + )), + } +} + +fn series_frame_output_schema(schema: &SeriesSchema) -> DataFusionResult { + let stream = + pack_evaluated_series(Box::pin(futures::stream::empty()), Arc::new(schema.clone())) + .map_err(|error| { + DataFusionError::Plan(format!("FusedSeries cannot declare frame schema: {error}")) + })?; + Ok(stream.schema()) +} + /// 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 @@ -186,6 +231,8 @@ pub struct FusedSeries { /// This row-state bridge is temporary: native final aggregation will consume series frames /// directly, removing state-row encoding and reconstruction. partial_fold: Option, + /// Fixed output representation; partial state export always remains PlainRows. + format: FusedOutputFormat, } impl FusedSeries { @@ -201,8 +248,39 @@ impl FusedSeries { binding: FusedSeriesBinding, children: Vec, output_schema: DFSchemaRef, + ) -> DataFusionResult { + Self::try_new_with_format( + promql, + tree, + binding, + children, + output_schema, + FusedOutputFormat::PlainRows, + ) + } + + /// Creates a node with an explicitly selected output format. + pub fn try_new_with_format( + promql: Arc, + tree: NativeTree, + binding: FusedSeriesBinding, + children: Vec, + output_schema: DFSchemaRef, + format: FusedOutputFormat, ) -> DataFusionResult { validate_binding(&tree, &binding, children.len())?; + if format == FusedOutputFormat::SeriesFrame && children.len() > 1 { + return Err(DataFusionError::Plan( + "FusedSeries SeriesFrame output does not yet support multiple sources".into(), + )); + } + let output_schema = if format == FusedOutputFormat::SeriesFrame { + Arc::new(datafusion::common::DFSchema::try_from( + series_frame_output_schema(evaluated_output(&tree)?.as_ref())?, + )?) + } else { + output_schema + }; Ok(Self { promql, tree, @@ -210,6 +288,7 @@ impl FusedSeries { children, output_schema, partial_fold: None, + format, }) } @@ -243,6 +322,7 @@ impl FusedSeries { children, output_schema, partial_fold: Some(op), + format: FusedOutputFormat::PlainRows, }) } @@ -296,6 +376,7 @@ impl FusedSeries { self.binding.time_index.clone(), self.binding.tags.clone(), self.partial_fold.clone(), + self.format, )?)) } } @@ -327,6 +408,9 @@ impl UserDefinedLogicalNodeCore for FusedSeries { self.promql, render_partial_fold(op) ), + None if self.format.as_str().is_some() => { + write!(f, "FusedSeries: {} (output: series_frame)", self.promql) + } None => write!(f, "FusedSeries: {}", self.promql), } } @@ -344,6 +428,7 @@ impl UserDefinedLogicalNodeCore for FusedSeries { children: inputs, output_schema: Arc::clone(&self.output_schema), partial_fold: self.partial_fold.clone(), + format: self.format, }) } } @@ -361,6 +446,7 @@ impl PartialEq for FusedSeries { && self.children == other.children && self.output_schema == other.output_schema && partial_key(self.partial_fold.as_ref()) == partial_key(other.partial_fold.as_ref()) + && self.format == other.format } } @@ -373,6 +459,7 @@ impl Hash for FusedSeries { self.children.hash(state); self.output_schema.hash(state); partial_key(self.partial_fold.as_ref()).hash(state); + self.format.hash(state); } } @@ -400,6 +487,10 @@ impl PartialOrd for FusedSeries { Some(Ordering::Equal) => {} ord => return ord, } + match self.format.cmp(&other.format) { + Ordering::Equal => {} + ord => return Some(ord), + } match tree_input_fields(&self.tree).cmp(&tree_input_fields(&other.tree)) { Ordering::Equal => {} ord => return Some(ord), @@ -783,6 +874,7 @@ pub struct FusedSeriesExec { metric: ExecutionPlanMetricsSet, /// Aggregate state exported by each partition, when this is the partial mode. partial_fold: Option, + format: FusedOutputFormat, /// This node's partitioning, emission and boundedness. properties: Arc, } @@ -807,9 +899,19 @@ impl FusedSeriesExec { time_column: String, tags: Vec, ) -> DataFusionResult { - Self::build(promql, inputs, tree, output_schema, time_column, tags, None) + Self::build( + promql, + inputs, + tree, + output_schema, + time_column, + tags, + None, + FusedOutputFormat::PlainRows, + ) } + #[allow(clippy::too_many_arguments)] fn build( promql: Arc, inputs: Vec<(Arc, RawLabelSource, SeriesOrdering)>, @@ -818,7 +920,13 @@ impl FusedSeriesExec { time_column: String, tags: Vec, partial_fold: Option, + format: FusedOutputFormat, ) -> DataFusionResult { + if partial_fold.is_none() && format == FusedOutputFormat::SeriesFrame && inputs.len() > 1 { + return Err(DataFusionError::Plan( + "FusedSeries SeriesFrame output does not yet support multiple sources".into(), + )); + } let raw_schemas = validate_physical_tree(&tree)?; if inputs.len() != raw_schemas.len() { return Err(DataFusionError::Plan(format!( @@ -853,7 +961,29 @@ impl FusedSeriesExec { } } } - let classification = if partial_fold.is_some() { + // A scalar root answers one value per grid point and names no series, so it is not something + // a frame of evaluated samples can carry. Refusing it here names the shape the caller asked + // for; `evaluated_output` would reach the same verdict later with a narrower message. + if partial_fold.is_none() && format == FusedOutputFormat::SeriesFrame && scalar_root { + return Err(DataFusionError::Plan( + "FusedSeries frame output requires evaluated samples, but this tree has a scalar \ + root and no raw source" + .into(), + )); + } + let frame_schema = if partial_fold.is_none() && format == FusedOutputFormat::SeriesFrame { + let declared = series_frame_output_schema(evaluated_output(&tree)?.as_ref())?; + if declared != output_schema { + return Err(DataFusionError::Plan( + "FusedSeriesExec output schema does not match the evaluated frame layout" + .into(), + )); + } + true + } else { + false + }; + let classification = if partial_fold.is_some() || frame_schema { Vec::new() } else if scalar_root { let scalar_fields = scalar_root_fields(&output_schema, &output_tags, &time_column)?; @@ -931,6 +1061,7 @@ impl FusedSeriesExec { metric: ExecutionPlanMetricsSet::new(), properties, partial_fold, + format, }) } @@ -953,6 +1084,8 @@ impl DisplayAs for FusedSeriesExec { write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree)?; if let Some(op) = &self.partial_fold { write!(f, "\npartial state: {}", render_partial_fold(op))?; + } else if self.format.as_str().is_some() { + write!(f, "\noutput: series_frame")?; } Ok(()) } @@ -1028,6 +1161,7 @@ impl ExecutionPlan for FusedSeriesExec { self.time_column.clone(), self.tags.to_vec(), self.partial_fold.clone(), + self.format, )?)) } @@ -1098,6 +1232,30 @@ impl ExecutionPlan for FusedSeriesExec { if streams.is_empty() { executor.end_of_input()?; } + // A frame is encoded from the tree's own evaluated output, and only when the node answers + // values rather than state: partial state is a third output the format does not reach. + let frame_mode = + self.partial_fold.is_none() && self.format == FusedOutputFormat::SeriesFrame; + let eval_schema = if frame_mode { + Some(evaluated_output(&self.tree)?) + } else { + None + }; + let frame_builder = eval_schema + .as_ref() + .map(|schema| { + SeriesFrameBuilder::try_new( + schema, + FrameLabelMode::Explicit, + None, + SeriesFrameBuilder::DEFAULT_MAX_ROWS, + SeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .map_err(|error| { + DataFusionError::Plan(format!("cannot create evaluated frame builder: {error}")) + }) + }) + .transpose()?; Ok(Box::pin(FusedSeriesStream { inputs: streams, executor: Some(executor), @@ -1106,6 +1264,8 @@ impl ExecutionPlan for FusedSeriesExec { tags: Arc::clone(&self.tags), output_schema: self.output_schema.clone(), partial_fold, + frame_builder, + eval_schema, _export_budget: Arc::clone(&export_budget), ready: VecDeque::new(), done: false, @@ -1226,6 +1386,8 @@ pub struct FusedSeriesStream { output_schema: SchemaRef, /// Per-partition aggregate state fold, if this stream exports a partial result. partial_fold: Option, + frame_builder: Option, + eval_schema: Option>, /// Keeps exported state batches charged for as long as this stream holds them. _export_budget: Arc>>, /// The batches the executor settled since the last poll of the input. @@ -1300,6 +1462,7 @@ impl FusedSeriesStream { self.done = true; self.executor = None; self.partial_fold = None; + self.frame_builder = None; self.inputs.clear(); // This reservation belongs only to this execution. Release it even if the consumer keeps // the failed stream and the plan alive. @@ -1320,13 +1483,44 @@ impl FusedSeriesStream { return Ok(()); } for chunk in &chunks { - let batch = series_to_record_batch( - &self.classification, - &self.tags, - &self.output_schema, - chunk, - )?; - self.ready.push_back(batch); + if let (Some(schema), Some(builder)) = + (self.eval_schema.as_ref(), self.frame_builder.as_mut()) + { + let mut cursor = SeriesFrameCursor::try_new_evaluated( + schema, + chunk.series.clone(), + chunk.is_last, + ) + .map_err(|error| { + DataFusionError::Execution(format!("cannot frame evaluated chunk: {error}")) + })?; + while let Some(batch) = cursor.append_to(builder).map_err(|error| { + DataFusionError::Execution(format!("cannot append evaluated frame: {error}")) + })? { + self.ready + .push_back(batch.with_schema(self.output_schema.clone())?); + } + } else { + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + chunk, + )?; + self.ready.push_back(batch); + } + } + Ok(()) + } + + fn flush_frames(&mut self) -> DataFusionResult<()> { + if let Some(builder) = self.frame_builder.as_mut() + && let Some(batch) = builder.flush().map_err(|error| { + DataFusionError::Execution(format!("cannot flush evaluated frame: {error}")) + })? + { + self.ready + .push_back(batch.with_schema(self.output_schema.clone())?); } Ok(()) } @@ -1372,7 +1566,7 @@ impl Stream for FusedSeriesStream { } } Ok(None) => { - if let Err(error) = this.finish_fold() { + if let Err(error) = this.finish_fold().and_then(|()| this.flush_frames()) { this.fail(); return Poll::Ready(Some(Err(error))); } @@ -1388,7 +1582,16 @@ impl Stream for FusedSeriesStream { }; let (source, polled) = (*source, input.as_mut().poll_next(cx)); match polled { - Poll::Pending => return Poll::Pending, + Poll::Pending => { + if let Err(error) = this.flush_frames() { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + if let Some(batch) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(batch))); + } + return Poll::Pending; + } Poll::Ready(None) => { this.inputs.pop_front(); if this.inputs.is_empty() diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 00df6f05aeb..910f587c5a8 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -36,6 +36,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::task::{Context, Poll}; +use common_query::series_frame::unpack_evaluated_series; use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; @@ -65,8 +66,9 @@ use datafusion::physical_plan::{ use datafusion::prelude::SessionContext; use datatypes::prelude::ConcreteDataType; use datatypes::series::frame::{ - FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, LabelDecodeSpec, - LabelDecodeTag, pack_raw_series_chunk, raw_series_frame_schema, series_frame_schema, + FrameContent, FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, + LabelDecodeSpec, LabelDecodeTag, SeriesFrameDeclaration, pack_raw_series_chunk, + raw_series_frame_schema, series_frame_schema, }; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{FieldSpec, Samples, Series, SeriesFields, SeriesSchema}; @@ -75,7 +77,10 @@ use datatypes::series::{Label, SeriesLabels, SeriesOrdering}; use futures::{Stream, StreamExt, TryStreamExt}; use store_api::region_engine::RawLabelSource; -use super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, series_to_record_batch}; +use super::{ + FusedOutputFormat, FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, + series_to_record_batch, +}; use crate::evaluated::EvaluatedChunk; use crate::extension_plan::aggregate_exchange::state_schema; use crate::native::execute::{InputId, NativeExecutionState, NativeExecutor}; @@ -85,6 +90,7 @@ use crate::native::plan::{ }; use crate::range_eval::{RangeFunction, RangeWindow}; use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels}; +use common_query::series_frame::series_frame_ordering; /// The query every tree here is compiled from. const PROMQL: &str = "rate(x[5m])"; @@ -3731,3 +3737,417 @@ async fn a_failed_input_exports_no_partial_state() { "refused for another reason: {error}" ); } + +/// Frame output is planned from the evaluated root and is emitted without flattening to rows. +#[tokio::test] +async fn series_frame_output_uses_evaluated_codec_schema() { + let tree = rate_tree(); + let frame_node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .expect("evaluated samples have a frame schema"); + assert!( + frame_node + .schema() + .fields() + .iter() + .all(|field| field.name() != "ts"), + "frame schema comes from the codec rather than the row output" + ); + assert!(format!("{frame_node:?}").contains("FusedSeries")); + let exec = plan_of(frame_node, memory_exec(one_series_frames())).unwrap(); + assert_eq!(1, exec.properties().partitioning.partition_count()); + let batches = collect(exec.clone(), SessionContext::default().task_ctx()) + .await + .unwrap(); + assert!(!batches.is_empty()); + assert!(batches.iter().all(|batch| batch.schema() == exec.schema())); + assert!(batches.iter().any(|batch| batch.num_rows() > 0)); +} + +/// Frame output refuses the currently unsupported multi-source physical shape at planning time. +#[test] +fn series_frame_rejects_multiple_sources_during_planning() { + let tree = binary_tree_over_sources(BinaryOp::Add, raw_schema(), raw_schema(), 0, 1); + let err = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(frame_schema()), input_plan(frame_schema())], + binary_output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .unwrap_err(); + assert!(err.to_string().contains("multiple sources")); +} + +/// A real four-partition input remains four independent frame-output executions. +#[tokio::test] +async fn series_frame_preserves_four_input_partitions() { + let partitions = (0..4) + .map(|index| { + frames(&[series_chunk( + &format!("part-{index}"), + &series_samples((index + 1) as f64), + true, + )]) + }) + .collect::>(); + let tree = rate_tree(); + let node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let exec = plan_of(node, memory_exec_partitioned(partitions)).unwrap(); + assert_eq!(4, exec.properties().partitioning.partition_count()); + let schema = exec.schema(); + let mut closes = 0; + for partition in 0..4 { + let stream = exec + .execute(partition, SessionContext::default().task_ctx()) + .unwrap(); + let decoded = unpack_evaluated_series(stream, None) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + 1, + decoded.len(), + "one chain closes in partition {partition}" + ); + assert!(decoded[0].1); + closes += decoded.iter().filter(|(_, last)| *last).count(); + let values = decoded[0].0.fields.get(0).unwrap(); + let value = values.iter().next().unwrap(); + let datatypes::series::samples::SampleRef::Float(value) = value else { + panic!("float sample expected") + }; + assert!((value - (partition + 1) as f64).abs() < 1e-9); + assert_eq!(schema, exec.schema()); + } + assert_eq!(4, closes); +} + +/// PlainRows and SeriesFrame encode precisely the same rate values, timestamps and labels. +#[tokio::test] +async fn series_frame_rate_matches_plain_rows_bitwise() { + let input = memory_exec(two_series_frames()); + let plain = plan(rate_tree(), Arc::clone(&input)).unwrap(); + let plain_rows = rows( + &collect(plain, SessionContext::default().task_ctx()) + .await + .unwrap(), + ); + let tree = rate_tree(); + let node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let exec = plan_of(node, input).unwrap(); + let batches = collect(exec.clone(), SessionContext::default().task_ctx()) + .await + .unwrap(); + assert!(batches.iter().all(|batch| batch.schema() == exec.schema())); + let stream = Box::pin(RecordBatchStreamAdapter::new( + exec.schema(), + futures::stream::iter(batches.into_iter().map(Ok)), + )); + let decoded = unpack_evaluated_series(stream, None) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut frame_rows = Vec::new(); + for (series, _) in decoded { + let timestamps = series.time.timestamps().unwrap(); + let samples = series.fields.get(0).unwrap(); + for (row, timestamp) in timestamps.iter().enumerate() { + let value = samples.iter().nth(row).unwrap(); + let datatypes::series::samples::SampleRef::Float(value) = value else { + panic!("float sample expected") + }; + frame_rows.push(( + *timestamp, + value.to_bits(), + series.metric.get(HOST).map(str::to_owned), + series.metric.name().map(str::to_owned), + )); + } + } + let mut expected = plain_rows + .into_iter() + .map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name)) + .collect::>(); + frame_rows.sort(); + expected.sort(); + assert_eq!( + expected, + frame_rows + .into_iter() + .map(|(time, bits, host, name)| (time, Some(bits), host, name)) + .collect::>() + ); +} + +/// The partial-state node keeps its state schema and explain annotation, independently of output format. +#[tokio::test] +async fn partial_fold_keeps_plain_rows_state_export() { + let node = partial_node(AggregateKind::Sum); + assert_eq!(FusedOutputFormat::PlainRows, node.format); + let exec = plan_of(node, memory_exec(one_series_frames())).unwrap(); + let state = state_schema(&[HOST.to_string()]); + assert_eq!(state.as_ref(), exec.schema().as_ref()); + assert!( + displayable(exec.as_ref()) + .indent(false) + .to_string() + .contains("partial state:") + ); + let batch = state_of(&exec, 0, &state).await; + assert_eq!(state.as_ref(), batch.schema().as_ref()); +} + +/// `EXPLAIN` says nothing extra about the row table and names the frame output. +/// +/// The row table is what a query returns, so its node line has to stay exactly what it was; a frame +/// output answers something else, and the line is the only place a reader can see that. +#[test] +fn the_explain_text_names_the_frame_format_only() { + let rows = fused(rate_tree()); + assert_eq!( + "FusedSeries: rate(x[5m])", + Explain(&rows).to_string(), + "the row-table line is what it always was" + ); + let framed = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + rows.tree().clone(), + rows.binding().clone(), + rows.inputs().into_iter().cloned().collect(), + Arc::clone(rows.schema()), + FusedOutputFormat::SeriesFrame, + ) + .expect("a rate tree is evaluated on a grid"); + assert!( + Explain(&framed) + .to_string() + .contains("(output: series_frame)"), + "a frame output names its format: {}", + Explain(&framed) + ); + + let plain = plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes the tree"); + let text = displayable(plain.as_ref()).indent(false).to_string(); + assert!( + text.starts_with("FusedSeriesExec: rate(x[5m])\n"), + "the row-table node line is unchanged: {text}" + ); +} + +/// The logical schema follows the format, and a frame schema is the codec's own declaration. +/// +/// The frame layout is not assembled here: it is read back through the declaration parser, so what +/// the plan claims is what a consumer of the frame will decode. +#[test] +fn the_logical_schema_follows_the_output_format() { + let rows = fused(rate_tree()); + assert_eq!(FusedOutputFormat::PlainRows, rows.format); + assert_eq!(&output_schema(), rows.schema().inner()); + + let framed = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + rows.tree().clone(), + rows.binding().clone(), + rows.inputs().into_iter().cloned().collect(), + Arc::clone(rows.schema()), + FusedOutputFormat::SeriesFrame, + ) + .expect("a rate tree is evaluated on a grid"); + assert_eq!(FusedOutputFormat::SeriesFrame, framed.format); + let declaration = + SeriesFrameDeclaration::parse(framed.schema().inner().as_ref()).expect("a frame schema"); + assert_eq!(FrameContent::Evaluated, declaration.content()); + assert_eq!(FrameLabelMode::Explicit, declaration.labels()); + assert_eq!( + SeriesOrdering::Contiguous, + series_frame_ordering(framed.schema().inner().as_ref()) + .expect("the frame declares its ordering"), + "computed values promise per-chain contiguity and nothing stronger" + ); +} + +/// Rebuilding a frame node keeps the format it was built with. +/// +/// The format decides the node's schema, so a rebuild that dropped it would answer rows where the +/// plan said frame. `with_exprs_and_inputs` is what the optimizer uses. +#[test] +fn rebuilding_a_frame_node_keeps_its_output_format() { + let rows = fused(rate_tree()); + let framed = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + rows.tree().clone(), + rows.binding().clone(), + rows.inputs().into_iter().cloned().collect(), + Arc::clone(rows.schema()), + FusedOutputFormat::SeriesFrame, + ) + .expect("a rate tree is evaluated on a grid"); + let rebuilt = framed + .with_exprs_and_inputs( + Vec::new(), + framed.inputs().into_iter().cloned().collect::>(), + ) + .expect("the node rebuilds over its own input"); + assert_eq!(FusedOutputFormat::SeriesFrame, rebuilt.format); + assert_eq!(framed.schema(), rebuilt.schema()); + assert_eq!(Explain(&framed).to_string(), Explain(&rebuilt).to_string()); +} + +/// A tree whose root is not evaluated samples cannot be framed, and says so when the plan is built. +/// +/// A scalar has no series to name, so it is not a series of float samples; an aggregate-state root is +/// refused by the same check for the same reason, and `partial_node` cannot even be built with a +/// frame format. Both refusals happen while the plan is built, because a plan that is built is a +/// plan that runs. +#[test] +fn a_series_frame_refuses_a_root_it_cannot_carry() { + let scalar = scalar_tree(grid(), 1.0); + // A scalar tree binds no raw source, so it is built over no child at all: the refusal below is + // then about the root, not about a source count that was wrong first. + let binding = FusedSeriesBinding { + grid: grid(), + metric_name: Some(Arc::from(METRIC)), + time_index: TIME.to_string(), + tags: vec![HOST.to_string(), NAME.to_string()], + }; + let error = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + scalar, + binding, + Vec::new(), + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .expect_err("a scalar is not a series to frame"); + assert!( + matches!(error, DataFusionError::Plan(_)), + "refused with {error}" + ); + assert!( + error.to_string().contains("scalar"), + "the refusal says what the root is: {error}" + ); + + // The partial mode is a different output entirely, so a frame format cannot even be asked for. + assert_eq!( + FusedOutputFormat::PlainRows, + partial_node(AggregateKind::Sum).format + ); +} + +/// A pending input hands the rows a frame already holds over before it reports pending. +/// +/// The builder buffers up to a batch, and a closing row only exists once its chunk is encoded. If a +/// pending input could hold those rows, a consumer waiting for a series to close would wait for +/// input that may never arrive. +#[tokio::test] +async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() { + let schema = frame_schema(); + let batches = one_series_frames(); + let staged_schema = Arc::clone(&schema); + let source = staged(Arc::clone(&schema), move || { + Box::pin( + StagedStream::new(Arc::clone(&staged_schema), batches.clone()) + // One batch is in the decoder's hands when the stream reports pending. + .pending_after(1), + ) + }); + let tree = rate_tree(); + let node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(schema)], + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .expect("a rate tree is evaluated on a grid"); + let exec = plan_of(node, source).expect("the runtime executes the tree"); + let declared = exec.schema(); + let mut stream = exec + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let first = stream + .next() + .await + .expect("the rows the node already encoded come out before the pending poll") + .expect("the frame is not an error"); + assert_eq!( + declared.as_ref(), + first.schema().as_ref(), + "a batch is declared as the plan declared it" + ); + + let mut answered = vec![first]; + answered.extend( + stream + .try_collect::>() + .await + .expect("the stream resumes where the pending poll left it"), + ); + let stream = Box::pin(RecordBatchStreamAdapter::new( + declared, + futures::stream::iter(answered.into_iter().map(Ok)), + )); + let decoded = unpack_evaluated_series(stream, None) + .expect("the node declares the frame it yields") + .try_collect::>() + .await + .expect("every batch of the frame decodes"); + assert_eq!( + 1, + decoded.iter().filter(|(_, last)| *last).count(), + "one series closes exactly once, across the pending boundary" + ); + // What the row table answers for the same input is the reference: `rate` decides which grid + // points carry a value, so computing the expected timestamps here would guess at it. + let plain = + plan(rate_tree(), memory_exec(one_series_frames())).expect("the runtime executes a tree"); + let mut expected = rows( + &collect(plain, SessionContext::default().task_ctx()) + .await + .expect("the row table runs"), + ) + .into_iter() + .map(|(time, _, _, _)| time) + .collect::>(); + expected.sort_unstable(); + let mut times = Vec::new(); + for (series, _) in &decoded { + for timestamp in series.time.timestamps().expect("a grid axis").iter() { + times.push(*timestamp); + } + } + times.sort_unstable(); + assert_eq!( + expected, times, + "every row the row table answers is in the frame, none twice" + ); +}