diff --git a/src/promql/src/extension_plan/aggregate_finalize.rs b/src/promql/src/extension_plan/aggregate_finalize.rs index 0cc2dfc802e..7540b46e147 100644 --- a/src/promql/src/extension_plan/aggregate_finalize.rs +++ b/src/promql/src/extension_plan/aggregate_finalize.rs @@ -501,6 +501,17 @@ fn split(node: &NativeNode, fields: &[FieldSpec], time_index: &str) -> DfResult< _ => Err(declined("an aggregate under a set operation")), } } + // The suffix above a split aggregate reads the merged state alone, so it has no source for + // the info series. + NativeExpr::Info { base, info, .. } => { + match ( + split(base, fields, time_index)?, + split(info, fields, time_index)?, + ) { + (None, None) => Ok(None), + _ => Err(declined("an aggregate under info")), + } + } } } @@ -527,9 +538,13 @@ fn prefix_fields(tree: &NativeTree) -> DfResult> { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => raw_fields(input), - NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { - raw_fields(left).or_else(|| raw_fields(right)) - } + NativeExpr::Binary { left, right, .. } + | NativeExpr::Set { left, right, .. } + | NativeExpr::Info { + base: left, + info: right, + .. + } => raw_fields(left).or_else(|| raw_fields(right)), } } diff --git a/src/promql/src/extension_plan/aggregate_finalize/tests.rs b/src/promql/src/extension_plan/aggregate_finalize/tests.rs index 329a8ea0467..c051f276678 100644 --- a/src/promql/src/extension_plan/aggregate_finalize/tests.rs +++ b/src/promql/src/extension_plan/aggregate_finalize/tests.rs @@ -91,7 +91,7 @@ fn instant() -> NativeNode { lookback_ms: 300_000, offset_ms: 0, at_ms: None, - timestamp: false, + value: crate::native::plan::InstantValue::Sample, output_grid: grid(), }, }, @@ -619,7 +619,13 @@ fn the_prefix_and_the_suffix_each_name_their_own_source_zero() { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => inputs(input, found), - NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { + NativeExpr::Binary { left, right, .. } + | NativeExpr::Set { left, right, .. } + | NativeExpr::Info { + base: left, + info: right, + .. + } => { inputs(left, found); inputs(right, found); } diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index 0b7fba4632c..cc593f84263 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -655,7 +655,13 @@ fn tree_inputs(tree: &NativeTree) -> Vec<&crate::native::plan::BoundInput> { | NativeExpr::Timestamp { input } | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => collect(input, inputs), - NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { + NativeExpr::Binary { left, right, .. } + | NativeExpr::Set { left, right, .. } + | NativeExpr::Info { + base: left, + info: right, + .. + } => { collect(left, inputs); collect(right, inputs); } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 5e877cf94ed..33d75bcc41f 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -89,7 +89,7 @@ use crate::map_label::LabelRewrite; use crate::native::execute::{InputId, NativeExecutionState, NativeExecutor}; use crate::native::plan::{ AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundWindow, Grouping, MapOp, NativeExpr, NativeNode, NativeTree, NativeType, + BoundWindow, Grouping, InstantValue, MapOp, NativeExpr, NativeNode, NativeTree, NativeType, }; use crate::range_eval::{RangeFunction, RangeWindow}; use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels}; @@ -249,7 +249,7 @@ fn binary_tree_over_sources( lookback_ms: 0, offset_ms: 0, at_ms: None, - timestamp: false, + value: InstantValue::Sample, output_grid: Arc::clone(&grid), }, }, @@ -269,7 +269,7 @@ fn binary_tree_over_sources( lookback_ms: 0, offset_ms: 0, at_ms: None, - timestamp: false, + value: InstantValue::Sample, output_grid: Arc::clone(&grid), }, }, @@ -321,7 +321,7 @@ fn instant_of(source_index: usize, grid: &Arc) -> NativeNode { lookback_ms: 0, offset_ms: 0, at_ms: None, - timestamp: false, + value: InstantValue::Sample, output_grid: Arc::clone(grid), }, }, @@ -362,7 +362,13 @@ fn declaring_orderings(tree: &NativeTree, orderings: &[SeriesOrdering]) -> Nativ | NativeExpr::Vector { input, .. } => { **input = rewrite(input, orderings, seen); } - NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => { + NativeExpr::Binary { left, right, .. } + | NativeExpr::Set { left, right, .. } + | NativeExpr::Info { + base: left, + info: right, + .. + } => { **left = rewrite(left, orderings, seen); **right = rewrite(right, orderings, seen); } 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 3e29af7e53e..ee7f59402e2 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 @@ -15,7 +15,7 @@ use super::*; use crate::map::MathOp; use crate::map_label::LabelRewrite; -use crate::native::plan::{BoundMap, MapOp}; +use crate::native::plan::{BoundMap, InstantValue, MapOp}; fn instant_tree(schema: Arc, grid: Arc) -> NativeTree { let output = NativeType::Samples(Arc::new(SeriesSchema { @@ -36,7 +36,7 @@ fn instant_tree(schema: Arc, grid: Arc) -> NativeTree { lookback_ms: 0, offset_ms: 0, at_ms: None, - timestamp: false, + value: InstantValue::Sample, output_grid: grid, }, }, diff --git a/src/promql/src/info.rs b/src/promql/src/info.rs index b0210416e6d..bfbafd92c8e 100644 --- a/src/promql/src/info.rs +++ b/src/promql/src/info.rs @@ -29,14 +29,15 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; -use datatypes::series::SeriesLabels; use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; +use datatypes::series::{Label, SeriesLabels}; use regex::Regex; use crate::grid_series::{ GridSample, GridSeriesBuilder, LabelSet, external, label_set, owned, render_label_set, }; use crate::select::MatchOp; +use crate::vector_match::{Budget, HISTOGRAM_ROW_BYTES, owned_histogram_bytes}; /// The labels an info series is joined on, as in Prometheus. const IDENTIFYING_LABELS: [&str; 2] = ["instance", "job"]; @@ -70,6 +71,33 @@ pub fn info( base: &[Series], info: &[Series], selector: &InfoSelector, +) -> Result> { + info_inner(schema, base, info, selector, None) +} + +/// [`info()`], handing the live workspace it holds to `reserve` before it allocates it. +/// +/// What the index of the info series and the rows of the results keep is held until the operation +/// returns; the merge of one label set's rows is a phase, and the builder's columns join the phase +/// of the output they are finished as. A total `reserve` refuses ends the operation with that error. +pub(crate) fn info_with_reservation( + schema: &Arc, + base: &[Series], + info: &[Series], + selector: &InfoSelector, + reserve: &mut dyn FnMut(usize) -> Result<()>, +) -> Result> { + let mut budget = Budget::new(reserve); + info_inner(schema, base, info, selector, Some(&mut budget)) +} + +/// The one info kernel, whether or not a caller accounts for its workspace. +fn info_inner( + schema: &Arc, + base: &[Series], + info: &[Series], + selector: &InfoSelector, + mut budget: Option<&mut Budget<'_>>, ) -> Result> { let grid = schema .eval_grid @@ -96,13 +124,21 @@ pub fn info( } let name_matchers = effective_name_matchers(&selector.name_matchers)?; let is_info_metric = |name: &str| name_matchers.iter().all(|matcher| matcher.matches(name)); + let width = schema.fields.len(); + // The presence of the base at each point and the identifying values of every base series + // stay live until the operation returns. + if let Some(budget) = budget.as_deref_mut() { + budget.hold(points.saturating_add(base.len().saturating_mul(PATTERN_BYTES)))?; + } let mut base_present = vec![false; points]; - for series in base { + let mut sampled = vec![false; base.len()]; + for (series, sampled) in base.iter().zip(&mut sampled) { for field in series.fields.iter() { for ((point, _), sample) in series.time.positions().zip(field.iter()) { if !matches!(sample, SampleRef::Missing) { base_present[point] = true; + *sampled = true; } } } @@ -114,8 +150,10 @@ pub fn info( // base series it enriches: a present label matches any value the pattern's base series have, // an absent one matches only absence. An info series no pattern admits is never fetched, so // it neither enriches nor conflicts. + // A base series without a sample is not in Prometheus' base matrix at all, so it selects no + // info series. let mut patterns: HashMap, Vec>>> = HashMap::new(); - for series in base { + for (series, _) in base.iter().zip(&sampled).filter(|(_, sampled)| **sampled) { if is_info_metric(series.metric.name().unwrap_or("")) { continue; } @@ -146,6 +184,17 @@ pub fn info( if !fetched { continue; } + // The entry of a signature and the newest series it keeps at each point the series + // holds a sample, charged before they are allocated. + if let Some(budget) = budget.as_deref_mut() { + budget.hold( + series + .time + .len() + .saturating_mul(KEPT_POINT_BYTES) + .saturating_add(SIGNATURE_BYTES), + )?; + } let name: Arc = Arc::from(series.metric.name().unwrap_or("")); info_names.insert(name.clone()); let kept = by_signature.entry((name, values)).or_default(); @@ -185,7 +234,6 @@ pub fn info( } } - let width = schema.fields.len(); let mut result_index: HashMap = HashMap::new(); let mut results: Vec<(SeriesLabels, Vec)> = Vec::new(); let mut added: Vec<(Arc, Arc)> = Vec::new(); @@ -205,13 +253,21 @@ pub fn info( }; // The result of the previous point, reused while the same info series are kept. let mut previous: Option = None; + // A row of this series is kept by the result it goes to, so every row it can have is held + // before the walk; a histogram is held as it is copied. + if let Some(budget) = budget.as_deref_mut() { + budget.hold(series.time.len().saturating_mul(row_bytes(width)))?; + } let mut cursors: Vec<_> = series.fields.iter().map(|field| field.cursor()).collect(); for (point, row) in series.time.positions() { let mut values = Vec::with_capacity(width); for cursor in &mut cursors { - values.push(owned( - cursor.at(row).expect("the axis and the fields are aligned"), - )?); + let sample = cursor.at(row).expect("the axis and the fields are aligned"); + if let (Some(budget), SampleRef::Histogram(view)) = (budget.as_deref_mut(), &sample) + { + budget.hold(owned_histogram_bytes(view))?; + } + values.push(owned(sample)?); } if values.iter().all(Option::is_none) { continue; @@ -246,6 +302,9 @@ pub fn info( Some(match result_index.entry(label_set(&metric)) { std::collections::hash_map::Entry::Occupied(entry) => *entry.get(), std::collections::hash_map::Entry::Vacant(entry) => { + if let Some(budget) = budget.as_deref_mut() { + budget.hold(result_set_bytes(&metric))?; + } results.push((metric, Vec::new())); *entry.insert(results.len() - 1) } @@ -270,6 +329,10 @@ pub fn info( // Rows of one label set arrive per base series; sorted by point, a field held twice at one // point is a collision. for (metric, rows) in &mut results { + // The sort's buffer and the merged rows, while the rows they replace are still alive. + if let Some(budget) = budget.as_deref_mut() { + budget.phase(rows.len().saturating_mul(2 * std::mem::size_of::()))?; + } rows.sort_by_key(|(point, _)| *point); let mut merged: Vec = Vec::with_capacity(rows.len()); for (point, values) in rows.drain(..) { @@ -301,8 +364,17 @@ pub fn info( let mut builder = GridSeriesBuilder::try_new(schema.clone())?; let mut output = Vec::with_capacity(results.len()); + if let Some(budget) = budget.as_deref_mut() { + budget.phase(0)?; + } for (metric, rows) in results { builder.begin_series(); + // This series' columns, reserved before its first row. They stay live inside the series + // it is finished as, so the phase keeps them. + if let Some(budget) = budget.as_deref_mut() { + budget.grow(column_bytes(&rows, width))?; + } + builder.reserve_rows(rows.len())?; for (point, row) in rows { builder.push_samples(point, &row)?; } @@ -314,6 +386,66 @@ pub fn info( /// The samples of one result series at one grid point, one per field. type Row = (usize, Vec>); +/// The identifying values one base series adds to its presence pattern, charged twice for the +/// spare room of the sets they go to. +const PATTERN_BYTES: usize = 2 * IDENTIFYING_LABELS.len() * std::mem::size_of::>(); + +/// The entry of one info series' signature: its key, its map of kept series, and the name the +/// info metric list keeps, charged twice for the spare room of the tables. +const SIGNATURE_BYTES: usize = 2 + * (std::mem::size_of::<(Signature, BTreeMap)>() + + IDENTIFYING_LABELS.len() * std::mem::size_of::>>() + + std::mem::size_of::>()); + +/// The series one signature keeps at one point, charged twice for the B-tree's spare room. +const KEPT_POINT_BYTES: usize = 2 * std::mem::size_of::<(usize, (f64, usize))>(); + +/// One row a result keeps: its entry in the result's row list, charged twice for the list's +/// growth, and its samples, allocated at their exact count. +fn row_bytes(width: usize) -> usize { + (2 * std::mem::size_of::()) + .saturating_add(width.saturating_mul(std::mem::size_of::>())) +} + +/// The bytes a result label set keeps: its metric, rebuilt by each label it is given, the label set +/// it is filed under with the copy of the name that set makes, and its entries in the result list +/// and index. +fn result_set_bytes(metric: &SeriesLabels) -> usize { + let labels = metric.tags().len().saturating_add(1); + labels + .saturating_mul( + 2 * std::mem::size_of::