From 2f173c58bdc9f032e900ce28d47c355b2170faf7 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:14:53 +0800 Subject: [PATCH] feat(promql): add a series frame output to the fused native node The fused node answers one way: a row table with the grid timestamps, one column per series field and one per declared tag. That is what a query returns, but a consumer of the native path needs to tell series and positions apart, and the row shape flattens both away - which series a value belongs to, and which grid positions that series actually holds. This adds a second output format, fixed when the plan is built rather than chosen at run time: PlainRows the existing row table, unchanged and still the only one a query reaches SeriesFrame the same evaluated chunks handed over as the unified series frame The format is part of the plan because it decides the node's output schema, so it takes part in the node's equality, its EXPLAIN text and its rebuild over new children. A row-table plan prints exactly what it printed before; a frame plan says so on its own line. A frame is encoded straight from each evaluated chunk, never by flattening it into rows and packing it again, so a sparse series keeps its runs and a closing chunk with no position is still one row. The output keeps one partition per input partition: a frame is not a reason to coalesce, and no exchange is inserted in either format. The partial-state output that a complete aggregate reads is a third thing again, and the format does not reach it. When a node exports partition-local state, the format is `PlainRows` by construction, `convert` hands every chunk to the fold before any format is consulted, and the EXPLAIN line keeps its `partial state:` suffix. Nothing in the aggregate path is rewritten to make room for frames. A tree whose root is not evaluated samples is refused while the plan is built, because aggregate state is a contribution to a later merge and a scalar has no series to name. So is a frame output over more than one raw source, which this node does not run in parallel yet. Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/extension_plan.rs | 4 +- src/promql/src/extension_plan/fused_series.rs | 227 +++++++++- .../src/extension_plan/fused_series/tests.rs | 426 +++++++++++++++++- 3 files changed, 641 insertions(+), 16 deletions(-) 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" + ); +}