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.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..aebc42ef3e4 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,39 @@ 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, +} + +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 +222,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 +239,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 +279,7 @@ impl FusedSeries { children, output_schema, partial_fold: None, + format, }) } @@ -243,6 +313,7 @@ impl FusedSeries { children, output_schema, partial_fold: Some(op), + format: FusedOutputFormat::PlainRows, }) } @@ -296,6 +367,7 @@ impl FusedSeries { self.binding.time_index.clone(), self.binding.tags.clone(), self.partial_fold.clone(), + self.format, )?)) } } @@ -327,6 +399,9 @@ impl UserDefinedLogicalNodeCore for FusedSeries { self.promql, render_partial_fold(op) ), + None if self.format == FusedOutputFormat::SeriesFrame => { + write!(f, "FusedSeries: {} (output: series_frame)", self.promql) + } None => write!(f, "FusedSeries: {}", self.promql), } } @@ -344,6 +419,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 +437,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 +450,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 +478,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 +865,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 +890,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 +911,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 +952,19 @@ impl FusedSeriesExec { } } } - let classification = if partial_fold.is_some() { + 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 +1042,7 @@ impl FusedSeriesExec { metric: ExecutionPlanMetricsSet::new(), properties, partial_fold, + format, }) } @@ -953,6 +1065,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 == FusedOutputFormat::SeriesFrame { + write!(f, "\noutput: series_frame")?; } Ok(()) } @@ -1028,6 +1142,7 @@ impl ExecutionPlan for FusedSeriesExec { self.time_column.clone(), self.tags.to_vec(), self.partial_fold.clone(), + self.format, )?)) } @@ -1098,6 +1213,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 +1245,9 @@ impl ExecutionPlan for FusedSeriesExec { tags: Arc::clone(&self.tags), output_schema: self.output_schema.clone(), partial_fold, + frame_builder, + eval_schema, + frame_cursor: None, _export_budget: Arc::clone(&export_budget), ready: VecDeque::new(), done: false, @@ -1226,6 +1368,10 @@ 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>, + /// 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. @@ -1249,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. @@ -1300,6 +1453,8 @@ impl FusedSeriesStream { self.done = true; 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. @@ -1311,26 +1466,65 @@ impl FusedSeriesStream { self.ready.clear(); } - /// Converts settled chunks into query rows or folds them into partial aggregate state. - 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) = 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, + &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| { + DataFusionError::Execution(format!("cannot flush evaluated frame: {error}")) + })? + { + self.ready + .push_back(batch.with_schema(self.output_schema.clone())?); + } + Ok(()) + } + /// Publishes the partition's fold only after the input and native execution have drained. fn finish_fold(&mut self) -> DataFusionResult<()> { if let Some(mut fold) = self.partial_fold.take() { @@ -1357,22 +1551,56 @@ 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))); } } 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 +1616,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..8ebc9b035b3 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -32,10 +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, 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}; @@ -65,17 +66,23 @@ 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, 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}; -use super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, series_to_record_batch}; +mod frame_output; + +use super::{ + FusedOutputFormat, FusedSeries, FusedSeriesBinding, FusedSeriesStream, 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 +92,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])"; @@ -2496,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 @@ -3731,3 +3948,619 @@ 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. 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); + // 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 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. +/// +/// 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()) + // 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(); + 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"); + + // 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(), + answered[0].schema().as_ref(), + "a batch is declared as the plan declared it" + ); + 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" + ); +} + +/// 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:?}" + ); +} + +/// 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 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", + ); + } +} 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..0e09ebafc5c --- /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_f64), ("web-2", 2.0)] { + for position in 0..GRID_LEN { + exact.push(( + GRID_START + position as i64 * GRID_STEP, + Some(scale.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..93284ddeb3a 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,211 @@ 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(); + 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); + 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 +2416,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",