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 1/4] 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" + ); +} From e3c6193e5736fea64f964c0f7c651e22a43a5eaf Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 03:30:36 +0800 Subject: [PATCH 2/4] test(promql): pin the frame output's boundaries and its pending hand-over The frame output was covered for its happy path and for the plan-level decisions, but three things the spec asks about had no test of their own at this boundary, and one test did not prove what it claimed. The pending test is the one that mattered. It awaited the stream, and `StagedStream` wakes itself right after reporting pending, so the await ran straight through the boundary: every later assertion held whether or not the buffered rows were handed over at the pending poll. It now polls that single poll by hand and requires it to answer the rows it already encoded. Deleting the `flush_frames` call from the pending branch makes it fail, which it did not before. The three boundaries now have their own tests: a chain that closes with no position survives as one row where the row table answers nothing at all; a sparse series carries exactly the positions the row table answers and no filled hole; and two chains that close on one label set are both kept, with values that differ so a lost chain would show. Also drops the `as_str` helper, which only asked whether to show a format and then spelled the same string twice, and corrects two comments: `convert` produces three outputs rather than two, and a built plan is not a promise that execution cannot fail. Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/extension_plan/fused_series.rs | 16 +- .../src/extension_plan/fused_series/tests.rs | 302 +++++++++++++++++- 2 files changed, 291 insertions(+), 27 deletions(-) diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index f35a5d4b01b..057f20be5d1 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -176,15 +176,6 @@ pub enum FusedOutputFormat { 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)), @@ -408,7 +399,7 @@ impl UserDefinedLogicalNodeCore for FusedSeries { self.promql, render_partial_fold(op) ), - None if self.format.as_str().is_some() => { + None if self.format == FusedOutputFormat::SeriesFrame => { write!(f, "FusedSeries: {} (output: series_frame)", self.promql) } None => write!(f, "FusedSeries: {}", self.promql), @@ -1084,7 +1075,7 @@ 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() { + } else if self.format == FusedOutputFormat::SeriesFrame { write!(f, "\noutput: series_frame")?; } Ok(()) @@ -1474,7 +1465,8 @@ impl FusedSeriesStream { self.ready.clear(); } - /// Converts settled chunks into query rows or folds them into partial aggregate state. + /// Converts settled chunks into the node's output: query rows, an evaluated series frame, or + /// partition-local state for a later aggregate merge. fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { if let Some(fold) = self.partial_fold.as_mut() { for chunk in chunks { diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 910f587c5a8..37699189fc8 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -4023,8 +4023,8 @@ fn rebuilding_a_frame_node_keeps_its_output_format() { /// /// 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. +/// frame format. The refusal is a planning error, so a caller learns about it before the query runs +/// rather than from a partition that fails later. #[test] fn a_series_frame_refuses_a_root_it_cannot_carry() { let scalar = scalar_tree(grid(), 1.0); @@ -4061,21 +4061,32 @@ fn a_series_frame_refuses_a_root_it_cannot_carry() { ); } -/// A pending input hands the rows a frame already holds over before it reports pending. +/// A pending input hands the rows a frame already encoded 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. +/// pending input could hold those rows, a consumer waiting for a series to close would wait for input +/// that may never arrive. +/// +/// This polls by hand rather than awaiting the stream: `StagedStream` reports pending once and then +/// wakes itself, so an `await` would run straight through the boundary and satisfy every later +/// assertion whether or not the rows were handed over at the right moment. What is checked instead is +/// the poll that *meets* the pending input - it has to return rows, and the input has to have been +/// read only once by then. #[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 handed = Arc::new(AtomicUsize::new(0)); let staged_schema = Arc::clone(&schema); + let staged_handed = Arc::clone(&handed); 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), + // The third read is the one that reports pending. By then the first two chunks have + // filled an evaluation window, so the encoder is holding rows: the boundary is met + // with something buffered, which is the case worth checking. + .pending_after(2) + .counted(&staged_handed), ) }); let tree = rate_tree(); @@ -4094,18 +4105,31 @@ async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() { .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"); + // Poll by hand so the boundary is observable: `await` runs through the pending poll, and every + // later assertion would then hold whether or not the rows came out at the right moment. + // + // The pending poll is the one under test. `StagedStream` wakes itself straight after reporting + // pending, so a stream that answered nothing there would still be collected correctly; the only + // way to see the difference is to look at what that single poll returned. Two chunks have been + // read by then and an evaluation window is full, so the encoder is holding rows: it must answer + // them, not pending. + let next = stream.next(); + futures::pin_mut!(next); + let mut answered = match futures::poll!(next) { + std::task::Poll::Ready(Some(Ok(batch))) => vec![batch], + std::task::Poll::Ready(Some(Err(error))) => panic!("the frame answered {error}"), + std::task::Poll::Ready(None) => panic!("the frame ended before answering a row"), + std::task::Poll::Pending => panic!( + "the pending input held the rows the frame already encoded: the input handed over {} \ + batches", + handed.load(Ordering::SeqCst) + ), + }; assert_eq!( declared.as_ref(), - first.schema().as_ref(), + answered[0].schema().as_ref(), "a batch is declared as the plan declared it" ); - - let mut answered = vec![first]; answered.extend( stream .try_collect::>() @@ -4151,3 +4175,251 @@ async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() { "every row the row table answers is in the frame, none twice" ); } + +/// A closed chain that holds no value still reaches the consumer of a frame. +/// +/// The last sample sits past the grid, so no evaluation point has a sample in its window and the +/// window closes the series with a chunk that carries no position. Dropping it would lose the only +/// signal that says this series is done - the row table has no row for it at all, which is exactly +/// why the frame has to keep it. +#[tokio::test] +async fn a_frame_carries_the_closing_chunk_a_row_table_drops() { + let schema = raw_schema(); + // Every sample is newer than the last grid point, so `rate` decides nothing at any of them. + let late = [(GRID_START + RANGE_MS + SAMPLE_STEP, 1.0)]; + let chunk = chunk_of( + &schema, + "late", + metric(&[(HOST, "late")], Some(METRIC)), + &late.iter().map(|(time, _)| *time).collect::>(), + vec![late.iter().map(|(_, value)| Some(*value)).collect()], + true, + ); + 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, + ) + .expect("a rate tree is evaluated on a grid"); + let exec = plan_of(node, memory_exec(vec![frame_of(&schema, &chunk)])) + .expect("the runtime executes the tree"); + let declared = exec.schema(); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the frame runs"); + assert!( + batches + .iter() + .all(|batch| batch.schema().as_ref() == declared.as_ref()), + "every batch carries the declared schema" + ); + let stream = Box::pin(RecordBatchStreamAdapter::new( + declared, + futures::stream::iter(batches.into_iter().map(Ok)), + )); + let decoded = unpack_evaluated_series(stream, None) + .expect("the node declares the frame it yields") + .try_collect::>() + .await + .expect("the frame decodes"); + + assert_eq!( + 1, + decoded.len(), + "the closing chunk is one frame row, not a dropped chunk: {decoded:?}" + ); + let (series, is_last) = &decoded[0]; + assert!(is_last, "that row is the one that closes the chain"); + assert!( + series.time.is_empty(), + "it carries no grid position: {:?}", + series.time.positions().collect::>() + ); + + // The row table answers nothing for the same input, which is what makes this chunk worth keeping. + let plain = plan(rate_tree(), memory_exec(vec![frame_of(&schema, &chunk)])) + .expect("the runtime executes a tree"); + let rows = rows( + &collect(plain, SessionContext::default().task_ctx()) + .await + .expect("the row table runs"), + ); + assert!( + rows.is_empty(), + "the row table drops the closing chunk entirely: {rows:?}" + ); +} + +/// A sparse series keeps its own runs through the output boundary, holes and all. +/// +/// `rate` over samples that only cover the first part of the grid decides values at the first few +/// points and nothing at the rest. The frame must carry exactly those positions: filling a hole +/// would answer a value the query never computed. +#[tokio::test] +async fn a_frame_keeps_a_sparse_series_runs() { + let schema = raw_schema(); + // Two samples one step apart: a single window has a value, and every later point has none. + let samples = [(0, 1.0), (SAMPLE_STEP, 2.0)]; + let chunk = chunk_of( + &schema, + "sparse", + metric(&[(HOST, "sparse")], Some(METRIC)), + &samples.iter().map(|(time, _)| *time).collect::>(), + vec![samples.iter().map(|(_, value)| Some(*value)).collect()], + true, + ); + let batches = vec![frame_of(&schema, &chunk)]; + + 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, + ) + .expect("a rate tree is evaluated on a grid"); + let exec = plan_of(node, memory_exec(batches.clone())).expect("the runtime executes the tree"); + let declared = exec.schema(); + let decoded = unpack_evaluated_series( + Box::pin(RecordBatchStreamAdapter::new( + declared, + futures::stream::iter( + collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the frame runs") + .into_iter() + .map(Ok), + ), + )), + None, + ) + .expect("the node declares the frame it yields") + .try_collect::>() + .await + .expect("the frame decodes"); + + let mut positions = Vec::new(); + for (series, _) in &decoded { + positions.extend(series.time.positions().map(|(position, _)| position)); + } + // The same input through the row table says which points carry a value; the frame has to agree, + // and it must not have invented a position for any point the row table leaves empty. + let plain = plan(rate_tree(), memory_exec(batches)).expect("the runtime executes a rate 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 expected_positions = expected + .iter() + .map(|time| usize::try_from((*time - GRID_START) / GRID_STEP).expect("a grid position")) + .collect::>(); + expected_positions.sort_unstable(); + positions.sort_unstable(); + assert_eq!( + expected_positions, positions, + "the frame carries exactly the points the row table answers, no filled hole" + ); + assert!( + positions.len() < GRID_LEN as usize, + "the fixture has to be sparse for this to prove anything: {positions:?}" + ); +} + +/// Two chains that close on one label set are both emitted, each closing itself. +/// +/// A range function drops the metric name, so two different raw series can reach the output under the +/// same complete labels. A frame carries one chain per chunk and the later one is a further +/// contribution, not a repeat of the first: refusing it as a duplicate identity would drop values the +/// query answers with. +#[tokio::test] +async fn two_chains_that_close_on_one_label_set_each_keep_their_frame() { + // `rate` drops `__name__`, and neither series declares any other tag, so both collapse to the + // same empty label set. + let schema = raw_schema(); + let chunk_of_host = |host: &str, scale: f64| { + let samples = series_samples(scale); + chunk_of( + &schema, + host, + metric(&[], None), + &samples.iter().map(|(time, _)| *time).collect::>(), + vec![samples.iter().map(|(_, value)| Some(*value)).collect()], + true, + ) + }; + let partitions = vec![ + vec![frame_of(&schema, &chunk_of_host("web-1", 1.0))], + vec![frame_of(&schema, &chunk_of_host("web-2", 2.0))], + ]; + 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, + ) + .expect("a rate tree is evaluated on a grid"); + let exec = + plan_of(node, memory_exec_partitioned(partitions)).expect("the runtime executes the tree"); + assert_eq!( + 2, + exec.properties().partitioning.partition_count(), + "the frame keeps the input's partition count" + ); + let declared = exec.schema(); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("both chains are legal"); + let decoded = unpack_evaluated_series( + Box::pin(RecordBatchStreamAdapter::new( + declared, + futures::stream::iter(batches.into_iter().map(Ok)), + )), + None, + ) + .expect("the node declares the frame it yields") + .try_collect::>() + .await + .expect("the frame decodes"); + + assert_eq!( + 2, + decoded.iter().filter(|(_, last)| *last).count(), + "each chain closes once, in its own partition" + ); + assert_eq!( + decoded[0].0.metric, decoded[1].0.metric, + "the two chains carry the one label set" + ); + let mut values = Vec::new(); + for (series, _) in &decoded { + for sample in series.fields.get(0).expect("one field").iter() { + if let datatypes::series::samples::SampleRef::Float(value) = sample { + values.push(value); + } + } + } + values.sort_by(f64::total_cmp); + assert_eq!( + 2 * GRID_LEN as usize, + values.len(), + "both chains' values survive, so neither was mistaken for a duplicate: {values:?}" + ); + assert!( + values.windows(2).any(|pair| pair[0] != pair[1]), + "the two chains answered different values, so a lost chain would show: {values:?}" + ); +} From 949eeeff385e3bcde8896e528322c5018d662205 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:02:25 +0800 Subject: [PATCH 3/4] fix(promql): stream evaluated frames with batch-level backpressure Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/promql/src/extension_plan/fused_series.rs | 138 ++++-- .../src/extension_plan/fused_series/tests.rs | 351 ++++++++++----- .../fused_series/tests/frame_output.rs | 412 ++++++++++++++++++ tests-integration/Cargo.toml | 1 + tests-integration/tests/http.rs | 217 +++++++++ 6 files changed, 967 insertions(+), 153 deletions(-) create mode 100644 src/promql/src/extension_plan/fused_series/tests/frame_output.rs diff --git a/Cargo.lock b/Cargo.lock index 606322e47b3..3f3d9e14f59 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15119,6 +15119,7 @@ dependencies = [ "paste", "pipeline", "plugins", + "promql", "prost 0.14.1", "query", "rand 0.9.4", diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index 057f20be5d1..aebc42ef3e4 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -952,16 +952,6 @@ impl FusedSeriesExec { } } } - // 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 { @@ -1257,6 +1247,7 @@ impl ExecutionPlan for FusedSeriesExec { partial_fold, frame_builder, eval_schema, + frame_cursor: None, _export_budget: Arc::clone(&export_budget), ready: VecDeque::new(), done: false, @@ -1379,6 +1370,8 @@ pub struct FusedSeriesStream { partial_fold: Option, frame_builder: Option, eval_schema: Option>, + /// The evaluated chunk whose frame is still being appended across downstream polls. + frame_cursor: 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. @@ -1402,11 +1395,18 @@ impl FusedSeriesStream { })? { executor.push_raw(*input_id, chunk.clone())?; } - let mut settled = Vec::new(); - while let Some(chunk) = executor.next_chunk()? { - settled.push(chunk); + if self.frame_builder.is_some() { + Ok(()) + } else { + let mut settled = Vec::new(); + while let Some(chunk) = executor.next_chunk()? { + settled.push(chunk); + } + for chunk in settled { + self.convert(chunk)?; + } + Ok(()) } - self.convert(settled) } /// Ends the input, so the tree settles what it was holding. @@ -1454,6 +1454,7 @@ impl FusedSeriesStream { self.executor = None; self.partial_fold = None; self.frame_builder = None; + self.frame_cursor = 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. @@ -1465,46 +1466,53 @@ impl FusedSeriesStream { self.ready.clear(); } - /// Converts settled chunks into the node's output: query rows, an evaluated series frame, or - /// partition-local state for a later aggregate merge. - fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { + /// Converts one settled chunk into the node's output: query rows, an evaluated series frame, + /// or partition-local state for a later aggregate merge. + fn convert(&mut self, chunk: EvaluatedChunk) -> DataFusionResult<()> { if let Some(fold) = self.partial_fold.as_mut() { - for chunk in chunks { - fold.update(chunk)?; - } + fold.update(chunk)?; return Ok(()); } - for chunk in &chunks { - 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); - } + if let Some(schema) = self.eval_schema.as_ref() + && self.frame_builder.is_some() + { + self.frame_cursor = Some( + SeriesFrameCursor::try_new_evaluated(schema, chunk.series, chunk.is_last).map_err( + |error| { + DataFusionError::Execution(format!("cannot frame evaluated chunk: {error}")) + }, + )?, + ); + } else { + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + &chunk, + )?; + self.ready.push_back(batch); } Ok(()) } + /// Appends only until the current evaluated chunk yields one batch. + fn advance_frame_cursor(&mut self) -> DataFusionResult> { + let (Some(cursor), Some(builder)) = + (self.frame_cursor.as_mut(), self.frame_builder.as_mut()) + else { + return Ok(None); + }; + match cursor.append_to(builder).map_err(|error| { + DataFusionError::Execution(format!("cannot append evaluated frame: {error}")) + })? { + Some(batch) => Ok(Some(batch.with_schema(self.output_schema.clone())?)), + None => { + self.frame_cursor = None; + Ok(None) + } + } + } + fn flush_frames(&mut self) -> DataFusionResult<()> { if let Some(builder) = self.frame_builder.as_mut() && let Some(batch) = builder.flush().map_err(|error| { @@ -1543,16 +1551,50 @@ impl Stream for FusedSeriesStream { // 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 this.frame_cursor.is_some() { + match this.advance_frame_cursor() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => {} + Err(error) => { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } if let Some(batch) = this.ready.pop_front() { return Poll::Ready(Some(Ok(batch))); } + if !this.done && this.frame_builder.is_some() && this.frame_cursor.is_none() { + let next = this + .executor + .as_mut() + .map_or(Ok(None), NativeExecutor::next_chunk); + match next { + Ok(Some(chunk)) => { + if let Err(error) = this.convert(chunk) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + continue; + } + Ok(None) => { + if this.inputs.is_empty() { + this.executor = None; + } + } + Err(error) => { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } if this.done { return Poll::Ready(None); } let Some((source, input)) = this.inputs.front_mut() else { match this.pull_one() { Ok(Some(chunk)) => { - if let Err(error) = this.convert(vec![chunk]) { + if let Err(error) = this.convert(chunk) { this.fail(); return Poll::Ready(Some(Err(error))); } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 37699189fc8..8ebc9b035b3 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -32,11 +32,11 @@ 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::sync::{Arc, Mutex}; use std::task::{Context, Poll}; -use common_query::series_frame::unpack_evaluated_series; +use common_query::series_frame::{unpack_evaluated_series, unpack_raw_series}; use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; @@ -67,19 +67,21 @@ use datafusion::prelude::SessionContext; use datatypes::prelude::ConcreteDataType; use datatypes::series::frame::{ FrameContent, FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, - LabelDecodeSpec, LabelDecodeTag, SeriesFrameDeclaration, pack_raw_series_chunk, - raw_series_frame_schema, series_frame_schema, + LabelDecodeSpec, LabelDecodeTag, SeriesFrameBuilder, 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}; use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; use datatypes::series::{Label, SeriesLabels, SeriesOrdering}; use futures::{Stream, StreamExt, TryStreamExt}; -use store_api::region_engine::RawLabelSource; +use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; + +mod frame_output; use super::{ - FusedOutputFormat, FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, - series_to_record_batch, + FusedOutputFormat, FusedSeries, FusedSeriesBinding, FusedSeriesStream, OutputColumn, + POLL_BUDGET, series_to_record_batch, }; use crate::evaluated::EvaluatedChunk; use crate::extension_plan::aggregate_exchange::state_schema; @@ -2502,6 +2504,215 @@ async fn a_ready_batch_is_returned_before_later_input_is_read() { assert!(total > 1, "the fixture has to outlast the first answer"); } +fn direct_frame_stream( + schema: Arc, + input: Option, +) -> FusedSeriesStream { + let output_schema = series_frame_schema(&schema, FrameLabelMode::Explicit, None).unwrap(); + let builder = SeriesFrameBuilder::try_new( + &schema, + FrameLabelMode::Explicit, + None, + SeriesFrameBuilder::DEFAULT_MAX_ROWS, + SeriesFrameBuilder::DEFAULT_MAX_BYTES, + ) + .unwrap(); + FusedSeriesStream { + inputs: input.into_iter().map(|input| (0, input)).collect(), + executor: None, + input_ids: Vec::new(), + classification: Arc::from([]), + tags: Arc::from([]), + output_schema, + partial_fold: None, + frame_builder: Some(builder), + eval_schema: Some(schema), + frame_cursor: None, + _export_budget: Arc::new(Mutex::new(None)), + ready: VecDeque::new(), + done: false, + } +} + +fn sparse_eval_chunk(rows: usize, value: f64) -> (Arc, EvaluatedChunk) { + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1, (rows * 2) as u64).unwrap()).unwrap()); + let schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: VALUE.into() }]), + eval_grid: Some(Arc::clone(&grid)), + }); + let axis = + TimeAxis::from_grid_spans(grid, (0..rows).map(|i| GridSpan::new(i * 2, 1)).collect()) + .unwrap(); + let chunk = EvaluatedChunk::try_new( + Series { + metric: SeriesLabels::try_new(Some(Arc::from(METRIC)), Vec::new()).unwrap(), + time: axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![value; rows]))), + }, + true, + ) + .unwrap(); + (schema, chunk) +} + +/// A large evaluated chunk is encoded over successive polls rather than queued as all of its +/// physical frame batches at once. +#[tokio::test] +async fn a_large_evaluated_chunk_keeps_a_cursor_between_polls() { + const ROWS: usize = 800; + let value = f64::from_bits(0x4009_21fb_5444_2d18); + let (schema, chunk) = sparse_eval_chunk(ROWS, value); + let mut stream = direct_frame_stream(Arc::clone(&schema), None); + stream + .convert(chunk) + .expect("the evaluated chunk becomes a cursor"); + assert!( + stream.ready.is_empty(), + "conversion must not eagerly queue frame batches" + ); + assert!( + stream.frame_cursor.is_some(), + "unencoded rows remain owned by the cursor" + ); + + let first = stream.next().await.unwrap().unwrap(); + assert!(first.num_rows() <= SeriesFrameBuilder::DEFAULT_MAX_ROWS); + assert!( + stream.frame_cursor.is_some(), + "the first batch cannot drain 800 rows" + ); + assert!(stream.ready.is_empty()); + + let mut batches = vec![first]; + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + assert!(batches.len() > 1); + let decoded = unpack_evaluated_series( + Box::pin(RecordBatchStreamAdapter::new( + batches[0].schema(), + futures::stream::iter(batches.into_iter().map(Ok)), + )), + None, + ) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut decoded_rows = Vec::new(); + let mut closing_chunks = 0; + for (series, last) in &decoded { + if *last { + closing_chunks += 1; + } + assert_eq!(Some(METRIC), series.metric.name()); + for (position, row) in series.time.positions() { + let value = match series.fields.as_slice()[0].sample(row) { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => panic!("the decoded field remains a present Float64"), + }; + decoded_rows.push((position, value.to_bits())); + } + } + decoded_rows.sort_unstable_by_key(|(position, _)| *position); + assert_eq!(1, closing_chunks, "the chunk closes exactly once"); + assert_eq!( + (0..ROWS) + .map(|i| (i * 2, 0x4009_21fb_5444_2d18)) + .collect::>(), + decoded_rows, + "every sparse position carries the exact expected float bits", + ); + + assert!(stream.next().await.is_none(), "EOF remains stable"); + assert!(stream.next().await.is_none(), "repeated EOF remains stable"); +} + +/// The cursor and its input are both released when the consumer drops an unfinished frame. +#[tokio::test] +async fn dropping_a_partly_emitted_frame_releases_its_input() { + const ROWS: usize = 800; + let (schema, mut chunk) = sparse_eval_chunk(ROWS, 3.0); + let name = Arc::::from(METRIC); + let name_weak = Arc::downgrade(&name); + chunk + .series + .metric + .set("__name__", Arc::clone(&name)) + .unwrap(); + drop(name); + let raw = raw_schema(); + let raw_chunk = series_chunk("cursor-source", &[(0, 1.0)], true); + let raw_frame = frame_of(&raw, &raw_chunk); + let dropped = Arc::new(AtomicBool::new(false)); + let raw_input = Box::pin(StagedStream::new(frame_schema(), vec![raw_frame]).dropped(&dropped)); + let input = unpack_raw_series(raw_input, Some(Arc::clone(&raw))).unwrap(); + let mut stream = direct_frame_stream(Arc::clone(&schema), Some(input)); + stream.convert(chunk).unwrap(); + assert!(stream.next().await.unwrap().unwrap().num_rows() > 0); + assert!(stream.frame_cursor.is_some()); + assert!( + name_weak.upgrade().is_some(), + "the cursor retains its unencoded label identity" + ); + assert!(!dropped.load(Ordering::SeqCst)); + assert_eq!(1, stream.inputs.len()); + drop(stream); + assert!(dropped.load(Ordering::SeqCst)); + assert!( + name_weak.upgrade().is_none(), + "dropping the cursor releases its label identity" + ); +} + +/// An upstream error clears rows buffered in the frame builder before any batch is delivered. +#[tokio::test] +async fn frame_upstream_error_discards_buffered_rows_and_terminates() { + let (schema, chunk) = sparse_eval_chunk(2, 5.0); + let raw = raw_schema(); + let dropped = Arc::new(AtomicBool::new(false)); + let input = unpack_raw_series( + Box::pin( + StagedStream::new(frame_schema(), Vec::new()) + .error("frame input failed") + .dropped(&dropped), + ), + Some(Arc::clone(&raw)), + ) + .unwrap(); + let mut stream = direct_frame_stream(schema, Some(input)); + stream.convert(chunk).unwrap(); + assert!(stream.advance_frame_cursor().unwrap().is_none()); + assert!( + stream.frame_cursor.is_none(), + "small rows sit buffered in the builder" + ); + assert!(stream.ready.is_empty()); + + let error = stream + .next() + .await + .expect("the upstream error is yielded") + .expect_err("buffered frame rows cannot hide the source error"); + assert!(error.to_string().contains("frame input failed")); + assert!( + stream.next().await.is_none(), + "the failed stream terminates once" + ); + assert!( + stream.next().await.is_none(), + "EOF after failure remains stable" + ); + assert!( + dropped.load(Ordering::SeqCst), + "failure releases the source" + ); + assert!(stream.frame_builder.is_none()); + assert!(stream.frame_cursor.is_none()); + assert!(stream.ready.is_empty()); +} + /// 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 @@ -4254,88 +4465,6 @@ async fn a_frame_carries_the_closing_chunk_a_row_table_drops() { ); } -/// A sparse series keeps its own runs through the output boundary, holes and all. -/// -/// `rate` over samples that only cover the first part of the grid decides values at the first few -/// points and nothing at the rest. The frame must carry exactly those positions: filling a hole -/// would answer a value the query never computed. -#[tokio::test] -async fn a_frame_keeps_a_sparse_series_runs() { - let schema = raw_schema(); - // Two samples one step apart: a single window has a value, and every later point has none. - let samples = [(0, 1.0), (SAMPLE_STEP, 2.0)]; - let chunk = chunk_of( - &schema, - "sparse", - metric(&[(HOST, "sparse")], Some(METRIC)), - &samples.iter().map(|(time, _)| *time).collect::>(), - vec![samples.iter().map(|(_, value)| Some(*value)).collect()], - true, - ); - let batches = vec![frame_of(&schema, &chunk)]; - - 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, - ) - .expect("a rate tree is evaluated on a grid"); - let exec = plan_of(node, memory_exec(batches.clone())).expect("the runtime executes the tree"); - let declared = exec.schema(); - let decoded = unpack_evaluated_series( - Box::pin(RecordBatchStreamAdapter::new( - declared, - futures::stream::iter( - collect(exec, SessionContext::default().task_ctx()) - .await - .expect("the frame runs") - .into_iter() - .map(Ok), - ), - )), - None, - ) - .expect("the node declares the frame it yields") - .try_collect::>() - .await - .expect("the frame decodes"); - - let mut positions = Vec::new(); - for (series, _) in &decoded { - positions.extend(series.time.positions().map(|(position, _)| position)); - } - // The same input through the row table says which points carry a value; the frame has to agree, - // and it must not have invented a position for any point the row table leaves empty. - let plain = plan(rate_tree(), memory_exec(batches)).expect("the runtime executes a rate 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 expected_positions = expected - .iter() - .map(|time| usize::try_from((*time - GRID_START) / GRID_STEP).expect("a grid position")) - .collect::>(); - expected_positions.sort_unstable(); - positions.sort_unstable(); - assert_eq!( - expected_positions, positions, - "the frame carries exactly the points the row table answers, no filled hole" - ); - assert!( - positions.len() < GRID_LEN as usize, - "the fixture has to be sparse for this to prove anything: {positions:?}" - ); -} - /// Two chains that close on one label set are both emitted, each closing itself. /// /// A range function drops the metric name, so two different raw series can reach the output under the @@ -4404,22 +4533,34 @@ async fn two_chains_that_close_on_one_label_set_each_keep_their_frame() { decoded[0].0.metric, decoded[1].0.metric, "the two chains carry the one label set" ); - let mut values = Vec::new(); - for (series, _) in &decoded { - for sample in series.fields.get(0).expect("one field").iter() { - if let datatypes::series::samples::SampleRef::Float(value) = sample { - values.push(value); - } + let mut per_chain = Vec::new(); + for (series, last) in &decoded { + assert!(*last, "each partition's single chain closes exactly once"); + assert_eq!(None, series.metric.name(), "rate drops the metric name"); + assert!( + series.metric.tags().is_empty(), + "both chains have empty labels" + ); + let mut rows = Vec::new(); + for (position, row) in series.time.positions() { + let value = match series.fields.as_slice()[0].sample(row) { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => panic!("each rate grid position has a present value"), + }; + rows.push((position, value.to_bits())); } + rows.sort_unstable_by_key(|(position, _)| *position); + per_chain.push(rows); + } + assert_eq!(2, per_chain.len(), "one closing chain from each partition"); + per_chain.sort(); + for (rows, value) in per_chain.iter().zip([1.0f64, 2.0]) { + assert_eq!( + (0..GRID_LEN as usize) + .map(|i| (i, value.to_bits())) + .collect::>(), + *rows, + "each chain preserves every exact grid position and value", + ); } - values.sort_by(f64::total_cmp); - assert_eq!( - 2 * GRID_LEN as usize, - values.len(), - "both chains' values survive, so neither was mistaken for a duplicate: {values:?}" - ); - assert!( - values.windows(2).any(|pair| pair[0] != pair[1]), - "the two chains answered different values, so a lost chain would show: {values:?}" - ); } diff --git a/src/promql/src/extension_plan/fused_series/tests/frame_output.rs b/src/promql/src/extension_plan/fused_series/tests/frame_output.rs new file mode 100644 index 00000000000..dd830045e7f --- /dev/null +++ b/src/promql/src/extension_plan/fused_series/tests/frame_output.rs @@ -0,0 +1,412 @@ +// 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. + +use super::*; +use crate::map::MathOp; +use crate::map_label::LabelRewrite; +use crate::native::plan::{BoundMap, MapOp}; + +fn instant_tree(schema: Arc, grid: Arc) -> NativeTree { + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields: schema.fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: 0, + output: NativeType::Samples(Arc::clone(&schema)), + ordering: SeriesOrdering::Contiguous, + }), + output: NativeType::Samples(Arc::clone(&schema)), + }), + op: BoundInstant { + lookback_ms: 0, + offset_ms: 0, + at_ms: None, + timestamp: false, + output_grid: grid, + }, + }, + output, + }) + .expect("instant selector tree is valid") +} + +async fn frame_decode( + exec: Arc, +) -> Vec<(datatypes::series::samples::Series, bool)> { + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("frame output executes"); + let stream = Box::pin(RecordBatchStreamAdapter::new( + exec.schema(), + futures::stream::iter(batches.into_iter().map(Ok)), + )); + unpack_evaluated_series(stream, None) + .expect("evaluated frame schema decodes") + .try_collect::>() + .await + .expect("evaluated frame decodes") +} + +fn compact_float_rows( + decoded: &[(datatypes::series::samples::Series, bool)], +) -> Vec<(i64, Vec>)> { + let mut rows = Vec::new(); + for (series, _) in decoded { + let timestamps = series.time.timestamps().expect("grid timestamps"); + for (row, timestamp) in timestamps.iter().enumerate() { + let values = series + .fields + .as_slice() + .iter() + .map( + |field| match field.iter().nth(row).expect("sample at timestamp") { + datatypes::series::samples::SampleRef::Float(value) => { + Some(value.to_bits()) + } + datatypes::series::samples::SampleRef::Missing => None, + _ => panic!("float or null sample expected"), + }, + ) + .collect(); + rows.push((*timestamp, values)); + } + } + rows.sort_by_key(|(timestamp, _)| *timestamp); + rows +} + +fn two_field_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(OTHER, 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 plain_float_rows(batches: &[RecordBatch], fields: &[&str]) -> Vec<(i64, Vec>)> { + let mut rows = Vec::new(); + for batch in batches { + let time = batch + .column_by_name(TIME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let columns = fields + .iter() + .map(|field| { + batch + .column_by_name(field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + }) + .collect::>(); + for row in 0..batch.num_rows() { + rows.push(( + time.value(row), + columns + .iter() + .map(|column| column.is_valid(row).then(|| column.value(row).to_bits())) + .collect(), + )); + } + } + rows.sort_by_key(|(timestamp, _)| *timestamp); + rows +} + +#[tokio::test] +async fn sparse_instant_frame_preserves_positions_nulls_and_nan_across_formats() { + let schema = two_field_schema(); + let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 6).unwrap()).unwrap()); + let tree = instant_tree(Arc::clone(&schema), grid); + let chunk = chunk_of( + &schema, + "sparse", + metric(&[(HOST, "sparse")], Some(METRIC)), + &[0, 2_000, 5_000], + vec![ + vec![Some(1.0), None, Some(f64::from_bits(0x7ff8_0000_0000_0042))], + vec![None, Some(2.0), None], + ], + true, + ); + let batches = vec![frame_of(&schema, &chunk)]; + let binding = binding_of(&tree, Arc::clone(&schema), vec![HOST.into(), NAME.into()]); + let make_node = |format| { + FusedSeries::try_new_with_format( + Arc::from("instant(x)"), + tree.clone(), + binding.clone(), + vec![input_plan(frame_schema_of(&schema))], + two_field_output_schema().to_dfschema_ref().unwrap(), + format, + ) + .expect("instant tree has the declared frame output") + }; + + let frame_exec = plan_of( + make_node(FusedOutputFormat::SeriesFrame), + memory_exec_with(frame_schema_of(&schema), batches.clone()), + ) + .unwrap(); + let decoded = frame_decode(frame_exec).await; + assert_eq!(1, decoded.iter().filter(|(_, last)| *last).count()); + let expected = vec![ + (0, vec![Some(1.0f64.to_bits()), None]), + (2_000, vec![None, Some(2.0f64.to_bits())]), + (5_000, vec![Some(0x7ff8_0000_0000_0042), None]), + ]; + assert_eq!(expected, compact_float_rows(&decoded)); + + let plain_exec = plan_of( + make_node(FusedOutputFormat::PlainRows), + memory_exec_with(frame_schema_of(&schema), batches), + ) + .unwrap(); + let plain_batches = collect(plain_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(expected, plain_float_rows(&plain_batches, &[VALUE, OTHER])); + assert!(decoded.iter().all(|(series, _)| { + series.metric.get(HOST) == Some("sparse") && series.metric.name() == Some(METRIC) + })); +} + +#[tokio::test] +async fn abs_rate_frame_matches_plain_rows_with_exact_values_and_labels() { + let rate = rate_tree(); + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Map { + input: Box::new(rate.root().clone()), + op: BoundMap { + ops: Arc::from([MapOp::Math(MathOp::Abs)]), + }, + }, + output: rate.root().output.clone(), + }) + .unwrap(); + let batches = two_series_frames(); + let binding = binding(&tree); + let make_node = |format| { + FusedSeries::try_new_with_format( + Arc::from("abs(rate(x[5m]))"), + tree.clone(), + binding.clone(), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + format, + ) + .unwrap() + }; + let plain_exec = plan_of( + make_node(FusedOutputFormat::PlainRows), + memory_exec(batches.clone()), + ) + .unwrap(); + let plain_batches = collect(plain_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + let mut expected = rows(&plain_batches) + .into_iter() + .map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name)) + .collect::>(); + expected.sort(); + let mut exact = Vec::new(); + for (host, scale) in [("web-1", 1.0), ("web-2", 2.0)] { + for position in 0..GRID_LEN { + exact.push(( + GRID_START + position as i64 * GRID_STEP, + Some((scale as f64).to_bits()), + Some(host.to_string()), + None, + )); + } + } + exact.sort(); + assert_eq!(exact, expected); + + let frame_exec = plan_of( + make_node(FusedOutputFormat::SeriesFrame), + memory_exec(batches), + ) + .unwrap(); + let decoded = frame_decode(frame_exec).await; + assert_eq!(2, decoded.iter().filter(|(_, last)| *last).count()); + let mut actual = Vec::new(); + for (series, _) in &decoded { + assert_eq!(None, series.metric.name()); + let host = series.metric.get(HOST).unwrap().to_string(); + for (time, values) in compact_float_rows(&[(series.clone(), true)]) { + actual.push((time, values[0], Some(host.clone()), None)); + } + } + actual.sort(); + assert_eq!(exact, actual); +} + +#[tokio::test] +async fn rewritten_identical_labels_keep_disjoint_chains_in_one_encoder() { + let schema = raw_schema(); + let semantic_grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 6).unwrap()).unwrap()); + let base = instant_tree(Arc::clone(&schema), semantic_grid); + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Map { + input: Box::new(base.root().clone()), + op: BoundMap { + ops: Arc::from([ + MapOp::Label(Arc::new( + LabelRewrite::replace("host", "shared", "host", ".*").unwrap(), + )), + MapOp::Label(Arc::new( + LabelRewrite::replace("zone", "edge", "host", ".*").unwrap(), + )), + ]), + }, + }, + output: base.root().output.clone(), + }) + .unwrap(); + let chunks = vec![ + chunk_of( + &schema, + "a", + metric(&[(HOST, "a")], Some(METRIC)), + &[0, 2_000], + vec![vec![Some(10.0), Some(12.0)]], + true, + ), + chunk_of( + &schema, + "b", + metric(&[(HOST, "b")], Some(METRIC)), + &[1_000, 5_000], + vec![vec![Some(21.0), Some(25.0)]], + true, + ), + ]; + let batches = frames_of(&schema, &chunks); + let binding = binding_of(&tree, Arc::clone(&schema), vec![HOST.into(), NAME.into()]); + let make_node = |format| { + FusedSeries::try_new_with_format( + Arc::from("rewritten instant"), + tree.clone(), + binding.clone(), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + format, + ) + .unwrap() + }; + + let exec = plan_of( + make_node(FusedOutputFormat::SeriesFrame), + memory_exec(batches.clone()), + ) + .unwrap(); + let decoded = frame_decode(exec).await; + assert_eq!( + 2, + decoded.iter().filter(|(_, last)| *last).count(), + "each chain closes once" + ); + let mut groups = Vec::new(); + let mut group_rows = Vec::new(); + let mut closes = 0; + for (series, last) in &decoded { + assert_eq!(Some("shared"), series.metric.get(HOST)); + assert_eq!(Some("edge"), series.metric.get("zone")); + assert_eq!(Some(METRIC), series.metric.name()); + let timestamps = series.time.timestamps().unwrap(); + let values = series.fields.get(0).unwrap(); + for (i, time) in timestamps.iter().enumerate() { + let value = match values.iter().nth(i).unwrap() { + datatypes::series::samples::SampleRef::Float(value) => value.to_bits(), + _ => panic!("float sample expected"), + }; + group_rows.push((*time, value)); + } + if *last { + closes += 1; + group_rows.sort_by_key(|(time, _)| *time); + groups.push(std::mem::take(&mut group_rows)); + } + } + assert_eq!(2, closes, "each chain closes exactly once"); + assert!(group_rows.is_empty(), "all output chains close"); + groups.sort(); + assert_eq!( + vec![ + vec![(0, 10.0f64.to_bits()), (2_000, 12.0f64.to_bits())], + vec![(1_000, 21.0f64.to_bits()), (5_000, 25.0f64.to_bits())], + ], + groups + ); + + let plain_exec = plan_of( + make_node(FusedOutputFormat::PlainRows), + memory_exec_with(frame_schema(), batches), + ) + .unwrap(); + let plain = collect(plain_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + let mut plain = rows(&plain) + .into_iter() + .map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name)) + .collect::>(); + plain.sort(); + assert_eq!( + vec![ + ( + 0, + Some(10.0f64.to_bits()), + Some("shared".into()), + Some(METRIC.into()) + ), + ( + 1_000, + Some(21.0f64.to_bits()), + Some("shared".into()), + Some(METRIC.into()) + ), + ( + 2_000, + Some(12.0f64.to_bits()), + Some("shared".into()), + Some(METRIC.into()) + ), + ( + 5_000, + Some(25.0f64.to_bits()), + Some("shared".into()), + Some(METRIC.into()) + ), + ], + plain + ); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 79df85475d9..6830158ec0d 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -109,6 +109,7 @@ otel-arrow-rust.workspace = true paste.workspace = true pipeline.workspace = true plugins.workspace = true +promql.workspace = true prost.workspace = true rand.workspace = true rskafka.workspace = true diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 972ad254257..ad9cc2331b7 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -46,15 +46,20 @@ use common_frontend::slow_query_event::{ }; use common_memory_manager::OnExhaustedPolicy; use common_options::plugin_options::StandaloneFlag; +use datafusion::common::{DFSchema, ToDFSchema}; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datatypes::arrow::array::{ Array as _, Float64Array, StringArray, StringViewArray, TimestampMillisecondArray, }; use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datatypes::series::samples::{SampleRef, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpec}; use flate2::Compression; use flate2::write::GzEncoder; use frontend::instance::Instance; +use futures::TryStreamExt; use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter}; use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; @@ -1855,6 +1860,216 @@ fn fixture_rows(value: impl Fn(f64) -> f64) -> Vec { .collect() } +/// Replays the captured raw scan through both public FusedSeries output formats. +async fn assert_raw_scan_fused_formats(observer: &NativeFusedObserver, frontend: &Arc) { + let table_id = frontend + .catalog_manager() + .table("greptime", "public", "native_demo", None) + .await + .unwrap() + .expect("fixture table exists") + .table_info() + .table_id(); + let (_, route) = frontend + .table_metadata_manager() + .table_route_manager() + .get_physical_table_route(table_id) + .await + .unwrap(); + assert_eq!( + route.region_routes.len(), + 1, + "fixture has exactly one physical region" + ); + + observer.clear(); + let (_, row_schema, _) = promql_output(frontend, "native_demo", true, "0", "2").await; + let origin = observer.take("native_demo"); + let children = origin.children(); + assert_eq!(children.len(), 1, "the fused origin has exactly one child"); + let mut scan_plan = children[0].clone(); + if scan_plan + .downcast_ref::() + .is_none() + { + let wrapper_children = scan_plan.children(); + assert_eq!( + wrapper_children.len(), + 1, + "the optional scan wrapper has one child" + ); + scan_plan = wrapper_children[0].clone(); + } + let scan = scan_plan + .downcast_ref::() + .expect("the fused input is a real SeriesScanExec"); + assert_eq!(scan.properties().output_partitioning().partition_count(), 4); + let raw_schema = Arc::clone(scan.series_schema()); + let label_source = scan.label_source().clone(); + let ordering = scan.raw_series_ordering(); + let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 3).unwrap()).unwrap()); + let output_schema = Arc::new(SeriesSchema { + fields: Arc::clone(&raw_schema.fields), + eval_grid: Some(Arc::clone(&grid)), + }); + let raw_type = promql::native::plan::NativeType::Samples(Arc::clone(&raw_schema)); + let root = promql::native::plan::NativeNode { + expr: promql::native::plan::NativeExpr::Instant { + input: Box::new(promql::native::plan::NativeNode { + expr: promql::native::plan::NativeExpr::Input(promql::native::plan::BoundInput { + source_index: 0, + output: raw_type.clone(), + ordering, + }), + output: raw_type, + }), + op: promql::native::plan::BoundInstant { + lookback_ms: 300_000, + offset_ms: 0, + at_ms: None, + timestamp: false, + output_grid: Arc::clone(&grid), + }, + }, + output: promql::native::plan::NativeType::Samples(Arc::clone(&output_schema)), + }; + let tree = promql::native::plan::NativeTree::try_new(root).unwrap(); + let binding = promql::extension_plan::FusedSeriesBinding { + grid, + metric_name: Some(Arc::from("native_demo")), + time_index: row_schema + .fields() + .iter() + .find(|field| { + matches!( + field.data_type(), + DataType::Timestamp(TimeUnit::Millisecond, _) + ) + }) + .expect("PromQL result has a millisecond timestamp field") + .name() + .clone(), + tags: vec!["host".to_string()], + }; + let logical_child = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchema::try_from(scan.schema()).unwrap().into(), + }); + let expected = fixture_rows(|value| value); + for format in [ + promql::extension_plan::FusedOutputFormat::PlainRows, + promql::extension_plan::FusedOutputFormat::SeriesFrame, + ] { + scan_plan.clone().reset_state().unwrap(); + let exec = promql::extension_plan::FusedSeries::try_new_with_format( + Arc::from("native_demo"), + tree.clone(), + binding.clone(), + vec![logical_child.clone()], + row_schema.clone().to_dfschema_ref().unwrap(), + format, + ) + .unwrap() + .to_execution_plan(vec![(scan_plan.clone(), label_source.clone(), ordering)]) + .unwrap(); + println!( + "{}", + datafusion::physical_plan::displayable(exec.as_ref()).indent(true) + ); + let context = "standalone raw SeriesScanExec replay"; + let mut actual = Vec::new(); + let mut closed = 0; + let mut partition_counts = Vec::new(); + for partition in 0..4 { + let stream = exec + .execute( + partition, + datafusion::prelude::SessionContext::new().task_ctx(), + ) + .unwrap(); + if format == promql::extension_plan::FusedOutputFormat::SeriesFrame { + let series = common_query::series_frame::unpack_evaluated_series(stream, None) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!( + !series.is_empty(), + "partition {partition} emitted no series frames" + ); + partition_counts.push( + series + .iter() + .map(|(s, _)| s.time.timestamps().unwrap().len()) + .sum::(), + ); + for (series, last) in series { + assert_eq!(series.metric.name(), Some("native_demo")); + assert_eq!(series.metric.tags().len(), 1); + assert_eq!(series.fields.len(), 1); + let host = series.metric.get("host").unwrap().to_string(); + let timestamps = series.time.timestamps().unwrap(); + let samples = series.fields.get(0).unwrap(); + for (i, timestamp) in timestamps.iter().enumerate() { + let SampleRef::Float(value) = samples.iter().nth(i).unwrap() else { + panic!("expected float sample") + }; + actual.push((vec![("host".to_string(), host.clone())], *timestamp, value)); + } + closed += usize::from(last); + } + } else { + let batches = stream.try_collect::>().await.unwrap(); + let mut count = 0; + for batch in batches { + let ts = batch + .column_by_name(&binding.time_index) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let host = batch + .column_by_name("host") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("val") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + actual.push(( + vec![("host".to_string(), host.value(i).to_string())], + ts.value(i), + values.value(i), + )); + count += 1; + } + } + assert_eq!(row_schema.fields().len(), 3); + assert!(count > 0, "partition {partition} emitted no rows"); + partition_counts.push(count); + } + } + assert_eq!(partition_counts.len(), 4); + assert!( + partition_counts.iter().all(|count| *count > 0), + "all four partitions emitted output: {partition_counts:?}" + ); + assert_eq!(partition_counts.iter().sum::(), 24); + if format == promql::extension_plan::FusedOutputFormat::SeriesFrame { + assert_eq!(closed, 8, "one closing frame per host series"); + } + actual.sort_unstable_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); + assert_rows_eq(&actual, &expected, context); + println!("{context}: {format:?} partition output counts: {partition_counts:?}"); + assert_native_partition_output(&exec, Some(4), 24, context); + } +} + /// Runs one PromQL query through the public frontend handler and returns the plan it executed, the /// schema it answered with and every row it reported. async fn promql_output( @@ -2206,6 +2421,8 @@ pub async fn test_native_promql_route_plans_the_query(store_type: StorageType) { .await; } + assert_raw_scan_fused_formats(&observer, frontend).await; + for (query, start, end, expected, scan_partitions) in [ ( "1 + 2", From f2041d8d9de87379896a30d34d60ad43fed0c13f Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:40:38 +0800 Subject: [PATCH 4/4] test(promql): keep the frame output tests lint-clean The workspace denies the default Clippy warnings, and the new tests tripped three of them without changing what any of them proves: - an expected value was built through `scale as f64` even though the tuple already carries an `f64`, so the conversion is spelled in the literal and the cast is gone; - the standalone replay printed the executed plan and the per-partition counts even though the surrounding assertions already require four nonempty partitions, 24 rows in total and eight closing frames, so the diagnostics are gone rather than exempted from the workspace lint. No production code, assertion or fixture changes. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/extension_plan/fused_series/tests/frame_output.rs | 4 ++-- tests-integration/tests/http.rs | 5 ----- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/promql/src/extension_plan/fused_series/tests/frame_output.rs b/src/promql/src/extension_plan/fused_series/tests/frame_output.rs index dd830045e7f..0e09ebafc5c 100644 --- a/src/promql/src/extension_plan/fused_series/tests/frame_output.rs +++ b/src/promql/src/extension_plan/fused_series/tests/frame_output.rs @@ -236,11 +236,11 @@ async fn abs_rate_frame_matches_plain_rows_with_exact_values_and_labels() { .collect::>(); expected.sort(); let mut exact = Vec::new(); - for (host, scale) in [("web-1", 1.0), ("web-2", 2.0)] { + for (host, scale) in [("web-1", 1.0_f64), ("web-2", 2.0)] { for position in 0..GRID_LEN { exact.push(( GRID_START + position as i64 * GRID_STEP, - Some((scale as f64).to_bits()), + Some(scale.to_bits()), Some(host.to_string()), None, )); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index ad9cc2331b7..93284ddeb3a 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1972,10 +1972,6 @@ async fn assert_raw_scan_fused_formats(observer: &NativeFusedObserver, frontend: .unwrap() .to_execution_plan(vec![(scan_plan.clone(), label_source.clone(), ordering)]) .unwrap(); - println!( - "{}", - datafusion::physical_plan::displayable(exec.as_ref()).indent(true) - ); let context = "standalone raw SeriesScanExec replay"; let mut actual = Vec::new(); let mut closed = 0; @@ -2065,7 +2061,6 @@ async fn assert_raw_scan_fused_formats(observer: &NativeFusedObserver, frontend: } actual.sort_unstable_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1))); assert_rows_eq(&actual, &expected, context); - println!("{context}: {format:?} partition output counts: {partition_counts:?}"); assert_native_partition_output(&exec, Some(4), 24, context); } }