diff --git a/src/common/query/src/series_frame/tests/evaluated.rs b/src/common/query/src/series_frame/tests/evaluated.rs index 5a2336a4707..dba22e8d702 100644 --- a/src/common/query/src/series_frame/tests/evaluated.rs +++ b/src/common/query/src/series_frame/tests/evaluated.rs @@ -27,7 +27,7 @@ use std::io::Cursor; use std::sync::Arc; use std::sync::atomic::Ordering; -use common_error::ext::{BoxedError, ErrorExt}; +use common_error::ext::{BoxedError, ErrorExt, PlainError}; use common_error::status_code::StatusCode; use common_recordbatch::DfSendableRecordBatchStream; use datatypes::arrow::array::{ @@ -68,8 +68,31 @@ fn labels(host: &str) -> SeriesLabels { /// One validated evaluated chunk: `values` on the runs `spans` of the shared test grid. fn series(spans: &[(usize, usize)], values: Vec>, metric: SeriesLabels) -> Series { + series_on(&evaluated_schema(), spans, values, metric) +} + +/// An evaluated schema large enough for multi-batch lifecycle fixtures. +fn large_evaluated_schema() -> Arc { + Arc::new(SeriesSchema { + fields: evaluated_schema().fields.clone(), + eval_grid: Some(Arc::new( + datatypes::series::time::EvalGrid::try_new( + datatypes::series::time::GridSpec::try_new(0, 1000, 1600).unwrap(), + ) + .unwrap(), + )), + }) +} + +/// One validated evaluated chunk on the supplied schema's grid. +fn series_on( + schema: &SeriesSchema, + spans: &[(usize, usize)], + values: Vec>, + metric: SeriesLabels, +) -> Series { let time = TimeAxis::from_grid_spans( - evaluated_grid(), + schema.eval_grid.as_ref().unwrap().clone(), spans .iter() .map(|(start, len)| GridSpan::new(*start, *len)) @@ -77,7 +100,7 @@ fn series(spans: &[(usize, usize)], values: Vec>, metric: SeriesLabe ) .unwrap(); let fields = SeriesFields::one(Samples::from_float(Float64Array::from(values))); - Series::try_new(&evaluated_schema(), metric, time, fields).unwrap() + Series::try_new(schema, metric, time, fields).unwrap() } /// An evaluated producer stream over the supplied chunks and their finality. @@ -87,6 +110,32 @@ fn source(items: Vec<(Series, bool)>) -> SendableEvaluatedSeriesStream { )) } +/// A counted source whose guard records when the packer releases its upstream stream. +fn observed_source( + items: Vec>, + polls: Arc, + dropped: Arc, +) -> SendableEvaluatedSeriesStream { + struct DropGuard(Arc); + impl Drop for DropGuard { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + Box::pin(futures_util::stream::unfold( + (items.into_iter(), DropGuard(dropped)), + move |(mut items, guard)| { + let polls = polls.clone(); + async move { + polls.fetch_add(1, Ordering::SeqCst); + let item = items.next()?; + Some((item, (items, guard))) + } + }, + )) +} + /// Packs the supplied chunks, keeping the stream so a test can observe its errors. fn pack(items: Vec<(Series, bool)>) -> DfSendableRecordBatchStream { pack_evaluated_series(source(items), evaluated_schema()).unwrap() @@ -330,6 +379,102 @@ fn single_position_spans(start: usize, count: usize) -> Vec<(usize, usize)> { (0..count).map(|offset| (start + offset, 1)).collect() } +/// The evaluated packer drains one large sparse chunk in bounded batches without re-polling upstream. +#[tokio::test] +async fn evaluated_pack_drains_large_sparse_chunk_before_polling_again() { + let schema = large_evaluated_schema(); + let positions: Vec = (0..1600).step_by(2).collect(); + let spans: Vec<(usize, usize)> = positions.iter().map(|position| (*position, 1)).collect(); + let expected: Vec> = (0..positions.len()) + .map(|index| Some(index as f64)) + .collect(); + let chunk = series_on(&schema, &spans, expected.clone(), labels("large")); + let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let upstream = observed_source(vec![Ok((chunk, true))], polls.clone(), dropped.clone()); + let mut stream = pack_evaluated_series(upstream, schema.clone()).unwrap(); + let mut batches = Vec::new(); + let mut batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), SeriesFrameBuilder::DEFAULT_MAX_ROWS); + batches.push(batch); + assert_eq!(polls.load(Ordering::SeqCst), 1); + for _ in 0..2 { + batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), SeriesFrameBuilder::DEFAULT_MAX_ROWS); + batches.push(batch); + assert_eq!(polls.load(Ordering::SeqCst), 1); + } + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + assert!(batches.len() > 1); + assert!(dropped.load(Ordering::SeqCst)); + assert!(stream.next().await.is_none()); + assert!(stream.next().await.is_none()); + let declared = batches[0].schema(); + let (source, _, _) = df_source(batches, declared); + let rows = drain(unpack_evaluated_series(source, Some(schema)).unwrap()).await; + let (actual_positions, actual_values) = flat(&rows); + assert_eq!(actual_positions, positions); + assert_float_payload_eq(&actual_values, &expected, "large sparse values"); + assert_eq!(rows.iter().filter(|(_, last)| *last).count(), 1); + assert!(rows.last().unwrap().1); +} + +/// Dropping a partially drained output releases both upstream and buffered label ownership. +#[tokio::test] +async fn evaluated_pack_drop_releases_unfinished_large_chunk() { + let identity: Arc = Arc::from("drop-aware"); + let weak = Arc::downgrade(&identity); + let metric = SeriesLabels::try_new(Some(identity.clone()), vec![]).unwrap(); + let schema = large_evaluated_schema(); + let positions: Vec = (0..1600).step_by(2).collect(); + let spans: Vec<(usize, usize)> = positions.iter().map(|position| (*position, 1)).collect(); + let chunk = series_on(&schema, &spans, vec![Some(1.0); spans.len()], metric); + drop(identity); + let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let upstream = observed_source(vec![Ok((chunk, true))], polls, dropped.clone()); + let mut stream = pack_evaluated_series(upstream, schema).unwrap(); + assert!(stream.next().await.unwrap().unwrap().num_rows() > 0); + assert!(!dropped.load(Ordering::SeqCst), "output is unfinished"); + assert!(weak.upgrade().is_some(), "buffered output retains labels"); + drop(stream); + assert!(dropped.load(Ordering::SeqCst)); + assert!( + weak.upgrade().is_none(), + "drop releases buffered label owner" + ); +} + +/// A source error discards buffered success rather than publishing it before the error. +#[tokio::test] +async fn evaluated_pack_error_discards_buffer_and_releases_source() { + let identity: Arc = Arc::from("error-owner"); + let weak = Arc::downgrade(&identity); + let metric = SeriesLabels::try_new(Some(identity.clone()), vec![]).unwrap(); + let chunk = series(&[(0, 1)], vec![Some(1.0)], metric); + drop(identity); + let error = BoxedError::new(PlainError::new( + "source failed".into(), + StatusCode::Unsupported, + )); + let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let upstream = observed_source( + vec![Ok((chunk, false)), Err(error)], + polls.clone(), + dropped.clone(), + ); + let mut stream = pack_evaluated_series(upstream, evaluated_schema()).unwrap(); + assert!(stream.next().await.unwrap().is_err()); + assert!(dropped.load(Ordering::SeqCst)); + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert!(weak.upgrade().is_none()); + assert!(stream.next().await.is_none()); + assert!(stream.next().await.is_none()); +} + /// Rows split across batches keep their order and finality: every position appears once and in /// order, and only the last row of the source's final chunk closes the chain. #[tokio::test] @@ -411,6 +556,75 @@ async fn consecutive_chains_with_equal_labels_each_close_independently() { assert_eq!(values(&rows[1].0), vec![Some(3.0), Some(4.0)]); } +/// Malformed chunks are rejected before buffered rows escape, and terminal failures release the +/// producer and any label owners held by its chunks. +#[tokio::test] +async fn evaluated_pack_rejects_malformed_sequences_before_publishing() { + let cases = [ + ( + "identity changed", + vec![ + (series(&[(0, 1)], vec![Some(1.0)], labels("first")), false), + (series(&[(1, 1)], vec![Some(2.0)], labels("other")), true), + ], + "identity changed", + ), + ( + "overlap", + vec![ + ( + series(&[(1, 2)], vec![Some(1.0), Some(2.0)], labels("same")), + false, + ), + (series(&[(2, 1)], vec![Some(3.0)], labels("same")), true), + ], + "positions are not increasing", + ), + ( + "backward", + vec![ + (series(&[(3, 1)], vec![Some(1.0)], labels("same")), false), + (series(&[(1, 1)], vec![Some(2.0)], labels("same")), true), + ], + "positions are not increasing", + ), + ]; + for (name, mut items, expected_error) in cases { + let owner: Arc = Arc::from("tracked-label-owner"); + let weak = Arc::downgrade(&owner); + if name != "identity changed" { + let metric_name = items[0].0.metric.name().map(Arc::::from); + let tracked = SeriesLabels::try_new( + metric_name, + vec![Label { + name: Arc::from("tracked"), + value: owner.clone(), + }], + ) + .unwrap(); + for (chunk, _) in &mut items { + chunk.metric = tracked.clone(); + } + } + drop(owner); + let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let upstream = observed_source(items.into_iter().map(Ok).collect(), polls, dropped.clone()); + let mut stream = pack_evaluated_series(upstream, evaluated_schema()).unwrap(); + let error = stream.next().await.unwrap().unwrap_err(); + assert!( + error.to_string().contains(expected_error), + "{name}: {error}" + ); + assert!(dropped.load(Ordering::SeqCst), "{name}: source dropped"); + if name != "identity changed" { + assert!(weak.upgrade().is_none(), "{name}: label owner released"); + } + assert!(stream.next().await.is_none(), "{name}: stable EOF"); + assert!(stream.next().await.is_none(), "{name}: repeated EOF"); + } +} + /// A chain still open at EOF is an error on both sides: a producer that never closed a series cannot /// finish, and a receiver cannot pretend the stream ended cleanly. #[tokio::test] diff --git a/src/promql/src/extension_plan/aggregate_finalize.rs b/src/promql/src/extension_plan/aggregate_finalize.rs index 8dcb7394aa1..69034776aeb 100644 --- a/src/promql/src/extension_plan/aggregate_finalize.rs +++ b/src/promql/src/extension_plan/aggregate_finalize.rs @@ -187,7 +187,8 @@ impl AggregateFinalize { /// its binding: its `Utf8` columns are the labels it writes, and its `Float64` columns name the /// merged fields by position, since the merge folds the cut aggregate's fields in the order the /// statement declares them. Building rejects an input that is not state, a suffix the runtime - /// does not execute and a schema [`classify_output`] cannot convert. + /// does not execute (the supported suffix operations are maps, aggregates, TopK and + /// constant-sided binary operations) and a schema [`classify_output`] cannot convert. pub fn to_execution_plan( &self, input: Arc, @@ -346,12 +347,18 @@ struct Split { /// Splits the innermost aggregate out of `node`. /// -/// `Ok(None)` is a subtree holding no aggregate. Only a map, an aggregate and a binary operation with -/// a constant on one side may sit above the cut: a map is point-wise and so is the same operation over -/// merged values, an aggregate runs over the complete merge, and a constant is broadcast over whatever -/// the other side answers - while an instant selection, a window, a vector-wide operation or a binary -/// operation between two vectors would read evaluated values, the whole vector or a matched pair, -/// shapes the merge does not produce. +/// `Ok(None)` is a subtree holding no aggregate. The currently verified suffix operations are maps, +/// aggregates, TopK/BottomK and binary operations with a constant on one side. TopK/BottomK run once +/// in the suffix over the fully merged aggregate series, per evaluation timestamp, and never in the +/// partition-local producer. A map is point-wise and an aggregate runs over the complete merge. +/// LimitK is declined because it selects the first `k` series in input order +/// ([`crate::select::limitk`]), while this merge does not preserve the storage input order. Other +/// vector operations have not yet been verified/supported by this split. +/// +/// An instant selection or window above the cut cannot be applied to the merged values under the +/// tree's current binding: the cut output is evaluated, whereas these nodes consume raw samples. +/// A binary operation between two vectors would require matching the cut's merged output against a +/// vector sibling, but this merge only gathers the cut's state. fn split(node: &NativeNode, fields: &[FieldSpec], time_index: &str) -> DfResult> { match &node.expr { NativeExpr::Input(_) | NativeExpr::Scalar(_) => Ok(None), @@ -412,6 +419,22 @@ fn split(node: &NativeNode, fields: &[FieldSpec], time_index: &str) -> DfResult< } } } + NativeExpr::Vector { + input, + op: op @ crate::native::plan::BoundVectorOp::TopK { .. }, + } => { + let Some(mut inner) = split(input, fields, time_index)? else { + return Ok(None); + }; + inner.suffix = NativeNode { + expr: NativeExpr::Vector { + input: Box::new(inner.suffix), + op: op.clone(), + }, + output: node.output.clone(), + }; + Ok(Some(inner)) + } NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Timestamp { input } @@ -620,9 +643,9 @@ const STRING_CAPACITY: usize = 1024; /// first item rather than a row the error invalidates after it. What the answer holds is charged to /// that stream, which is the only owner it has: the plan answers once. /// -/// The cut is the innermost aggregate of a tree of inputs, maps and aggregates, so the suffix the -/// runtime is asked to build is one of those; [`ExecutionPlan::execute`] builds its state and is -/// where a shape the runtime cannot run is refused. +/// The suffix above the cut is built from the aggregate's output and may contain maps, aggregates, +/// TopK/BottomK or a binary operation with a constant on one side. [`ExecutionPlan::execute`] +/// builds its state and is where a shape the runtime cannot run is refused. /// /// The plain state-batch input is temporary. The final design consumes series frames directly /// from FusedSeries instead of decoding and rebuilding this intermediate state-row layout. diff --git a/src/promql/src/extension_plan/aggregate_finalize/tests.rs b/src/promql/src/extension_plan/aggregate_finalize/tests.rs index 3554e6e1572..241677217d4 100644 --- a/src/promql/src/extension_plan/aggregate_finalize/tests.rs +++ b/src/promql/src/extension_plan/aggregate_finalize/tests.rs @@ -24,6 +24,7 @@ use datafusion::common::DFSchema; use datatypes::series::samples::FieldSpec; use datatypes::series::time::{EvalGrid, GridSpec}; +use crate::native::plan::BoundVectorOp; use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels}; /// The statement's grid, which every evaluated chunk is addressed on. @@ -337,14 +338,86 @@ fn a_constant_operand_keeps_the_cut_below_it() { "the merged values enter where the aggregate was: {}", node.suffix ); + assert!(!node.suffix.to_string().contains("Window")); + } +} + +fn vector_over(input: NativeNode, op: BoundVectorOp) -> NativeNode { + let output = input.output.clone(); + NativeNode { + expr: NativeExpr::Vector { + input: Box::new(input), + op, + }, + output, + } +} + +#[test] +fn topk_and_bottomk_run_after_the_aggregate_cut() { + for top in [true, false] { + let aggregate = complete(AggregateKind::Sum, by(&["k0"]))(window()); + let root = vector_over( + aggregate, + BoundVectorOp::TopK { + k: 10.0, + top, + group: by(&[]), + }, + ); + let tree = NativeTree::try_new(root).unwrap(); + assert!(AggregateFinalize::can_split(&tree, &raw_fields(), "ts")); + let node = finalize(tree.root().clone(), output_schema(&["sum(v0)"], &["k0"])).unwrap(); + let NativeExpr::Vector { + input, + op: BoundVectorOp::TopK { + top: actual_top, .. + }, + } = &node.suffix.root().expr + else { + panic!("{}", node.suffix); + }; + assert_eq!(top, *actual_top); assert!( - !node.suffix.to_string().contains("Window"), - "{}", - node.suffix + matches!(&input.expr, NativeExpr::Input(bound) if bound.source_index == 0 && bound.output.samples().is_some_and(|schema| schema.fields.len() == 1 && schema.fields[0].name.as_ref() == "v0" && schema.eval_grid.as_ref().is_some_and(|actual| actual.spec() == grid().spec()))) ); } } +#[test] +fn limitk_above_the_cut_is_declined_because_merge_does_not_preserve_input_order() { + let error = finalize( + vector_over( + complete(AggregateKind::Sum, by(&["k0"]))(window()), + BoundVectorOp::LimitK { + k: 1.0, + group: by(&[]), + }, + ), + output_schema(&["sum(v0)"], &["k0"]), + ) + .unwrap_err(); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "{error}" + ); +} + +#[test] +fn topk_without_aggregate_is_not_split() { + let tree = NativeTree::try_new(vector_over( + window(), + BoundVectorOp::TopK { + k: 10.0, + top: true, + group: by(&[]), + }, + )) + .unwrap(); + assert!(!AggregateFinalize::can_split(&tree, &raw_fields(), "ts")); + assert!(split(tree.root(), &raw_fields(), "ts").unwrap().is_none()); +} + /// A binary operation above the cut is only the operation over merged values when its other side is a /// constant: two aggregates are two cuts, a vector sibling would have to be matched with the merged /// values by label, and a prefix that pairs two vectors reads the whole vector rather than one series. diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index 69b11a57956..2ccffbc5393 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -45,7 +45,8 @@ use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use common_query::series_frame::{ - ensure_frame_matches_label_source, pack_evaluated_series, unpack_raw_series, + SendableEvaluatedSeriesStream, ensure_frame_matches_label_source, pack_evaluated_series, + unpack_raw_series, }; use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::ScalarBuffer; @@ -66,7 +67,6 @@ 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; @@ -270,13 +270,15 @@ impl FusedSeries { "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 - }; + if format == FusedOutputFormat::SeriesFrame { + let expected = series_frame_output_schema(evaluated_output(&tree)?.as_ref())?; + if output_schema.as_arrow() != expected.as_ref() { + return Err(DataFusionError::Plan( + "FusedSeries SeriesFrame output schema does not match the evaluated frame layout" + .into(), + )); + } + } Ok(Self { promql, tree, @@ -1245,27 +1247,7 @@ impl ExecutionPlan for FusedSeriesExec { // 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 { + let stream = FusedSeriesStream { inputs, executor: Some(executor), input_ids, @@ -1273,13 +1255,19 @@ impl ExecutionPlan for FusedSeriesExec { tags: Arc::clone(&self.tags), output_schema: self.output_schema.clone(), partial_fold, - frame_builder, - eval_schema, - frame_cursor: None, + frame_mode, _export_budget: Arc::clone(&export_budget), ready: VecDeque::new(), done: false, - })) + }; + if frame_mode { + let schema = evaluated_output(&self.tree)?; + let source: SendableEvaluatedSeriesStream = + Box::pin(FusedEvaluatedStream { source: stream }); + Ok(pack_evaluated_series(source, schema)?) + } else { + Ok(Box::pin(stream)) + } } fn metrics(&self) -> Option { @@ -1396,10 +1384,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>, - /// The evaluated chunk whose frame is still being appended across downstream polls. - frame_cursor: Option, + /// Whether evaluated chunks are emitted as series frames rather than converted to rows. + frame_mode: bool, /// 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. @@ -1423,7 +1409,7 @@ impl FusedSeriesStream { })? { executor.push_raw(*input_id, chunk.clone())?; } - if self.frame_builder.is_some() { + if self.frame_mode { Ok(()) } else { let mut settled = Vec::new(); @@ -1431,7 +1417,7 @@ impl FusedSeriesStream { settled.push(chunk); } for chunk in settled { - self.convert(chunk)?; + self.convert_rows(chunk)?; } Ok(()) } @@ -1481,8 +1467,6 @@ impl FusedSeriesStream { self.done = true; self.executor = None; self.partial_fold = None; - self.frame_builder = None; - self.frame_cursor = None; self.inputs = None; // This reservation belongs only to this execution. Release it even if the consumer keeps // the failed stream and the plan alive. @@ -1494,62 +1478,19 @@ impl FusedSeriesStream { self.ready.clear(); } - /// 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<()> { + /// Converts one settled chunk into a row batch. Evaluated frames leave this driver as chunks. + fn convert_rows(&mut self, chunk: EvaluatedChunk) -> DataFusionResult<()> { if let Some(fold) = self.partial_fold.as_mut() { fold.update(chunk)?; return Ok(()); } - 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| { - DataFusionError::Execution(format!("cannot flush evaluated frame: {error}")) - })? - { - self.ready - .push_back(batch.with_schema(self.output_schema.clone())?); - } + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + &chunk, + )?; + self.ready.push_back(batch); Ok(()) } @@ -1562,6 +1503,87 @@ impl FusedSeriesStream { } } +enum FusedOutput { + Rows(RecordBatch), + Evaluated(EvaluatedChunk), +} + +impl FusedSeriesStream { + fn poll_output(&mut self, cx: &mut Context<'_>) -> Poll>> { + for _ in 0..POLL_BUDGET { + if let Some(batch) = self.ready.pop_front() { + return Poll::Ready(Some(Ok(FusedOutput::Rows(batch)))); + } + if self.frame_mode && !self.done { + match self + .executor + .as_mut() + .map_or(Ok(None), NativeExecutor::next_chunk) + { + Ok(Some(chunk)) => return Poll::Ready(Some(Ok(FusedOutput::Evaluated(chunk)))), + Ok(None) if self.inputs.is_none() => self.executor = None, + Ok(None) => {} + Err(error) => { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + if self.done { + return Poll::Ready(None); + } + let Some(inputs) = self.inputs.as_mut() else { + match self.pull_one() { + Ok(Some(chunk)) if self.frame_mode => { + return Poll::Ready(Some(Ok(FusedOutput::Evaluated(chunk)))); + } + Ok(Some(chunk)) => { + if let Err(error) = self.convert_rows(chunk) { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Ok(None) => { + if let Err(error) = self.finish_fold() { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + self.done = true; + continue; + } + Err(error) => { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + } + continue; + }; + match Pin::new(inputs).poll_next(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + self.inputs = None; + if let Err(error) = self.close_input() { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + self.fail(); + return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); + } + Poll::Ready(Some(Ok((source, chunk)))) => { + if let Err(error) = self.push(source, chunk) { + self.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + } + cx.waker().wake_by_ref(); + Poll::Pending + } +} + impl RecordBatchStream for FusedSeriesStream { fn schema(&self) -> SchemaRef { self.output_schema.clone() @@ -1571,110 +1593,47 @@ impl RecordBatchStream for FusedSeriesStream { impl Stream for FusedSeriesStream { type Item = DataFusionResult; - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll>> { - let this = &mut *self; - // A batch that is already settled is handed over before anything is read, so a bounded - // poll never delays one that is ready behind more input. - for _ in 0..POLL_BUDGET { - if 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_none() { - this.executor = None; - } - } - Err(error) => { - this.fail(); - return Poll::Ready(Some(Err(error))); - } - } - } - if this.done { - return Poll::Ready(None); - } - let Some(inputs) = this.inputs.as_mut() else { - match this.pull_one() { - Ok(Some(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().and_then(|()| this.flush_frames()) { - this.fail(); - return Poll::Ready(Some(Err(error))); - } - this.done = true; - continue; - } - Err(error) => { - this.fail(); - return Poll::Ready(Some(Err(error))); - } - } - continue; - }; - match Pin::new(inputs).poll_next(cx) { - 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 = None; - if let Err(error) = this.close_input() { - this.fail(); - return Poll::Ready(Some(Err(error))); - } - } - Poll::Ready(Some(Err(error))) => { - this.fail(); - return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); - } - Poll::Ready(Some(Ok((source, chunk)))) => { - if let Err(error) = this.push(source, chunk) { - this.fail(); - return Poll::Ready(Some(Err(error))); - } - } + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.poll_output(cx) { + Poll::Ready(Some(Ok(FusedOutput::Rows(batch)))) => Poll::Ready(Some(Ok(batch))), + Poll::Ready(Some(Ok(FusedOutput::Evaluated(_)))) => { + self.fail(); + Poll::Ready(Some(Err(DataFusionError::Internal( + "evaluated chunk reached the row stream adapter".into(), + )))) } + Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))), + Poll::Ready(None) => Poll::Ready(None), + Poll::Pending => Poll::Pending, + } + } +} + +struct FusedEvaluatedStream { + source: FusedSeriesStream, +} + +impl Stream for FusedEvaluatedStream { + type Item = Result<(datatypes::series::samples::Series, bool), common_error::ext::BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match Pin::new(&mut self.source).poll_output(cx) { + Poll::Ready(Some(Ok(FusedOutput::Evaluated(chunk)))) => { + Poll::Ready(Some(Ok((chunk.series, chunk.is_last)))) + } + Poll::Ready(Some(Ok(FusedOutput::Rows(_)))) => { + self.source.fail(); + Poll::Ready(Some(Err(common_error::ext::BoxedError::new( + common_query::error::Error::from(DataFusionError::Internal( + "row batch reached the evaluated stream adapter".into(), + )), + )))) + } + Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err( + common_error::ext::BoxedError::new(common_query::error::Error::from(error)), + ))), + Poll::Ready(None) => Poll::Ready(None), + Poll::Pending => Poll::Pending, } - // The budget ran out on input that settled nothing. Everything read so far is kept, so the - // next poll continues where this one stopped; waking first is what asks for that poll. - cx.waker().wake_by_ref(); - Poll::Pending } } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 2599aeaa678..41627bb0ae4 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -36,7 +36,7 @@ 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 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}; @@ -48,9 +48,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; -use datafusion::execution::memory_pool::{ - GreedyMemoryPool, MemoryConsumer, MemoryPool, UnboundedMemoryPool, -}; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; @@ -69,25 +67,24 @@ use datafusion::prelude::SessionContext; use datatypes::prelude::ConcreteDataType; use datatypes::series::frame::{ FrameContent, FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, - LabelDecodeSpec, LabelDecodeTag, SeriesFrameBuilder, SeriesFrameDeclaration, - pack_raw_series_chunk, raw_series_frame_schema, series_frame_schema, + 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}; 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, SendableRawSeriesStream}; +use store_api::region_engine::RawLabelSource; mod frame_output; use super::{ - FusedOutputFormat, FusedSeries, FusedSeriesBinding, FusedSeriesStream, OutputColumn, - POLL_BUDGET, series_to_record_batch, + FusedOutputFormat, FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, + series_to_record_batch, }; use crate::evaluated::EvaluatedChunk; use crate::extension_plan::aggregate_exchange::state_schema; -use crate::extension_plan::series_gather::SeriesGather; use crate::native::execute::{InputId, NativeExecutionState, NativeExecutor}; use crate::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundScalar, @@ -535,6 +532,12 @@ fn keyed_frame_schema_of(schema: &SeriesSchema) -> SchemaRef { .expect("key labels are a frame") } +fn canonical_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef { + let evaluated = super::evaluated_output(tree).unwrap(); + let schema = super::series_frame_output_schema(evaluated.as_ref()).unwrap(); + Arc::new(datafusion::common::DFSchema::try_from(schema).unwrap()) +} + /// The columns this node produces: the field, the tags and the grid timestamps. fn output_schema() -> SchemaRef { Arc::new(Schema::new(vec![ @@ -2573,221 +2576,6 @@ 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.map(|input| { - SeriesGather::new( - vec![(0, input)], - MemoryConsumer::new("test") - .register(&(Arc::new(UnboundedMemoryPool::default()) as Arc)), - ) - }), - 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!(stream.inputs.is_some()); - 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 @@ -4033,7 +3821,7 @@ async fn series_frame_output_uses_evaluated_codec_schema() { tree.clone(), binding(&tree), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .expect("evaluated samples have a frame schema"); @@ -4045,6 +3833,16 @@ async fn series_frame_output_uses_evaluated_codec_schema() { .all(|field| field.name() != "ts"), "frame schema comes from the codec rather than the row output" ); + let wrong_schema = 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_err("the caller's incorrect output schema is rejected"); + assert!(wrong_schema.to_string().contains("schema does not match")); 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()); @@ -4065,7 +3863,7 @@ fn series_frame_rejects_multiple_sources_during_planning() { tree.clone(), binding(&tree), vec![input_plan(frame_schema()), input_plan(frame_schema())], - binary_output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .unwrap_err(); @@ -4090,7 +3888,7 @@ async fn series_frame_preserves_four_input_partitions() { tree.clone(), binding(&tree), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .unwrap(); @@ -4141,7 +3939,7 @@ async fn series_frame_rate_matches_plain_rows_bitwise() { tree.clone(), binding(&tree), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .unwrap(); @@ -4226,7 +4024,7 @@ fn the_explain_text_names_the_frame_format_only() { rows.tree().clone(), rows.binding().clone(), rows.inputs().into_iter().cloned().collect(), - Arc::clone(rows.schema()), + canonical_frame_schema(rows.tree()), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); @@ -4261,7 +4059,7 @@ fn the_logical_schema_follows_the_output_format() { rows.tree().clone(), rows.binding().clone(), rows.inputs().into_iter().cloned().collect(), - Arc::clone(rows.schema()), + canonical_frame_schema(rows.tree()), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); @@ -4290,7 +4088,7 @@ fn rebuilding_a_frame_node_keeps_its_output_format() { rows.tree().clone(), rows.binding().clone(), rows.inputs().into_iter().cloned().collect(), - Arc::clone(rows.schema()), + canonical_frame_schema(rows.tree()), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); @@ -4381,7 +4179,7 @@ async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() { tree.clone(), binding(&tree), vec![input_plan(schema)], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); @@ -4462,6 +4260,56 @@ async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() { ); } +/// Errors from a fused frame input terminate once and release its upstream immediately. +#[tokio::test] +async fn fused_frame_input_error_terminates_and_releases_input() { + let schema = frame_schema(); + let partial = frame(&series_chunk( + "web-1", + &[(GRID_START + RANGE_MS + SAMPLE_STEP, 1.0)], + false, + )); + let dropped = Arc::new(AtomicBool::new(false)); + let staged_schema = Arc::clone(&schema); + let staged_dropped = Arc::clone(&dropped); + let source = staged(Arc::clone(&schema), move || { + Box::pin( + StagedStream::new(Arc::clone(&staged_schema), vec![partial.clone()]) + .error("fused frame input failed") + .dropped(&staged_dropped), + ) + }); + let tree = rate_tree(); + let node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree.clone(), + binding(&tree), + vec![input_plan(schema)], + canonical_frame_schema(&tree), + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let exec = plan_of(node, source).unwrap(); + let mut output = exec + .execute(0, SessionContext::default().task_ctx()) + .unwrap(); + let error = match output.next().await { + Some(Err(error)) => error, + Some(Ok(_)) => panic!("unexpected output before the upstream failure"), + None => panic!("the upstream failure was swallowed"), + }; + assert!(error.to_string().contains("fused frame input failed")); + assert!( + dropped.load(Ordering::SeqCst), + "the failed adapter releases input" + ); + assert!(output.next().await.is_none(), "failure is yielded once"); + assert!( + output.next().await.is_none(), + "EOF remains stable after failure" + ); +} + /// 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 @@ -4487,7 +4335,7 @@ async fn a_frame_carries_the_closing_chunk_a_row_table_drops() { tree.clone(), binding(&tree), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); @@ -4572,7 +4420,7 @@ async fn two_chains_that_close_on_one_label_set_each_keep_their_frame() { tree.clone(), binding(&tree), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + canonical_frame_schema(&tree), FusedOutputFormat::SeriesFrame, ) .expect("a rate tree is evaluated on a grid"); 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 0e09ebafc5c..3e29af7e53e 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 @@ -162,7 +162,11 @@ async fn sparse_instant_frame_preserves_positions_nulls_and_nan_across_formats() tree.clone(), binding.clone(), vec![input_plan(frame_schema_of(&schema))], - two_field_output_schema().to_dfschema_ref().unwrap(), + if format == FusedOutputFormat::SeriesFrame { + canonical_frame_schema(&tree) + } else { + two_field_output_schema().to_dfschema_ref().unwrap() + }, format, ) .expect("instant tree has the declared frame output") @@ -217,7 +221,11 @@ async fn abs_rate_frame_matches_plain_rows_with_exact_values_and_labels() { tree.clone(), binding.clone(), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + if format == FusedOutputFormat::SeriesFrame { + canonical_frame_schema(&tree) + } else { + output_schema().to_dfschema_ref().unwrap() + }, format, ) .unwrap() @@ -317,7 +325,11 @@ async fn rewritten_identical_labels_keep_disjoint_chains_in_one_encoder() { tree.clone(), binding.clone(), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + if format == FusedOutputFormat::SeriesFrame { + canonical_frame_schema(&tree) + } else { + output_schema().to_dfschema_ref().unwrap() + }, format, ) .unwrap() diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index f58a9ee3ff7..5300cb0b002 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -287,8 +287,9 @@ impl DistPlannerAnalyzer { let plan = plan.transform_with_subqueries(&|plan| { Self::inspect_plan_with_subquery(plan, in_process_region_transport) })?; - let mut rewriter = PlanRewriter::new(&plan.data, in_process_region_transport); - let result = plan.data.rewrite(&mut rewriter)?.data; + let plan = rewrite_hybrid_binary_operands(plan.data, in_process_region_transport)?; + let mut rewriter = PlanRewriter::new(&plan, in_process_region_transport); + let result = plan.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } @@ -508,6 +509,56 @@ struct PlanRewriter { new_child_plan: Option, } +fn native_operand_wrapper_child(plan: &LogicalPlan) -> Option<&LogicalPlan> { + match plan { + LogicalPlan::SubqueryAlias(alias) => Some(alias.input.as_ref()), + LogicalPlan::Projection(projection) => Some(projection.input.as_ref()), + LogicalPlan::Filter(filter) => Some(filter.input.as_ref()), + LogicalPlan::Window(window) => Some(window.input.as_ref()), + _ => None, + } +} + +fn native_operand_fused(plan: &LogicalPlan) -> Option<&LogicalPlan> { + let mut node = plan; + while let Some(child) = native_operand_wrapper_child(node) { + node = child; + } + (is_fused_series_plan(node) && matches!(node.inputs().as_slice(), [LogicalPlan::TableScan(_)])) + .then_some(node) +} + +fn is_native_operand_subtree(plan: &LogicalPlan) -> bool { + native_operand_fused(plan).is_some() +} + +fn rewrite_hybrid_binary_operands(plan: LogicalPlan, transport: bool) -> DfResult { + plan.transform_up(|plan| { + let LogicalPlan::Join(join) = &plan else { + return Ok(Transformed::no(plan)); + }; + if !is_native_operand_subtree(&join.left) || !is_native_operand_subtree(&join.right) { + return Ok(Transformed::no(plan)); + } + fn rewrite_side(side: &LogicalPlan, transport: bool) -> DfResult { + if let Some(child) = native_operand_wrapper_child(side) { + let rewritten_child = rewrite_side(child, transport)?; + return side + .clone() + .with_new_exprs(side.expressions(), vec![rewritten_child]); + } + let mut rw = PlanRewriter::new(side, transport); + Ok(side.clone().rewrite(&mut rw)?.data) + } + let left = rewrite_side(&join.left, transport)?; + let right = rewrite_side(&join.right, transport)?; + let mut rebuilt = LogicalPlan::Join(join.clone()); + rebuilt = rebuilt.with_new_exprs(rebuilt.expressions(), vec![left, right])?; + Ok(Transformed::yes(rebuilt)) + }) + .map(|transformed| transformed.data) +} + fn is_admitted_pair_plan(plan: &LogicalPlan) -> bool { let LogicalPlan::Extension(extension) = plan else { return false; diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index b512655dd2c..a899c824679 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -16,6 +16,7 @@ use std::collections::{BTreeSet, HashSet}; use std::pin::Pin; use std::sync::Arc; +use crate::promql::native_compile::{compile, statement_grid}; use arrow::datatypes::{DataType, IntervalDayTime, TimeUnit}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MITO_ENGINE}; use common_error::ext::BoxedError; @@ -43,10 +44,14 @@ use datafusion_functions::datetime::expr_fn::now; use datafusion_functions::unicode::expr_fn::substring; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; +use datatypes::series::SeriesOrdering; +use datatypes::series::samples::{FieldSpec, SeriesSchema}; use futures::Stream; use futures::task::{Context, Poll}; use pretty_assertions::assert_eq; -use promql::extension_plan::{HistogramFold, HistogramFoldOperation}; +use promql::extension_plan::{ + FusedSeries, FusedSeriesBinding, HistogramFold, HistogramFoldOperation, +}; use regex::Regex; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; @@ -322,6 +327,134 @@ impl Stream for EmptyStream { } } +fn hybrid_fused(scan: LogicalPlan, metric: &str) -> LogicalPlan { + let stmt = promql_parser::parser::EvalStmt { + expr: promql_parser::parser::parse(metric).unwrap(), + start: std::time::UNIX_EPOCH, + end: std::time::UNIX_EPOCH + std::time::Duration::from_secs(10), + interval: std::time::Duration::from_secs(1), + lookback_delta: std::time::Duration::from_secs(1), + }; + let series = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile(&stmt, series, SeriesOrdering::KeyThenTime).unwrap(); + let fused = FusedSeries::try_new( + Arc::from(metric), + tree, + FusedSeriesBinding { + grid: statement_grid(&stmt).unwrap(), + time_index: "ts".into(), + tags: vec![], + metric_name: Some(Arc::from(metric)), + }, + vec![scan], + Arc::new( + datafusion_common::DFSchema::try_from(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("v0", DataType::Float64, true), + ])) + .unwrap(), + ), + ) + .unwrap(); + LogicalPlan::Extension(Extension { + node: Arc::new(fused), + }) +} + +fn count_merge_scans(plan: &LogicalPlan) -> usize { + usize::from( + matches!(plan, LogicalPlan::Extension(extension) if extension.node.as_any().is::()), + ) + plan + .inputs() + .iter() + .map(|input| count_merge_scans(input)) + .sum::() +} + +#[test] +fn hybrid_binary_join_rewrites_each_native_operand_boundary() { + let left_table = TestTable::table_with_name(0, "left_metric".into()); + let right_table = TestTable::table_with_name(1, "right_metric".into()); + let scan = |table: TableRef| { + LogicalPlanBuilder::scan( + table.table_info().name.clone(), + Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))), + None, + ) + .unwrap() + .build() + .unwrap() + }; + let left = LogicalPlanBuilder::from(hybrid_fused(scan(left_table), "left_metric")) + .alias(TableReference::bare("left_metric")) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(hybrid_fused(scan(right_table), "right_metric")) + .alias(TableReference::bare("right_metric")) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .join( + LogicalPlanBuilder::from(right).build().unwrap(), + JoinType::Inner, + (vec![Column::from_name("v0")], vec![Column::from_name("v0")]), + None, + ) + .unwrap() + .build() + .unwrap(); + let rewritten = DistPlannerAnalyzer::new(true).try_push_down(plan).unwrap(); + let LogicalPlan::Join(join) = &rewritten else { + panic!("Join must remain frontend: {rewritten}"); + }; + for (side, expected_table) in [(&join.left, "left_metric"), (&join.right, "right_metric")] { + assert_eq!(1, count_merge_scans(side), "{rewritten}"); + assert_no_fused_series_outside_merge_scan(side); + + let merge_scan = find_merge_scan(side).expect("each join operand has a MergeScan"); + let LogicalPlan::Extension(extension) = merge_scan.input() else { + panic!("MergeScan remote input must be a FusedSeries root"); + }; + let fused = extension + .node + .as_any() + .downcast_ref::() + .expect("MergeScan remote input must be a FusedSeries root"); + assert_eq!(expected_table, fused.promql()); + assert_eq!(1, fused.children().len()); + let [LogicalPlan::TableScan(table_scan)] = fused.children() else { + panic!("FusedSeries operand must directly consume a TableScan: {fused:?}"); + }; + assert_eq!(expected_table, table_scan.table_name.to_string()); + } + assert_eq!( + 2, + count_merge_scans(&rewritten), + "Join must be outside both MergeScans: {rewritten}" + ); +} + +fn assert_no_fused_series_outside_merge_scan(plan: &LogicalPlan) { + if let LogicalPlan::Extension(extension) = plan { + if extension.node.as_any().is::() { + return; + } + assert!( + !extension.node.as_any().is::(), + "FusedSeries must be contained in a MergeScan remote input: {plan}" + ); + } + for input in plan.inputs() { + assert_no_fused_series_outside_merge_scan(input); + } +} + fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { if let LogicalPlan::Extension(extension) = plan && let Some(merge_scan) = extension diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c6c96a33156..09397fbbc7e 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -586,7 +586,11 @@ impl PromPlanner { // schema. self.ctx.start = end; self.ctx.end = end; - let plan = self.prom_expr_to_plan(expr, query_engine_state).await?; + let native_fused = self.native_fused; + self.native_fused = false; + let plan_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; + let plan = plan_result?; return Ok(LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, schema: plan.schema().clone(), @@ -594,7 +598,11 @@ impl PromPlanner { }; self.ctx.start = start; self.ctx.end = end; - self.prom_expr_to_plan(expr, query_engine_state).await + let native_fused = self.native_fused; + self.native_fused = false; + let plan_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; + plan_result } async fn prom_subquery_expr_to_plan( @@ -646,10 +654,14 @@ impl PromPlanner { self.ctx.start = child_start; self.ctx.end = child_end; } - let input = self.prom_expr_to_plan(expr, query_engine_state).await?; + let native_fused = self.native_fused; + self.native_fused = false; + let input_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; self.ctx.interval = current_interval; self.ctx.start = current_start; self.ctx.end = current_end; + let input = input_result?; // Return no rows with the child schema so enclosing fallbacks still work. let input = match child_window { Some(_) => input, @@ -1365,32 +1377,57 @@ impl PromPlanner { } // both are columns. join them on time index (None, None) => { - let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?; - let left_field_columns = self.ctx.field_columns.clone(); - let left_time_index_column = self.ctx.time_index_column.clone(); + let native_pair = + self.native_fused && native::hybrid_binary_admitted(*op, modifier, lhs, rhs); + let native = if native_pair { + match ( + self.try_plan_native_operand(lhs, query_engine_state) + .await?, + self.try_plan_native_operand(rhs, query_engine_state) + .await?, + ) { + (Some(left), Some(right)) => Some((left, right)), + _ => None, + } + } else { + None + }; + let used_native = native.is_some(); + let (mut left_input, mut left_context, mut right_input, mut right_context) = + if let Some(((left, left_ctx), (right, right_ctx))) = native { + (left, left_ctx, right, right_ctx) + } else { + let left = self.prom_expr_to_plan(lhs, query_engine_state).await?; + let left_ctx = self.ctx.clone(); + let right = self.prom_expr_to_plan(rhs, query_engine_state).await?; + let right_ctx = self.ctx.clone(); + (left, left_ctx, right, right_ctx) + }; + let left_field_columns = left_context.field_columns.clone(); + let left_time_index_column = left_context.time_index_column.clone(); + self.ctx = left_context.clone(); let mut left_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let mut left_context = self.ctx.clone(); - - let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?; - let right_field_columns = self.ctx.field_columns.clone(); - let right_time_index_column = self.ctx.time_index_column.clone(); + let right_field_columns = right_context.field_columns.clone(); + let right_time_index_column = right_context.time_index_column.clone(); + self.ctx = right_context.clone(); let mut right_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let mut right_context = self.ctx.clone(); // Both operands are planned first because only the planned contexts tell a tag // from a value field. The rewrite merely adds matchers to a selector, so the // table reference, time index and field columns captured above stay valid. - if let Some(rewritten) = matching_filters::propagate( - binary_expr, - &left_context.tag_columns, - &right_context.tag_columns, - &left_context.aggregation_field_labels, - &right_context.aggregation_field_labels, - ) { + if !used_native + && let Some(rewritten) = matching_filters::propagate( + binary_expr, + &left_context.tag_columns, + &right_context.tag_columns, + &left_context.aggregation_field_labels, + &right_context.aggregation_field_labels, + ) + { // A copied matcher belongs to the scan, not to the operand's identity: // `absent()` turns `selector_matcher` into the labels it reports, so the // re-planned context keeps the matchers the operand was written with. diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index e35712c0237..90dab4bae3c 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -40,7 +40,7 @@ use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; -use datafusion::logical_expr::{Extension, LogicalPlan, TableScanBuilder}; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder, TableScanBuilder}; use datafusion::prelude::Expr as DfExpr; use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; @@ -51,7 +51,8 @@ use promql::vector_match::{Cardinality, MatchingLabels, SetOp}; use promql_parser::label::Matchers; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ - EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, + BinModifier, EvalStmt, Expr as PromExpr, LabelModifier, MatrixSelector, ParenExpr, + VectorMatchCardinality, VectorSelector, token, }; use session::context::QueryContextRef; use snafu::ResultExt; @@ -109,7 +110,70 @@ type PreparedInput = ( Vec, ); +pub(crate) fn hybrid_binary_admitted( + op: token::TokenType, + modifier: &Option, + lhs: &PromExpr, + rhs: &PromExpr, +) -> bool { + let supported_op = matches!( + op.id(), + token::T_ADD + | token::T_SUB + | token::T_MUL + | token::T_DIV + | token::T_MOD + | token::T_POW + | token::T_ATAN2 + ); + let supported_matching = match modifier { + None => true, + Some(modifier) => { + matches!(&modifier.card, VectorMatchCardinality::OneToOne) + && (matches!(&modifier.matching, Some(LabelModifier::Include(labels)) if !labels.labels.is_empty()) + || matches!(&modifier.matching, Some(LabelModifier::Exclude(labels)) if labels.labels.is_empty())) + } + }; + supported_op + && supported_matching + && lhs.value_type() == ValueType::Vector + && rhs.value_type() == ValueType::Vector + && admitted_selector(lhs).is_some() + && admitted_selector(rhs).is_some() +} + impl PromPlanner { + pub(super) async fn try_plan_native_operand( + &mut self, + expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result> { + if admitted_selector(expr).is_none() { + return Ok(None); + } + let saved = self.ctx.clone(); + let (plan, operand_ctx) = match self + .try_plan_native_fused_with_context(expr, query_engine_state) + .await? + { + Some((plan, Some(operand_ctx))) => (plan, operand_ctx), + _ => { + self.ctx = saved; + return Ok(None); + } + }; + self.ctx = operand_ctx.clone(); + let table_ref = self.table_ref(); + self.ctx = saved; + let table_ref = table_ref?; + let plan = LogicalPlanBuilder::from(plan) + .alias(table_ref) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + Ok(Some((plan, operand_ctx))) + } + /// Plans `prom_expr` as one native node, or returns `None` to keep the existing plan. /// /// `query_engine_state` is only read to build the *existing* plan of an aggregate query, whose @@ -133,6 +197,17 @@ impl PromPlanner { prom_expr: &PromExpr, query_engine_state: &QueryEngineState, ) -> Result> { + Ok(self + .try_plan_native_fused_with_context(prom_expr, query_engine_state) + .await? + .map(|(plan, _)| plan)) + } + + async fn try_plan_native_fused_with_context( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result)>> { if !self.native_fused { return Ok(None); } @@ -188,13 +263,17 @@ impl PromPlanner { output_schema, ) .context(DataFusionPlanningSnafu)?; - return Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))); + return Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + None, + ))); } let original_context = self.ctx.clone(); let preparation = async { let mut prepared: Vec = Vec::new(); + let mut single_input_context = None; for selector in &selectors { self.ctx = original_context.clone(); // These matchers are alternatives the one scan below the node cannot represent, so the @@ -210,6 +289,9 @@ impl PromPlanner { if self.setup_context().await?.is_some() { return Ok(None); } + if selectors.len() == 1 { + single_input_context = Some(self.ctx.clone()); + } let (Some(metric), Some(time)) = ( self.ctx.table_name.clone(), self.ctx.time_index_column.clone(), @@ -360,11 +442,11 @@ impl PromPlanner { } prepared.push((input, series_schema, Some(Arc::from(metric)), time, tags)); } - Ok(Some(prepared)) + Ok(Some((prepared, single_input_context))) } .await; self.ctx = original_context; - let Some(mut prepared) = preparation? else { + let Some((mut prepared, single_input_context)) = preparation? else { return Ok(None); }; // Two selectors over one table share one region, so the node runs there over both scans. @@ -536,9 +618,12 @@ impl PromPlanner { Err(DataFusionError::NotImplemented(_)) => return Ok(None), Err(error) => return Err(error).context(DataFusionPlanningSnafu), }; - return Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))); + return Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + single_input_context, + ))); } let tree = compiled.tree; let root_fields = tree @@ -593,9 +678,12 @@ impl PromPlanner { output_schema, ) .context(DataFusionPlanningSnafu)?; - Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))) + Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + single_input_context, + ))) } /// The statement being planned, rebuilt from the evaluation numbers the planner holds. @@ -972,8 +1060,8 @@ mod tests { use super::{ NATIVE_PROMQL_EXTENSION_KEY, admitted_selector, admitted_single_selector, - collect_selectors, fused_output_tags, independent_scan_pair, native_promql_requested, - scan_time_filter, + collect_selectors, fused_output_tags, hybrid_binary_admitted, independent_scan_pair, + native_promql_requested, scan_time_filter, }; fn compiled_test_tree(query: &str) -> NativeTree { @@ -1146,6 +1234,32 @@ mod tests { } } + fn hybrid_admission(query: &str) -> bool { + let expression = promql_parser::parser::parse(query).unwrap(); + let promql_parser::parser::Expr::Binary(binary) = expression else { + panic!("expected binary expression: {query}"); + }; + hybrid_binary_admitted(binary.op, &binary.modifier, &binary.lhs, &binary.rhs) + } + + #[test] + fn hybrid_binary_admission_is_narrow() { + for query in ["A / B", "rate(A[1m]) / rate(B[1m])"] { + assert!(hybrid_admission(query), "{query}"); + } + assert!(hybrid_admission("A + on(job) B")); + for query in [ + "A and B", + "A + on() B", + "A + on(job) group_left B", + "sum(A) / B", + "100 / B", + "A / (B + C)", + ] { + assert!(!hybrid_admission(query), "{query}"); + } + } + #[test] fn millisecond_scan_bounds_preserve_open_closed_and_maximum_edges() { let inclusive = scan_time_filter(0, 10, true, "ts").to_string(); diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index c6b30c24821..76fc2fac8c7 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2662,16 +2662,26 @@ async fn native_fused_cross_table_pair_reads_each_table_through_its_own_source() } } -/// Two selectors over tables split into regions keep the existing plan, across tables as on one. +/// Two selectors over tables split into regions are not one two-input node, across tables as on +/// one; the hybrid route may still plan each operand natively under the existing join. #[tokio::test] async fn native_fused_cross_table_pair_on_partitioned_tables_falls_back() { + fn has_pair_node(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(fused) = node.as_any().downcast_ref::() + && fused.inputs().len() > 1 + { + return true; + } + plan.inputs().into_iter().any(has_pair_node) + } let plan = build_native_fused_plan_with_provider( "some_metric / other_metric", true, native_two_table_provider(vec![0]).await, ) .await; - assert!(find_fused_series(&plan).is_none(), "{plan}"); + assert!(!has_pair_node(&plan), "{plan}"); } /// `some_metric` as [`native_two_table_provider`] builds it, and `other_metric` whose tags are diff --git a/tests-integration/src/tests/native_aggregate_test.rs b/tests-integration/src/tests/native_aggregate_test.rs index ab8617dde98..4b476aeb143 100644 --- a/tests-integration/src/tests/native_aggregate_test.rs +++ b/tests-integration/src/tests/native_aggregate_test.rs @@ -67,23 +67,32 @@ //! existing plan. [`check_fallback_case`] pins that the computed scalars do, and that a fold over a //! second source and a subquery keep it too. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt; use std::sync::{Arc, Mutex}; use common_base::Plugins; use common_query::OutputData; use common_recordbatch::RecordBatches; use common_recordbatch::util::collect_batches; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::coop::CooperativeExec; use datafusion::physical_plan::displayable; use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; use datatypes::arrow::array::{Array, ArrayRef, Float64Array, Int64Array, StringArray}; use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; +use datatypes::arrow::record_batch::RecordBatch; use frontend::instance::Instance; +use futures::StreamExt; use query::datafusion::QUERY_PARALLELISM_HINT; use query::dist_plan::MergeScanExec; use query::options::QueryOptions; @@ -116,8 +125,17 @@ const EVAL_MILLIS: i64 = EVAL_SECS * 1000; const GAUGE_TABLE: &str = "native_agg_gauge"; /// The table the `rate` queries read: the same five series as [`GAUGE_TABLE`], holding a counter. const COUNTER_TABLE: &str = "native_agg_counter"; +/// The right operand of the hybrid binary cases. It has one region but a partition column, so a +/// pair reading it is declined by the two-scan native route and planned by the hybrid one. +const COUNTER_RHS_TABLE: &str = "native_agg_counter_rhs"; /// The table the range-window smoke test reads. const RANGE_TABLE: &str = "native_agg_range"; +const HYBRID_MATCH_LEFT: &str = "native_hybrid_match_left"; +/// Partitioned like [`COUNTER_RHS_TABLE`], so the hybrid route plans a pair reading it. +const HYBRID_MATCH_RIGHT: &str = "native_hybrid_match_right"; +/// The same rows as [`HYBRID_MATCH_RIGHT`] without a partition column, so the two-scan route plans +/// a pair reading it. +const HYBRID_MATCH_RIGHT_UNPARTITIONED: &str = "native_hybrid_match_right_unpartitioned"; /// The tolerance of the value comparisons: relative, with an absolute floor. const FLOAT_TOLERANCE: f64 = 1e-10; @@ -182,6 +200,8 @@ const COUNTER_EAST_REGION: &str = "us-east-1"; #[derive(Default)] struct NativePlanObserver { seen: Mutex>>, + record_original_streams: bool, + partial_batches: Arc>>, } impl NativePlanObserver { @@ -194,6 +214,10 @@ impl NativePlanObserver { fn since(&self, mark: usize) -> Vec> { self.seen.lock().unwrap()[mark..].to_vec() } + + fn partial_batches(&self) -> Vec<(usize, RecordBatch)> { + std::mem::take(&mut *self.partial_batches.lock().unwrap()) + } } impl PhysicalPlanWrapper for NativePlanObserver { @@ -202,8 +226,84 @@ impl PhysicalPlanWrapper for NativePlanObserver { origin: Arc, _ctx: QueryContextRef, ) -> Arc { - self.seen.lock().unwrap().push(Arc::clone(&origin)); - origin + if origin.name() == FUSED_EXEC_NAME && self.record_original_streams { + let recorder = Arc::new(PartialStateRecorder { + input: Arc::clone(&origin), + properties: Arc::clone(origin.properties()), + batches: Arc::clone(&self.partial_batches), + }); + self.seen + .lock() + .unwrap() + .push(Arc::clone(&recorder) as Arc); + recorder + } else { + self.seen.lock().unwrap().push(Arc::clone(&origin)); + origin + } + } +} + +/// Transparent test wrapper that records the fused producer's original partition streams. +#[derive(Debug)] +struct PartialStateRecorder { + input: Arc, + properties: Arc, + batches: Arc>>, +} + +impl DisplayAs for PartialStateRecorder { + fn fmt_as(&self, format: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.input.fmt_as(format, f) + } +} + +impl ExecutionPlan for PartialStateRecorder { + fn name(&self) -> &str { + "PartialStateRecorder" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> datafusion::common::Result, + ) -> datafusion::common::Result { + self.input.apply_expressions(f) + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> datafusion::common::Result> { + if children.len() != 1 { + return Err(datafusion::error::DataFusionError::Internal( + "PartialStateRecorder expects one child".into(), + )); + } + Ok(Arc::new(Self { + input: children.remove(0), + properties: Arc::clone(&self.properties), + batches: Arc::clone(&self.batches), + })) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion::common::Result { + let stream = self.input.execute(partition, context)?; + let batches = Arc::clone(&self.batches); + let schema = stream.schema(); + let recorded = stream.map(move |result| { + if let Ok(batch) = &result { + batches.lock().unwrap().push((partition, batch.clone())); + } + result + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, recorded))) } } @@ -434,6 +534,463 @@ async fn run_query( ) } +#[tokio::test(flavor = "multi_thread")] +async fn native_hybrid_binary_05_06_08() { + common_telemetry::init_default_ut_logging(); + let observer = Arc::new(NativePlanObserver::default()); + let plugins = Plugins::new(); + plugins.insert(Arc::clone(&observer) as PhysicalPlanWrapperRef); + let standalone = GreptimeDbStandaloneBuilder::new("native_hybrid_binary") + .with_query_options(QueryOptions { + experimental_native_promql: true, + ..Default::default() + }) + .with_plugin(plugins) + .build() + .await; + let instance = Arc::clone(standalone.fe_instance()); + setup_fixture(&instance).await; + execute_sql(&instance, &format!("CREATE TABLE {HYBRID_MATCH_LEFT} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance))")).await; + execute_sql(&instance, &format!("CREATE TABLE {HYBRID_MATCH_RIGHT} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance)) PARTITION ON COLUMNS (job) ()")).await; + execute_sql(&instance, &format!("CREATE TABLE {HYBRID_MATCH_RIGHT_UNPARTITIONED} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance))")).await; + execute_sql( + &instance, + &format!("INSERT INTO {HYBRID_MATCH_LEFT} VALUES ({EVAL_MILLIS}, 'api', 'a1', 10)"), + ) + .await; + for table in [HYBRID_MATCH_RIGHT, HYBRID_MATCH_RIGHT_UNPARTITIONED] { + execute_sql(&instance, &format!("INSERT INTO {table} VALUES ({EVAL_MILLIS}, 'api', 'b1', 2), ({EVAL_MILLIS}, 'other', 'b2', 3), ({EVAL_MILLIS}, 'other', 'b3', 4)" )).await; + } + let propagated = format!("{HYBRID_MATCH_LEFT}{{job=\"api\"}} / on(job) {HYBRID_MATCH_RIGHT}"); + let (legacy, _) = run_query(&instance, &propagated, false, &observer).await; + assert_eq!( + vec![5.0], + legacy + .rows + .iter() + .map(|row| row.value.unwrap()) + .collect::>() + ); + let mark = observer.mark(); + let (hybrid, capture) = run_query(&instance, &propagated, true, &observer).await; + assert_eq!( + vec![5.0], + hybrid + .rows + .iter() + .map(|row| row.value.unwrap()) + .collect::>() + ); + assert_eq!( + 2, + capture.fused.len(), + "both binary operands use native execution" + ); + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .expect("frontend Join plan"); + let join = find_binary_join(frontend).unwrap(); + let uniqueness = join + .children() + .iter() + .map(|branch| uniqueness_operators_above_merge(branch)) + .fold((false, false), |(window, filter), found| { + (window || found.0, filter || found.1) + }); + assert!( + uniqueness.0 && uniqueness.1, + "on(job) FE uniqueness operators: {}", + displayable(frontend.as_ref()).indent(false) + ); + + // On the two-scan route the right side's two `job="other"` series are a duplicate match group, + // which Prometheus refuses whether or not the left side has a series to match them with + // (`VectorBinop` in engine.go). The 5 above comes from the matcher `{job="api"}` the existing + // plan copies to the right side before matching. + let unpartitioned = + format!("{HYBRID_MATCH_LEFT}{{job=\"api\"}} / on(job) {HYBRID_MATCH_RIGHT_UNPARTITIONED}"); + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, "true".to_string()); + let output = instance + .do_query(&tql_eval(&unpartitioned), Arc::new(query_ctx)) + .await + .remove(0) + .expect("the query plans"); + let error = match output.data { + OutputData::Stream(stream) => collect_batches(stream) + .await + .expect_err("a duplicate match group fails"), + other => panic!("`{unpartitioned}` was expected to stream: {other:?}"), + }; + assert!( + error.to_string().contains( + "found duplicate series for the match group {job=\"other\"} on the right hand-side" + ), + "{unpartitioned}: {error}" + ); + + check_hybrid_binary_cases(&instance, &observer).await; + let subquery = format!( + "max_over_time((rate({COUNTER_TABLE}[5m]) / on(job, instance) rate({COUNTER_RHS_TABLE}[5m]))[2m:1m])" + ); + let (legacy, legacy_capture) = run_query(&instance, &subquery, false, &observer).await; + let (native, native_capture) = run_query(&instance, &subquery, true, &observer).await; + assert!( + legacy_capture.fused.is_empty() && native_capture.fused.is_empty(), + "hybrid must be disabled inside subquery: {subquery}" + ); + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "subquery legacy/native row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("subquery fallback `{subquery}`"), + ); + check_reduced_label_matching(&instance, &observer).await; +} + +async fn check_reduced_label_matching(instance: &Arc, observer: &NativePlanObserver) { + let unique_query = format!( + "{GAUGE_TABLE}{{instance=~\"h1|h4\"}} / on(job) {COUNTER_RHS_TABLE}{{instance=~\"h1|h4\"}}" + ); + let (legacy, legacy_capture) = run_query(instance, &unique_query, false, observer).await; + assert!(legacy_capture.fused.is_empty()); + let mark = observer.mark(); + let (native, capture) = run_query(instance, &unique_query, true, observer).await; + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "{unique_query}: row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("on(job) native == legacy `{unique_query}`"), + ); + assert_eq!( + 1, + native.rows.len(), + "the matching instance filters leave one matching job" + ); + assert_eq!(2, capture.fused.len(), "{unique_query}"); + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .expect("frontend Join plan"); + let join = find_binary_join(frontend).unwrap(); + let uniqueness = join + .children() + .iter() + .map(|branch| uniqueness_operators_above_merge(branch)) + .fold((false, false), |(window, assertion_filter), found| { + (window || found.0, assertion_filter || found.1) + }); + assert!( + uniqueness.0, + "on(job) requires a frontend uniqueness WindowAggExec above MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + uniqueness.1, + "on(job) requires a frontend prom_assert_unique_match_group FilterExec above MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + for branch in join.children() { + assert!( + contains_execution_node(branch, "MergeScanExec"), + "on(job) FE branch must retain remote boundary: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + !contains_execution_node(branch, FUSED_EXEC_NAME), + "on(job) fused node must stay in region" + ); + } + for fused in &capture.fused { + assert_fused_reads_scan_directly(fused, &unique_query); + } + + let duplicate_query = format!("{GAUGE_TABLE} / on(job) {COUNTER_RHS_TABLE}"); + for enabled in [false, true] { + let sql = tql_eval(&duplicate_query); + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, SCAN_PARTITIONS.to_string()); + query_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, enabled.to_string()); + let output = instance + .do_query(&sql, Arc::new(query_ctx)) + .await + .remove(0) + .expect("duplicate query should produce an execution stream"); + let error = match output.data { + OutputData::Stream(stream) => collect_batches(stream) + .await + .expect_err("duplicate on(job) matches must fail"), + OutputData::RecordBatches(_) => panic!("duplicate on(job) matches must fail"), + other => panic!("unexpected duplicate-query output: {other:?}"), + }; + let message = error.to_string().to_ascii_lowercase(); + assert!( + message.contains("many-to-many") + || message.contains("multiple matches") + || message.contains("duplicate"), + "`{sql}` should fail for duplicate matching groups, got: {error}" + ); + } +} + +async fn check_hybrid_binary_cases(instance: &Arc, observer: &NativePlanObserver) { + let cases = [ + ( + format!("{GAUGE_TABLE} / on(job, instance) {COUNTER_RHS_TABLE}{{instance!=\"h5\"}}"), + 4usize, + ), + ( + format!( + "100 * (1 - ({GAUGE_TABLE} / on(job, instance) {COUNTER_RHS_TABLE}{{instance!=\"h5\"}}))" + ), + 4usize, + ), + ( + format!("rate({COUNTER_TABLE}[5m]) / on(job, instance) rate({COUNTER_RHS_TABLE}[5m])"), + 5usize, + ), + ]; + for (index, (promql, expected_rows)) in cases.into_iter().enumerate() { + let (legacy, legacy_capture) = run_query(instance, &promql, false, observer).await; + assert!(legacy_capture.fused.is_empty()); + let mark = observer.mark(); + let (native, capture) = run_query(instance, &promql, true, observer).await; + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "{promql}: full row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("native == legacy `{promql}`"), + ); + assert_eq!(expected_rows, native.rows.len(), "{promql}"); + assert_eq!( + legacy + .columns + .iter() + .filter(|column| column.ends_with("ts")) + .count(), + 1, + "{promql}: legacy timestamp column" + ); + assert_eq!( + native + .columns + .iter() + .filter(|column| column.ends_with("ts")) + .count(), + 1, + "{promql}: native timestamp column" + ); + if index == 2 { + let expected = [0.5; 5]; + assert!( + native.rows.iter().zip(expected).all(|(row, expected)| row + .value + .is_some_and(|value| close_enough(value, expected))), + "{promql}: independently scaled RHS rates must have expected ratios: {:#?}", + native.rows + ); + } + if index < 2 { + assert!( + native + .rows + .iter() + .all(|row| row.labels.get("instance").is_none_or(|v| v != "h5")), + "one-to-one match must drop unmatched series" + ); + } + assert_eq!(2, capture.fused.len(), "{promql}: plans since {mark}"); + let region_plans = observer + .since(mark) + .into_iter() + .filter(|plan| contains_execution_node(plan, "SeriesScanExec")) + .collect::>(); + let scan_tables = region_plans + .iter() + .map(|plan| { + let text = displayable(plan.as_ref()).indent(false).to_string(); + if text.contains("native_agg_gauge") { + GAUGE_TABLE + } else if text.contains(COUNTER_RHS_TABLE) { + COUNTER_RHS_TABLE + } else if text.contains(COUNTER_TABLE) { + COUNTER_TABLE + } else { + panic!("{promql}: unable to identify scan table: {text}") + } + }) + .collect::>(); + assert!( + scan_tables.contains(&GAUGE_TABLE) || scan_tables.contains(&COUNTER_TABLE), + "{promql}: gauge/counter source missing: {scan_tables:?}" + ); + assert!( + scan_tables.contains(&COUNTER_RHS_TABLE), + "{promql}: independent RHS scan missing: {scan_tables:?}" + ); + for fused in &capture.fused { + assert_fused_reads_scan_directly(fused, &format!("hybrid `{promql}`")); + } + for fused in &capture.fused { + let per_fused = NativeCapture { + fused: vec![Arc::clone(fused)], + finalize: Vec::new(), + }; + let counts = + fused_output_values_per_partition(&per_fused, &format!("hybrid `{promql}`")); + assert!( + counts.iter().all(|count| count.is_some_and(|n| n > 0)), + "{promql}: each scan partition must emit values, got {counts:?}" + ); + } + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .unwrap_or_else(|| panic!("{promql}: FE JoinExec missing; plans: {plans:#?}")); + let join = find_binary_join(frontend).unwrap(); + assert_eq!( + 2, + join.children().len(), + "{promql}: FE join branches: {}", + displayable(frontend.as_ref()).indent(false) + ); + for branch in join.children() { + assert!( + contains_execution_node(branch, "MergeScanExec"), + "{promql}: FE Join branch lacks MergeScanExec: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + !contains_execution_node(branch, FUSED_EXEC_NAME), + "{promql}: FusedSeriesExec leaked onto frontend: {}", + displayable(frontend.as_ref()).indent(false) + ); + } + assert!( + !is_nested_under_merge(frontend, "HashJoinExec"), + "{promql}: FE Join is under MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + if index == 1 { + assert!( + has_projection_parent_of_join(frontend), + "06 outer scalar arithmetic must project above JoinExec: {}", + displayable(frontend.as_ref()).indent(false) + ); + } + } +} + +fn find_binary_join(plan: &Arc) -> Option<&Arc> { + if matches!( + plan.name(), + "HashJoinExec" | "NestedLoopJoinExec" | "CrossJoinExec" + ) { + return Some(plan); + } + plan.children() + .iter() + .find_map(|child| find_binary_join(child)) +} + +fn find_execution_node<'a>( + plan: &'a Arc, + name: &str, +) -> Option<&'a Arc> { + if plan.name() == name { + return Some(plan); + } + plan.children() + .iter() + .find_map(|child| find_execution_node(child, name)) +} + +fn contains_execution_node(plan: &Arc, name: &str) -> bool { + find_execution_node(plan, name).is_some() +} + +/// Finds the PromQL uniqueness validation operators on the frontend side of a remote scan. +/// Traversal stops at each MergeScan so it cannot accidentally match region-side execution. +fn uniqueness_operators_above_merge(plan: &Arc) -> (bool, bool) { + if plan.name() == "MergeScanExec" { + return (false, false); + } + + let mut has_window = matches!(plan.name(), "WindowAggExec" | "BoundedWindowAggExec"); + let mut has_assertion_filter = plan.name() == "FilterExec" + && format!("{plan:?}").contains("prom_assert_unique_match_group"); + for child in plan.children() { + let (child_window, child_filter) = uniqueness_operators_above_merge(child); + has_window |= child_window; + has_assertion_filter |= child_filter; + } + (has_window, has_assertion_filter) +} + +fn is_nested_under_merge(plan: &Arc, name: &str) -> bool { + fn walk(plan: &Arc, under_merge: bool, name: &str) -> bool { + if under_merge && plan.name() == name { + return true; + } + plan.children() + .iter() + .any(|child| walk(child, under_merge || plan.name() == "MergeScanExec", name)) + } + walk(plan, false, name) +} + +fn has_projection_parent_of_join(plan: &Arc) -> bool { + plan.name() == "ProjectionExec" + && plan + .children() + .iter() + .any(|child| find_binary_join(child).is_some()) + || plan + .children() + .iter() + .any(|child| has_projection_parent_of_join(child)) +} + /// Executes one statement and drops its result. async fn execute_sql(instance: &Arc, sql: &str) { for result in instance.do_query(sql, QueryContext::arc()).await { @@ -499,6 +1056,7 @@ async fn setup_fixture(instance: &Arc) { } execute_sql(instance, &counter_table_sql()).await; + execute_sql(instance, &format!("CREATE TABLE {COUNTER_RHS_TABLE} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance)) PARTITION ON COLUMNS (job) () WITH ('compaction.type' = 'twcs', 'compaction.twcs.trigger_file_num' = '100')")).await; let (counter_first, counter_second) = ( COUNTER_FIRST_SECS * 1000, (COUNTER_FIRST_SECS + COUNTER_STEP_SECS) * 1000, @@ -521,8 +1079,15 @@ async fn setup_fixture(instance: &Arc) { ), ) .await; + let rhs_second_value = COUNTER_FIRST_VALUE + 2.0 * increase; + execute_sql(instance, &format!("INSERT INTO {COUNTER_RHS_TABLE} (ts, job, instance, v) VALUES ({counter_first}, '{job}', '{host}', {COUNTER_FIRST_VALUE:?}), ({counter_second}, '{job}', '{host}', {rhs_second_value:?})")).await; if index + 1 < COUNTER_SERIES.len() { execute_sql(instance, &format!("ADMIN FLUSH_TABLE('{COUNTER_TABLE}')")).await; + execute_sql( + instance, + &format!("ADMIN FLUSH_TABLE('{COUNTER_RHS_TABLE}')"), + ) + .await; } } @@ -950,6 +1515,197 @@ async fn check_case(instance: &Arc, observer: &NativePlanObserver, cas } } +#[tokio::test(flavor = "multi_thread")] +async fn native_aggregate_topk_over_global_partitions() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_topk_counters"; + const T1: i64 = BASE_SECS + 180; + const T2: i64 = BASE_SECS + 600; + let observer = Arc::new(NativePlanObserver { + record_original_streams: true, + ..Default::default() + }); + let plugins = Plugins::new(); + plugins.insert(Arc::clone(&observer) as PhysicalPlanWrapperRef); + let standalone = GreptimeDbStandaloneBuilder::new("native_aggregate_topk") + .with_query_options(QueryOptions { + experimental_native_promql: true, + ..Default::default() + }) + .with_plugin(plugins) + .build() + .await; + let instance = Arc::clone(standalone.fe_instance()); + execute_sql(&instance, &format!( + "CREATE TABLE {TABLE} (ts TIMESTAMP(3) TIME INDEX, hostname STRING, device STRING, \"region\" STRING, v DOUBLE, PRIMARY KEY (hostname, device, \"region\")) WITH ('compaction.type' = 'twcs', 'compaction.twcs.trigger_file_num' = '100')" + )).await; + // Flush host00/dev0 first and host00/dev1 last; flush each of the 28 series independently. + let order = std::iter::once((0usize, 0usize)) + .chain((1..14).flat_map(|host| (0..2).map(move |device| (host, device)))) + .chain(std::iter::once((0usize, 1usize))); + for (host, device) in order { + let slope_before = 140 - 10 * host as i64; + let slope_after = 10 + 10 * host as i64; + let mut values = Vec::new(); + for elapsed in (0i64..=600).step_by(30) { + let value = + 100_000 + slope_before * elapsed.min(300) + slope_after * (elapsed - 300).max(0); + values.push(format!( + "({}, 'host{host:02}', 'dev{device}', 'us-west-2', {value}.0)", + (BASE_SECS + elapsed) * 1000 + )); + } + execute_sql( + &instance, + &format!( + "INSERT INTO {TABLE} (ts, hostname, device, \"region\", v) VALUES {}", + values.join(",") + ), + ) + .await; + execute_sql(&instance, &format!("ADMIN FLUSH_TABLE('{TABLE}')")).await; + } + let promql = format!( + "topk(10, sum by(hostname)(rate({TABLE}{{device!=\"lo\",region=\"us-west-2\"}}[5m])))" + ); + let sql = format!("TQL EVAL ({T1}, {T2}, '7m') {promql}"); + let run = async |native: bool| { + let mark = observer.mark(); + let mut ctx = QueryContext::with_db_name(None); + ctx.set_extension(QUERY_PARALLELISM_HINT, SCAN_PARTITIONS.to_string()); + ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, native.to_string()); + let output = instance + .do_query(&sql, Arc::new(ctx)) + .await + .remove(0) + .unwrap(); + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + other => panic!("unexpected output: {other:?}"), + }; + ( + QueryResult::from_batches(&batches, &sql), + NativeCapture::from_plans(&observer.since(mark)), + ) + }; + let (legacy, legacy_capture) = run(false).await; + assert!(legacy_capture.fused.is_empty() && legacy_capture.finalize.is_empty()); + let (native, capture) = run(true).await; + let partial_batches = observer.partial_batches(); + assert_same_columns(&native, &legacy, &sql); + assert_rows( + &native.rows, + &legacy.rows, + &format!("native vs legacy `{sql}`"), + ); + assert_eq!(20, native.rows.len(), "both timestamps return ten rows"); + let distinct_timestamps = native + .rows + .iter() + .map(|row| row.ts_millis) + .collect::>(); + assert_eq!( + [T1 * 1000, T2 * 1000] + .into_iter() + .collect::>(), + distinct_timestamps + ); + let mut expected = Vec::new(); + for (ts, hosts, post) in [(T1 * 1000, 0..10, false), (T2 * 1000, 4..14, true)] { + for host in hosts { + let slope = if post { + 10 + 10 * host + } else { + 140 - 10 * host + }; + let scale = if post { 2.0 } else { 1.3 }; + expected.push(ResultRow { + ts_millis: ts, + labels: [("hostname".to_string(), format!("host{host:02}"))].into(), + value: Some(slope as f64 * scale), + }); + } + } + expected.sort_by(|a, b| (a.ts_millis, &a.labels).cmp(&(b.ts_millis, &b.labels))); + assert_rows(&native.rows, &expected, &format!("expected `{sql}`")); + assert_eq!(1, capture.fused.len()); + assert_eq!(1, capture.finalize.len()); + let fused_plan = displayable(capture.fused[0].as_ref()) + .indent(false) + .to_string(); + assert!( + fused_plan.contains("partial state: sum by (hostname) Complete"), + "{fused_plan}" + ); + let producer_tree = fused_plan + .split_once("NativeTree") + .map(|(_, tree)| tree.split("partial state:").next().unwrap_or(tree)) + .unwrap_or(&fused_plan); + assert!( + !producer_tree.to_ascii_lowercase().contains("topk"), + "{fused_plan}" + ); + let finalize_plan = displayable(capture.finalize[0].as_ref()) + .indent(false) + .to_string(); + assert!( + finalize_plan.to_ascii_lowercase().contains("topk"), + "{finalize_plan}" + ); + let per_partition = fused_output_values_per_partition(&capture, &sql); + assert!( + per_partition + .iter() + .all(|count| count.is_some_and(|count| count > 0)), + "{per_partition:?}" + ); + // Prove the partition-local states overlap without decoding or summing their values: at each + // evaluation point, at least one selected hostname/point/field was emitted by distinct original + // producer partitions. + let mut owners = BTreeMap::<(String, u64, u32), BTreeSet>::new(); + for (partition, batch) in &partial_batches { + let schema = batch.schema(); + // `state_schema` names its fixed columns and labels `metric_name`, `tag_N`, `point` and + // `field`; use those names rather than baking their positions into this assertion. + let hostname = batch + .column(schema.index_of("tag_0").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let batch_point = batch + .column(schema.index_of("point").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let field = batch + .column(schema.index_of("field").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + owners + .entry(( + hostname.value(row).to_string(), + batch_point.value(row), + field.value(row), + )) + .or_default() + .insert(*partition); + } + } + for (point, expected_hosts) in [(0u64, 0..10), (1u64, 4..14)] { + assert!( + expected_hosts.into_iter().any(|host| { + owners + .get(&(format!("host{host:02}"), point, 0)) + .is_some_and(|partitions| partitions.len() >= 2) + }), + "no selected hostname has rows from distinct original producer partitions at point {point}: {owners:?}" + ); + } +} + #[tokio::test(flavor = "multi_thread")] async fn native_aggregate_complete_over_four_nonempty_mito_partitions() { common_telemetry::init_default_ut_logging(); diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index e24f81115f1..c1df8a9afb3 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -3467,8 +3467,9 @@ async fn native_promql_cross_table_comparisons() { } } -/// A two-selector query over a table split into regions keeps the existing plan: the two-scan -/// native node has no placement there, and the query used to fail in the dist analyzer. +/// A two-selector query over a table split into regions is not planned as one two-scan native node: +/// that node has no placement there, and the query used to fail in the dist analyzer. The hybrid +/// route may still run each operand natively under the existing join. #[tokio::test(flavor = "multi_thread")] async fn native_promql_two_selectors_on_a_partitioned_table_keep_the_legacy_plan() { let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( @@ -3506,7 +3507,12 @@ async fn native_promql_two_selectors_on_a_partitioned_table_keep_the_legacy_plan let sql = format!("TQL EVAL (5, 5, '1s') {query}"); observer.clear(); let native = native_result_rows(native_tql_output(&instance, &sql, true).await).await; - assert!(observer.fused().is_none(), "`{query}` was planned natively"); + assert!( + observer + .fused() + .is_none_or(|plan| plan.children().len() < 2), + "`{query}` was planned as one two-scan native node" + ); assert_eq!(native, expected, "{query}"); let legacy = native_result_rows(native_tql_output(&instance, &sql, false).await).await; assert_eq!(native, legacy, "{query}"); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 93284ddeb3a..987ae089e2a 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1960,13 +1960,22 @@ async fn assert_raw_scan_fused_formats(observer: &NativeFusedObserver, frontend: promql::extension_plan::FusedOutputFormat::PlainRows, promql::extension_plan::FusedOutputFormat::SeriesFrame, ] { - scan_plan.clone().reset_state().unwrap(); + let scan_plan = 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(), + if format == promql::extension_plan::FusedOutputFormat::SeriesFrame { + let packed = common_query::series_frame::pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&output_schema), + ) + .unwrap(); + Arc::new(DFSchema::try_from(packed.schema()).unwrap()) + } else { + row_schema.clone().to_dfschema_ref().unwrap() + }, format, ) .unwrap()