diff --git a/src/promql/src/engine/aggregate.rs b/src/promql/src/engine/aggregate.rs index 9fbf4b056c..a1516b3278 100644 --- a/src/promql/src/engine/aggregate.rs +++ b/src/promql/src/engine/aggregate.rs @@ -46,10 +46,12 @@ impl Engine { return fused::fused_range_agg( modifier, range_input, - shape.func.as_ref(), + shape.func, shape.op, &self.eval_ctx, - ); + self.ctx.query_ctx.timeout, + ) + .await; } let input = self.exec_expr(expr).await?; diff --git a/src/promql/src/fused/eval.rs b/src/promql/src/fused/eval.rs index d1c05206ab..c7d4acfc47 100644 --- a/src/promql/src/fused/eval.rs +++ b/src/promql/src/fused/eval.rs @@ -13,38 +13,38 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::time::Duration; +use std::sync::Arc; use config::meta::promql::{ NAME_LABEL, - value::{CounterSeries, EvalContext, ExtrapolationKind, RangeValue, Sample, Value}, + value::{EvalContext, Value}, }; use datafusion::error::{DataFusionError, Result}; use promql_parser::parser::LabelModifier; use rayon::prelude::*; -use super::{accumulator::FusedAccumulator, op::FusedAggOp}; +use super::{ + fold::{FoldParams, fold_partition, merge_folds, run_folds}, + op::FusedAggOp, +}; use crate::{ - aggregations::{group_series_by_labels, projected_labels}, - functions::{KEEP_METRIC_NAME_FUNC, RangeFunc, advance_sample_window}, - micros, + functions::{KEEP_METRIC_NAME_FUNC, RangeFunc}, + series_stream::matrix::matrix_sources, }; -/// Series per parallel fold chunk; smaller than the generic `AGG_PARALLEL_CHUNK` -/// because a fused chunk also pays the range-function evaluation. -const FUSED_PARALLEL_CHUNK: usize = 1024; - -/// Evaluates a range function and folds its values straight into aggregation -/// groups, skipping the generic path's per-series sample materialization. +/// Evaluates a range function over an already-materialized matrix and folds +/// its values straight into aggregation groups, through the same fold the +/// streaming path uses. /// /// The engine selects this only for the exact `agg(range_func(...))` shape; /// everything else stays on the generic evaluator, the correctness reference. -pub(crate) fn fused_range_agg( +pub(crate) async fn fused_range_agg( param: &Option, data: Value, - func: &dyn RangeFunc, + func: Arc, op: FusedAggOp, eval_ctx: &EvalContext, + timeout: u64, ) -> Result { let func_name = func.name(); let mut matrix = match data { @@ -78,127 +78,58 @@ pub(crate) fn fused_range_agg( }); } - let timestamps = eval_ctx.timestamps(); - let groups = group_series_by_labels(&matrix, param); - let counter_kind = func.counter_extrapolation(); - - let results = groups - .par_iter() - .filter_map(|(_, series_indices)| { - let labels = projected_labels(param, &matrix[series_indices[0]].labels); - // Series order first, timestamps second — the generic accumulation order. - let fold_chunk = |chunk: &[usize]| { - let mut acc = FusedAccumulator::new(op, timestamps.len()); - for &series_idx in chunk { - let metric = &matrix[series_idx]; - let range = metric - .time_window - .as_ref() - .expect("range function input must have a time window") - .range; - fold_series( - &mut acc, - &metric.samples, - range, - func, - counter_kind, - eval_ctx, - ×tamps, - ); - } - acc - }; - - // An ungrouped `sum(rate(...))` puts every series in one group; fold large - // groups in parallel chunks, merging partials in chunk order for determinism. - let acc = if series_indices.len() >= 2 * FUSED_PARALLEL_CHUNK { - let mut chunks = series_indices - .par_chunks(FUSED_PARALLEL_CHUNK) - .map(fold_chunk) - .collect::>() - .into_iter(); - let mut acc = chunks.next().expect("group has at least one chunk"); - for chunk in chunks { - acc.merge(chunk); - } - acc - } else { - fold_chunk(series_indices) - }; - - let samples = acc.into_samples(×tamps); - // The generic range function drops no-output series, and their groups with them. - if samples.is_empty() { - return None; - } - Some(RangeValue { - labels, - samples, - exemplars: None, - time_window: None, - }) + // the load applied one window to every series, so the first speaks for all + let range = matrix[0] + .time_window + .as_ref() + .expect("range function input must have a time window") + .range; + let params = Arc::new(FoldParams { + op, + func: func.clone(), + counter_kind: func.counter_extrapolation(), + range, + eval_ctx: eval_ctx.clone(), + timestamps: eval_ctx.timestamps(), + }); + let (matrix, sources) = matrix_sources(matrix, param, config::get_config().limit.cpu_num); + let folds = sources + .into_iter() + .map(|source| { + let params = params.clone(); + async move { fold_partition(source, params).await } }) .collect::>(); + let folds = run_folds(folds, timeout).await?; + let value = merge_folds( + folds.into_iter().map(|(groups, _)| groups).collect(), + ¶ms.timestamps, + ); - // Free the per-series allocations on the rayon pool; dropping them single-threaded is slow. - matrix.into_par_iter().for_each(drop); + if let Ok(matrix) = Arc::try_unwrap(matrix) { + // Free the per-series allocations on the rayon pool; dropping them single-threaded is slow. + matrix.into_par_iter().for_each(drop); + } log::info!( "[trace_id: {trace_id}] [PromQL Timing] fused {}({func_name}) completed in {:?}, folded {input_series} series into {} series", op.name(), start.elapsed(), - results.len() + match &value { + Value::Matrix(matrix) => matrix.len(), + _ => 0, + }, ); - if results.is_empty() { - return Ok(Value::None); - } - Ok(Value::Matrix(results)) -} - -/// Evaluates `func` over one series' time-ordered samples and pushes each -/// produced value into `acc` at its evaluation slot. -pub(super) fn fold_series( - acc: &mut FusedAccumulator, - samples: &[Sample], - range: Duration, - func: &dyn RangeFunc, - counter_kind: Option, - eval_ctx: &EvalContext, - timestamps: &[i64], -) { - let range_micros = micros(range); - let mut start_index = 0; - let mut end_index = 0; - let counter = CounterSeries::try_new(samples, counter_kind, eval_ctx, range_micros); - - for (slot, &eval_ts) in timestamps.iter().enumerate() { - let window_samples = advance_sample_window( - samples, - eval_ts - range_micros, - eval_ts, - &mut start_index, - &mut end_index, - ); - if window_samples.is_empty() { - continue; - } - let value = match &counter { - Some(counter) => counter.extrapolate(start_index, end_index, eval_ts, range), - None => func.exec(window_samples, eval_ts, &range), - }; - if let Some(value) = value { - acc.push(slot, value); - } - } + Ok(value) } #[cfg(test)] mod tests { use std::{sync::Arc, time::Duration}; - use config::meta::promql::value::{Label, Sample, TimeWindow}; + use config::meta::promql::value::{Label, RangeValue, Sample, TimeWindow}; use super::*; - use crate::{aggregations, functions}; + use crate::{aggregations, functions, series_stream::matrix::MATRIX_PARTITION_CHUNK}; type CanonicalSeries = (Vec<(String, String)>, Vec<(i64, u64)>); type GenericAgg = fn(&Option, Value, &EvalContext) -> Result; @@ -302,8 +233,19 @@ mod tests { })) } - #[test] - fn test_fused_range_agg_matches_generic_for_all_pairs() { + async fn run_fused( + modifier: &Option, + matrix: Vec, + func_name: &str, + op: FusedAggOp, + eval_ctx: &EvalContext, + ) -> Result { + let func: Arc = Arc::from(functions::fusable_range_func(func_name).unwrap()); + fused_range_agg(modifier, Value::Matrix(matrix), func, op, eval_ctx, 30).await + } + + #[tokio::test] + async fn test_fused_range_agg_matches_generic_for_all_pairs() { let agg_cases: [(FusedAggOp, GenericAgg); 8] = [ (FusedAggOp::Avg, aggregations::avg), (FusedAggOp::Count, aggregations::count), @@ -348,15 +290,9 @@ mod tests { generic_range(Value::Matrix(matrix.clone()), &eval_ctx).unwrap(); let expected = generic_agg(modifier, generic_input, &eval_ctx).unwrap(); - let func = functions::fusable_range_func(func_name).unwrap(); - let actual = fused_range_agg( - modifier, - Value::Matrix(matrix.clone()), - func.as_ref(), - op, - &eval_ctx, - ) - .unwrap(); + let actual = run_fused(modifier, matrix.clone(), func_name, op, &eval_ctx) + .await + .unwrap(); assert_eq!( canonical_matrix(expected), @@ -369,12 +305,12 @@ mod tests { } } - #[test] - fn test_fused_chunked_large_group_matches_generic_and_is_deterministic() { + #[tokio::test] + async fn test_fused_chunked_large_group_matches_generic_and_is_deterministic() { let eval_ctx = eval_ctx(); - // Ungrouped and per-`by(path)` folds both cross the chunk threshold; + // Ungrouped and per-`by(path)` folds both cross the partition threshold; // integer values keep float ops exact, so results must match generic bits. - let series_count = 2 * (2 * FUSED_PARALLEL_CHUNK) + 100; + let series_count = 2 * (2 * MATRIX_PARTITION_CHUNK) + 100; let matrix = (0..series_count) .map(|i| { let value = (i % 97) as f64; @@ -402,32 +338,30 @@ mod tests { (FusedAggOp::Stdvar, aggregations::stdvar), (FusedAggOp::Sum, aggregations::sum), ]; - let func = functions::fusable_range_func("sum_over_time").unwrap(); for (op, generic_agg) in agg_cases { for modifier in [None, by(&["path"])] { let generic_input = functions::sum_over_time(Value::Matrix(matrix.clone()), &eval_ctx).unwrap(); let expected = generic_agg(&modifier, generic_input, &eval_ctx).unwrap(); - let run = || { - fused_range_agg( - &modifier, - Value::Matrix(matrix.clone()), - func.as_ref(), - op, - &eval_ctx, - ) - .unwrap() - }; - let first = canonical_matrix(run()); + let first = canonical_matrix( + run_fused(&modifier, matrix.clone(), "sum_over_time", op, &eval_ctx) + .await + .unwrap(), + ); assert_eq!( canonical_matrix(expected), first, "chunked fused {}(sum_over_time) diverged from generic (modifier: {modifier:?})", op.name(), ); + let second = canonical_matrix( + run_fused(&modifier, matrix.clone(), "sum_over_time", op, &eval_ctx) + .await + .unwrap(), + ); assert_eq!( first, - canonical_matrix(run()), + second, "chunked fused {}(sum_over_time) must be deterministic", op.name(), ); @@ -435,13 +369,13 @@ mod tests { } } - #[test] - fn test_fused_chunked_float_values_stay_within_epsilon_of_generic() { + #[tokio::test] + async fn test_fused_chunked_float_values_stay_within_epsilon_of_generic() { let eval_ctx = eval_ctx(); - // In [chunk threshold, 65536) fused folds in chunks while generic is - // sequential; Kahan is non-associative, so fractional sums may drift + // Above the partition threshold fused folds in partitions while generic + // is sequential; Kahan is non-associative, so fractional sums may drift // in the last bits but never materially. - let series_count = 2 * FUSED_PARALLEL_CHUNK + 500; + let series_count = 2 * MATRIX_PARTITION_CHUNK + 500; let matrix = (0..series_count) .map(|i| { let base = (i % 97) as f64 * 0.1; @@ -464,16 +398,10 @@ mod tests { ] { let generic_input = functions::rate(Value::Matrix(matrix.clone()), &eval_ctx).unwrap(); let expected = canonical_matrix(generic_agg(&None, generic_input, &eval_ctx).unwrap()); - let func = functions::fusable_range_func("rate").unwrap(); let actual = canonical_matrix( - fused_range_agg( - &None, - Value::Matrix(matrix.clone()), - func.as_ref(), - op, - &eval_ctx, - ) - .unwrap(), + run_fused(&None, matrix.clone(), "rate", op, &eval_ctx) + .await + .unwrap(), ); assert_eq!(expected.len(), actual.len()); @@ -498,36 +426,42 @@ mod tests { } } - #[test] - fn test_fused_range_agg_none_and_invalid_input() { + #[tokio::test] + async fn test_fused_range_agg_none_and_invalid_input() { let eval_ctx = eval_ctx(); - let func = functions::fusable_range_func("rate").unwrap(); + let func: Arc = Arc::from(functions::fusable_range_func("rate").unwrap()); let result = fused_range_agg( &None, Value::None, - func.as_ref(), + func.clone(), FusedAggOp::Sum, &eval_ctx, + 30, ) + .await .unwrap(); assert!(matches!(result, Value::None)); let result = fused_range_agg( &None, Value::Float(1.0), - func.as_ref(), + func.clone(), FusedAggOp::Sum, &eval_ctx, - ); + 30, + ) + .await; assert!(result.is_err()); let result = fused_range_agg( &None, Value::Matrix(vec![]), - func.as_ref(), + func, FusedAggOp::Sum, &eval_ctx, + 30, ) + .await .unwrap(); assert!(matches!(result, Value::None)); } diff --git a/src/promql/src/fused/fold.rs b/src/promql/src/fused/fold.rs new file mode 100644 index 0000000000..795641628c --- /dev/null +++ b/src/promql/src/fused/fold.rs @@ -0,0 +1,191 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! The single fused fold: every series source partition folds its series into +//! per-group accumulators, and the partitions merge in order at the end. The +//! producers only decide how series arrive; the aggregation lives here once. + +use std::{sync::Arc, time::Duration}; + +use config::meta::promql::value::{ + CounterSeries, EvalContext, ExtrapolationKind, Labels, RangeValue, Sample, Value, +}; +use datafusion::error::{DataFusionError, Result}; +use hashbrown::{HashMap, hash_map::Entry}; +use infra::errors::{Error, ErrorCodes}; + +use super::{accumulator::FusedAccumulator, op::FusedAggOp}; +use crate::{ + functions::{RangeFunc, advance_sample_window}, + micros, + series_stream::SeriesSource, +}; + +pub(super) type GroupAccs = HashMap; + +pub(super) struct FoldParams { + pub(super) op: FusedAggOp, + pub(super) func: Arc, + pub(super) counter_kind: Option, + pub(super) range: Duration, + pub(super) eval_ctx: EvalContext, + pub(super) timestamps: Vec, +} + +pub(super) struct GroupEntry { + labels: Labels, + acc: FusedAccumulator, +} + +/// Folds one partition's series into its group accumulators, dropping each +/// series as soon as it is folded. +pub(super) async fn fold_partition( + mut source: S, + params: Arc, +) -> Result<(GroupAccs, usize)> { + let mut groups = GroupAccs::new(); + let mut series_count = 0; + while let Some(sig) = source.advance().await? { + let entry = match groups.entry(sig) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(GroupEntry { + labels: source.labels(), + acc: FusedAccumulator::new(params.op, params.timestamps.len()), + }), + }; + let samples = source.consume().await?; + fold_series( + &mut entry.acc, + samples, + params.range, + params.func.as_ref(), + params.counter_kind, + ¶ms.eval_ctx, + ¶ms.timestamps, + ); + series_count += 1; + } + Ok((groups, series_count)) +} + +/// Evaluates `func` over one series' time-ordered samples and pushes each +/// produced value into `acc` at its evaluation slot. +fn fold_series( + acc: &mut FusedAccumulator, + samples: &[Sample], + range: Duration, + func: &dyn RangeFunc, + counter_kind: Option, + eval_ctx: &EvalContext, + timestamps: &[i64], +) { + let range_micros = micros(range); + let mut start_index = 0; + let mut end_index = 0; + let counter = CounterSeries::try_new(samples, counter_kind, eval_ctx, range_micros); + + for (slot, &eval_ts) in timestamps.iter().enumerate() { + let window_samples = advance_sample_window( + samples, + eval_ts - range_micros, + eval_ts, + &mut start_index, + &mut end_index, + ); + if window_samples.is_empty() { + continue; + } + let value = match &counter { + Some(counter) => counter.extrapolate(start_index, end_index, eval_ts, range), + None => func.exec(window_samples, eval_ts, &range), + }; + if let Some(value) = value { + acc.push(slot, value); + } + } +} + +/// Aborts the partition folds when `timeout` elapses before they all finish. +pub(super) async fn run_folds(folds: Vec, timeout: u64) -> Result> +where + Fut: Future> + Send + 'static, +{ + let mut tasks = Vec::with_capacity(folds.len()); + let mut abort_handles = Vec::with_capacity(folds.len()); + for fold in folds { + let task = tokio::task::spawn(fold); + abort_handles.push(task.abort_handle()); + tasks.push(task); + } + tokio::select! { + joined = futures::future::try_join_all(tasks) => { + joined + .map_err(|e| DataFusionError::Execution(e.to_string()))? + .into_iter() + .collect() + } + _ = tokio::time::sleep(Duration::from_secs(timeout)) => { + for handle in abort_handles { + handle.abort(); + } + Err(DataFusionError::Plan( + Error::ErrorCode(ErrorCodes::SearchTimeout( + "[PromQL] fused agg timeout".to_string(), + )) + .to_string(), + )) + } + } +} + +/// Merges the partition-local groups in partition order and materializes the +/// result; groups whose accumulator produced no samples are dropped like the +/// generic path drops no-output series. +pub(super) fn merge_folds(folds: Vec, timestamps: &[i64]) -> Value { + let mut folds = folds.into_iter(); + let Some(mut merged) = folds.next() else { + return Value::None; + }; + for fold in folds { + for (sig, entry) in fold { + match merged.entry(sig) { + Entry::Occupied(mut occupied) => occupied.get_mut().acc.merge(entry.acc), + Entry::Vacant(vacant) => { + vacant.insert(entry); + } + } + } + } + let results: Vec = merged + .into_values() + .filter_map(|entry| { + let samples = entry.acc.into_samples(timestamps); + if samples.is_empty() { + return None; + } + Some(RangeValue { + labels: entry.labels, + samples, + exemplars: None, + time_window: None, + }) + }) + .collect(); + if results.is_empty() { + Value::None + } else { + Value::Matrix(results) + } +} diff --git a/src/promql/src/fused/mod.rs b/src/promql/src/fused/mod.rs index 0cdea306ed..157cb74f6d 100644 --- a/src/promql/src/fused/mod.rs +++ b/src/promql/src/fused/mod.rs @@ -19,6 +19,7 @@ mod accumulator; mod eval; +mod fold; mod op; pub(crate) use eval::fused_range_agg; diff --git a/src/promql/src/lib.rs b/src/promql/src/lib.rs index 8b9d121abb..df60ffb318 100644 --- a/src/promql/src/lib.rs +++ b/src/promql/src/lib.rs @@ -40,6 +40,7 @@ mod functions; mod fused; pub mod load_series; pub mod promql; +mod series_stream; pub mod utils; pub const DEFAULT_LOOKBACK: Duration = Duration::from_secs(300); // 5m diff --git a/src/promql/src/series_stream/matrix.rs b/src/promql/src/series_stream/matrix.rs new file mode 100644 index 0000000000..2fd3db397d --- /dev/null +++ b/src/promql/src/series_stream/matrix.rs @@ -0,0 +1,96 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! The matrix producer: an already-materialized matrix served through the +//! series-source contract, so layouts that cannot stream still share the +//! streaming consumers. + +use std::sync::Arc; + +use config::meta::promql::value::{Labels, RangeValue, Sample}; +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use super::SeriesSource; +use crate::aggregations::{group_series_by_labels, projected_labels}; + +/// Series per partition; matches the fused fold's historical chunk size. +pub(crate) const MATRIX_PARTITION_CHUNK: usize = 1024; + +/// One partition of a materialized matrix, iterated group-contiguously. +pub(crate) struct MatrixSource { + matrix: Arc>, + modifier: Option, + /// `(group signature, matrix index)` of this partition's series. + order: Vec<(u64, u32)>, + pos: usize, +} + +impl SeriesSource for MatrixSource { + async fn advance(&mut self) -> Result> { + Ok(self.order.get(self.pos).map(|&(sig, _)| sig)) + } + + fn labels(&self) -> Labels { + let (_, index) = self.order[self.pos]; + projected_labels(&self.modifier, &self.matrix[index as usize].labels) + } + + async fn consume(&mut self) -> Result<&[Sample]> { + let (_, index) = self.order[self.pos]; + self.pos += 1; + Ok(&self.matrix[index as usize].samples) + } +} + +/// Splits a matrix into group-contiguous partitions behind shared ownership; +/// partition boundaries depend only on the series count, so folds merge +/// deterministically. +pub(crate) fn matrix_sources( + matrix: Vec, + modifier: &Option, + max_partitions: usize, +) -> (Arc>, Vec) { + let mut groups: Vec<(u64, Vec)> = group_series_by_labels(&matrix, modifier) + .into_iter() + .collect(); + groups.sort_unstable_by_key(|(sig, _)| *sig); + let order: Vec<(u64, u32)> = groups + .into_iter() + .flat_map(|(sig, indices)| indices.into_iter().map(move |index| (sig, index as u32))) + .collect(); + + // small folds stay sequential, keeping them bit-identical to the generic path + let partitions = if order.len() < 2 * MATRIX_PARTITION_CHUNK { + 1 + } else { + max_partitions + .max(1) + .min(order.len().div_ceil(MATRIX_PARTITION_CHUNK)) + }; + let chunk = order.len().div_ceil(partitions).max(1); + + let matrix = Arc::new(matrix); + let sources = order + .chunks(chunk) + .map(|chunk| MatrixSource { + matrix: matrix.clone(), + modifier: modifier.clone(), + order: chunk.to_vec(), + pos: 0, + }) + .collect(); + (matrix, sources) +} diff --git a/src/promql/src/series_stream/mod.rs b/src/promql/src/series_stream/mod.rs new file mode 100644 index 0000000000..daf14adc6d --- /dev/null +++ b/src/promql/src/series_stream/mod.rs @@ -0,0 +1,36 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Sources that deliver a query's series one at a time, so a consumer can +//! evaluate and drop each series without holding the full set: the matrix +//! producer adapts an already-materialized matrix, and ordered-scan producers +//! plug in behind the same contract. + +pub(crate) mod matrix; + +use config::meta::promql::value::{Labels, Sample}; +use datafusion::error::Result; + +/// One partition's series, delivered whole and grouped under a signature. +/// +/// Protocol per series: `advance` yields the group signature, `labels` may be +/// read while the series is current, `consume` takes its samples and moves on. +pub(crate) trait SeriesSource: Send { + fn advance(&mut self) -> impl Future>> + Send; + /// Group labels of the current series; valid only before `consume`. + fn labels(&self) -> Labels; + /// Time-ordered samples of the current series. + fn consume(&mut self) -> impl Future> + Send; +}