diff --git a/src/promql/src/aggregate.rs b/src/promql/src/aggregate.rs index 228f67ba4bc..0ab1e26a291 100644 --- a/src/promql/src/aggregate.rs +++ b/src/promql/src/aggregate.rs @@ -27,11 +27,14 @@ pub use state::{ use std::cmp::Reverse; use std::collections::hash_map::Entry; use std::collections::{BTreeMap, BinaryHeap, HashMap}; +use std::mem::size_of; use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; +use datatypes::arrow::array::Float64Array; use datatypes::series::MetricName; -use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; +use datatypes::series::samples::{SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::{GridSpan, TimeAxis}; use crate::go_format::format_float_f; use crate::grid_series::{GridSeriesBuilder, external}; @@ -97,6 +100,7 @@ pub struct FloatAggregator { labels: Vec>, group_index: HashMap, groups: Vec, + retained: usize, } impl FloatAggregator { @@ -104,25 +108,78 @@ impl FloatAggregator { schema: &Arc, op: AggregateOp, grouping: Grouping<'_>, + ) -> Result { + Self::try_new_budgeted(schema, op, grouping, |_| Ok(())) + } + + pub(crate) fn try_new_budgeted( + schema: &Arc, + op: AggregateOp, + grouping: Grouping<'_>, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result { validate_input(schema, &[], grouping)?; let (by, labels) = match grouping { Grouping::By(labels) => (true, labels), Grouping::Without(labels) => (false, labels), }; + let label_storage = + grouping_bytes(grouping).saturating_add(labels_capacity_bytes(grouping)); + budget(label_storage)?; + let mut owned_labels: Vec> = Vec::new(); + owned_labels + .try_reserve_exact(labels.len()) + .map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate grouping labels: {error}" + )) + })?; + owned_labels.extend(labels.iter().map(|label| Arc::from(*label))); + let retained = owned_labels + .capacity() + .saturating_mul(size_of::>()) + .saturating_add(owned_labels.iter().fold(0usize, |bytes, label| { + bytes.saturating_add(label.len() + 2 * size_of::()) + })); Ok(Self { schema: schema.clone(), op, by, - labels: labels.iter().map(|label| Arc::from(*label)).collect(), + labels: owned_labels, group_index: HashMap::new(), groups: Vec::new(), + retained, }) } /// Folds `series` into its group. A series that fails validation leaves the state unchanged. + /// + /// A budget refusal is not a validation failure: the fold may have opened a group or filled + /// blocks before the allocation it could not admit, so a failed [`Self::update_budgeted`] + /// leaves the state half folded and its retained figure stale. The runtime stops the execution + /// and drops the aggregator on that error, which is why no rollback is kept; a caller holding + /// one must drop it rather than fold into it again. pub fn update(&mut self, series: &Series) -> Result<()> { + self.update_budgeted(series, |_| Ok(())) + } + + pub(crate) fn update_budgeted( + &mut self, + series: &Series, + mut budget: impl FnMut(usize) -> Result<()>, + ) -> Result<()> { validate_float_series(&self.schema, series)?; + let existing_bytes = self.state_bytes(); + let grouping_workspace = grouping_workspace_bytes(&series.metric); + let mut workspace_bytes = + grouping_workspace.saturating_add(pending_bytes(self.schema.fields.len())); + budget(existing_bytes.saturating_add(workspace_bytes))?; + workspace_bytes = workspace_bytes.saturating_add( + self.labels + .len() + .saturating_mul(std::mem::size_of::<&str>()), + ); + budget(existing_bytes.saturating_add(workspace_bytes))?; let labels: Vec<&str> = self.labels.iter().map(|label| &**label).collect(); let grouping = if self.by { Grouping::By(&labels) @@ -130,39 +187,210 @@ impl FloatAggregator { Grouping::Without(&labels) }; let metric = grouping.apply(&series.metric)?; - let index = match self.group_index.entry(label_key(&metric)) { + let key_bytes = metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()); + workspace_bytes = workspace_bytes + .saturating_add(grouping_workspace_bytes(&metric)) + .saturating_add(key_bytes); + budget(existing_bytes.saturating_add(workspace_bytes))?; + let key = label_key(&metric); + let key_retained = key + .capacity() + .saturating_mul(size_of::<(Arc, Arc)>()); + let new_group = !self.group_index.contains_key(&key); + let group_capacity = self.groups.capacity(); + let map_capacity = self.group_index.capacity(); + let old_map_backing = hash_map_backing_bytes(map_capacity); + if new_group { + let next_group_capacity = if self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + let group_realloc_bytes = if next_group_capacity > group_capacity { + next_group_capacity.saturating_mul(size_of::()) + } else { + 0 + }; + budget( + existing_bytes + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_bytes) + // A realloc that cannot grow in place keeps both buffers alive while it + // copies, so a growing step admits the old buffer beside the new one -- + // the same reason the map below admits both tables. The old buffer alone + // is already in `existing_bytes`; what is added here is the new one. + .saturating_add(group_realloc_bytes) + .saturating_add(if self.group_index.len() == map_capacity { + hash_map_backing_bytes(next_hash_capacity(map_capacity)) + } else { + 0 + }) + .saturating_add(workspace_bytes), + )?; + } + let group_capacity_growth = if new_group && self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + if new_group && group_capacity_growth > group_capacity { + self.groups.reserve(1); + } + if new_group && self.group_index.len() == self.group_index.capacity() { + self.group_index.reserve(1); + } + let index = match self.group_index.entry(key) { Entry::Occupied(entry) => *entry.get(), Entry::Vacant(entry) => { self.groups.push(Group { - metric, + metric: metric.clone(), blocks: BTreeMap::new(), }); - *entry.insert(self.groups.len() - 1) + let index = self.groups.len() - 1; + entry.insert(index); + index } }; - self.groups[index].accumulate(series, self.op); + if new_group { + self.retained = self + .retained + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_retained) + .saturating_add( + self.groups + .capacity() + .saturating_sub(group_capacity) + .saturating_mul(size_of::()), + ) + .saturating_add( + hash_map_backing_bytes(self.group_index.capacity()) + .saturating_sub(old_map_backing), + ); + } + drop(labels); + drop(metric); + self.retained = self.groups[index].accumulate( + series, + self.op, + &mut budget, + self.retained, + self.schema.fields.len(), + )?; Ok(()) } + /// Number of groups currently retained by this fold. + pub(crate) fn groups(&self) -> usize { + self.groups.len() + } + + /// Bytes reserved for this fold, using bounded allowances for its maps and group payloads. + pub(crate) fn state_bytes(&self) -> usize { + self.retained + } + + #[cfg(test)] + pub(crate) fn recounted_state_bytes(&self) -> usize { + let groups = self.groups.iter().fold(0usize, |total, group| { + total + .saturating_add(metric_bytes(&group.metric)) + .saturating_add( + group + .blocks + .len() + .saturating_mul(BLOCK_POINTS) + .saturating_mul(self.schema.fields.len()) + .saturating_mul(std::mem::size_of::()), + ) + .saturating_add(btree_backing_bytes(group.blocks.len())) + }); + groups + .saturating_add( + self.labels + .capacity() + .saturating_mul(std::mem::size_of::>()), + ) + .saturating_add(hash_map_backing_bytes(self.group_index.capacity())) + .saturating_add(self.group_index.keys().fold(0usize, |bytes, key| { + bytes.saturating_add( + key.capacity() + .saturating_mul(size_of::<(Arc, Arc)>()), + ) + })) + .saturating_add( + self.labels + .iter() + .map(|label| label.len() + 2 * std::mem::size_of::()) + .sum::(), + ) + .saturating_add( + self.groups + .capacity() + .saturating_mul(std::mem::size_of::()), + ) + } + /// The result of every group, in the order the groups first appeared. pub fn finish(self) -> Result> { - finish_groups(self.schema, self.op, self.groups) + self.finish_budgeted(|_| Ok(())) + } + + pub(crate) fn finish_budgeted( + self, + budget: impl FnMut(usize) -> Result<()>, + ) -> Result> { + let state = self.state_bytes(); + finish_groups_budgeted(self.schema, self.op, self.groups, state, budget) } } -/// Turns the state of `groups` into result series, in group order. -fn finish_groups( +/// Turns the state of `groups` into result series, in group order, admitting as it builds. +fn finish_groups_budgeted( schema: Arc, op: AggregateOp, groups: Vec, + state: usize, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result> { let width = schema.fields.len(); + let group_count = groups.len(); + if group_count == 0 { + return Ok(Vec::new()); + } + let groups = groups.into_iter(); let mut builder = GridSeriesBuilder::try_new(schema)?; - let mut output = Vec::with_capacity(groups.len()); - let mut row = Vec::with_capacity(width); + let output_capacity_bytes = group_count.saturating_mul(std::mem::size_of::()); + let row_capacity_bytes = width.saturating_mul(std::mem::size_of::>()); + let base = state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes); + budget(base)?; + let mut output = Vec::new(); + output.try_reserve_exact(group_count).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate outputs: {error}")) + })?; + let mut row = Vec::new(); + row.try_reserve_exact(width).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate row: {error}")) + })?; + let mut completed_output_bytes = 0usize; for group in groups { // One builder serves every group, so each result series starts explicitly. builder.begin_series(); + let (rows, runs) = group_rows_runs(&group, width); + let finish_bytes = finish_working_bytes(rows, runs, width); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes) + .saturating_add(finish_bytes), + )?; + builder.reserve_rows(rows)?; for (&block, states) in &group.blocks { // The last block can extend past the grid; only points with a sample are pushed. for (offset, point) in states.chunks_exact(width).enumerate() { @@ -174,11 +402,89 @@ fn finish_groups( builder.push(block * BLOCK_POINTS + offset, &row)?; } } - output.extend(builder.finish(group.metric)?); + let finished = builder.finish(group.metric)?; + completed_output_bytes = completed_output_bytes.saturating_add( + finished + .iter() + .map(crate::grid_series::series_bytes) + .sum::(), + ); + output.extend(finished); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes), + )?; } Ok(output) } +/// The result of every group, without a budget: the entry the partial-state merger takes, which +/// does not hold a query pool to admit against. +fn finish_groups( + schema: Arc, + op: AggregateOp, + groups: Vec, +) -> Result> { + finish_groups_budgeted(schema, op, groups, 0, |_| Ok(())) +} + +fn group_rows_runs(group: &Group, width: usize) -> (usize, usize) { + if width == 0 { + return (0, 0); + } + let mut rows = 0usize; + let mut runs = 0usize; + let mut previous: Option = None; + for (&block, states) in &group.blocks { + for (offset, point) in states.chunks_exact(width).enumerate() { + if point.iter().any(|state| state.seen) { + let index = block.saturating_mul(BLOCK_POINTS).saturating_add(offset); + rows = rows.saturating_add(1); + if previous.is_none_or(|previous| previous.checked_add(1) != Some(index)) { + runs = runs.saturating_add(1); + } + previous = Some(index); + } + } + } + (rows, runs) +} + +fn finish_working_bytes(rows: usize, runs: usize, width: usize) -> usize { + let values = rows + .saturating_mul(width) + .saturating_mul(std::mem::size_of::()); + let validity_bytes = rows + .saturating_add(7) + .checked_div(8) + .unwrap_or(0) + .div_ceil(64) + .saturating_mul(64); + let fields = crate::grid_series::field_builder_bytes(width).saturating_add( + width.saturating_mul( + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(std::mem::size_of::()), + ), + ); + let indices = rows.saturating_mul(std::mem::size_of::()); + let spans = runs.saturating_mul(std::mem::size_of::()); + let axis_conversion = runs + .saturating_mul( + 2 * std::mem::size_of::() + std::mem::size_of::>(), + ) + .saturating_add(rows.saturating_mul(std::mem::size_of::())); + values + .saturating_add(width.saturating_mul(validity_bytes)) + .saturating_add(fields) + .saturating_add(indices) + .saturating_add(spans) + .saturating_add(axis_conversion) + .saturating_add(std::mem::size_of::() + std::mem::size_of::()) +} + /// `quantile(q, v)`, with `q` read at each grid point from an unlabelled single-field scalar on /// the same grid, as Prometheus evaluates the parameter at every step. A point where `q` has no /// value has no result; `q` outside `[0, 1]` or NaN gives what Prometheus' `quantile` gives. @@ -438,6 +744,105 @@ impl Grouping<'_> { type LabelKey = Vec<(Arc, Arc)>; +fn metric_bytes(metric: &MetricName) -> usize { + metric + .name() + .map_or(0, str::len) + .saturating_add( + metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) + .saturating_add( + metric + .tags() + .len() + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +fn grouping_workspace_bytes(metric: &MetricName) -> usize { + metric_bytes(metric) + .saturating_add( + metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +/// Conservative backing estimate for Rust 1.98.1's insert-only `HashMap` (hashbrown 0.17.1). +/// Bucket count follows its 7/8 load factor; entries, one control byte per bucket and the +/// 16-byte control-group tail are counted. This models requested table storage, not allocator slack. +/// An upper bound on the bytes the group index holds for `capacity` entries. +/// +/// The table keeps its entries in a power-of-two number of buckets below `16/7` of the usable +/// capacity, with one control byte beside each bucket. Three buckets per entry therefore covers the +/// table whatever exact growth the map uses, which is what keeps this a bound rather than a model of +/// one implementation: a toolchain that changes how the map grows cannot make it under-count. +fn hash_map_backing_bytes(capacity: usize) -> usize { + capacity + .saturating_mul(3) + .saturating_mul(size_of::<(LabelKey, usize)>() + 1) +} + +/// An upper bound on the capacity a map holding `capacity` entries grows to. +fn next_hash_capacity(capacity: usize) -> usize { + capacity.saturating_mul(2).max(7) +} + +/// An upper bound on the bytes one B-tree node holds. +/// +/// A node keeps at most eleven key/value pairs and twelve child pointers, plus a header of a few +/// words; the leaf form is smaller than that. The header allowance is generous on purpose, so the +/// bound does not depend on how a toolchain lays the node out. +fn btree_node_bytes() -> usize { + 11 * size_of::<(usize, Box<[FloatState]>)>() + 16 * size_of::() +} + +/// An upper bound on the bytes a B-tree of `entries` holds. +/// +/// Every node but the root keeps at least five entries, so `entries / 5` nodes is enough to hold +/// them; a tree of one entry is one node. +fn btree_backing_bytes(entries: usize) -> usize { + if entries == 0 { + return 0; + } + entries.div_ceil(5).saturating_mul(btree_node_bytes()) +} + +fn pending_bytes(fields: usize) -> usize { + working_set_bytes(fields, BLOCK_POINTS) +} + +fn working_set_bytes(fields: usize, pending_capacity: usize) -> usize { + pending_capacity + .saturating_mul(std::mem::size_of::<(usize, f64)>()) + .saturating_add(fields.saturating_mul(std::mem::size_of::>() + 1)) +} + +fn labels_capacity_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels.len().saturating_mul(std::mem::size_of::>()) +} + +fn grouping_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels + .iter() + .map(|label| label.len() + std::mem::size_of::>()) + .sum::() +} + fn label_key(metric: &MetricName) -> LabelKey { let mut labels = Vec::with_capacity(metric.tags().len() + 1); if let Some(name) = metric.name() { @@ -463,9 +868,29 @@ struct Group { blocks: BTreeMap>, } +/// Which field of which layout a fold is writing, and how it folds it. +/// +/// One value carried through the per-block work instead of three separate arguments: the three +/// never vary independently within one series, and the block loop is the only place that reads +/// them. +#[derive(Clone, Copy)] +struct FieldFold { + width: usize, + field_index: usize, + op: AggregateOp, +} + impl Group { - fn accumulate(&mut self, series: &Series, op: AggregateOp) { - let width = series.fields.len(); + fn accumulate( + &mut self, + series: &Series, + op: AggregateOp, + budget: &mut impl FnMut(usize) -> Result<()>, + mut retained: usize, + width: usize, + ) -> Result { + let pending_capacity_bytes = BLOCK_POINTS.saturating_mul(size_of::<(usize, f64)>()); + budget(retained.saturating_add(pending_capacity_bytes))?; let mut pending = Vec::with_capacity(BLOCK_POINTS); for (field_index, field) in series.fields.iter().enumerate() { // The axis says which absolute positions the field holds and the field iterator @@ -475,35 +900,81 @@ impl Group { let mut block = usize::MAX; for ((index, _), sample) in series.time.positions().zip(field.iter()) { if index / BLOCK_POINTS != block { - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; block = index / BLOCK_POINTS; } if let SampleRef::Float(sample) = sample { pending.push((index % BLOCK_POINTS, sample)); } } - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; } + Ok(retained) } /// Applies one field's samples within a block, so a block costs one lookup per field run. fn update_block( &mut self, block: usize, - width: usize, - field_index: usize, - op: AggregateOp, + fold: FieldFold, pending: &mut Vec<(usize, f64)>, - ) { + budget: &mut impl FnMut(usize) -> Result<()>, + retained: usize, + ) -> Result { if pending.is_empty() { - return; + return Ok(retained); + } + let new_block = !self.blocks.contains_key(&block); + let block_bytes = if new_block { + BLOCK_POINTS + .saturating_mul(fold.width) + .saturating_mul(std::mem::size_of::()) + } else { + 0 + }; + let backing_delta = if new_block { + btree_backing_bytes(self.blocks.len() + 1) + .saturating_sub(btree_backing_bytes(self.blocks.len())) + } else { + 0 + }; + if new_block { + budget( + retained + .saturating_add(backing_delta) + .saturating_add(block_bytes) + .saturating_add(pending.capacity().saturating_mul(size_of::<(usize, f64)>())), + )?; } let states = self.blocks.entry(block).or_insert_with(|| { - vec![FloatState::default(); BLOCK_POINTS * width].into_boxed_slice() + vec![FloatState::default(); BLOCK_POINTS * fold.width].into_boxed_slice() }); for (offset, sample) in pending.drain(..) { - states[offset * width + field_index].update(op, sample); + states[offset * fold.width + fold.field_index].update(fold.op, sample); } + Ok(retained + .saturating_add(backing_delta) + .saturating_add(block_bytes)) } } diff --git a/src/promql/src/aggregate/tests.rs b/src/promql/src/aggregate/tests.rs index 31082bd784a..0fd86561fdd 100644 --- a/src/promql/src/aggregate/tests.rs +++ b/src/promql/src/aggregate/tests.rs @@ -14,6 +14,7 @@ use std::collections::BTreeMap; +use datafusion::error::DataFusionError; use datatypes::arrow::array::{Float64Array, UInt32Array}; use datatypes::series::Label; use datatypes::series::histogram::HistogramBatchBuilder; @@ -345,6 +346,27 @@ fn avg_compensation_matches_prometheus_bits() { } } +#[test] +fn one_field_missing_and_another_nan_keep_row_nullability() { + let schema = schema(&["missing", "value"]); + let series = input( + &schema, + metric("m", &[]), + 0, + &[ + vec![None, None, None], + vec![Some(1.0), Some(f64::NAN), Some(2.0)], + ], + ); + let output = aggregate_float(&schema, &[series], AggregateOp::Sum, Grouping::By(&[])).unwrap(); + assert_eq!(output.len(), 1); + assert_eq!(values(&output[0], 0).1, vec![None, None, None]); + let values = values(&output[0], 1).1; + assert_eq!(values[0], Some(1.0)); + assert!(values[1].unwrap().is_nan()); + assert_eq!(values[2], Some(2.0)); +} + #[test] fn fields_are_aggregated_independently() { // A field without input values at a row is null there, even for count and group; the row @@ -476,6 +498,301 @@ fn fields_are_aggregated_independently() { } } +#[test] +fn incremental_retained_bytes_match_full_recount_after_updates() { + let schema = schema_with_grid_len(&["value", "other"], 256); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["zone", "rack"])) + .unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "initial state" + ); + let updates = [ + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(1.0)], vec![Some(2.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 64, + &[vec![Some(3.0)], vec![Some(4.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(5.0)], vec![Some(6.0)]], + ), + input( + &schema, + metric("m", &[("zone", "b"), ("rack", "2")]), + 128, + &[vec![Some(7.0)], vec![Some(8.0)]], + ), + ]; + for (index, series) in updates.iter().enumerate() { + let previous = aggregator.state_bytes(); + aggregator.update(series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "update {index}" + ); + if index == 2 { + assert_eq!(aggregator.state_bytes(), previous); + } + } +} + +#[test] +fn incremental_retained_bytes_match_full_recount_for_many_groups() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["id"])).unwrap(); + for id in 0..512 { + let id = id.to_string(); + let series = input( + &schema, + metric("m", &[("id", &id)]), + id.parse::().unwrap() % 64, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "group {id}" + ); + } +} + +/// The unbudgeted entry point is the budgeted one with a no-op budget, so this compares one path +/// with itself: it is a smoke check that the wrapper still runs, not evidence that the kernels are +/// right. The numbers are checked against Prometheus in `aggregates_match_prometheus` and +/// `stat_aggregates_match_prometheus`. +#[test] +fn every_supported_kind_runs_through_the_unbudgeted_entry_point() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0), Some(3.0)]]); + for op in [ + AggregateOp::Sum, + AggregateOp::Count, + AggregateOp::Avg, + AggregateOp::Min, + AggregateOp::Max, + AggregateOp::Group, + AggregateOp::Stddev, + AggregateOp::Stdvar, + ] { + let mut plain = FloatAggregator::try_new(&schema, op, Grouping::By(&[])).unwrap(); + plain.update(&series).unwrap(); + let plain = plain.finish().unwrap(); + let mut budgeted = + FloatAggregator::try_new_budgeted(&schema, op, Grouping::By(&[]), |_| Ok(())).unwrap(); + budgeted.update_budgeted(&series, |_| Ok(())).unwrap(); + let budgeted = budgeted.finish_budgeted(|_| Ok(())).unwrap(); + assert_eq!( + values_of(&plain, 0, 128), + values_of(&budgeted, 0, 128), + "{op:?}" + ); + } +} + +#[test] +fn repeated_updates_in_one_group_and_block_keep_one_block_charge() { + let schema = schema(&["value"]); + let first = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + let second = input(&schema, metric("m", &[]), 1, &[vec![Some(3.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + aggregator.update(&first).unwrap(); + let retained = aggregator.state_bytes(); + aggregator.update(&second).unwrap(); + assert_eq!(aggregator.state_bytes(), retained); +} + +#[test] +fn adjacent_blocks_are_admitted_and_accounted_separately() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 63, &[vec![Some(1.0), Some(2.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut admissions = Vec::new(); + aggregator + .update_budgeted(&series, |bytes| { + admissions.push(bytes); + Ok(()) + }) + .unwrap(); + assert_eq!(aggregator.groups(), 1); + assert_eq!(aggregator.groups[0].blocks.len(), 2); + assert_eq!( + admissions.last().copied(), + Some( + aggregator + .state_bytes() + .saturating_add(std::mem::size_of::<(usize, f64)>() * BLOCK_POINTS,) + ) + ); +} + +#[test] +fn a_rejected_block_allocation_does_not_insert_a_block() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + // The figure a new block's admission adds: the block's states and its pending buffer, on + // top of whatever the fold already holds. Located by the value, not by which call reaches it, + // so reordering the earlier admissions cannot make this test pass without testing anything. + let block_states = BLOCK_POINTS + .saturating_mul(schema.fields.len()) + .saturating_mul(std::mem::size_of::()); + let pending = BLOCK_POINTS * std::mem::size_of::<(usize, f64)>(); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut refused = None; + let result = aggregator.update_budgeted(&series, |bytes| { + // Everything up to the block's own admission is admitted; only the block is refused. + if bytes >= block_states.saturating_add(pending) { + refused = Some(bytes); + return Err(DataFusionError::Execution("block allocation budget".into())); + } + Ok(()) + }); + assert!( + result.is_err(), + "the block was admitted with no budget left" + ); + let refused = refused.expect("the refusal did not come from a block admission"); + assert!( + refused >= block_states.saturating_add(pending), + "the refusal came from an admission smaller than a block" + ); + assert_eq!(aggregator.groups(), 1); + assert!( + aggregator.groups[0].blocks.is_empty(), + "a refused block was inserted anyway" + ); +} + +#[test] +fn a_map_growth_rejection_does_not_insert_the_new_group() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + let first = input(&schema, metric("m", &[("i", "0")]), 0, &[vec![Some(1.0)]]); + aggregator.update(&first).unwrap(); + let mut index = 1; + while aggregator.group_index.len() < aggregator.group_index.capacity() { + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + index += 1; + } + let group_count = aggregator.groups(); + let old_state = aggregator.state_bytes(); + assert_eq!(aggregator.group_index.capacity(), group_count); + + // A resize keeps the old table while it fills the new one, so the figure the growth has to + // admit is the whole new table on top of the old state -- not the difference between them. + // That is the property asserted here: the refusal threshold is the state plus a whole new + // table, so the test fails if the admission ever drops to the delta. + let new_table = hash_map_backing_bytes(next_hash_capacity(group_count)); + let threshold = old_state.saturating_add(new_table); + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + let mut peak = 0; + let result = aggregator.update_budgeted(&series, |bytes| { + peak = peak.max(bytes); + if bytes >= threshold { + Err(DataFusionError::Execution("old+new map budget".into())) + } else { + Ok(()) + } + }); + assert!( + result.is_err(), + "the map grew with no budget for both tables" + ); + assert!( + peak >= threshold, + "the growth admitted {peak}, less than the old state plus a new table ({threshold})" + ); + assert_eq!(aggregator.group_index.capacity(), group_count); + assert_eq!(aggregator.groups(), group_count); + assert_eq!( + aggregator + .groups + .iter() + .map(|group| group.blocks.len()) + .sum::(), + group_count + ); +} + +#[test] +fn finish_rejection_after_first_group_is_budgeted_before_the_next() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + for (label, value) in [("a", 1.0), ("b", 2.0)] { + let series = input( + &schema, + metric("m", &[("i", label)]), + 0, + &[vec![Some(value)]], + ); + aggregator.update(&series).unwrap(); + } + let mut calls = Vec::new(); + let result = aggregator.finish_budgeted(|bytes| { + calls.push(bytes); + if calls.len() == 4 { + Err(DataFusionError::Execution("budget".into())) + } else { + Ok(()) + } + }); + assert!(result.is_err()); + assert_eq!(calls.len(), 4); + assert!(calls[1] > calls[0], "first builder peak"); + assert!(calls[2] < calls[1], "first group settle"); + assert!(calls[3] > calls[2], "second builder peak"); + + // The base already includes the output Vec backing; later callbacks add only the + // builder workspace or completed Series allocation. + let builder_workspace = finish_working_bytes(1, 1, 1); + let expected_output = aggregate_float( + &schema, + &[ + input(&schema, metric("m", &[("i", "a")]), 0, &[vec![Some(1.0)]]), + input(&schema, metric("m", &[("i", "b")]), 0, &[vec![Some(2.0)]]), + ], + AggregateOp::Sum, + Grouping::By(&["i"]), + ) + .unwrap(); + let completed_series = crate::grid_series::series_bytes(&expected_output[0]); + assert_eq!(calls[1] - calls[0], builder_workspace); + assert_eq!(calls[2] - calls[0], completed_series); + assert_eq!(calls[3] - calls[2], builder_workspace); +} + #[test] fn state_blocks_keep_absolute_indices_across_boundaries_and_gaps() { // A 512-point grid spans eight state blocks; inputs straddle a block edge, leave whole diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 750ee41b5e8..baf57a275c8 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -14,6 +14,7 @@ mod absent; mod empty_metric; +mod fused_series; mod histogram_fold; mod instant_manipulate; mod normalize; @@ -37,6 +38,7 @@ use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; +pub use fused_series::{FusedSeries, FusedSeriesBinding, FusedSeriesExec, FusedSeriesStream}; pub use histogram_fold::{ HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, }; diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs new file mode 100644 index 00000000000..d957851902a --- /dev/null +++ b/src/promql/src/extension_plan/fused_series.rs @@ -0,0 +1,1555 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! One node that runs a compiled native tree over raw series. +//! +//! [`FusedSeries`] is the logical half: it holds the PromQL text the tree was compiled from, the +//! tree itself and the binding its input was compiled against. [`FusedSeriesExec`] is the physical +//! half: it owns one execution of that tree per partition and decodes its input frames into raw +//! series chunks and converts its evaluated chunks back into batches. +//! +//! The two halves show different things in `EXPLAIN`, which is the point of having both. The +//! logical node prints the query a reader wrote ([`FusedSeries`]'s `fmt_for_explain`), and the +//! physical node prints the tree it will run ([`FusedSeriesExec`]'s `fmt_as`, which embeds +//! [`NativeTree`]'s own display rather than a second renderer). A reader comparing the two plans +//! can therefore see whether a query was compiled the way they expected without decoding node +//! names. +//! +//! Scope of this layer: the input is a stream of native raw series frames +//! ([`raw_series_frame_schema`](datatypes::series::frame::raw_series_frame_schema)) — one row per +//! chunk, carrying the series' encoded key, its finality, its timestamps and its fields, plus the +//! complete labels when the frame is explicit — with each series' rows contiguous and its +//! timestamps ascending, which is what the caller's [`SeriesOrdering`] names. The output is one +//! batch per evaluated chunk, with the statement's grid timestamps in the declared time column. +//! Everything the tree cannot do is rejected when the execution plan is built, not when the first +//! batch arrives: a tree the native runtime does not execute, a frame whose schema the tree's input +//! was not compiled against, or an output schema the conversion cannot classify. + +use std::cmp::Ordering; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, VecDeque}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::buffer::ScalarBuffer; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::DFSchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::EvalGrid; +use futures::Stream; +use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; + +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::{LabelSet, label_set}; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::metrics::{self, NodePath}; +use crate::native::plan::{NativeExpr, NativeTree}; +use crate::native::runs::{Run, absorb}; +use crate::source::materialize_raw_labels; + +#[cfg(test)] +mod tests; + +/// The label a metric name is stored under, which is not a tag column of the input. +const METRIC_NAME: &str = "__name__"; + +/// How many input steps one poll of [`FusedSeriesStream`] runs before it wakes itself. +/// +/// A step polls the input once and hands what it settled to the executor. Settling is not +/// guaranteed: raw samples before the first evaluation point settle nothing, so an input of many +/// such chunks would otherwise be read to its end inside one poll, holding the task's thread for +/// an unbounded time and delaying both other tasks on it and the cancellation the consumer asked +/// for. The budget makes that a poll that yields instead, and the waker makes sure it is polled +/// again, so nothing is lost by stopping early. +const POLL_BUDGET: usize = 64; + +/// What the tree's input was compiled against, and how its output columns are named. +/// +/// The binding travels with the tree: a tree is legal only against the schema it was bound to, and +/// a reader of the output needs to know which grid a chunk's absolute positions are on. Keeping it +/// in one value is what lets the logical node be built from the same description, and the physical +/// node reads the same facts back off the tree rather than trusting a second copy of them. +/// +/// What the logical node checks is that the schema the tree's input node declares is the one this +/// binding names. What it cannot check is whether the caller meant *some other* schema with the +/// same field names, or whether [`time_index`](Self::time_index) is the column the query reports +/// positions in: from here both are indistinguishable from a correct binding, so picking them is +/// the caller's obligation. The physical node reads neither the metric name nor a time column from +/// the input — a raw frame carries the labels of each series and its own timestamps. +#[derive(Clone, Debug)] +pub struct FusedSeriesBinding { + /// The raw series schema the tree's input node was bound to, without an evaluation grid. + pub series_schema: Arc, + /// The statement's grid, which is the grid every evaluated chunk is addressed on. + pub grid: Arc, + /// The metric name every input series is bound to, when the input reads one named metric. + /// + /// The logical node requires it when [`tags`](Self::tags) declares [`METRIC_NAME`]: a reader of + /// the output has to be able to tell that a range function dropped the name, which needs a name + /// to drop. The physical node takes the names from the frames, which carry each series' own. + pub metric_name: Option>, + /// The output column holding the grid timestamps. + pub time_index: String, + /// The logical tag names the output carries, read from each series' labels. + /// + /// [`METRIC_NAME`] is special: it is resolved to the series' metric name instead of a column, + /// which is how a reader sees that a range function dropped it. + pub tags: Vec, +} + +impl FusedSeriesBinding { + /// Whether `other` binds the same fields under the same names. + fn same_fields(&self, other: &SeriesSchema) -> bool { + self.series_schema.fields.len() == other.fields.len() + && self + .series_schema + .fields + .iter() + .zip(other.fields.iter()) + .all(|(mine, theirs)| mine.name == theirs.name) + } +} + +impl PartialEq for FusedSeriesBinding { + fn eq(&self, other: &Self) -> bool { + // The schemas are compared by the field names they declare, and the grids by their spec: + // neither type has structural equality, and neither identity nor allocation is meant here. + self.metric_name == other.metric_name + && self.time_index == other.time_index + && self.tags == other.tags + && self.grid.spec() == other.grid.spec() + && self.same_fields(&other.series_schema) + } +} + +impl Eq for FusedSeriesBinding {} + +impl Hash for FusedSeriesBinding { + fn hash(&self, state: &mut H) { + self.metric_name.hash(state); + self.time_index.hash(state); + self.tags.hash(state); + let spec = self.grid.spec(); + (spec.start_ms, spec.step_ms, spec.len).hash(state); + for field in self.series_schema.fields.iter() { + field.name.hash(state); + } + } +} + +impl PartialOrd for FusedSeriesBinding { + fn partial_cmp(&self, other: &Self) -> Option { + // A total order over the same fields equality reads, so equal bindings compare equal. + match self + .metric_name + .as_deref() + .partial_cmp(&other.metric_name.as_deref()) + { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.time_index.partial_cmp(&other.time_index) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tags.partial_cmp(&other.tags) { + Some(Ordering::Equal) => {} + ord => return ord, + } + let (mine, theirs) = (self.grid.spec(), other.grid.spec()); + match (mine.start_ms, mine.step_ms, mine.len).partial_cmp(&( + theirs.start_ms, + theirs.step_ms, + theirs.len, + )) { + Some(Ordering::Equal) => {} + ord => return ord, + } + Some( + self.series_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .cmp( + other + .series_schema + .fields + .iter() + .map(|field| field.name.as_ref()), + ), + ) + } +} + +/// The logical plan of one node that runs a compiled native tree. +/// +/// The node adds nothing to the tree: it records the query the tree was compiled from, so +/// `EXPLAIN` shows the reader their own query rather than the shape it compiled to, and it carries +/// the binding the tree's input was compiled against so the execution plan can be built from it. +#[derive(Debug)] +pub struct FusedSeries { + /// The PromQL text this tree was compiled from. + promql: Arc, + /// The compiled tree to run. + tree: NativeTree, + /// What the tree's input was compiled against. + binding: FusedSeriesBinding, + /// The plan producing the input series. + input: LogicalPlan, + /// The columns this node produces. + output_schema: DFSchemaRef, +} + +impl FusedSeries { + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeries" + } + + /// Creates the node, checking that `tree` is legal against `binding`. + pub fn try_new( + promql: Arc, + tree: NativeTree, + binding: FusedSeriesBinding, + input: LogicalPlan, + output_schema: DFSchemaRef, + ) -> DataFusionResult { + validate_binding(&tree, &binding)?; + Ok(Self { + promql, + tree, + binding, + input, + output_schema, + }) + } + + /// The PromQL text this node was compiled from. + pub fn promql(&self) -> &str { + &self.promql + } + + /// The compiled tree this node runs. + pub fn tree(&self) -> &NativeTree { + &self.tree + } + + /// What the tree's input was compiled against. + pub fn binding(&self) -> &FusedSeriesBinding { + &self.binding + } + + /// The physical node that runs `exec_input` as this node's input. + /// + /// `label_source` and `ordering` are the binding of the frames `exec_input` produces: the + /// frame carries the encoded key and, in explicit mode, the complete labels, but no schema can + /// say which of the two a producer meant or how it ordered the series, so the caller that built + /// the frame passes both on. Building the node rejects a tree the runtime cannot execute, a + /// frame schema the tree's input was not compiled against, and a schema the conversion cannot + /// classify, so a plan that is built is a plan that runs. + pub fn to_execution_plan( + &self, + exec_input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + ) -> DataFusionResult> { + let output_schema: SchemaRef = self.output_schema.inner().clone(); + Ok(Arc::new(FusedSeriesExec::try_new( + Arc::clone(&self.promql), + exec_input, + label_source, + ordering, + Arc::new(self.tree.clone()), + output_schema, + self.binding.time_index.clone(), + self.binding.tags.clone(), + )?)) + } +} + +impl UserDefinedLogicalNodeCore for FusedSeries { + fn name(&self) -> &str { + Self::name() + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + // The tree's operands are bound inside it; there is no expression for an optimizer to + // rewrite separately from the tree that owns it. + Vec::new() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeries: {}", self.promql) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DataFusionResult { + if inputs.len() != 1 { + return Err(DataFusionError::Internal( + "FusedSeries should have at exact one input".to_string(), + )); + } + + let input = inputs.pop().unwrap(); + Ok(Self { + promql: Arc::clone(&self.promql), + tree: self.tree.clone(), + binding: self.binding.clone(), + input, + output_schema: Arc::clone(&self.output_schema), + }) + } +} + +impl PartialEq for FusedSeries { + fn eq(&self, other: &Self) -> bool { + // `NativeTree` carries no structural equality, so the rendered tree stands in for it. The + // renderer names every node with its bound operation and grids, which is what makes two + // nodes the same computation; comparing the text alone would call two different trees + // equal whenever their PromQL happens to read the same. + self.promql == other.promql + && self.tree.to_string() == other.tree.to_string() + && self.binding == other.binding + && self.input == other.input + && self.output_schema == other.output_schema + } +} + +impl Eq for FusedSeries {} + +impl Hash for FusedSeries { + fn hash(&self, state: &mut H) { + self.promql.hash(state); + self.binding.hash(state); + self.input.hash(state); + self.output_schema.hash(state); + } +} + +impl PartialOrd for FusedSeries { + fn partial_cmp(&self, other: &Self) -> Option { + match self.promql.partial_cmp(&other.promql) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.binding.partial_cmp(&other.binding) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.input.partial_cmp(&other.input) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tree.to_string().partial_cmp(&other.tree.to_string()) { + Some(Ordering::Equal) => {} + ord => return ord, + } + if self.output_schema == other.output_schema { + Some(Ordering::Equal) + } else { + None + } + } +} + +/// What the tree's input node reads, read off the tree itself. +/// +/// The tree is the one authority on the schema its input was compiled against: a caller's binding +/// can only agree or disagree with it, and the frame this node reads has to be decoded against the +/// same schema, so both halves read it here instead of each carrying a copy. +struct TreeInput { + /// The position the tree's raw chunks enter at: a position is how the runtime addresses a node, + /// so nothing on the node has to carry an identity that the tree itself does not. + path: NodePath, + /// The raw series schema the tree's input node was compiled against. + schema: Arc, +} + +/// Reads what the tree's input node was compiled against, rejecting a tree this node cannot feed. +/// +/// A node feeds one raw input, and its output is addressed on a grid; a tree that does not have +/// that shape is one this layer cannot hand chunks to or read positions from. +fn tree_input(tree: &NativeTree) -> DataFusionResult { + let mut sources = 0; + let mut input_path = None; + let mut schema = None; + metrics::walk(tree.root(), &mut |node, path| { + let NativeExpr::Input(input) = &node.expr else { + return; + }; + sources += 1; + input_path.get_or_insert_with(|| path.clone()); + // Raw, not evaluated: this node feeds the tree from a scan, so a tree whose input was + // compiled against an already-evaluated schema would read the wrong thing. + schema = input + .output + .samples() + .filter(|schema| schema.eval_grid.is_none()) + .map(Arc::clone); + }); + if sources != 1 { + return Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree with {sources} inputs, and it needs exactly one" + ))); + } + let (Some(path), Some(schema)) = (input_path, schema) else { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose input is not raw series".to_string(), + )); + }; + let evaluated = tree + .root() + .output + .samples() + .is_some_and(|schema| schema.eval_grid.is_some()); + if !evaluated { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on a grid".to_string(), + )); + } + Ok(TreeInput { path, schema }) +} + +/// Refuses a tree whose logical series can be split across the partitions this node runs. +/// +/// Keeping the metric name is not what keeps two raw series apart. An empty tag value and an absent +/// tag are the same canonical label set, so `{host="h", zone=""}` and `{host="h"}` are one series to +/// Prometheus while the scan encodes them as two raw keys. A window that evaluated its input without +/// merging it therefore passes such a pair through as two series under one label set, exactly as the +/// functions that drop `__name__` do by merging them, and a map drops `__name__` outright. +/// +/// A stateful node above one of those needs both halves together, and the runtime does provide that +/// within one execution: a window whose input can hold two series under one label set collects the +/// fragments, merges each label set once, and only then evaluates. That merge is per execution, and +/// this node runs one execution per partition. Two raw series that become one label set can land in +/// different partitions - the scan keeps a *series* in one partition, which says nothing about two +/// different series that later become indistinguishable - so each execution merges only the half it +/// was given. The consumer above then answers from one half and reports nothing for the other, which +/// is an empty result rather than an error. +/// +/// Whether a window merged its input is runtime state this check does not have: it runs over a bare +/// tree, where a window's input is only something that produces samples. A window below another +/// window is therefore treated as one that did not merge, whatever function it runs, which refuses +/// `rate` under `rate` (which drops the name) and `last_over_time` under `rate` (which keeps it) +/// alike. +/// +/// Reading such a tree needs an exchange above this node that merges a logical series across +/// partitions, which does not exist yet. Until it does the tree is refused where it can hurt: over +/// one partition there is nothing to split, so only a tree run over more than one is checked. +/// +/// The compiler cannot build a nested window today - [`compile`](crate::promql::native_compile::compile) +/// requires a range function to apply directly to a selector, `([])`, +/// and a selector is raw samples - so no tree it produces has a window above a window, and this +/// check costs no executable query. It is checked because [`FusedSeriesExec::try_new`] accepts any +/// well-formed tree, and a tree it accepts is a tree it must answer correctly. +fn reject_a_split_logical_series( + tree: &NativeTree, + partition_count: usize, +) -> DataFusionResult<()> { + if partition_count == 1 { + return Ok(()); + } + // The walk is outermost first, so it meets nodes in the reverse of the order data flows: a + // consumer met before a rewriter is one that the rewriter feeds. + let mut consumer: Option = None; + let mut rejection: Option<(String, String)> = None; + metrics::walk(tree.root(), &mut |node, path| { + if rejection.is_some() { + return; + } + // A node can leave two raw series on one canonical label set. A map always can: it drops + // `__name__` and every empty-valued tag. An instant selection preserves its input labels, + // rather than dropping them, but can still leave two raw series on one canonical label set: + // `{zone=""}` and the absent `zone` are one label set and two raw keys. A window can + // whenever it did not merge its input, which a bare tree says nothing about: a check on the + // function alone accepts `last_over_time`, whose output keeps the name and can still hold + // both raw series. Every map, instant selection, and window is therefore taken to be a + // possible rewriter, and the refusal below is what narrows that to the shape that costs + // correctness - a window with a rewriter below it, a window being the only consumer here. + let rewrites = matches!( + node.expr, + NativeExpr::Map { .. } | NativeExpr::Instant { .. } | NativeExpr::Window { .. } + ); + if rewrites && let Some(above) = &consumer { + rejection = Some((path.as_str().to_string(), above.clone())); + return; + } + let consumes = matches!(node.expr, NativeExpr::Window { .. }); + if consumes && consumer.is_none() { + consumer = Some(path.as_str().to_string()); + } + }); + let Some((rewriter, above)) = rejection else { + return Ok(()); + }; + Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree where `{rewriter}` can leave two raw series on one canonical label \ + set and `{above}` above it reads a whole series, over {partition_count} partitions: an \ + empty tag value and an absent tag are the same label set, the runtime merges such \ + fragments once per execution, and this node runs one execution per partition, so `{above}` \ + would read one half and answer an empty result. Reading this tree needs an exchange between \ + `{rewriter}` and `{above}` that merges a logical series across partitions, which is not \ + implemented" + ))) +} + +/// Checks that `tree` reads the schema `binding` declares, on the grid it declares. +/// +/// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule +/// [`NativeTree`]'s own validation and the kernels apply, and a chunk's timestamps already come +/// from its absolute grid positions, so a separately allocated grid with the same spec is the same +/// grid to everything downstream. +/// +/// Answers with what the tree's input reads, which is what a caller names when it hands a raw chunk +/// over. +fn validate_binding( + tree: &NativeTree, + binding: &FusedSeriesBinding, +) -> DataFusionResult { + if binding.series_schema.eval_grid.is_some() { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema, which cannot have an evaluation grid" + .to_string(), + )); + } + if binding.tags.iter().any(|tag| tag == METRIC_NAME) && binding.metric_name.is_none() { + return Err(DataFusionError::Plan(format!( + "FusedSeries declares a {METRIC_NAME:?} tag column but binds no metric name" + ))); + } + + // The tree's own input node says which schema it reads, so the binding is checked against the + // tree rather than trusted: a mismatch would evaluate one schema and label the result as + // another. + let input = tree_input(tree)?; + if !binding.same_fields(&input.schema) { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema the tree's input was not compiled against" + .to_string(), + )); + } + match tree.root().output.samples() { + Some(schema) + if schema + .eval_grid + .as_ref() + .is_some_and(|grid| grid.spec() == binding.grid.spec()) => + { + Ok(input) + } + _ => Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on the statement's grid" + .to_string(), + )), + } +} + +/// Where one output column of [`FusedSeriesExec`] comes from. +#[derive(Clone, Copy, Debug)] +enum OutputColumn { + /// The evaluation grid timestamps of the current chunk. + Time, + /// Series field by index in the bound series schema. + Field(usize), + /// Declared logical tag by index in the caller-supplied tag list. + Tag(usize), +} + +/// Classifies every output column, rejecting anything the conversion cannot produce. +fn classify_output( + output_schema: &SchemaRef, + series_schema: &SeriesSchema, + tags: &[String], + time_column: &str, +) -> DataFusionResult> { + for (index, tag) in tags.iter().enumerate() { + if tags[..index].contains(tag) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {tag:?} in FusedSeriesExec output schema" + ))); + } + if tag == time_column { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec cannot be the time column" + ))); + } + if series_schema + .fields + .iter() + .any(|field| field.name.as_ref() == tag) + { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec clashes with a series field" + ))); + } + } + + let mut time_seen = false; + let mut field_seen = vec![false; series_schema.fields.len()]; + let mut tag_seen = vec![false; tags.len()]; + let mut classification = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { + let name = field.name(); + if name == time_column { + if std::mem::replace(&mut time_seen, true) { + return Err(DataFusionError::Plan(format!( + "duplicate time column {name:?} in FusedSeriesExec output schema" + ))); + } + match field.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, None) => {} + other => { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must be Timestamp(Millisecond, \ + None), got {other:?}" + ))); + } + } + if field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must not be nullable" + ))); + } + classification.push(OutputColumn::Time); + continue; + } + if let Some(index) = series_schema + .fields + .iter() + .position(|field| field.name.as_ref() == name) + { + if std::mem::replace(&mut field_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate field column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be Float64, got {:?}", + field.data_type() + ))); + } + // A grid point inside the statement's interval can have no value in this field. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be nullable, because a grid \ + point can be missing" + ))); + } + classification.push(OutputColumn::Field(index)); + continue; + } + if let Some(index) = tags.iter().position(|tag| tag == name) { + if std::mem::replace(&mut tag_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Utf8 { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be Utf8, got {:?}", + field.data_type() + ))); + } + // `rate` and the map kernels drop `__name__`, so a tag can be absent from a row. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be nullable, because range \ + functions and map kernels drop the metric name" + ))); + } + classification.push(OutputColumn::Tag(index)); + continue; + } + return Err(DataFusionError::Plan(format!( + "unknown column {name:?} in FusedSeriesExec output schema" + ))); + } + + if !time_seen { + return Err(DataFusionError::Plan(format!( + "missing time column {time_column:?} in FusedSeriesExec output schema" + ))); + } + for (index, seen) in field_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing field column {:?} in FusedSeriesExec output schema", + series_schema.fields[index].name + ))); + } + } + for (index, seen) in tag_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing tag column {:?} in FusedSeriesExec output schema", + tags[index] + ))); + } + } + Ok(classification) +} + +/// The physical node that runs a compiled native tree over its input's raw series. +pub struct FusedSeriesExec { + /// The PromQL text the tree was compiled from, which is what the plan line shows. + promql: Arc, + /// The tree this node runs, shared by every partition's execution. + tree: Arc, + /// The raw series schema the tree's input was compiled against, which every frame is decoded + /// against. + series_schema: Arc, + /// How the input frames name their series: labels encoded in the key, or carried along. + label_source: RawLabelSource, + /// The ordering the input frames guarantee. + ordering: SeriesOrdering, + /// The position of the tree's input node, where every raw chunk a partition reads enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The output column holding the grid timestamps, kept so this node can be rebuilt over a new + /// child without losing which column its time came from. + time_column: String, + /// The tag column names, indexed by [`OutputColumn::Tag`]. + tags: Arc<[String]>, + /// The plan producing the raw series frames. + input: Arc, + /// The columns this node produces. + output_schema: SchemaRef, + /// The grid positions every label set in this plan's output already occupies, and what they + /// keep charged on the memory pool. + /// + /// Plan-level rather than per-partition: see [`LabelSets`]. + labels: SharedLabelSets, + /// Per-node metrics, registered by each execution of the tree. + metric: ExecutionPlanMetricsSet, + /// This node's partitioning, emission and boundedness. + properties: Arc, +} + +impl FusedSeriesExec { + /// Creates the node, rejecting whatever this node cannot execute or convert. + /// + /// `input` produces native raw series frames: the series' encoded key, its finality, its + /// timestamps and its fields, plus the complete labels when `label_source` is + /// [`RawLabelSource::Explicit`]. `time_column` and `tags` name the columns of the output, which + /// is a row table rather than a frame. + /// + /// The parameters stay separate because each is a fact its own caller has: the query, the plan + /// below, the frame's label binding and ordering, the tree, and the three things the output + /// schema is classified by. A struct over them would be the same list under another name. + #[allow(clippy::too_many_arguments)] + pub fn try_new( + promql: Arc, + input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + tree: Arc, + output_schema: SchemaRef, + time_column: String, + tags: Vec, + ) -> DataFusionResult { + let tree_input = tree_input(&tree)?; + // Building a state is what proves the runtime executes this tree: a node kind it does not + // run yet is refused here, so a plan that is built is a plan that starts. The throwaway + // state is dropped at once, returning its (unbounded) reservations. + let _ = NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + )?; + let classification = + classify_output(&output_schema, &tree_input.schema, &tags, &time_column)?; + // Read here rather than below, because whether the split below can happen decides whether + // the tree shape is one this node may run at all. + let input_properties = input.properties(); + let partition_count = input_properties.partitioning.partition_count(); + if partition_count == 0 { + return Err(DataFusionError::Plan( + "FusedSeriesExec input has no partitions".into(), + )); + } + reject_a_split_logical_series(&tree, partition_count)?; + // The frame states its own label layout and series ordering; this plan's binding is an + // expectation checked against that declaration, not an input to decoding. + ensure_frame_matches_label_source(input.schema().as_ref(), &label_source)?; + // The frame is checked against what the tree's input reads, so a frame this node cannot + // decode is reported when the plan is built. The shared decoder is the authority on that + // layout, so it is asked rather than copied: handing it an empty stream of this schema makes + // it validate the declared frame with no batch to decode, which is the check + // `unpack_raw_series` makes before its first poll too. + unpack_raw_series( + Box::pin(RecordBatchStreamAdapter::new( + input.schema(), + futures::stream::empty(), + )), + Some(Arc::clone(&tree_input.schema)), + )?; + // One evaluator per input partition, and the input's own partitioning kept as it is. + // + // A window holds the state of the series it is reading and closes it on the chunk carrying + // `is_last`, so it can answer a series only if the whole series arrives in the partition it + // reads. The input guarantees that: a raw series frame is produced by `SeriesScan`, whose + // contract is that the same primary key always reaches the same partition, and the frame + // decoder holds each series' chunks contiguous within that partition. + // + // What is *not* asked of the optimizer is a distribution. A `KeyPartitioned` or + // `HashPartitioned` requirement is classed as a hash requirement, and a hash requirement + // makes a hash exchange necessary whenever the child has more than one partition - even if + // the child already satisfies it, and even when the partition counts match. A hash exchange + // does not preserve order, so it would interleave the chunks of one series and the decoder + // would close a series it has not finished. `UnspecifiedDistribution` leaves the child's + // layout alone, and `benefits_from_input_partitioning` is false so no round-robin is added + // to raise parallelism either. Between them no exchange can appear between the scan and this + // node, and the scan's guarantee arrives intact. + // As many partitions as the input hands over: the node runs one evaluator per partition and + // answers each as it settles. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + input_properties.emission_type, + input_properties.boundedness, + )); + // Shared by every partition's stream: a label set colliding is a property of the whole + // output vector, not of one partition. Building it needs no pool; the owner's reservation is + // attached in `execute`, which is the first point that has one. + let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default())); + Ok(Self { + promql, + tree, + series_schema: tree_input.schema, + label_source, + ordering, + input_path: tree_input.path, + classification: classification.into(), + time_column, + tags: tags.into(), + input, + output_schema, + labels, + metric: ExecutionPlanMetricsSet::new(), + properties, + }) + } + + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeriesExec" + } +} + +impl std::fmt::Debug for FusedSeriesExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("FusedSeriesExec").finish() + } +} + +impl DisplayAs for FusedSeriesExec { + /// The query, then the tree it compiles to: the tree is rendered by [`NativeTree`]'s own + /// `Display`, so `EXPLAIN` and the node's metrics name the same nodes. + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree) + } +} + +impl ExecutionPlan for FusedSeriesExec { + fn name(&self) -> &str { + Self::name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + /// Nothing is asked of the input's distribution, and nothing may be inserted above it. + /// + /// A hash requirement would make a hash exchange necessary as soon as the child has more than + /// one partition - even where the child already satisfies the requirement - and a hash exchange + /// does not preserve order. The input's own layout is already the one this node needs, because + /// the scan it reads from guarantees a series stays in one partition, so the requirement says + /// only that the layout is left alone. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution]) + } + + /// Splitting the input further would not help: a second partition is a second set of series + /// answered by a second tree, which this node already runs, and raising the count is what an + /// optimizer does by adding a round-robin exchange that reorders nothing it can rely on. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + // The tree's operands are bound inside it; this node owns no expression of its own. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + let [child] = <[Arc; 1]>::try_from(children).map_err(|children| { + DataFusionError::Plan(format!( + "FusedSeriesExec requires exactly one child, got {}", + children.len() + )) + })?; + // Rebuilt rather than patched: the new child's schema is what the frame is validated + // against, so a child that changed shape is refused instead of decoded with a stale layout. + Ok(Arc::new(Self::try_new( + Arc::clone(&self.promql), + child, + self.label_source.clone(), + self.ordering, + Arc::clone(&self.tree), + self.output_schema.clone(), + self.time_column.clone(), + self.tags.to_vec(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + // The label-set state is the plan's, so the reservation that pays for it is attached exactly + // once, to the pool of the first execution. Every partition calls this and only the first + // takes effect: re-registering per stream would move a charge the other partitions keep + // growing, and would give it back when that stream ends while the entries stay. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .attach(&context.runtime_env().memory_pool); + let input = self.input.execute(partition, Arc::clone(&context))?; + // Batches become raw series chunks through the shared frame decoder, which is the inverse of + // what the producer packed: one chunk per frame row, each series closed exactly once by the + // row carrying `is_last`. + let input = unpack_raw_series(input, Some(Arc::clone(&self.series_schema)))?; + // A source bound in key mode carries its identity encoded and no decoded labels, and a chunk + // handed to the tree has to name the series it holds. + let input = + materialize_raw_labels(input).map_err(|e| DataFusionError::External(Box::new(e)))?; + let state = NativeExecutionState::try_new( + &self.tree, + Arc::clone(&context.runtime_env().memory_pool), + &self.metric, + partition, + )?; + // What the label sets hold is charged to the reservation attached above, so nothing is + // reserved per stream: the entries outlive every stream and are released with the plan. + Ok(Box::pin(FusedSeriesStream { + input: Some(input), + executor: Some(NativeExecutor::new(state)), + input_path: self.input_path.clone(), + classification: Arc::clone(&self.classification), + tags: Arc::clone(&self.tags), + output_schema: self.output_schema.clone(), + labels: Arc::clone(&self.labels), + ready: VecDeque::new(), + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metric.clone_inner()) + } +} + +/// Converts one evaluated chunk into one output batch. +/// +/// The time column is the absolute grid positions the chunk carries, missing samples are nulls and +/// ordinary NaN stays a value. +fn series_to_record_batch( + classification: &[OutputColumn], + tags: &[String], + output_schema: &SchemaRef, + chunk: &EvaluatedChunk, +) -> DataFusionResult { + let series = &chunk.series; + // A chunk that forms one run is read from its range, which lends the grid's buffers and so + // allocates nothing. A chunk whose rows form several runs has no single range to read + // ([`EvaluatedChunk::range`] answers `None` for it), and the positions between its runs are + // holes rather than rows: the axis is walked for its carried rows instead, which is the one + // traversal that visits exactly those. + let timestamps = match chunk.range() { + Some(range) => range.timestamps(), + None => series + .time + .timestamps() + .map(|timestamps| ScalarBuffer::from(timestamps.into_owned())), + } + .map_err(|error| DataFusionError::Execution(error.to_string()))?; + let rows = timestamps.len(); + let mut columns: Vec = Vec::with_capacity(classification.len()); + for (position, column) in classification.iter().enumerate() { + match column { + OutputColumn::Time => columns.push(Arc::new(TimestampMillisecondArray::new( + timestamps.clone(), + None, + ))), + OutputColumn::Field(index) => { + let samples = series.fields.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "evaluated series has no field at index {index}" + )) + })?; + if samples.len() != rows { + return Err(DataFusionError::Execution(format!( + "evaluated series field {index} has {} rows, time axis has {rows}", + samples.len() + ))); + } + columns.push(field_column(samples)?); + } + OutputColumn::Tag(index) => { + // `__name__` resolves to the metric name, every other tag through the label set. + let name = tags.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "output column {position} declares unavailable tag {index}" + )) + })?; + let value = series.metric.get(name); + columns.push(Arc::new(StringArray::from(vec![value; rows]))); + } + } + } + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) +} + +/// Converts one evaluated series field into one `Float64` column. +/// +/// Dense float samples are reused as they are: the array already carries exactly the right nulls +/// and keeps NaN a value, so copying it through `Samples::iter()` would only cost an allocation +/// and a re-encode. Chunked samples go through the iterator, and histograms have no float column. +fn field_column(samples: &Samples) -> DataFusionResult { + match samples.dense() { + Some(DenseSamplesRef::Float(values)) => Ok(Arc::new(values.clone())), + Some(DenseSamplesRef::Histogram(_)) | Some(DenseSamplesRef::Mixed(_)) => { + Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )) + } + None => { + let values = samples + .iter() + .map(|sample| match sample { + SampleRef::Float(value) => Ok(Some(value)), + SampleRef::Missing => Ok(None), + SampleRef::Histogram(_) => Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )), + }) + .collect::>>()?; + Ok(Arc::new(Float64Array::from(values))) + } + } +} + +/// The bytes one label set holds, for the memory reservation. +/// +/// It counts each label's name and value with the pair that refers to them, which is the same +/// convention [`crate::native::execute`]'s own label-set table keeps. +fn labels_bytes(labels: &LabelSet) -> usize { + labels + .iter() + .map(|(name, value)| size_of::<(Arc, Arc)>() + name.len() + value.len()) + .sum() +} + +/// Renders the label set Prometheus compares the way it prints one, for the duplicate-label-set +/// error. +fn render_labels(labels: &LabelSet) -> String { + let parts = labels + .iter() + .map(|(name, value)| format!("{name}={value:?}")) + .collect::>(); + format!("{{{}}}", parts.join(", ")) +} + +/// The grid positions each final label set already occupies, per field and as merged, sorted runs. +/// +/// Shared by every partition's stream, because a label-set collision is a property of the whole +/// output vector and not of one partition. Two series that carry the same labels after the range +/// function and its maps drop a label land in different partitions, and the runtime's own per-node +/// `Emitted` table never sees both: it belongs to one execution. Checking only within a partition +/// would miss exactly those collisions. +/// +/// The key is [`crate::grid_series::label_set`]'s normalization of the final labels rather than +/// [`datatypes::series::MetricName`] itself: a label with an empty value is no label to +/// Prometheus, so `{host="h", zone=""}` and `{host="h"}` are one label set and have to be compared +/// as one. Encoding the tags as they are would file them apart and let one point hold two values. +/// +/// Occupancy is per field and per position, not per chunk: two series whose label sets became equal +/// while holding values in *different* fields merge cleanly in Prometheus, so only the same field +/// at the same position collides. That is the rule +/// [`crate::grid_series::merge_same_labelset`] applies. +/// +/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed +/// for the partitions still running, so the entries are only released with the plan itself. The +/// charge has the same lifetime, which is what makes the bytes the map holds - one entry per label +/// set of the output, kept until the plan dies - visible to the memory pool. Charging each stream +/// for what it inserted instead returned those bytes when that stream ended, while the entries the +/// later partitions are checked against stayed: everything the map held after the first partition +/// was unaccounted for. +#[derive(Default)] +struct LabelSets { + /// The grid positions each final label set already occupies, one run list per field of the + /// series schema. Each list is sorted, disjoint and has the runs that touch joined, which is + /// what lets a claim look a position up instead of walking the history. + seen: HashMap>>, + /// What `seen` keeps charged on the task's memory pool. `None` until the first execution, + /// since construction has no `TaskContext` and therefore no pool to register against. + reservation: Option, +} + +impl LabelSets { + /// Attaches the reservation to the pool of the first execution. + /// + /// Later executions keep the first reservation: it already pays for what `seen` holds, and + /// moving it to another pool would drop the bytes other partitions are still growing. + fn attach(&mut self, pool: &Arc) { + if self.reservation.is_none() { + self.reservation = + Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool)); + } + } + + /// Claims the occupied grid positions of one chunk of a final series for its label set. + /// + /// A collision is an error only when both series hold a sample of *one field* at the same + /// evaluation timestamp, which is Prometheus' range-query rule. Series with disjoint points are + /// legal, and so are series that fill the same points field by field: Prometheus merges both + /// into one output series, while this node emits them as separate batches and defers that merge + /// to matrix assembly, so every sample survives. + fn claim(&mut self, series: &Series) -> DataFusionResult<()> { + let runs = occupied_runs(series)?; + // A chunk holding no value occupies nothing, so there is no collision to remember and no + // entry to write. Such a chunk is not rare: a selection over a series that staleness or the + // range left empty still closes it, and hands over a closing chunk with no rows. Recording + // one entry per such series would grow the table - and its charge, which the plan holds for + // its whole life - for a query that never answers a value, until an empty result exhausted + // the pool. The executor's own `Emitted::occupy` takes the same early return. + if runs.iter().all(Vec::is_empty) { + return Ok(()); + } + let key = label_set(&series.metric); + let Self { seen, reservation } = self; + // Nothing reaches a claim without an `execute` first, and an execution that ran attached the + // reservation. Claiming without one would leave the map growing uncharged, so it is refused + // rather than silently skipped. + let Some(reservation) = reservation.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec label-set state is not attached to a memory pool".to_string(), + )); + }; + match seen.entry(key) { + Entry::Occupied(entry) => { + if let Some(field) = colliding_field(entry.get(), &runs) { + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset: {} (field \ + {field})", + render_labels(entry.key()) + ))); + } + let stored = entry.into_mut(); + // An entry recorded from a chunk of fewer fields is padded, so every field of this + // chunk has a slot of its own: a slot is a field index of the schema, so both sides + // count the same field the same way. A chunk of *fewer* fields than the entry + // leaves the remaining slots to the chunks that fill them. + // Stage only the affected run lists, and reserve before committing. A forward + // append whose existing capacity suffices needs no copy of its stored history. + let mut growth = 0usize; + let mut shrink = 0usize; + let mut merged_runs = Vec::with_capacity(runs.len()); + for (index, incoming) in runs.iter().enumerate() { + let Some(old) = stored.get(index) else { + let copy = incoming.clone(); + growth += copy.capacity(); + merged_runs.push(Some(copy)); + continue; + }; + let append_fits = incoming.is_empty() + || (old.last().is_none_or(|last| incoming[0].start >= last.end) + && old.capacity() + >= old.len() + incoming.len() + - usize::from(old.last().is_some_and(|last| { + incoming + .first() + .is_some_and(|first| first.start == last.end) + }))); + if append_fits { + merged_runs.push(None); + } else { + let mut merged = old.clone(); + absorb(&mut merged, incoming); + if merged.capacity() >= old.capacity() { + growth += merged.capacity() - old.capacity(); + } else { + shrink += old.capacity() - merged.capacity(); + } + merged_runs.push(Some(merged)); + } + } + reservation.try_grow(growth * size_of::())?; + for (index, (incoming, merged)) in runs.iter().zip(merged_runs).enumerate() { + if let Some(merged) = merged { + if index < stored.len() { + stored[index] = merged; + } else { + stored.push(merged); + } + } else if index < stored.len() { + absorb(&mut stored[index], incoming); + } + } + reservation.try_shrink(shrink * size_of::())?; + } + Entry::Vacant(entry) => { + // Only the label-set state is this node's own cross-poll state that grows with + // the output, so it is the only thing reserved here. The runs are charged by the + // capacity they were gathered with rather than by their length, so that a later + // claim appending into that slack is charged nothing for it. + let bytes = labels_bytes(entry.key()) + + runs.iter().map(|runs| runs.capacity()).sum::() * size_of::(); + reservation.try_grow(bytes)?; + entry.insert(runs); + } + } + Ok(()) + } +} + +type SharedLabelSets = Arc>; + +/// The grid positions the rows of one chunk hold a value at, one run list per field. +/// +/// Each [`Run`] is a half-open range of absolute grid indices, and the runs of one field are sorted +/// and disjoint. Occupancy is kept per field because it is what a collision is: two series whose +/// label sets became equal merge field by field in Prometheus, so only two values of the same field +/// at one position collide. Marking the whole position as taken as soon as *any* field held a value +/// would refuse the complementary pair that merges cleanly. +fn occupied_runs(series: &Series) -> DataFusionResult>> { + // Rows are addressed on the statement's grid, so a row sits at the absolute position two + // series can collide on. An explicit axis has no such position, and a chunk never carries + // one: this is the refusal `EvaluatedChunk` already makes when the chunk is built. + if series.time.grid().is_none() { + return Err(DataFusionError::Execution( + "FusedSeriesExec requires evaluated series on a time grid".to_string(), + )); + } + let rows = series.time.len(); + let range = series.time.as_grid_range(); + let mut runs: Vec> = vec![Vec::new(); series.fields.len()]; + // A field that is one dense, null-free float run holds a value on every row the chunk carries, + // so it occupies exactly the range the chunk reports and no row of it has to be inspected. This + // is the shape the fused evaluator produces for a range function over ordinary float samples. + // Every other field is walked below, and a field filled here is skipped there. + let mut walked: Vec<(usize, &Samples)> = Vec::new(); + for (index, samples) in series.fields.iter().enumerate() { + match (range, samples.dense()) { + (Some(range), Some(DenseSamplesRef::Float(values))) + if rows > 0 && values.len() == rows && values.null_count() == 0 => + { + runs[index] = vec![Run { + start: range.start(), + end: range.start() + rows, + }]; + } + _ => walked.push((index, samples)), + } + } + if walked.is_empty() { + return Ok(runs); + } + // `positions` walks a chunk that forms one run and one whose rows form several the same way, + // in position order, so the holes between two runs are never read as rows. + let mut cursors: Vec<_> = walked.iter().map(|(_, samples)| samples.cursor()).collect(); + for (position, row) in series.time.positions() { + for ((index, _), cursor) in walked.iter().zip(&mut cursors) { + if !matches!(cursor.at(row), Some(sample) if !matches!(sample, SampleRef::Missing)) { + continue; + } + match runs[*index].last_mut() { + Some(run) if run.end == position => run.end += 1, + _ => runs[*index].push(Run { + start: position, + end: position + 1, + }), + } + } + } + Ok(runs) +} + +/// The field of `runs` that holds a value where `stored` already holds one, if any. +/// +/// A collision is two values of one field at one position, so each field is checked on its own. +/// Each incoming run is looked up rather than the stored history walked: the stored runs are sorted +/// and disjoint, so the only run that can share a position with `run` is the first one that does not +/// end before it, which is the look-up the runtime's own `Emitted::occupy` makes. Walking the +/// history instead costs a sparse series one pass over everything it stored per chunk - a query +/// whose chunks move forward one position at a time, a hole apart, is legal and would cost the +/// square of its chunk count under the lock every partition shares. +fn colliding_field(stored: &[Vec], runs: &[Vec]) -> Option { + for (field, (stored, new)) in stored.iter().zip(runs).enumerate() { + for run in new { + let at = stored.partition_point(|stored| stored.end <= run.start); + if stored.get(at).is_some_and(|stored| stored.overlaps(*run)) { + return Some(field); + } + } + } + None +} + +/// The batches one execution of [`FusedSeriesExec`] yields. +pub struct FusedSeriesStream { + /// The input frames, read as raw series chunks, or `None` once the stream ended in an error. + /// + /// The decoder behind it holds its current batch, the chunks it decoded from it and its own + /// source, so a terminated stream that kept it alive would hold upstream resources until the + /// consumer drops the stream. + input: Option, + /// The execution of this node's tree, or `None` once it has been closed. + executor: Option, + /// The position of the tree's input node, which is where every chunk read enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The tag column names, indexed by `OutputColumn::Tag`. + tags: Arc<[String]>, + /// The schema of every batch this stream yields. + output_schema: SchemaRef, + /// The label-set occupancy shared with the other partitions' streams. + /// + /// What it holds is charged to the plan's own reservation, attached by the first execution. + labels: SharedLabelSets, + /// The batches the executor settled since the last poll of the input. + ready: VecDeque, + /// Whether the input ended and everything it settled has been handed out. + done: bool, +} + +impl FusedSeriesStream { + /// Hands one raw chunk over and queues what it settled. + fn push(&mut self, chunk: RawSeriesChunk) -> DataFusionResult<()> { + let Some(executor) = self.executor.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec stream pushed a chunk after its execution closed".to_string(), + )); + }; + // The chunk enters at the tree's input node by position: a position is how the runtime + // addresses a node, and a tree that gains a second input names them the same way. + executor.push_raw_at(&self.input_path, chunk)?; + let mut settled = Vec::new(); + while let Some(chunk) = executor.next_chunk()? { + settled.push(chunk); + } + self.convert(settled) + } + + /// Ends the input, so the tree settles what it was holding. + /// + /// `finish()` would be the same thing plus a drain, but the drain collects: it holds the whole + /// answer in a `Vec` before anything is handed out, so a tree that answers many chunks after its + /// input ends would sit on all of them while the consumer waits for the first. Only the input is + /// ended here, and the answers are pulled one per call below, so each chunk leaves as the tree + /// settles it and what has not been pulled yet stays charged to the node holding it. + fn close_input(&mut self) -> DataFusionResult<()> { + match self.executor.as_mut() { + Some(executor) => executor.end_of_input(), + // Every answer was pulled and the execution closed itself. + None => Ok(()), + } + } + + /// Pulls the next answer the tree settles after its input ended, or `None` when it has none. + /// + /// One chunk per call is what keeps this a stream: collecting them first would hold the whole + /// answer while the consumer waits for the first, and would charge the pool for all of it at + /// once. `None` is also when the execution closes and returns its reservations, so it is + /// dropped here rather than staying alive until the stream is. + fn pull_one(&mut self) -> DataFusionResult> { + let Some(executor) = self.executor.as_mut() else { + return Ok(None); + }; + match executor.next_chunk() { + Ok(Some(chunk)) => Ok(Some(chunk)), + Ok(None) => { + self.executor = None; + Ok(None) + } + Err(error) => Err(error), + } + } + + /// Ends the stream in an error, releasing the execution and the input at once. + /// + /// An ordinary drop of the stream releases all of this as well, so this is not about the last + /// release: it is about a consumer that keeps the terminated stream, which would otherwise + /// hold the evaluator, the decoder's held batch and chunks, and the plan upstream of it. + fn fail(&mut self) { + self.done = true; + self.executor = None; + self.input = None; + // The label-set map and its charge are the plan's, so a stream stopping early gives back + // nothing: the entries another partition is still checking against have to stay charged, and + // are released when the plan itself is dropped. + // A batch queued before the error was settled by a run that has now stopped. + self.ready.clear(); + } + + /// Converts the chunks the executor settled into batches, in the order it produced them. + fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { + for chunk in &chunks { + // The label set is claimed before its batch is queued, so a collision surfaces at the + // chunk that reached it and nothing this plan would refuse is handed out after it. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .claim(&chunk.series)?; + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + chunk, + )?; + self.ready.push_back(batch); + } + Ok(()) + } +} + +impl RecordBatchStream for FusedSeriesStream { + fn schema(&self) -> SchemaRef { + self.output_schema.clone() + } +} + +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 let Some(batch) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(batch))); + } + if this.done { + return Poll::Ready(None); + } + // An error is the only thing that clears the input, and it ends the stream with it. + let Some(input) = this.input.as_mut() else { + // The input ended, so what the tree answers now comes from what it held. One chunk + // per iteration keeps it a stream, and `None` is both "nothing more" and the point + // where the execution releases its reservations. + match this.pull_one() { + Ok(Some(chunk)) => { + if let Err(error) = this.convert(vec![chunk]) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Ok(None) => { + this.done = true; + continue; + } + Err(error) => { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + continue; + }; + let polled = input.as_mut().poll_next(cx); + match polled { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + // The input ended, so every series it opened has to be closed. The tree settles + // what it was holding, and the answers are pulled below in this same poll + // rather than collected first. + this.input = None; + if let Err(error) = this.close_input() { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + // An input error stops the execution: a stopped evaluator must not be reused. + this.fail(); + return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); + } + Poll::Ready(Some(Ok(chunk))) => { + if let Err(error) = this.push(chunk) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + } + // 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 new file mode 100644 index 00000000000..017f4ef6628 --- /dev/null +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -0,0 +1,2882 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! What this node has to get right: the two `EXPLAIN` lines it shows, the batches it runs end to +//! end, and what it refuses to run. +//! +//! The trees here are built node by node with the same bindings the PromQL compiler produces for +//! `rate(x[5m])`: the compiler itself lives in the `query` crate, which depends on this one, so a +//! test in this crate cannot call it. +//! +//! The input is what a scan hands over: native raw series frames, one batch per chunk here, packed +//! by the same code the scan uses ([`pack_raw_series_chunk`]). A frame carries the series' encoded +//! key, the chunk's finality, its timestamps and its fields, plus the complete labels, so a test +//! says which series a chunk belongs to and whether it closes it instead of relying on where a row +//! boundary happens to fall. +//! +//! Kernel numerics are not the subject: `rate` is checked against a counter that advances one +//! second per second, so the value it has to produce is exactly 1 (or 2) at every grid point and +//! any position or finalization mistake shows up as a wrong, missing or duplicated row. + +use std::collections::VecDeque; +use std::fmt::{self, Display, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::ToDFSchema; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +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, MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, collect, displayable, +}; +use datafusion::prelude::SessionContext; +use datatypes::series::frame::{RawLabelMode, pack_raw_series_chunk, raw_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, MetricName, SeriesOrdering}; +use futures::{Stream, StreamExt, TryStreamExt}; +use store_api::region_engine::RawLabelSource; + +use super::{ + FusedSeries, FusedSeriesBinding, LabelSets, OutputColumn, POLL_BUDGET, label_set, + occupied_runs, series_to_record_batch, +}; +use crate::evaluated::EvaluatedChunk; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundWindow, Grouping, + NativeExpr, NativeNode, NativeTree, NativeType, +}; +use crate::native::runs::absorb_rebuilds; +use crate::range_eval::{RangeFunction, RangeWindow}; + +/// The query every tree here is compiled from. +const PROMQL: &str = "rate(x[5m])"; +/// The range the statement's window selects, which is the `5m` of [`PROMQL`]. +const RANGE_MS: i64 = 300_000; +/// The statement's grid: `rate(x[5m])` evaluated every minute from five minutes in. +const GRID_START: i64 = 300_000; +const GRID_STEP: i64 = 60_000; +const GRID_LEN: u64 = 5; +/// One raw sample per minute, from 0 to 9 minutes: enough to fill every window on the grid. +const SAMPLE_STEP: i64 = 60_000; +const SAMPLES: i64 = 10; +/// How many one-sample chunks the quiet series below arrives in. +/// +/// Two bounds hold at once: the series reaches the end of the last grid point's window (540s, so +/// indices `0..=540`), which is what lets it rate at 1 over the whole grid, and it arrives in more +/// batches than one poll's budget, which is what makes a poll that ignored the budget visible. +const QUIET_SAMPLES: i64 = 541; +/// One second between those samples, so the series rates at 1 once it is read to its end. +/// +/// The step is a second and not a minute on purpose: at a minute the series would arrive in ten +/// batches, fewer than [`POLL_BUDGET`], and the budget the test above exists to check would never +/// be reached. +const QUIET_STEP: i64 = 1_000; + +const TIME: &str = "time"; +const VALUE: &str = "value"; +const OTHER: &str = "other"; +const HOST: &str = "host"; +const ZONE: &str = "zone"; +const NAME: &str = "__name__"; +/// A metadata label, which the window drops along with the metric name. +const TYPE: &str = "__type__"; +const METRIC: &str = "http_requests_total"; + +fn grid() -> Arc { + Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ) +} + +/// The raw schema the tree's input is bound to: one float field called [`VALUE`]. +fn raw_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: VALUE.into() }]), + eval_grid: None, + }) +} + +/// A raw schema of two float fields, so one series can fill one of them and another the other. +fn two_field_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([ + FieldSpec { name: VALUE.into() }, + FieldSpec { name: OTHER.into() }, + ]), + eval_grid: None, + }) +} + +/// The tree `rate(x[5m])` compiles to over `schema`: raw samples folded by one window onto the grid. +fn rate_tree_of(schema: Arc) -> NativeTree { + let grid = grid(); + // One schema shared by the window's binding and its output, which is what `NativeTree`'s + // validation compares by identity. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: schema.fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(schema); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + // The input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: grid, + }, + }, + output: evaluated, + }) + .unwrap() +} + +/// The tree `rate(x[5m])` compiles to over one field. +fn rate_tree() -> NativeTree { + rate_tree_of(raw_schema()) +} + +/// An instant selection feeding a range window. +fn instant_window_tree() -> NativeTree { + let instant = instant_tree().root().clone(); + let grid = instant + .output + .samples() + .expect("an instant selection produces samples") + .eval_grid + .clone() + .expect("an instant selection's output is evaluated on a grid"); + let evaluated = instant.output.clone(); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(instant), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + input_grid: Some(Arc::clone(&grid)), + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// Two windows stacked, `rate(last_over_time(x[1s])[2s])`, which no compiler produces. +/// +/// A range function has to apply directly to a selector, and a selector is raw samples, so a tree +/// with a window above a window is only reachable by writing it down by hand - which is a tree +/// [`FusedSeriesExec`](super::FusedSeriesExec) accepts, and therefore one it has to answer correctly. +/// +/// The inner window is the interesting one: `last_over_time` keeps `__name__`, so the pair it can +/// split is not two series whose names it dropped but `{host="h", zone=""}` and `{host="h"}`, one +/// canonical label set and two raw keys. Its fragments are not merged unless the node above it +/// merges them, and that merge happens once per execution: with one execution per partition each +/// partition sees one fragment, so the outer window reads one half of the pair. +fn nested_window_tree() -> NativeTree { + let grid = grid(); + // The inner window's output, which is also the outer window's output: a window reports one value + // per input field on its own grid. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::LastOverTime, + window: RangeWindow { + range_ms: 1_000, + offset_ms: 0, + anchor_ms: None, + }, + // The inner window's input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated.clone(), + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 2_000, + offset_ms: 0, + anchor_ms: None, + }, + // The outer window reads the inner one's evaluated output, on that output's grid, + // and reports one value per field there like the inner one does. + input_grid: Some(Arc::clone(&grid)), + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// The same tree asking for `quantile`, whose state the runtime does not keep yet. +fn quantile_aggregate_tree() -> NativeTree { + let complete = aggregate_tree(); + let root = complete.root(); + let NativeExpr::Aggregate { input, op } = &root.expr else { + unreachable!("the fixture is an aggregate over a window"); + }; + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: input.clone(), + op: BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: op.group.clone(), + stage: AggregateStage::Complete, + }, + }, + output: root.output.clone(), + }) + .expect("the bound tree is legal") +} + +/// The `rate(x[5m])` tree with a `sum by (job)` over it. +/// +/// A `Complete` aggregate: every contribution to each group lies inside one execution. +fn aggregate_tree() -> NativeTree { + let rate = rate_tree(); + let window = rate.root().clone(); + let grid = window + .output + .samples() + .expect("a window produces samples") + .eval_grid + .clone() + .expect("a window's output is evaluated on a grid"); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(window), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: Grouping::By(Arc::new(["job".into()])), + stage: AggregateStage::Complete, + }, + }, + output: NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(grid), + })), + }) + .expect("the bound tree is legal") +} + +/// The tree an instant selection compiles to: the last raw sample at or before each grid point. +/// +/// Unlike a window or a map, a selection rewrites no label: the series it answers carries the labels +/// its input did. +fn instant_tree() -> NativeTree { + let grid = grid(); + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundInstant { + lookback_ms: RANGE_MS, + offset_ms: 0, + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// What the tree's input was compiled against, for the input schema and output tags given. +fn binding_of( + tree: &NativeTree, + series_schema: Arc, + tags: Vec, +) -> FusedSeriesBinding { + FusedSeriesBinding { + series_schema, + grid: tree + .root() + .output + .samples() + .expect("this tree produces samples") + .eval_grid + .clone() + .expect("this tree's output is evaluated on a grid"), + metric_name: Some(Arc::from(METRIC)), + time_index: TIME.to_string(), + tags, + } +} + +/// What the tree's input was compiled against, read off the tree itself. +fn binding(tree: &NativeTree) -> FusedSeriesBinding { + binding_of(tree, raw_schema(), vec![HOST.to_string(), NAME.to_string()]) +} + +/// The schema of a frame carrying series of `schema` with their labels explicit. +fn frame_schema_of(schema: &SeriesSchema) -> SchemaRef { + raw_series_frame_schema(schema, RawLabelMode::Explicit).expect("explicit labels are a frame") +} + +/// The schema of a frame carrying these series, which is what the node declares as its input. +fn frame_schema() -> SchemaRef { + frame_schema_of(&raw_schema()) +} + +/// The columns this node produces: the field, the tags and the grid timestamps. +fn output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +fn input_plan(schema: SchemaRef) -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.to_dfschema_ref().unwrap(), + }) +} + +/// A physical input producing `batches` in one partition, with `schema` as its schema. +fn memory_exec_with(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) +} + +/// A physical input producing `batches` in one partition, on the frame schema of these tests. +fn memory_exec(batches: Vec) -> Arc { + memory_exec_with(frame_schema(), batches) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own, with +/// `schema` as its schema. +fn memory_exec_partitioned_with( + schema: SchemaRef, + partitions: Vec>, +) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&partitions, schema, None) + .expect("every partition holds batches of the given frame schema"), + ))) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own. +fn memory_exec_partitioned(partitions: Vec>) -> Arc { + memory_exec_partitioned_with(frame_schema(), partitions) +} + +/// A task context whose memory pool is `pool`, so what the node charges can be read afterwards. +fn task_context(pool: &Arc) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(pool) as Arc) + .build() + .expect("the runtime environment builds"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) +} + +/// A physical input whose one partition's stream `make` builds. +fn staged( + schema: SchemaRef, + make: impl Fn() -> SendableRecordBatchStream + Send + Sync + 'static, +) -> Arc { + Arc::new(StagedPlan { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + make: Box::new(make), + }) +} + +/// A physical input, so a test can stage what a memory source cannot: a pending poll, an error +/// after some rows, and a stream whose drop it can see. +struct StagedPlan { + make: Box SendableRecordBatchStream + Send + Sync>, + properties: Arc, +} + +impl std::fmt::Debug for StagedPlan { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("StagedPlan").finish() + } +} + +impl DisplayAs for StagedPlan { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { + Ok(()) + } +} + +impl ExecutionPlan for StagedPlan { + fn name(&self) -> &str { + "StagedPlan" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + // A scan answers the partitions it declares and no others, which is what makes an index + // past this plan's count an error the node above has to pass on. + if partition >= self.properties.partitioning.partition_count() { + return Err(DataFusionError::Execution(format!( + "StagedPlan has no partition {partition}" + ))); + } + Ok((self.make)()) + } +} + +/// The stream a [`StagedPlan`] hands over: `batches`, then an optional error, with a pending poll +/// and a drop a test can observe. +struct StagedStream { + schema: SchemaRef, + batches: VecDeque, + handed: usize, + /// Hand this many batches, then answer [`Poll::Pending`] once. + pending_after: Option, + /// The error to report once the batches are exhausted. + error: Option<&'static str>, + /// Set when this stream is dropped, so a test can see the input being released. + dropped: Option>, + /// Counts the batches this stream has handed over, so a test can see how far one poll read. + handed_out: Option>, +} + +impl StagedStream { + fn new(schema: SchemaRef, batches: Vec) -> Self { + Self { + schema, + batches: batches.into(), + handed: 0, + pending_after: None, + error: None, + dropped: None, + handed_out: None, + } + } + + fn pending_after(mut self, batches: usize) -> Self { + self.pending_after = Some(batches); + self + } + + fn error(mut self, message: &'static str) -> Self { + self.error = Some(message); + self + } + + fn dropped(mut self, flag: &Arc) -> Self { + self.dropped = Some(Arc::clone(flag)); + self + } + + fn counted(mut self, counter: &Arc) -> Self { + self.handed_out = Some(Arc::clone(counter)); + self + } +} + +impl Stream for StagedStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if this.pending_after == Some(this.handed) { + this.pending_after = None; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + match this.batches.pop_front() { + Some(batch) => { + this.handed += 1; + if let Some(counter) = &this.handed_out { + counter.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(Some(Ok(batch))) + } + None => match this.error.take() { + Some(message) => { + Poll::Ready(Some(Err(DataFusionError::Execution(message.to_string())))) + } + None => Poll::Ready(None), + }, + } + } +} + +impl RecordBatchStream for StagedStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Drop for StagedStream { + fn drop(&mut self) { + if let Some(dropped) = &self.dropped { + dropped.store(true, Ordering::SeqCst); + } + } +} + +fn fused(tree: NativeTree) -> FusedSeries { + let binding = binding(&tree); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The physical node of `node` over `input`, reading the complete labels every frame carries. +fn plan_of( + node: FusedSeries, + input: Arc, +) -> DataFusionResult> { + node.to_execution_plan(input, RawLabelSource::Explicit, SeriesOrdering::Contiguous) +} + +/// The physical node of `tree` over `input`. +fn plan( + tree: NativeTree, + input: Arc, +) -> DataFusionResult> { + plan_of(fused(tree), input) +} + +/// The timestamps of `rows`. +fn times_of(rows: &[(i64, f64)]) -> Vec { + rows.iter().map(|(time, _)| *time).collect() +} + +/// The values of `rows`, one sample per timestamp. +fn values_of(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|(_, value)| Some(*value)).collect() +} + +/// One missing sample per timestamp of `rows`, which is a field the series holds no value in. +fn no_values(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|_| None).collect() +} + +/// One chunk of `metric`, keyed `key`, holding `times` and one sample per field of `schema`. +fn chunk_of( + schema: &SeriesSchema, + key: &str, + metric: MetricName, + times: &[i64], + values: Vec>>, + last: bool, +) -> RawSeriesChunk { + assert_eq!( + schema.fields.len(), + values.len(), + "one value list per field of the schema" + ); + assert!( + values.iter().all(|values| values.len() == times.len()), + "one sample per timestamp per field" + ); + let fields = values + .into_iter() + .map(|values| Samples::from_float(Float64Array::from(values))) + .collect::>(); + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key.as_bytes()), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).expect("ascending timestamps"), + SeriesFields::try_from_vec(fields).expect("one sample per timestamp per field"), + last, + ) + .expect("one series' samples are a legal raw chunk") +} + +/// One chunk of the series `{__name__=METRIC, host=}`: `rows` as its samples, and the chunk +/// that closes the series when `last`. The host names the series, so it doubles as the key. +fn series_chunk(host: &str, rows: &[(i64, f64)], last: bool) -> RawSeriesChunk { + chunk_of( + &raw_schema(), + host, + metric(&[(HOST, host)], Some(METRIC)), + ×_of(rows), + vec![values_of(rows)], + last, + ) +} + +/// The metric these tests read, with `labels` on it and the name given. +fn metric(labels: &[(&str, &str)], name: Option<&str>) -> MetricName { + MetricName::try_new( + name.map(Arc::from), + labels + .iter() + .map(|(name, value)| Label { + name: Arc::from(*name), + value: Arc::from(*value), + }) + .collect(), + ) + .expect("the labels of these tests are legal") +} + +/// One frame batch carrying `chunk` of `schema`, packed by the scan's own encoder. +fn frame_of(schema: &SeriesSchema, chunk: &RawSeriesChunk) -> RecordBatch { + let mut batches = pack_raw_series_chunk(schema, RawLabelMode::Explicit, chunk) + .expect("a chunk of this schema packs into a frame"); + assert_eq!(1, batches.len(), "one small chunk is one frame batch"); + batches.remove(0) +} + +/// One frame batch carrying `chunk` of the one-field schema of these tests. +fn frame(chunk: &RawSeriesChunk) -> RecordBatch { + frame_of(&raw_schema(), chunk) +} + +/// The frames carrying `chunks` of `schema`, one batch each, so a test can say how much input one +/// poll reads. +fn frames_of(schema: &SeriesSchema, chunks: &[RawSeriesChunk]) -> Vec { + chunks.iter().map(|chunk| frame_of(schema, chunk)).collect() +} + +/// The frames carrying `chunks` of the one-field schema of these tests, one batch each. +fn frames(chunks: &[RawSeriesChunk]) -> Vec { + frames_of(&raw_schema(), chunks) +} + +/// One series' samples: one per minute from 0, `scale` seconds per second. +fn series_samples(scale: f64) -> Vec<(i64, f64)> { + (0..SAMPLES) + .map(|index| { + let time = index * SAMPLE_STEP; + (time, scale * time as f64 / 1000.0) + }) + .collect() +} + +/// Two series, each split over two chunks. A series' chunks stay adjacent, and the second of each +/// pair closes it, which is the input contract: a series' chunks arrive together. +fn two_series_frames() -> Vec { + let first = series_samples(1.0); + let second = series_samples(2.0); + frames(&[ + series_chunk("web-1", &first[..3], false), + series_chunk("web-1", &first[3..], true), + series_chunk("web-2", &second[..3], false), + series_chunk("web-2", &second[3..], true), + ]) +} + +/// One series split over three chunks, so its chunks have to be stitched. +fn one_series_frames() -> Vec { + let rows = series_samples(1.0); + frames(&[ + series_chunk("web-1", &rows[..4], false), + series_chunk("web-1", &rows[4..7], false), + series_chunk("web-1", &rows[7..], true), + ]) +} + +/// The rows of `batches`: timestamp, value, host and metric name. +type Row = (i64, Option, Option, Option); + +fn rows(batches: &[RecordBatch]) -> Vec { + let mut rows = Vec::new(); + for batch in batches { + let time = batch + .column_by_name(TIME) + .expect("every batch has a time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let value = batch + .column_by_name(VALUE) + .expect("every batch has the field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + let host = batch + .column_by_name(HOST) + .expect("every batch has the tag column") + .as_any() + .downcast_ref::() + .expect("the tag column is Utf8"); + let name = batch + .column_by_name(NAME) + .expect("every batch has the metric name column") + .as_any() + .downcast_ref::() + .expect("the metric name column is Utf8"); + for row in 0..batch.num_rows() { + rows.push(( + time.value(row), + value.is_valid(row).then(|| value.value(row)), + host.is_valid(row).then(|| host.value(row).to_string()), + name.is_valid(row).then(|| name.value(row).to_string()), + )); + } + } + rows +} + +/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. +fn zone_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// `rate(x[5m])` over frames that carry `zone` as well as `host`, with both in the output. +fn zone_node() -> FusedSeries { + let tree = rate_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The instant-selection node over frames that carry `zone` as well as `host`. +/// +/// Its output labels are the input's, so an empty-valued label reaches this node's own state as the +/// frame wrote it. +fn instant_node() -> FusedSeries { + let tree = instant_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The two chunks of the collision: one series once with an empty `zone` and once with no `zone`. +/// +/// They are one identity to Prometheus — an empty label value is no label — but two different raw +/// series: the scan encodes the empty value and the absent one differently, so they reach the +/// runtime as two series that close on the same label set. +fn zone_chunks() -> Vec { + let rows = series_samples(1.0); + vec![ + chunk_of( + &raw_schema(), + "zone-empty", + metric(&[(HOST, "web-1"), (ZONE, "")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + chunk_of( + &raw_schema(), + "zone-absent", + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + ] +} + +/// The canonical label set these tests collide on, as the runtime renders it in its refusal. +const DUPLICATE_LABEL_SET: &str = + "the node at root produces two series under the label set [(\"host\", \"web-1\")]"; + +/// Two raw series that are already one Prometheus identity are refused by the runtime that would +/// have to answer both. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one identity to Prometheus: an empty label +/// value is no label, so the two chunks are the same series and `rate` creates no collision here. +/// What the two *raw* keys differ in is the empty `zone` against the absent one, which the scan +/// encodes differently and the decoder therefore opens two input series for. `rate` drops the +/// metric name and the empty value, so the second raw series arrives on the canonical label set the +/// first one already occupies; the refusal names that label set rather than answering it, because +/// answering would report one point twice. +#[tokio::test] +async fn differing_raw_keys_that_close_on_one_label_set_are_refused() { + let exec = plan_of(zone_node(), memory_exec(frames(&zone_chunks()))) + .expect("the runtime executes a rate tree"); + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second raw key cannot be answered under the first one's label set"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// A runtime error after buffered rows is reported, and the terminated stream releases its input. +/// +/// The decoder holds the batch it is reading, the chunks it decoded from it and its own source, so a +/// stream that kept its input after the error would hold all three until the consumer drops it. The +/// collision above is the error: it happens on the second series, after the first one's rows have +/// been handed out. +#[tokio::test] +async fn a_runtime_error_after_rows_releases_the_input() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = frames(&zone_chunks()); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).dropped(&dropped)) + }); + let mut stream = plan_of(zone_node(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!( + "the collision was answered instead of refused: {:?}", + rows(&answered) + ), + } + }; + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); + // The first series' rows were handed out before the error, and the input is gone although the + // terminated stream is still alive here. + assert!( + !answered.is_empty(), + "the first series is answered before the error" + ); + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// Two series with *different* names collapse once `rate` drops the name. +/// +/// This is the collision the canonical label set exists for, at the layer that sees it: +/// `http_requests_total{host="web-1"}` and `http_errors_total{host="web-1"}` reach the window as two +/// identities and leave it as `{host="web-1"}`. [`FusedSeries`] reads one metric per scan, so this +/// collision cannot be staged through it; the runtime is driven directly, with the names on the raw +/// chunks. +#[test] +fn two_differently_named_series_collide_once_rate_drops_the_name() { + let tree = rate_tree(); + let mut execution = NativeExecutor::new( + NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + ) + .expect("the runtime executes a rate tree"), + ); + let samples = series_samples(1.0); + execution + .push_raw(chunk_of( + &raw_schema(), + METRIC, + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the first series is a legal raw chunk"); + assert!( + execution + .next_chunk() + .expect("the first series runs") + .is_some(), + "the first series is answered" + ); + execution + .push_raw(chunk_of( + &raw_schema(), + "http_errors_total", + metric(&[(HOST, "web-1")], Some("http_errors_total")), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the second series is a legal raw chunk"); + let error = execution + .next_chunk() + .expect_err("the second series puts a value where the first one already holds one"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// The grid timestamps of the statement, in order. +fn grid_timestamps() -> Vec { + (0..GRID_LEN) + .map(|index| GRID_START + index as i64 * GRID_STEP) + .collect() +} + +/// A waker a test can ask whether it was woken. +/// +/// A poll that stops on its own budget has to wake itself, and the default waker a `poll!` uses +/// swallows that, so the tests that check it poll with this one. +#[derive(Default)] +struct CountingWake(AtomicBool); + +impl futures::task::ArcWake for CountingWake { + fn wake_by_ref(woken: &Arc) { + woken.0.store(true, Ordering::SeqCst); + } +} + +impl CountingWake { + fn woken(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + +/// One series as many chunks that settle nothing: one sample per chunk, all before the grid starts. +/// +/// The counter advances one second per second, so `rate` answers 1 at every grid point once the +/// series is read to its end; until a sample reaches the grid's first point, a push decides +/// nothing and settles no chunk. +fn quiet_series_frames() -> Vec { + frames( + &(0..QUIET_SAMPLES) + .map(|index| { + series_chunk( + "web-1", + &[(index * QUIET_STEP, index as f64)], + index == QUIET_SAMPLES - 1, + ) + }) + .collect::>(), + ) +} + +/// Asserts that `rows` are `rates`' series, each over the whole grid in order. +fn assert_series_rates(rows: &[Row], rates: &[(f64, &str)]) { + let expected = rates + .iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, *rate, *host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } +} + +fn count(metrics: &MetricsSet, name: &str) -> usize { + match metrics.sum_by_name(name) { + Some(MetricValue::Count { count, .. }) => count.value(), + other => panic!("no count named {name}: {other:?}"), + } +} + +/// Prints `node` the way the logical `EXPLAIN` prints it. +struct Explain<'a>(&'a FusedSeries); + +impl Display for Explain<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.0.fmt_for_explain(f) + } +} + +/// The logical node prints the query a reader wrote, not the tree it compiled to. +#[test] +fn fmt_for_explain_prints_the_promql() { + assert_eq!( + "FusedSeries: rate(x[5m])", + Explain(&fused(rate_tree())).to_string() + ); +} + +/// The physical node prints the tree it was built from, rendered by the tree's own display. +#[test] +fn fmt_as_prints_the_native_tree() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let text = displayable(exec.as_ref()).indent(false).to_string(); + + assert!( + text.starts_with("FusedSeriesExec: rate(x[5m])"), + "the node line names the query first: {text}" + ); + // A distinctive line of the compiled tree: the window's bound range, offset, anchor and grid. + assert!( + text.contains("root Window rate [300000ms] offset 0 anchor - on grid 300000/60000x5"), + "the node line carries the tree it runs: {text}" + ); + assert!( + text.contains("input Input raw samples"), + "the node line carries the tree's source: {text}" + ); +} + +/// `rate(x[5m])` runs end to end over two raw series: values, tags and the dropped metric name. +#[tokio::test] +async fn runs_rate_over_two_series() { + let exec = plan(rate_tree(), memory_exec(two_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // A counter advancing one second per second rates at 1, and at 2 for the series that advances + // two: the value at each grid point says the right window was read for the right series. + let expected = [(1.0, "web-1"), (2.0, "web-2")] + .into_iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, rate, host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } + // `rate` drops `__name__`, so the tag column that reports it is empty for every row. + assert!( + rows.iter().all(|(_, _, _, name)| name.is_none()), + "rate drops the metric name: {rows:?}" + ); +} + +/// A series that spans several input chunks keeps its absolute grid positions and closes once. +#[tokio::test] +async fn a_series_spanning_chunks_keeps_absolute_positions_and_closes_once() { + let exec = plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // Every grid position exactly once, in order: a chunk's range is never read as a row number. + let times = rows.iter().map(|row| row.0).collect::>(); + assert_eq!(grid_timestamps(), times); + assert!( + rows.iter().all(|(_, value, ..)| value.is_some()), + "every grid point of this series has a value: {rows:?}" + ); + + // The source closed the series on exactly one chunk, and the window saw exactly that one + // closing: one `is_last`, not one per input batch. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(1, count(&metrics, "native.root.input.completed_series")); + assert_eq!(1, count(&metrics, "native.root.completed_series")); +} + +/// A tree the runtime does not execute is refused when the execution plan is built. +/// +/// The runtime folds the supported `Complete` kinds now, so the tree below asks for one it still +/// refuses on purpose: `quantile` needs the parameter and the member buffer that the other kinds +/// do not. What the test pins is the refusal happening while the plan is built, not the particular +/// kind of node that provokes it. +#[tokio::test] +async fn a_tree_the_runtime_cannot_run_is_rejected_when_the_exec_is_built() { + let node = fused(quantile_aggregate_tree()); + let error = plan_of(node, memory_exec(Vec::new())) + .expect_err("the runtime does not execute a quantile aggregate"); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "the runtime refuses the tree it cannot run: {error}" + ); + assert!(error.to_string().contains("Quantile"), "{error}"); +} + +/// A tree evaluated on another grid than the node declares is refused before it runs. +#[test] +fn a_tree_bound_to_another_grid_is_rejected() { + let tree = rate_tree(); + let other_grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, GRID_STEP, GRID_LEN).unwrap()).unwrap()); + let binding = FusedSeriesBinding { + grid: other_grid, + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect_err("the tree is not evaluated on the declared grid"); + assert!(error.to_string().contains("grid"), "{error}"); +} + +/// A grid allocated separately from the tree's, with the same spec, is the tree's grid. +/// +/// The comparison is by spec, which is the rule [`NativeTree`]'s own validation and the kernels +/// apply, and a chunk's timestamps come from its absolute grid positions rather than from the +/// allocation, so nothing downstream needs the `Arc` to be shared. This pairs with +/// [`a_tree_bound_to_another_grid_is_rejected`]: a different spec is still refused. +#[test] +fn a_separately_allocated_grid_with_an_equal_spec_is_accepted() { + let tree = rate_tree(); + let equal = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ); + assert!( + !Arc::ptr_eq(&equal, &binding(&tree).grid), + "the two grids are separate allocations" + ); + let binding = FusedSeriesBinding { + grid: equal, + ..binding(&tree) + }; + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("an equal grid spec is the same grid"); +} + +/// An input that ends without a series still runs the tree, and yields no rows. +#[tokio::test] +async fn an_empty_input_runs_the_tree_and_yields_no_rows() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("a tree with no series in it still closes"); + assert!(rows(&batches).is_empty(), "{:?}", rows(&batches)); + // The run closed on no series, so nothing was completed and nothing was decided. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(0, count(&metrics, "native.root.completed_series")); + assert_eq!(0, count(&metrics, "native.root.input.completed_series")); +} + +/// An upstream error after some rows is reported, and the terminated stream releases that input. +#[tokio::test] +async fn an_upstream_error_after_rows_is_reported() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + .error("the source failed") + .dropped(&dropped), + ) + } + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The error is reported rather than swallowed by the rows the input already handed over. + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!("the source failure was swallowed: {:?}", rows(&answered)), + } + }; + assert!( + error.to_string().contains("the source failed"), + "reported another error: {error}" + ); + // Still alive here: the release is the stream's, not the drop of a temporary. + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// A pending poll loses nothing when the next-item future is dropped and the stream polled again. +#[tokio::test] +async fn a_pending_boundary_loses_no_rows() { + let expected = rows( + &collect( + plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"), + SessionContext::default().task_ctx(), + ) + .await + .expect("the node runs"), + ); + assert!(!expected.is_empty(), "the fixture decides some grid points"); + + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + // One batch is in the decoder's hands when the stream reports pending. + .pending_after(1), + ) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The future is dropped at the end of this block, as a consumer that gives up waiting does. + { + let next = stream.next(); + futures::pin_mut!(next); + assert!( + futures::poll!(next.as_mut()).is_pending(), + "the input's pending poll reaches the consumer" + ); + } + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the pending poll left it"); + assert_eq!(expected, rows(&batches)); +} + +/// A frame whose labels are encoded in its key is refused when the plan is built. +/// +/// The tags this node declares are read off each chunk's labels, and a key-mode frame carries none: +/// decoding it with a binding that promises complete labels would produce chunks that name no +/// series. The frame's own declared layout says which of the two it is, so the mismatch is caught +/// before the first batch, not when the missing labels reach the output. +#[test] +fn a_frame_whose_labels_are_encoded_is_rejected_when_the_plan_is_built() { + let keyed = + raw_series_frame_schema(&raw_schema(), RawLabelMode::Key).expect("key labels are a frame"); + let error = plan(rate_tree(), memory_exec_with(keyed, Vec::new())) + .expect_err("the frames carry no labels to read tags from"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// A frame that carries an evaluation grid is refused when the plan is built. +/// +/// A frame holds raw series: its rows are samples, not grid positions, and the grid every output +/// row is addressed on is the tree's own. A schema that declares a grid is one this node would +/// read positions out of, so it is refused rather than decoded. The source below is one partition, +/// which is the shape the node accepts, so what refuses the plan here is the frame's own +/// declaration and not the node's partition count: the refusal has to name the grid. +#[test] +fn a_frame_with_an_evaluation_grid_is_rejected_when_the_plan_is_built() { + let mut metadata = frame_schema().metadata().clone(); + metadata.insert("greptime.series.grid".to_string(), "v1".to_string()); + let gridded = Arc::new(Schema::new_with_metadata( + frame_schema().fields().clone(), + metadata, + )); + // One empty partition: every check but the frame's passes. + let source = memory_exec_with(gridded, Vec::new()); + assert_eq!( + 1, + source.properties().output_partitioning().partition_count(), + "the node only looks at the frame when its input has a partition" + ); + let error = plan(rate_tree(), source).expect_err("the frame declares a grid"); + let message = error_chain(&error); + assert!( + message.contains("raw series schema cannot have an evaluation grid"), + "refused for another reason: {message}" + ); + assert!( + !message.contains("no partitions"), + "the partition count is not what refuses this plan: {message}" + ); +} + +/// A frame of another series schema is refused when the plan is built. +/// +/// The tree's input node was compiled against one schema, and a frame for a different one names +/// fields the tree does not read and omits ones it does: decoding it would hand the tree a column +/// it never bound. Nothing about the two schemas says which one was meant, so the frame is refused +/// instead of half-read. +#[test] +fn a_frame_of_another_series_schema_is_rejected_when_the_plan_is_built() { + let other = frame_schema_of(&two_field_schema()); + let error = plan(rate_tree(), memory_exec_with(other, Vec::new())) + .expect_err("the tree's input reads one field"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// An input with no partitions is refused when the plan is built, and no other count is. +/// +/// The node runs one evaluator per input partition and adds none of its own, so an input with no +/// partition is one that could answer nothing: the plan is refused rather than built as a node that +/// would report an empty result for data it never read. Any non-zero count is the input's to +/// choose: the node takes it and answers each partition's own series. +#[test] +fn an_input_with_no_partitions_is_rejected_when_the_plan_is_built() { + let schema = frame_schema(); + let empty = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[], Arc::clone(&schema), None) + .expect("no partitions is a legal source"), + ))); + let error = plan(rate_tree(), empty).expect_err("there is no partition to run"); + assert!( + error.to_string().contains("input has no partitions"), + "refused for another reason: {error}" + ); + // The only count refused is zero: two partitions build and answer, which is what the + // multi-partition tests below run. + let two = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("a two-partition input is not refused"); + assert_eq!( + 2, + two.properties().output_partitioning().partition_count(), + "the node takes the input's count" + ); +} + +/// Two windows stacked are refused over more than one partition, because the inner one's fragments +/// need not have been merged. +/// +/// The counterexample is legal PromQL once nesting is expressible: the inner `last_over_time` runs +/// over two raw series that are one canonical label set, `{host="h", zone=""}` holding `(1s, 1)` and +/// `{host="h"}` holding `(2s, 2)`, and the outer `rate` needs both in one window. In a single +/// partition the outer window merges the two fragments and answers `(2s, 1)`; in two partitions each +/// executor is handed one fragment and has one sample to rate, so the outer window answers nothing. +/// That is an empty result rather than an error, which is why no check further down can see it. +/// +/// `last_over_time` keeps `__name__`, so a check that looked only for a dropped metric name accepts +/// this tree. What makes it unsafe is that a window's output can hold two series under one label set +/// whenever it did not merge its input, which a bare tree cannot say. +/// +/// One partition stays accepted: there is nothing to split there, which is what makes the refusal +/// about the split rather than about the shape of the tree. +#[test] +fn a_nested_window_over_two_partitions_is_refused() { + let error = plan( + nested_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each of two partitions would merge only the fragment it was given"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the inner window: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the window above it: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); + + // The same tree over one partition is accepted, because one execution sees every fragment and + // merges them into one series before the window above it evaluates. + plan(nested_window_tree(), memory_exec(Vec::new())) + .expect("a single partition cannot split a logical series"); +} + +/// An instant selection can leave distinct raw labels on one canonical label set for a window. +#[test] +fn an_instant_then_window_over_two_partitions_is_refused() { + let error = plan( + instant_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each partition could hand the window only one fragment"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the instant selection: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the consuming window: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); +} + +/// A single partition cannot split a logical series, even when an instant feeds a window. +#[test] +fn an_instant_then_window_over_one_partition_is_accepted() { + plan(instant_window_tree(), memory_exec(Vec::new())).expect("one partition sees all fragments"); +} + +/// A window directly over raw input remains valid across partitions without a preceding rewriter. +#[test] +fn a_window_over_raw_input_on_two_partitions_is_accepted() { + plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("ordinary input-to-window trees do not split a logical series"); +} + +/// One poll reads at most [`POLL_BUDGET`] chunks of input, and stopping there loses nothing. +/// +/// `rate(x[5m])` decides nothing until a sample reaches the grid's first point, so a series can +/// arrive as many chunks that settle no batch at all: a poll that kept reading them would hold the +/// task's thread until the series ended, which is what stops other tasks on that thread and delays +/// the cancellation a consumer asked for. The first series here is [`QUIET_SAMPLES`] such chunks; +/// the second one covers the grid, so a poll that resumes properly still answers both. +#[tokio::test] +async fn one_poll_reads_at_most_the_budget_and_loses_nothing() { + let mut batches = quiet_series_frames(); + batches.push(frame(&series_chunk("web-2", &series_samples(2.0), true))); + let total = batches.len(); + assert!( + total > POLL_BUDGET, + "the fixture has to outlast one poll: {total} batches" + ); + + let pulled = Arc::new(AtomicUsize::new(0)); + let woken = Arc::new(CountingWake::default()); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The first poll: the quiet series settles nothing, so the budget is what ends this poll. + let waker = futures::task::waker(Arc::clone(&woken)); + let mut context = Context::from_waker(&waker); + assert!( + stream.as_mut().poll_next(&mut context).is_pending(), + "chunks that settle nothing cannot answer a poll" + ); + let read = pulled.load(Ordering::SeqCst); + assert!(read > 0, "the poll did read its input"); + // One step polls the decoder once, and the decoder reads at most one frame per poll: every + // frame here carries one chunk. The bound is therefore one frame per step, and the budget is + // what it is measured in. + assert!( + read <= POLL_BUDGET + 1, + "one poll read {read} of {total} batches, past its budget of {POLL_BUDGET} steps" + ); + assert!( + read < total, + "one poll read the whole input ({read} batches): the budget did not stop it" + ); + assert!( + woken.woken(), + "a poll that stopped on its budget has to wake itself" + ); + + // Nothing the budget stopped is lost: the stream resumes where the poll left it. + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the bounded poll left it"); + assert_series_rates(&rows(&batches), &[(1.0, "web-1"), (2.0, "web-2")]); +} + +/// A batch that is ready is handed over before any batch behind it is read. +/// +/// Each series below is one chunk, and a chunk that carries `is_last` closes its series, so the +/// first answer needs only the first chunk. A poll that read on would answer the same batch only +/// after the whole input, which is exactly what the budget above exists to prevent. +#[tokio::test] +async fn a_ready_batch_is_returned_before_later_input_is_read() { + let batches = (0..20) + .map(|index| { + frame(&series_chunk( + &format!("web-{index}"), + &series_samples(1.0), + true, + )) + }) + .collect::>(); + let total = batches.len(); + let pulled = Arc::new(AtomicUsize::new(0)); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let first = stream + .next() + .await + .expect("the first series is answered") + .expect("the first series is answered without an error"); + let answered = rows(&[first]); + assert_series_rates(&answered, &[(1.0, "web-0")]); + // One batch: the series' own chunk closes it, so nothing behind it has to be read. + assert_eq!( + 1, + pulled.load(Ordering::SeqCst), + "the answer waited for input behind the series it reports" + ); + assert!(total > 1, "the fixture has to outlast the first answer"); +} + +/// 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 +/// ([`EvaluatedChunk::range`] answers `None` for it). Reading the hole as rows would invent grid +/// points this series holds no value at, and reading positions as compact rows would shift every +/// row after the hole onto the wrong timestamp: the batch has to hold exactly the carried rows, at +/// their absolute grid times, with nulls nowhere else. +#[test] +fn a_chunk_of_several_runs_converts_its_own_rows_only() { + let grid = grid(); + let axis = TimeAxis::from_grid_spans( + Arc::clone(&grid), + // Grid positions 0..2 and 4..5: position 2 and 3 hold no row at all. + vec![GridSpan::new(0, 2), GridSpan::new(4, 1)], + ) + .expect("the spans are increasing and inside the grid"); + let chunk = EvaluatedChunk::try_new( + Series { + metric: MetricName::try_new(Some(Arc::from(METRIC)), Vec::new()) + .expect("one metric name is a legal identity"), + time: axis, + fields: SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from( + // One value per carried row, so the second run's value sits at grid position 4. + vec![7.0, 8.0, 11.0], + ))]) + .expect("one field is a legal field list"), + }, + true, + ) + .expect("a chunk on a grid"); + + // The columns of `output_schema`, in its order: the field, the two tags, then the time. + let batch = series_to_record_batch( + &[ + OutputColumn::Field(0), + OutputColumn::Tag(0), + OutputColumn::Tag(1), + OutputColumn::Time, + ], + &[HOST.to_string(), NAME.to_string()], + &output_schema(), + &chunk, + ) + .expect("the chunk converts"); + + let stamps = batch + .column_by_name(TIME) + .expect("the batch has its time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let values = batch + .column_by_name(VALUE) + .expect("the batch has its field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + assert_eq!(3, batch.num_rows(), "only the carried rows are rows"); + let grid_stamps = grid.timestamps().expect("the grid timestamps build"); + assert_eq!( + vec![grid_stamps[0], grid_stamps[1], grid_stamps[4]], + stamps.values()[..3].to_vec(), + "the carried rows keep their absolute grid positions" + ); + assert_eq!(vec![7.0, 8.0, 11.0], values.values()[..3].to_vec()); +} + +/// Every message in `error`'s chain, so an error a layer above wrapped is still read. +fn error_chain(error: &DataFusionError) -> String { + let mut message = error.to_string(); + let mut source = std::error::Error::source(error); + while let Some(cause) = source { + message.push_str(": "); + message.push_str(&cause.to_string()); + source = cause.source(); + } + message +} + +/// Two partitions, one series each, and both series leave the window on the same label set. +/// +/// `rate` drops `__name__`, so `http_requests_total{host="web-1"}` reads as `{host="web-1"}` in +/// the output. The key column keeps each of the two whole inside its own partition, so a check that +/// lived inside one execution would see one series each and report nothing: only state shared by +/// the partitions sees that the output vector holds the label set twice. +#[tokio::test] +async fn a_duplicate_label_set_across_partitions_is_refused() { + let series = frame(&series_chunk("web-1", &series_samples(1.0), true)); + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![vec![series.clone()], vec![series]]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second partition lands on the label set the first one claimed"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// The same label set on points that do not overlap is legal, and every point survives. +/// +/// A rate needs two samples in its window, so the first series only has a value at the first grid +/// point and the second only at the last: the two share a final label set but no evaluation +/// timestamp. Prometheus merges such series into one; this node emits both batches and leaves the +/// merge to matrix assembly, so the check has to allow what it would otherwise refuse. +#[tokio::test] +async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { + let first = vec![frame(&series_chunk( + "web-1", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-1", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("disjoint grid points are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + let mut rows = rows(&batches); + rows.sort_by_key(|row| row.0); + assert_eq!( + vec![GRID_START, last], + rows.iter().map(|row| row.0).collect::>(), + "both series must survive with the one point each holds" + ); + for ((_, value, host, _), rate) in rows.iter().zip([1.0, 1.0]) { + assert_eq!(Some("web-1"), host.as_deref()); + let value = value.expect("the point each series holds has a value"); + assert!((rate - value).abs() < 1e-9, "{value} is not {rate}"); + } +} + +/// What a label set holds stays charged after the stream that first claimed it ends, and what a +/// later partition adds to it is charged as well. +/// +/// The map belongs to the plan: it outlives every partition and is released with the plan. Charging +/// each stream for what it inserted returned the bytes of a label set when *that* stream ended, +/// while the entry the next partition is checked against stayed in the map, so everything the map +/// held after the first partition was invisible to the pool. What is charged here is therefore read +/// between the partitions: it has to survive the first stream, and the runs the second partition +/// merges into the entry have to cost what they hold. +#[tokio::test] +async fn a_label_set_stays_charged_after_its_first_stream_ends() { + // The first series holds a value only at the first grid point and the second only at the last, + // so the two share a final label set without sharing an evaluation timestamp: the second + // partition is legal and has to be answered, and it has to find the first one's entry there. + let first = vec![frame(&series_chunk( + "web-0", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-0", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + + // The first partition claims the label set, answers it and ends. + let mut stream = exec + .execute(0, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the first partition runs")); + } + assert_eq!(1, answered.len(), "one series over one grid point"); + drop(stream); + let charged = pool.reserved(); + assert!( + charged > 0, + "the first partition's label set is still in the plan's map, but nothing is charged for it" + ); + + // The second partition lands on that same label set on points the first one does not hold: it + // is answered, and the run it merges into the entry is what its charge grows by. + let mut stream = exec + .execute(1, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the second partition runs")); + } + assert_eq!(1, answered.len(), "disjoint points are legal"); + drop(stream); + assert!( + pool.reserved() > charged, + "the second partition's run grew the entry, but nothing was charged for it: {} is not \ + above {charged}", + pool.reserved() + ); + + // The entries live as long as the plan, so dropping it is what gives their bytes back. + drop(exec); + assert_eq!( + 0, + pool.reserved(), + "the plan's drop releases what its label sets held" + ); +} + +/// Each partition runs its own series and answers it, with no label set claimed twice. +#[tokio::test] +async fn every_partition_evaluates_its_own_series() { + let partitions = [1.0, 2.0, 3.0] + .into_iter() + .enumerate() + .map(|(index, scale)| { + vec![frame(&series_chunk( + &format!("web-{index}"), + &series_samples(scale), + true, + ))] + }) + .collect::>(); + let exec = plan(rate_tree(), memory_exec_partitioned(partitions)) + .expect("the runtime executes a rate tree"); + assert_eq!( + 3, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("each partition runs its own tree"); + assert_eq!(3, batches.len(), "one batch per partition"); + let rows = rows(&batches); + assert_eq!( + 3 * GRID_LEN as usize, + rows.len(), + "every partition reports its own series over the whole grid" + ); + // Batches arrive in whatever order the partitions settle, so the series are read apart. + for (index, scale) in [1.0, 2.0, 3.0].into_iter().enumerate() { + let host = format!("web-{index}"); + let mut series = rows + .iter() + .filter(|row| row.2.as_deref() == Some(host.as_str())) + .cloned() + .collect::>(); + series.sort_by_key(|row| row.0); + assert_eq!(GRID_LEN as usize, series.len(), "{host}: {series:?}"); + assert_series_rates(&series, &[(scale, host.as_str())]); + } +} + +/// The optimizer inserts no exchange between a two-partition scan and this node. +/// +/// This is the claim the whole arrangement rests on, so it is checked against the real rule rather +/// than against the flags that are supposed to satisfy it: a plan that reports the right +/// requirements can still be repartitioned, and that is exactly how the previous `KeyPartitioned` +/// requirement failed - a hash requirement is classed as necessary for any child with more than one +/// partition, so a hash exchange was inserted even though the child already satisfied it. +/// +/// The input here is a two-partition source, and `target_partitions` is raised above it because +/// that is when an optimizer is most willing to repartition. Both are the conditions under which the +/// old requirement broke, so a plan that survives them is one where the guarantee the scan makes +/// arrives at this node intact. +#[test] +fn the_optimizer_inserts_no_exchange_between_the_input_and_this_node() { + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![ + vec![frame(&series_chunk("web-0", &series_samples(1.0), true))], + vec![frame(&series_chunk("web-1", &series_samples(2.0), true))], + ]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs a child with more than one partition" + ); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 8; + let optimized = EnsureRequirements::new() + .optimize(Arc::clone(&exec), &config) + .expect("the rule accepts the plan"); + + let exchanges = exchanges_between(&optimized, &exec); + assert!( + exchanges.is_empty(), + "the optimizer inserted {exchanges:?} between the input and the node, and an exchange \ + reorders a series' chunks: the node would close a series it has not finished" + ); + // No exchange is not the same as no change: a rule that raised the count some other way would + // still split a series. The count is the input's, which is what the node reports and what it + // runs one evaluator per. + assert_eq!( + 2, + optimized + .properties() + .output_partitioning() + .partition_count(), + "the optimizer kept no exchange but changed how many partitions the node runs" + ); + assert_eq!( + 2, + optimized.children()[0] + .properties() + .output_partitioning() + .partition_count(), + "the input's own partition count changed under the node" + ); +} + +/// The `RepartitionExec`s in `plan` that sit above `stop` and below no other `RepartitionExec`. +/// +/// Named by the kind the rule chose, so a failure says whether it was a hash or a round-robin +/// exchange: the two are inserted for different reasons and only the hash one is the trap. +fn exchanges_between(plan: &Arc, stop: &Arc) -> Vec { + let mut found = Vec::new(); + let mut stack = vec![Arc::clone(plan)]; + while let Some(node) = stack.pop() { + if Arc::ptr_eq(&node, stop) { + continue; + } + if let Some(exchange) = + node.downcast_ref::() + { + found.push(format!( + "{}({:?})", + exchange.name(), + exchange.partitioning() + )); + } + stack.extend(node.children().into_iter().cloned()); + } + found +} + +/// The node asks the optimizer for nothing, so no exchange can be inserted under it. +/// +/// A window holds the state of the series it is reading and closes it on the chunk carrying +/// `is_last`, so it can answer a series only if the whole series arrives in the partition it reads. +/// What makes that true is the input: a raw series frame comes from `SeriesScan`, whose contract is +/// that the same series (primary key) always reaches the same partition +/// (`src/mito2/src/read/series_scan.rs`), and the frame decoder keeps each series' chunks contiguous +/// within that partition. The node therefore needs the input's own layout kept, not a particular +/// layout of its own. +/// +/// The trap is the requirement writing that as `KeyPartitioned`, which reads like the honest +/// description of the guarantee. A `Hash|KeyPartitioned` requirement is classed as a hash +/// requirement, and a hash requirement forces a hash exchange whenever the child has more than one +/// partition - even where the child already satisfies it, and even when the partition counts match. +/// A hash exchange does not preserve order, so it would interleave one series' chunks across +/// partitions and the decoder would close a series it has not finished. Asking for +/// [`Distribution::UnspecifiedDistribution`] is what leaves the child alone, and declining extra +/// partitioning is the second half of the same thing: no round-robin is added to raise parallelism +/// either. The node is not pinned to one partition — it runs one tree per input partition and keeps +/// the input's count. +#[test] +fn the_input_is_left_alone_so_no_exchange_can_be_inserted() { + let exec = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("the runtime executes a rate tree"); + let distributions = exec.input_distribution_requirements().into_per_child(); + // `Distribution` has no equality of its own, so the exact requirement is matched structurally: + // one child, and what it asks for is that the child's layout is left alone. + assert!( + matches!( + distributions.as_slice(), + [Distribution::UnspecifiedDistribution] + ), + "anything but leaving the child alone makes a hash exchange necessary: {distributions:?}" + ); + assert_eq!( + vec![false], + exec.benefits_from_input_partitioning(), + "raising the partition count would add a round-robin exchange" + ); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "the node's own count is its input's" + ); +} + +/// A partition index the input does not have is refused rather than answered. +/// +/// The node runs one evaluator per input partition and adds none of its own, so it forwards the +/// index to its input: the count it declares is the input's, and an index past that is the input's +/// refusal, not a series read from another partition. +#[tokio::test] +async fn a_partition_past_the_inputs_count_is_refused() { + let schema = frame_schema(); + let source = staged(schema.clone(), move || { + Box::pin(StagedStream::new(schema.clone(), Vec::new())) + }); + let exec = plan(rate_tree(), source).expect("the runtime executes a rate tree"); + assert_eq!( + 1, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + exec.execute(0, SessionContext::default().task_ctx()) + .expect("the one partition the input has"); + let error = match exec.execute(1, SessionContext::default().task_ctx()) { + Ok(_) => panic!("a partition the input does not have is answered"), + Err(error) => error, + }; + assert!( + error.to_string().contains("partition"), + "refused for another reason: {error}" + ); +} + +/// The `rate(x[5m])` tree over two fields, with both fields and the tags in the output. +fn two_field_node() -> FusedSeries { + let tree = rate_tree_of(two_field_schema()); + let binding = binding_of( + &tree, + two_field_schema(), + vec![HOST.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema_of(&two_field_schema())), + two_field_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The output of [`two_field_node`]: both fields, the tags and the grid times. +fn two_field_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(OTHER, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// One chunk of `{__name__=METRIC, host=, __type__=}` whose `field` holds `rows` and +/// whose other field is missing everywhere. +fn one_field_chunk(field: usize, host: &str, kind: &str, last: bool) -> RawSeriesChunk { + let rows = series_samples(1.0); + let (value, other) = if field == 0 { + (values_of(&rows), no_values(&rows)) + } else { + (no_values(&rows), values_of(&rows)) + }; + chunk_of( + &two_field_schema(), + &format!("{host}-{kind}"), + metric(&[(HOST, host), (TYPE, kind)], Some(METRIC)), + ×_of(&rows), + vec![value, other], + last, + ) +} + +/// Two series that close on one label set may still fill it, one field each. +/// +/// `rate` drops `__name__` and `__type__`, so the two series below leave the window as +/// `{host="web-1"}` and read as one label set. They hold values in *different* fields, which is the +/// pair [`crate::grid_series::merge_same_labelset`] merges: one value of one field at each position, +/// never two. Occupancy is therefore per field, and refusing this pair would refuse a series +/// Prometheus answers. +#[tokio::test] +async fn a_duplicate_label_set_of_complementary_fields_is_emitted() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(1, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("complementary fields are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + // Each batch answers one field of the shared label set and leaves the other missing: the two + // together are the merged series, with no value of one field holding a position twice. + let answered = batches + .iter() + .map(|batch| { + let host = batch + .column_by_name(HOST) + .and_then(|host| host.as_any().downcast_ref::()) + .expect("every batch has its tag column"); + assert!( + (0..host.len()).all(|row| host.value(row) == "web-1"), + "both batches answer the one label set" + ); + let present = [VALUE, OTHER] + .into_iter() + .filter(|field| { + let column = batch + .column_by_name(field) + .unwrap_or_else(|| panic!("every batch has a {field} column")) + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + column.null_count() < column.len() + }) + .collect::>(); + assert_eq!( + 1, + present.len(), + "each series fills exactly one field: {present:?}" + ); + present[0] + }) + .collect::>(); + let mut answered = answered; + answered.sort_unstable(); + assert_eq!(vec![OTHER, VALUE], answered, "the two fields are filled"); +} + +/// Two series that close on one label set may not hold the same field at the same position. +/// +/// This is the pair no merge resolves: at one grid position one field of the label set would carry +/// two values. It is refused even though the two series name different `__type__` labels, because +/// the window drops that label and what reaches the output is one label set twice. +#[tokio::test] +async fn a_duplicate_label_set_of_one_field_is_refused() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both series hold a value of one field at one position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// Two partitions whose label sets differ only in an empty value hold one label set. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one series to Prometheus: an empty label value +/// is no label. An instant selection rewrites nothing, so both label sets reach this node's own +/// state as their frames wrote them, and a state that compared the tags as they arrived would file +/// them apart. Neither partition sees the other's series, so nothing else in the plan can refuse +/// the pair: this node's state is what compares them, and what it has to compare is the label set +/// Prometheus compares. +#[tokio::test] +async fn partitions_whose_label_sets_differ_only_in_an_empty_value_are_refused() { + let [empty, absent] = <[_; 2]>::try_from(zone_chunks()).expect("the fixture has two chunks"); + let exec = plan_of( + instant_node(), + memory_exec_partitioned(vec![frames(&[empty]), frames(&[absent])]), + ) + .expect("the runtime executes an instant selection"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both partitions hold one value of one field at the same position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// A chunk holding no value records nothing, so a query that answers nothing holds nothing. +/// +/// A selection over a series that staleness or the range left empty still closes that series, and +/// hands over a closing chunk with no rows - [`crate::native::execute::no_row_chunk`] is where it is +/// built. Every chunk reaches [`LabelSets::claim`], so an implementation that wrote an entry per +/// chunk would record one per empty series and charge the plan's reservation for it, for a query +/// whose answer is empty: enough of them would exhaust a finite pool without a single row produced. +/// +/// The entries are checked rather than the charge, because a charge of zero is what a mistake here +/// would still show if the label bytes happened to be free: what must not grow is the table. +#[test] +fn a_chunk_holding_no_value_is_claimed_without_recording_it() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let empty_series = |index: usize| Series { + metric: metric(&[(HOST, &format!("web-{index}"))], Some(METRIC)), + time: axis.clone(), + // Five positions were never filled, so no field occupies one. + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }; + + for index in 0..64 { + labels + .claim(&empty_series(index)) + .expect("an empty chunk cannot collide with anything"); + } + + assert_eq!( + 0, + labels.seen.len(), + "64 empty series were recorded, and the table outlives the query: {:?}", + labels.seen.keys().collect::>() + ); + assert_eq!( + 0, + pool.reserved(), + "nothing is held, so nothing may be charged" + ); +} + +/// An empty chunk does not stop a later chunk of the same label set from being recorded. +/// +/// Skipping the empty chunk is only safe while it also records nothing a later one needs: a series +/// that answers nothing holds no position, so a second series under one label set has to be free to +/// claim the positions it really holds. +#[test] +fn a_later_chunk_still_claims_the_label_set_an_empty_one_skipped() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let empty_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let full_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: empty_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }) + .expect("an empty chunk cannot collide with anything"); + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: full_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }) + .expect("the first chunk to hold a value claims the positions it holds"); + + assert_eq!(1, labels.seen.len(), "the positions are recorded once"); + assert!(pool.reserved() > 0, "the recorded positions are charged"); +} + +/// The claim compares the label set Prometheus compares, so an empty value is no label. +/// +/// Driven directly at the layer that keeps the state: a window and a map both drop an empty-valued +/// label before a chunk reaches this node, so no tree this crate can build produces the pair below. +/// The state has to key on the same normalization they apply all the same, or the equality it +/// enforces would be one Prometheus does not have. +#[test] +fn a_claim_compares_label_sets_without_empty_values() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + let series_of = |metric: MetricName| Series { + metric, + time: axis.clone(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }; + + labels + .claim(&series_of(metric(&[(HOST, "h"), (ZONE, "")], Some(METRIC)))) + .expect("the first label set is claimed"); + let error = labels + .claim(&series_of(metric(&[(HOST, "h")], Some(METRIC)))) + .expect_err("an empty label value is no label"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(Arc::from(HOST), Arc::from("h"))], + label_set(&metric(&[(HOST, "h"), (ZONE, "")], None)), + "the empty value is not part of the label set" + ); +} + +/// A series whose chunks move forward, a hole apart, appends to its stored runs rather than +/// re-merging them. +/// +/// One grid point per chunk with a hole after it (`0, 2, 4, ...`) is a legal stream and the shape a +/// short `count_over_time` over sub-sampled, chunked input produces: the chunks of one label set +/// hold their values at disjoint positions, so they merge into one series and no position ever +/// carries two values. What it used to cost was quadratic - `claim` walked the whole stored history +/// to look for a collision and then rebuilt all of it with a merge, for each chunk - and all of it +/// ran under the label-set lock every partition shares, so a sparse series of `K` chunks made the +/// other partitions wait for `K(K+1)/2` run moves. +/// +/// The counter is the runtime's own, so this test measures the append path the native table takes +/// as well: one fix, not two that can drift apart again. +#[test] +fn a_forward_chunk_does_not_rebuild_the_stored_runs() { + // How many one-position chunks the series arrives in: enough that the quadratic cost it used to + // pay would be unmistakable, since that cost was the square of this figure. + const CHUNKS: usize = 64; + + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + // Two positions per chunk, so a hole follows every point the series holds. + let grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 2 * CHUNKS as u64).unwrap()) + .unwrap(), + ); + let chunk = |position: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(position, 1)]) + .expect("one position of the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + + let before = absorb_rebuilds(); + for index in 0..CHUNKS { + labels + .claim(&chunk(2 * index)) + .expect("the chunks of one label set hold values at disjoint positions"); + } + + assert_eq!( + 0, + absorb_rebuilds() - before, + "the {CHUNKS} forward chunks rebuilt the stored runs instead of appending to them" + ); + assert_eq!( + 1, + labels.seen.len(), + "every chunk belongs to one label set, so the history is one entry" + ); +} + +/// A backward disjoint claim after a spare-capacity run merge updates the exact memory charge. +#[test] +fn a_backward_disjoint_claim_after_spare_capacity_is_charged_correctly() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start, len| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + labels.claim(&chunk(10, 2)).unwrap(); + let runs = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + let old_capacity = runs.capacity(); + runs.reserve(16); + let new_capacity = runs.capacity(); + labels + .reservation + .as_mut() + .unwrap() + .try_grow((new_capacity - old_capacity) * std::mem::size_of::()) + .unwrap(); + let before = pool.reserved(); + + labels + .claim(&chunk(0, 1)) + .expect("the position is disjoint and earlier than stored history"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (10, 12)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + before - (new_capacity - runs.capacity()) * std::mem::size_of::(), + pool.reserved(), + "the accounting follows the replacement vector's actual capacity" + ); +} + +/// A backward claim that joins two spare-capacity runs charges a replacement vector's shrinkage. +#[test] +fn a_backward_claim_merging_spare_capacity_runs_accounts_for_shrinkage() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 3).unwrap()).unwrap()); + let chunk = |values: Vec>| Series { + metric: metric(&[(HOST, "shrink")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 3)]) + .expect("the chunk covers the small grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(values))), + }; + + labels + .claim(&chunk(vec![Some(1.0), None, Some(1.0)])) + .expect("the first claim stores runs at positions 0 and 2"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (2, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + let before_capacity = runs.capacity(); + let before_reserved = pool.reserved(); + + labels + .claim(&chunk(vec![None, Some(1.0), None])) + .expect("the disjoint middle position joins both existing runs"); + + let runs = labels.seen.values().next().unwrap().first().unwrap(); + let after_capacity = runs.capacity(); + assert_eq!( + vec![(0, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert!( + before_capacity > after_capacity, + "the regression requires replacement-vector shrinkage (before={before_capacity}, after={after_capacity})" + ); + let run_size = std::mem::size_of::(); + assert_eq!( + before_reserved - (before_capacity - after_capacity) * run_size, + pool.reserved(), + "the reservation tracks the replacement vector's actual capacity decrease" + ); +} + +/// A newly retained field is charged for the cloned run vector's capacity, not incoming slack. +#[test] +fn a_new_field_claim_is_charged_for_its_retained_capacity() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 4).unwrap()).unwrap()); + let first = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&first).unwrap(); + + let incoming = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![None, None, None, None])), + Samples::from_float(Float64Array::from(vec![Some(1.0), None, Some(2.0), None])), + ]) + .unwrap(), + }; + let incoming_runs = occupied_runs(&incoming).unwrap(); + assert_eq!(2, incoming_runs[1].len()); + assert!( + incoming_runs[1].capacity() > incoming_runs[1].len(), + "the incoming run vector has spare capacity" + ); + let before = pool.reserved(); + labels.claim(&incoming).unwrap(); + + let retained = &labels.seen.values().next().unwrap()[1]; + assert_eq!(2, retained.len()); + assert_eq!( + retained.capacity() * std::mem::size_of::(), + pool.reserved() - before, + "the reservation charges the retained cloned vector capacity" + ); + assert!( + retained.capacity() < incoming_runs[1].capacity(), + "the incoming spare capacity was not retained or charged" + ); +} + +/// Chunked fields are traversed sequentially while occupancy is converted into grid runs. +#[test] +fn chunked_samples_occupancy_preserves_sparse_rows() { + let values = Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), None])), + Samples::from_float(Float64Array::from(vec![None, Some(2.0)])), + ]) + .unwrap(); + let series = Series { + metric: metric(&[(HOST, "chunked")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::one(values), + }; + let runs = occupied_runs(&series).expect("chunked occupancy is read through a forward cursor"); + assert_eq!( + vec![(0, 1), (3, 4)], + runs[0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); +} + +/// Appending a run that touches stored history reuses its last slot instead of rebuilding it. +#[test] +fn a_touching_append_with_full_capacity_does_not_rebuild_history() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start| Series { + metric: metric(&[(HOST, "touch")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&chunk(0)).unwrap(); + let stored = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + stored.shrink_to_fit(); + let cap = stored.capacity(); + let allocation = stored.as_ptr(); + assert_eq!( + stored.len(), + cap, + "the test requires a full-capacity run list" + ); + let before = absorb_rebuilds(); + labels + .claim(&chunk(1)) + .expect("the touching position extends the same run"); + let stored = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 2)], + stored + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + cap, + stored.capacity(), + "extending a run did not allocate or copy its history" + ); + assert_eq!( + allocation, + stored.as_ptr(), + "extending a run retained the stored vector allocation" + ); + assert_eq!( + before, + absorb_rebuilds(), + "touching append uses the append path" + ); +} + +/// A failed claim does not mutate the shared entry, while its reservation owner remains alive. +#[test] +fn a_refused_growth_keeps_shared_label_entry_unchanged() { + use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; + let limit = 4096; + let pool = Arc::new(GreedyMemoryPool::new(limit)); + let erased: Arc = pool.clone(); + let mut labels = LabelSets::default(); + labels.attach(&erased); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |host: &str, start| Series { + metric: metric(&[(HOST, host)], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels + .claim(&chunk("shared", 0)) + .expect("initial label fits"); + let key = label_set(&chunk("shared", 0).metric); + let before_runs = labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(); + let before_reserved = pool.reserved(); + let blocker: MemoryReservation = MemoryConsumer::new("test blocker").register(&erased); + blocker + .try_grow(limit - before_reserved) + .expect("the blocker occupies remaining capacity"); + let error = labels + .claim(&chunk("shared", 2)) + .expect_err("growth exceeds the pool"); + assert!(!error.to_string().is_empty(), "pool refusal is reported"); + assert_eq!( + before_runs, + labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(), + "the persistent entry is unchanged" + ); + assert_eq!(limit, pool.reserved(), "both reservations remain live"); + drop(blocker); + assert_eq!( + before_reserved, + pool.reserved(), + "the plan-owned entry remains charged" + ); +} + +/// Binding ordering includes field names, while schema ordering refuses unequal DF schemas. +#[test] +fn unequal_binding_and_output_schema_are_not_ordered_equal() { + let tree = rate_tree(); + let left = binding(&tree); + let other_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: "different".into(), + }]), + eval_grid: None, + }); + let right = binding_of(&tree, other_schema, left.tags.clone()); + assert_ne!(left, right); + assert_ne!(Some(std::cmp::Ordering::Equal), left.partial_cmp(&right)); + + let left_node = fused(tree.clone()); + let nullability_changed = Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, false), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let right_node = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding(&rate_tree()), + input_plan(frame_schema()), + nullability_changed.to_dfschema_ref().unwrap(), + ) + .unwrap(); + assert_ne!(left_node, right_node); + assert_ne!( + Some(std::cmp::Ordering::Equal), + left_node.partial_cmp(&right_node) + ); +} + +/// A chunk reaching back before the last stored run still merges into the history. +/// +/// The forward path only takes a chunk that starts at or after the end of the last stored run; a +/// chunk that starts before it can interleave with what is stored, so it has to go through the full +/// merge, and the result still has to be sorted, disjoint and with the runs that touch joined - +/// which is what the look-up in [`LabelSets::claim`] relies on. +/// +/// A chunk reaching back over a stored run is the other case, and it stays refused: `[0, 11)` shares +/// position 10 with the stored `[10, 12)`, so it is a collision and the stored history is left +/// as it was. +#[test] +fn a_chunk_reaching_back_still_merges() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + // One chunk holding `len` positions from `start`, every one of them with a value. + let chunk = |start: usize, len: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .expect("one run over the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + // The runs stored under the one label set, which one field of the output holds. + let stored = |labels: &LabelSets| -> Vec<(usize, usize)> { + labels + .seen + .values() + .next() + .expect("the chunk was recorded under one label set") + .first() + .expect("the series has one field") + .iter() + .map(|run| (run.start, run.end)) + .collect() + }; + + labels + .claim(&chunk(10, 2)) + .expect("the first chunk claims the positions it holds"); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "the run `[10, 12)` is stored" + ); + + let error = labels + .claim(&chunk(0, 11)) + .expect_err("`[0, 11)` holds a value at position 10, where `[10, 12)` holds one"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "a refused claim left the stored history as it was" + ); + + let before = absorb_rebuilds(); + labels + .claim(&chunk(0, 10)) + .expect("`[0, 10)` reaches back before the last stored run without sharing a position"); + assert_eq!( + vec![(0, 12)], + stored(&labels), + "the run reaching back did not merge into the stored one and join it" + ); + assert_eq!( + 1, + absorb_rebuilds() - before, + "a chunk starting before the last stored run was appended rather than merged" + ); +} diff --git a/src/promql/src/grid_series.rs b/src/promql/src/grid_series.rs index c9d8bc19d1f..90cfa0f2e62 100644 --- a/src/promql/src/grid_series.rs +++ b/src/promql/src/grid_series.rs @@ -19,12 +19,12 @@ use std::collections::hash_map::Entry; use std::sync::Arc; use datafusion::arrow::array::builder::NullBufferBuilder; -use datafusion::arrow::array::{Float64Array, UInt32Array}; +use datafusion::arrow::array::{Array, Float64Array, UInt32Array}; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; -use datatypes::series::histogram::HistogramBatchBuilder; +use datatypes::series::histogram::{HistogramBatch, HistogramBatchBuilder}; use datatypes::series::samples::{ - HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, + DenseSamplesRef, HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, }; use datatypes::series::time::{EvalGrid, GridSpan, TimeAxis}; @@ -35,6 +35,70 @@ pub(crate) mod reference; #[cfg(test)] mod tests; +pub(crate) fn series_bytes(series: &Series) -> usize { + let runs = series.time.runs().count(); + let segments: usize = series + .fields + .iter() + .map(|field| field.segments().len()) + .sum(); + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) + .saturating_add( + segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), + ) + .saturating_add(fields_bytes(&series.fields)) + .saturating_add(series.metric.name().map_or(0, str::len)) + .saturating_add(std::mem::size_of_val(series.metric.tags())) + .saturating_add( + series + .metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) +} + +pub(crate) fn fields_bytes(fields: &SeriesFields) -> usize { + fields + .iter() + .map(|field| { + field.segments().iter().fold(0usize, |total, segment| { + total.saturating_add(segment_bytes(segment)) + }) + }) + .fold(0usize, usize::saturating_add) +} + +/// The bytes one dense segment of samples holds. +pub(crate) fn segment_bytes(segment: &Samples) -> usize { + match segment.dense() { + Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), + Some(DenseSamplesRef::Histogram(histograms)) => histograms + .indices() + .get_array_memory_size() + .saturating_add(table_bytes(histograms.batch())), + Some(DenseSamplesRef::Mixed(mixed)) => mixed + .floats() + .get_array_memory_size() + .saturating_add(mixed.histograms().indices().get_array_memory_size()) + .saturating_add(table_bytes(mixed.histograms().batch())), + // `segments` hands out dense segments only, so a nested chunk never reaches here. + None => 0, + } +} + +/// The bytes one histogram table holds: the batch struct and the arrays its columns view. +pub(crate) fn table_bytes(table: &Arc) -> usize { + std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) +} + +pub(crate) fn field_builder_bytes(width: usize) -> usize { + width.saturating_mul(std::mem::size_of::()) +} + /// Builds one grid series from rows pushed at absolute grid indices. /// /// A push names one position and the value each field holds there. Rows are accumulated @@ -114,6 +178,36 @@ impl GridSeriesBuilder { }) } + /// Reserves capacity for the exact number of valued rows in the next series. + /// + /// Call after `begin_series` and before its first push. Existing public callers may continue + /// to rely on normal growth; callers with a budget can reserve the concrete row capacity first. + pub(crate) fn reserve_rows(&mut self, rows: usize) -> Result<()> { + self.indices + .try_reserve_exact(rows.saturating_sub(self.indices.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid row indices: {error}")) + })?; + if rows == 0 { + return Ok(()); + } + if self.fields.is_empty() { + self.fields = (0..self.schema.fields.len()) + .map(|_| FieldBuilder::new(rows)) + .collect(); + } + for field in &mut self.fields { + field + .values + .try_reserve_exact(rows.saturating_sub(field.values.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid values: {error}")) + })?; + field.nulls = NullBufferBuilder::new(rows); + } + Ok(()) + } + /// Pushes the field values at absolute grid `index`; `None` is a missing value. /// /// Indices must be strictly increasing; skipped indices are positions that hold no value. @@ -176,7 +270,7 @@ impl GridSeriesBuilder { return Ok(false); } if self.fields.is_empty() { - self.fields = (0..len).map(|_| FieldBuilder::new()).collect(); + self.fields = (0..len).map(|_| FieldBuilder::new(0)).collect(); } self.indices.push(index); Ok(true) @@ -239,24 +333,28 @@ impl GridSeriesBuilder { } let fields = std::mem::take(&mut self.fields); let indices = std::mem::take(&mut self.indices); - let samples = fields - .into_iter() - .map(|mut field| { - let floats = Float64Array::new(field.values.into(), field.nulls.finish()); - let Some(column) = field.histograms else { - return Ok(Samples::from_float(floats)); - }; + let mut samples = Vec::new(); + samples.try_reserve_exact(fields.len()).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid field descriptions: {error}")) + })?; + for mut field in fields { + let floats = Float64Array::new(field.values.into(), field.nulls.finish()); + let sample = if let Some(column) = field.histograms { let batch = column.batch.finish().map_err(external)?; let histograms = HistogramSamples::try_new(UInt32Array::from(column.indices), batch) .map_err(external)?; - Samples::try_from_mixed(floats, histograms).map_err(external) - }) - .collect::>>()?; + Samples::try_from_mixed(floats, histograms).map_err(external)? + } else { + Samples::from_float(floats) + }; + samples.push(sample); + } let fields = SeriesFields::try_from_vec(samples).map_err(external)?; // The runs are all the axis has to remember about the gaps, and the rows are already // compact, so no position holding no value is visited or stored. - let spans = spans_of(&indices); + let run_count = count_runs(&indices); + let spans = spans_of(&indices, run_count)?; let time = TimeAxis::from_grid_spans(self.grid.clone(), spans).map_err(external)?; Series::try_new(&self.schema, metric, time, fields) .map(Some) @@ -265,8 +363,19 @@ impl GridSeriesBuilder { } /// Cuts the retained indices into maximal runs of consecutive positions. -fn spans_of(indices: &[usize]) -> Vec { +fn count_runs(indices: &[usize]) -> usize { + indices + .windows(2) + .filter(|pair| pair[1] != pair[0] + 1) + .count() + .saturating_add(usize::from(!indices.is_empty())) +} + +fn spans_of(indices: &[usize], runs: usize) -> Result> { let mut spans: Vec = Vec::new(); + spans.try_reserve_exact(runs).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid spans: {error}")) + })?; let mut start = 0; while let Some(&position) = indices.get(start) { let mut len = 1; @@ -276,14 +385,14 @@ fn spans_of(indices: &[usize]) -> Vec { spans.push(GridSpan::new(position, len)); start += len; } - spans + Ok(spans) } impl FieldBuilder { - fn new() -> Self { + fn new(rows: usize) -> Self { Self { values: Vec::new(), - nulls: NullBufferBuilder::new(0), + nulls: NullBufferBuilder::new(rows), histograms: None, } } @@ -454,7 +563,8 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R // The runs are all the axis remembers about the holes, so a sparse merge stores one span // per run and no row for the positions between them. let time = - TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions)).map_err(external)?; + TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions, count_runs(&positions))?) + .map_err(external)?; let fields = SeriesFields::try_from_vec(fields).map_err(external)?; // The first fragment names the output, as it always has: a label set ignores empty values, // so the metric of the first member is the one the merged series keeps. diff --git a/src/promql/src/native/aggregate.rs b/src/promql/src/native/aggregate.rs new file mode 100644 index 00000000000..975dfb6c3c2 --- /dev/null +++ b/src/promql/src/native/aggregate.rs @@ -0,0 +1,243 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The run-time state of an aggregate node. +//! +//! An aggregate folds a whole vector into one series per group, so its node is the one that cannot +//! answer as its input arrives: a group's value at a grid point is the fold of every value every +//! member of the group has there, and the last member to arrive is the one that decides it. What the +//! node keeps between chunks is therefore the fold's states and nothing else -- a chunk is folded as +//! it arrives and dropped -- and it produces once, when the input ends. +//! +//! That is also why this runtime cannot tell an incomplete group from a complete one: it folds what +//! it is handed. A [`AggregateStage::Complete`] aggregate is correct only if every contribution to +//! each group lies within the streams this execution reads, which is the physical plan's job -- the +//! caller that schedules scans and exchanges guarantees it, and a plan that ran a `Complete` +//! aggregate once per partition and then combined the answers would be combining values that were +//! each missing the other partitions' contributions. The stages that exist for that shape are +//! [`AggregateStage::Partial`] and [`AggregateStage::Final`], which this runtime does not execute +//! yet. + +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result}; +use datatypes::series::samples::{Series, SeriesSchema}; + +use crate::aggregate::{AggregateOp, FloatAggregator, Grouping}; +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::series_bytes; +use crate::map::validate_float_series; +use crate::native::execute::{NodeBuffer, evaluated_bytes}; +use crate::native::metrics::NodePath; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, Grouping as BoundGrouping, +}; + +/// One aggregate node's states, from the first chunk it is handed to the series it finishes with. +pub(crate) struct AggregateRuntime { + /// What the node computes, as the tree bound it: the kind it folds by and the stage it reports. + op: BoundAggregate, + /// The layout its input and output share: the fields every chunk has and the grid they are + /// addressed on. A `Complete` aggregate neither adds nor drops a field, and it evaluates on its + /// input's grid, so one schema describes both sides. + schema: Arc, + /// The fold. `None` once [`Self::compute`] has taken it, which is also what says the node has + /// nothing left to fold into. + accumulator: Option, + /// Whether the fold has been finished and its series handed over. A node computes once: a second + /// pull finds the series it queued rather than a fold that ran again. + /// + /// The driver reads this to decide whether the node settled: an aggregate answers only once its + /// fold has run, so a child with nothing left to hand over is not an answer. + pub(crate) computed: bool, +} + +impl AggregateRuntime { + /// What this node computes, for a rejection that has to say what it refused. + /// + /// It names the kind this node computes, which distinguishes it from another node in the tree. + fn describe(&self) -> String { + format!("{:?}", self.op.kind) + } + + /// The states of a `Complete` aggregate over `schema`. + /// + /// A stage other than `Complete` and a kind without a supported kernel are refused here, before a single + /// chunk is folded: a node this runtime cannot compute must not start, because a partial answer + /// would be one the caller never asked for. The rejection names what was found, so a reader can + /// tell a tree this runtime cannot run from one it cannot run correctly. + pub(crate) fn new( + op: BoundAggregate, + schema: Arc, + buffer: &mut NodeBuffer, + ) -> Result { + if op.stage != AggregateStage::Complete { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute a `{:?}` aggregate yet: a `Partial` aggregate \ + reports the state a later stage merges and a `Final` aggregate merges that state, \ + and only `Complete` -- every contribution inside one execution -- is computed here", + op.stage + ))); + } + let kind = match op.kind { + AggregateKind::Sum => AggregateOp::Sum, + AggregateKind::Count => AggregateOp::Count, + AggregateKind::Avg => AggregateOp::Avg, + AggregateKind::Min => AggregateOp::Min, + AggregateKind::Max => AggregateOp::Max, + AggregateKind::Group => AggregateOp::Group, + AggregateKind::Stddev => AggregateOp::Stddev, + AggregateKind::Stdvar => AggregateOp::Stdvar, + other => { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute the `{other:?}` aggregate yet: only the supported `Complete` aggregate kinds are computed here" + ))); + } + }; + let name_count = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => labels.len(), + }; + let name_bytes = name_count.saturating_mul(std::mem::size_of::<&str>()); + buffer.grow(name_bytes)?; + let names: Vec<&str> = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => { + labels.iter().map(|label| label.as_ref()).collect() + } + }; + let grouping = match &op.group { + BoundGrouping::By(_) => Grouping::By(&names), + BoundGrouping::Without(_) => Grouping::Without(&names), + }; + let name_bytes = names.len().saturating_mul(std::mem::size_of::<&str>()); + let initial = FloatAggregator::try_new_budgeted(&schema, kind, grouping, |bytes| { + buffer.resize(bytes.saturating_add(name_bytes)) + })?; + drop(names); + buffer.resize(initial.state_bytes())?; + Ok(Self { + op, + schema, + accumulator: Some(initial), + computed: false, + }) + } + + /// What folding `chunk` reserves while it is folded. + /// + /// The chunk is not retained: it is folded into the states and dropped at the end of the call, + /// so what this node holds afterwards is the fold and none of the input. The charge is therefore + /// the chunk's own bytes -- the figure [`evaluated_bytes`] reports, which is one pass over the + /// chunk and no state between chunks -- and not the `operand_chunk_bytes` figure a binary + /// operand keeps, which exists because such an operand holds a series for the whole execution + /// and so must count the slices of one buffer instead of the buffer once per chunk. + pub(crate) fn pending(&self, chunk: &EvaluatedChunk) -> usize { + evaluated_bytes(chunk) + } + + /// Folds one chunk into its group. + /// + /// A chunk that closes its series closes nothing here: a group is not a series, so the last + /// contribution of one member says nothing about the group. The fold keeps its states, which is + /// what makes a group's answer its whole group's answer once the input ends. + pub(crate) fn accumulate( + &mut self, + path: &NodePath, + chunk: EvaluatedChunk, + buffer: &mut NodeBuffer, + pending: usize, + ) -> Result<()> { + // `Series` keeps its fields public, so a chunk can be built by struct literal against + // another schema; the states are indexed by absolute position on this node's grid, so a + // chunk of another field count or another grid would be folded into the wrong places. + validate_float_series(&self.schema, &chunk.series) + .map_err(|error| folded_chunk(path, &self.describe(), error))?; + let accumulator = self.accumulator.as_mut().ok_or_else(|| no_fold(path))?; + // Keep the chunk covered while the fold decides whether it needs a new group or block. + accumulator.update_budgeted(&chunk.series, |fold_bytes| { + buffer.resize(pending.saturating_add(fold_bytes)) + })?; + drop(chunk); + buffer.resize(accumulator.state_bytes()) + } + + /// The number of retained groups, for budgeting the output queue backing storage. + pub(crate) fn groups(&self) -> usize { + self.accumulator.as_ref().map_or(0, FloatAggregator::groups) + } + + /// Finishes the fold and returns one series per group, or `None` if the fold already ran. + /// + /// `None` is what a node reports after it has queued its series: the caller reads them from the + /// queue, exactly as it reads the chunks a binary node queued, so a second pull does not finish + /// a fold whose states are gone. + pub(crate) fn compute( + &mut self, + path: &NodePath, + buffer: &mut NodeBuffer, + queue_bytes: usize, + ) -> Result>> { + if self.computed { + return Ok(None); + } + if self.accumulator.is_none() { + return Err(no_fold(path)); + } + self.computed = true; + let output = self + .accumulator + .take() + .ok_or_else(|| no_fold(path))? + .finish_budgeted(|total| buffer.resize(total.saturating_add(queue_bytes)))?; + let output_bytes = output.iter().fold(0usize, |bytes, series| { + bytes.saturating_add(series_bytes(series)) + }); + let output_vec_bytes = output + .capacity() + .saturating_mul(std::mem::size_of::()); + buffer.resize( + output_bytes + .saturating_add(output_vec_bytes) + .saturating_add(queue_bytes), + )?; + Ok(Some(output)) + } +} + +/// The rejection of a chunk that is not the layout this node folds. +/// +/// The fold indexes its states by absolute grid position and by field, so a chunk of another field +/// count or another grid would be folded into states that mean something else. `error` is what the +/// series check reported, which says which of the two disagreed. +/// Wraps the error of a chunk this node cannot fold. +/// +/// Only a chunk that does not fit the schema is the node's own refusal; anything else is the +/// kernel saying it does not compute this shape at all -- a native histogram in a float aggregate, +/// for instance -- and a caller that falls back to another planner on `NotImplemented` still can. +/// Wrapping that one in an `Execution` would deny it. +fn folded_chunk(path: &NodePath, aggregate: &str, error: DataFusionError) -> DataFusionError { + match error { + DataFusionError::NotImplemented(_) => error, + error => DataFusionError::Execution(format!( + "the `{aggregate}` aggregate at {path} folds chunks bound to its own schema, and this \ + one is not: {error}" + )), + } +} + +/// The error for a fold that has already been finished. +fn no_fold(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!( + "the aggregate at {path} already finished its fold, so it takes no more input" + )) +} diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index cc79b2b8fa2..8813deff82a 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -91,9 +91,9 @@ use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datatypes::arrow::array::{Array, Float64Array}; +use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::histogram::{HistogramBatch, HistogramView}; +use datatypes::series::histogram::HistogramView; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ DenseSamplesRef, RowCursor, SampleRef, Samples, Series, SeriesFields, SeriesSchema, @@ -103,7 +103,8 @@ use datatypes::series::{MetricName, SeriesOrdering}; use crate::aggregate; use crate::grid_series::{ - GridSample, GridSeriesBuilder, LabelSet, label_set, merge_same_labelset, owned, + GridSample, GridSeriesBuilder, LabelSet, fields_bytes, label_set, merge_same_labelset, owned, + series_bytes, table_bytes, }; use crate::histogram_eval::evaluate_histograms; use crate::instant_eval::select_instants; @@ -111,12 +112,16 @@ use crate::map::{ ArithmeticOp, Comparison, ComparisonOp, MathOp, ScalarOperand, drop_metric_name, map_arithmetic, map_compare, map_math, }; +use crate::native::aggregate::AggregateRuntime; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::{NodeMetrics, NodePath}; use crate::native::plan::{ BoundInstant, BoundScalar, BoundVectorOp, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; +#[cfg(test)] +pub(crate) use crate::native::runs::absorb_rebuilds; +use crate::native::runs::{Run, absorb}; use crate::range_eval::RangeEvaluator; use crate::select::{self, SortOrder}; use crate::vector_match::{ @@ -190,34 +195,11 @@ fn raw_bytes(chunk: &RawSeriesChunk) -> usize { /// The bytes an evaluated chunk holds: its struct and the Arrow buffers behind its fields. /// /// The grid is shared by every chunk addressed on it, so it is not counted per chunk. -fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { - std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) -} - -/// The bytes one evaluated series holds: its header, the axis that addresses it, the identity it -/// carries and the sample buffers behind its fields. /// -/// It is the figure a fragment of an affected window is charged with as it arrives. The axis's runs -/// and the segments of its fields are counted in one pass over the series, so the estimate costs -/// the fragment rather than a walk over everything the node holds. A payload a fragment shares with -/// another — a histogram table behind two slices of it, the grid behind every axis — is counted -/// once per fragment: that over-count is accepted here rather than a deduplication that would have -/// to keep state across fragments. -fn series_bytes(series: &Series) -> usize { - let runs = series.time.runs().count(); - let segments: usize = series - .fields - .iter() - .map(|field| field.segments().len()) - .sum(); - std::mem::size_of::() - .saturating_add(std::mem::size_of::()) - .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) - .saturating_add( - segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), - ) - .saturating_add(fields_bytes(&series.fields)) - .saturating_add(metric_bytes(&series.metric)) +/// It is what an aggregate node charges for a chunk it is folding: that node hands the chunk over +/// to its fold and drops it, so the chunk alone is what it holds for as long as it holds anything. +pub(super) fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { + std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) } /// The bytes one chunk a binary operand retains holds on its own. @@ -343,42 +325,6 @@ fn merge_workspace_bytes(rows: usize, fragments: usize, fields: usize) -> usize .saturating_add(output) } -/// The bytes the fields of a chunk hold: the sample buffers of every segment, and the histogram -/// tables those segments reference. -fn fields_bytes(fields: &SeriesFields) -> usize { - fields - .iter() - .map(|field| { - field.segments().iter().fold(0usize, |total, segment| { - total.saturating_add(segment_bytes(segment)) - }) - }) - .fold(0usize, usize::saturating_add) -} - -/// The bytes one dense segment of samples holds. -fn segment_bytes(segment: &Samples) -> usize { - match segment.dense() { - Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), - Some(DenseSamplesRef::Histogram(histograms)) => histograms - .indices() - .get_array_memory_size() - .saturating_add(table_bytes(histograms.batch())), - Some(DenseSamplesRef::Mixed(mixed)) => mixed - .floats() - .get_array_memory_size() - .saturating_add(mixed.histograms().indices().get_array_memory_size()) - .saturating_add(table_bytes(mixed.histograms().batch())), - // `segments` hands out dense segments only, so a nested chunk never reaches here. - None => 0, - } -} - -/// The bytes one histogram table holds: the batch struct and the arrays its columns view. -fn table_bytes(table: &Arc) -> usize { - std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) -} - /// The bytes the allocation `buffer` is a slice of. /// /// The capacity is the whole allocation, which is what a slice keeps alive: a key of a few bytes @@ -409,6 +355,8 @@ enum NodeTag { /// Matches two evaluated vectors and applies one operation to every matched pair of series, or /// broadcasts one constant over the chunks of one evaluated vector. Binary, + /// Folds every series of the input into one series per group, when the input ends. + Aggregate, /// Keeps some samples of its whole input vector at each grid point. Vector, } @@ -423,6 +371,7 @@ impl NodeTag { Self::Window => "window", Self::Map => "map", Self::Binary => "binary", + Self::Aggregate => "aggregate", Self::Vector => "vector operation", } } @@ -577,7 +526,9 @@ impl MergeInput { /// in the query; the metrics report the same figure so a slow node can be told apart from a wide /// one. The reservation is also the figure itself: a node holds nothing beside what it reserved, so /// a gauge reads its size rather than a copy that could drift from it. -struct NodeBuffer { +// Visible to the aggregate node so that its fold can reserve for what it is about to allocate; the +// reservation belongs to a node, not to a module. +pub(super) struct NodeBuffer { metrics: Arc, reservation: MemoryReservation, } @@ -593,7 +544,7 @@ impl NodeBuffer { } /// Reserves `bytes` beside what this node already holds. - fn grow(&mut self, bytes: usize) -> Result<()> { + pub(super) fn grow(&mut self, bytes: usize) -> Result<()> { self.reservation.try_grow(bytes)?; self.record(); Ok(()) @@ -604,7 +555,7 @@ impl NodeBuffer { /// What a node that consumes its input chunk by chunk holds afterwards is not the sum of what /// it consumed, so the reservation is set from the state that is left rather than grown per /// chunk. - fn resize(&mut self, bytes: usize) -> Result<()> { + pub(super) fn resize(&mut self, bytes: usize) -> Result<()> { self.reservation.try_resize(bytes)?; self.record(); Ok(()) @@ -629,20 +580,6 @@ impl NodeBuffer { } } -/// One run of grid positions, half-open, that hold a value in one field. -#[derive(Clone, Copy, Debug)] -struct Run { - start: usize, - end: usize, -} - -impl Run { - /// Whether two runs share a position. - fn overlaps(self, other: Self) -> bool { - self.start < other.end && other.start < self.end - } -} - /// The positions one node has handed out, per output label set and field. /// /// A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name @@ -877,109 +814,6 @@ fn misaligned_field( )) } -// How many times `absorb` rebuilt the stored runs instead of appending to them. -// -// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows -// only in what it costs, and a test that reads its own thread's figure is not disturbed by another -// test's chunks. -#[cfg(test)] -thread_local! { - static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; -} - -/// How many times this thread's [`absorb`] rebuilt the stored runs. -#[cfg(test)] -fn absorb_rebuilds() -> usize { - ABSORB_REBUILDS.with(std::cell::Cell::get) -} - -/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that -/// touch. -/// -/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and -/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square -/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values -/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves -/// for a single series. -/// -/// One pass per chunk is still a pass over everything the series has stored so far, so it is not -/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, -/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision -/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run -/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends -/// before anything this chunk holds — and it is appended instead, which is amortised constant per -/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk -/// that starts before the last stored run can split a run or reach back past it, so it still costs -/// the full merge. -/// -/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what -/// [`Emitted::occupy`] looks a position up by. -fn absorb(stored: &mut Vec, incoming: &[Run]) { - if incoming.is_empty() { - return; - } - if stored - .last() - .is_none_or(|last| incoming[0].start >= last.end) - { - append_runs(stored, incoming); - } else { - merge_runs(stored, incoming); - } -} - -/// Appends `incoming` to `stored`, all of which lies past the last stored run. -/// -/// Nothing already stored has to be placed again: the two sides are each in position order and -/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and -/// the vector's own growth is the only move — amortised over the runs appended. Where its first run -/// touches the end of the last stored run the two are one run, as they are after a merge, so it is -/// joined rather than left beside it. -fn append_runs(stored: &mut Vec, incoming: &[Run]) { - let mut incoming = incoming; - if let Some(last) = stored.last_mut() - && let Some(first) = incoming.first() - && first.start <= last.end - { - last.end = last.end.max(first.end); - incoming = &incoming[1..]; - } - stored.extend_from_slice(incoming); -} - -/// Merges `incoming` into `stored`, where the two can interleave. -/// -/// One pass is enough because each side is sorted and disjoint: whichever run starts first is -/// placed, and it joins the run before it where the two touch. -fn merge_runs(stored: &mut Vec, incoming: &[Run]) { - #[cfg(test)] - ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); - let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); - let (mut a, mut b) = (0, 0); - while a < stored.len() || b < incoming.len() { - let larger = match (stored.get(a), incoming.get(b)) { - (Some(left), Some(right)) if right.start < left.start => { - b += 1; - *right - } - (Some(left), _) => { - a += 1; - *left - } - (None, Some(right)) => { - b += 1; - *right - } - (None, None) => break, - }; - match merged.last_mut() { - Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), - _ => merged.push(larger), - } - } - *stored = merged; -} - /// What a map node applies. struct MapRuntime { ops: Arc<[MathOp]>, @@ -2293,6 +2127,14 @@ struct NodeRuntime { /// `None` for every other node, and for a binary node that broadcasts a constant instead: that /// one holds no operand at all and is described by [`NodeRuntime::broadcast`]. binary: Option, + /// What an aggregate node holds: the fold over what it was handed. + /// + /// Unlike a binary node's operands, nothing of the input survives a chunk: the fold keeps its + /// states and the chunk it read is dropped, so this is the one node whose reservation does not + /// follow what it was handed but what it accumulated. That is also why this node cannot report + /// its input complete from what it holds -- there is nothing to hold -- and reads the driver's + /// [`NativeExecutor::inputs_finished`] and its child's [`NativeExecutionState::settled`] instead. + aggregate: Option, /// What a vector node holds. vector: Option, /// What a binary node holds when one operand is a constant: the operation and the layout of the @@ -2702,6 +2544,102 @@ impl NodeRuntime { self.buffer.resize(queued) } + /// Folds one chunk into this node's aggregate. + /// + /// The chunk is reserved for before it is folded, by what the node reports it is worth, and the + /// reservation settles afterwards on the fold that is left. Unlike [`Self::retain_binary`], what + /// stops being held here is the chunk itself: the fold keeps its states, so a chunk that closes + /// its series leaves nothing behind and the node is empty of input again. + fn retain_aggregate(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { + let NativeChunk::Evaluated(chunk) = chunk else { + return Err(DataFusionError::Execution(format!( + "the aggregate at {path} consumes evaluated chunks, not raw samples" + ))); + }; + // The figure is read before the chunk is folded, because reserving after the fact would be + // reserving for memory the fold already read. + let pending = self + .aggregate + .as_ref() + .ok_or_else(|| no_aggregate(path))? + .pending(&chunk); + let positions = chunk.series.time.len(); + self.buffer.grow(pending)?; + let metrics = Arc::clone(&self.buffer.metrics); + { + let _timer = metrics.timer(); + let (aggregate, buffer) = (&mut self.aggregate, &mut self.buffer); + aggregate + .as_mut() + .ok_or_else(|| no_aggregate(path))? + .accumulate(path, chunk, buffer, pending)?; + } + metrics.record_input(1, positions); + // A chunk that closed its series closed it for the node below: a group is not a series, so + // the fold completing nothing on it is why nothing is counted here. The series this node + // closes are the ones it hands out, counted when it queues them. + Ok(()) + } + + /// Finishes this node's fold and queues one chunk per group it produced. + /// + /// Each group is one series, decided in full by the fold, so the one chunk it is handed over as + /// closes it: a consumer above reads "nothing more here" from that chunk rather than from another + /// one arriving. + fn compute_aggregate(&mut self, path: &NodePath) -> Result<()> { + let metrics = Arc::clone(&self.buffer.metrics); + let _timer = metrics.timer(); + let groups = self + .aggregate + .as_ref() + .map_or(0, |aggregate| aggregate.groups()); + let queue_capacity = self.queued.capacity(); + let queue_target = self.queued.len().saturating_add(groups); + let queue_growth = queue_target + .saturating_sub(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer.grow(queue_growth)?; + if groups > 0 { + self.queued.try_reserve_exact(groups).map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate output queue: {error}" + )) + })?; + } + let queue_bytes = queue_target + .max(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + let output = { + let aggregate = self.aggregate.as_mut().ok_or_else(|| no_aggregate(path))?; + aggregate.compute(path, &mut self.buffer, queue_bytes)? + }; + // The fold has returned its series; reserve their handover before appending them to the queue. + let Some(output) = output else { + return Ok(()); + }; + let mut closed = 0usize; + for series in output { + let bytes = series_bytes(&series); + let chunk = NativeChunk::Evaluated(EvaluatedChunk::try_new(series, true)?); + self.queued.push_back((chunk, bytes)); + closed += 1; + } + // A group is one series and this node closes each of them exactly once, on the one chunk it + // hands it over as. + self.buffer.metrics.record_completed_series(closed); + // The output Vec has been consumed; the node now holds only queue backing and chunks. + let queued_bytes = self + .queued + .iter() + .fold(0usize, |total, (_, bytes)| total.saturating_add(*bytes)); + let queue_backing = self + .queued + .capacity() + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer + .resize(queued_bytes.saturating_add(queue_backing)) + } + /// Feeds one chunk to this node's instant selection. /// /// `None` means the series this chunk belongs to has not ended: the selection reads one whole @@ -2904,6 +2842,8 @@ impl NodeRuntime { self.window = None; self.instant = None; self.binary = None; + // The fold is what an aggregate node holds, and dropping it is what gives its states back. + self.aggregate = None; self.vector = None; self.broadcast = None; self.source_schema = None; @@ -2996,15 +2936,14 @@ impl NativeExecutionState { // two vectors, or a constant broadcast over one -- is decided where the operands' // layouts are read. NativeExpr::Binary { .. } => NodeTag::Binary, + // An aggregate is built here and narrowed below, for the same reason: `validate` + // admits every kind and every stage as a legal tree, and the runtime computes one + // kind at one stage. + NativeExpr::Aggregate { .. } => NodeTag::Aggregate, NativeExpr::Vector { .. } => NodeTag::Vector, // A scalar is a value node: it produces the constant it was bound to, hands over no // chunk and holds no state, so it is built like any other node and pulled like none. NativeExpr::Scalar(_) => NodeTag::Scalar, - // An aggregate is legal in a bound tree, and `validate` stays the authority on that; - // this runtime simply does not execute it yet. Rejected before a single chunk is - // read: a tree the runtime cannot execute must not start, because a partial answer - // would be one the caller never asked for. - NativeExpr::Aggregate { .. } => return Err(not_implemented("Aggregate", &path)), }; let buffer = NodeBuffer::new(&path, handles, &pool); let mut runtime = NodeRuntime { @@ -3015,6 +2954,7 @@ impl NativeExecutionState { window: None, map: None, binary: None, + aggregate: None, vector: None, broadcast: None, source_schema: None, @@ -3114,6 +3054,29 @@ impl NativeExecutionState { // out. A broadcast that keeps the labels it reads is not where two series // become one either, so it keeps none; the one that drops them does, above. } + NativeExpr::Aggregate { input, op } => { + // A `Complete` aggregate neither adds nor drops a field and evaluates on its + // input's grid, so the node's own output is the layout it folds. What it does + // not execute -- another kind, or another stage -- is refused here, before the + // pass reaches the input that would feed it, and naming what it found is that + // rejection's job: a `Partial` aggregate's output is state and names no samples + // at all, so the input's layout stands in for it and the stage decides. + let schema = node + .output + .samples() + .or_else(|| input.output.samples()) + .cloned() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "the aggregate at {path} does not produce samples" + )) + })?; + runtime.aggregate = Some(AggregateRuntime::new( + op.clone(), + schema, + &mut runtime.buffer, + )?); + } NativeExpr::Vector { op, .. } => { let schema = node.output.samples().cloned().ok_or_else(|| { DataFusionError::Execution(format!( @@ -3136,10 +3099,9 @@ impl NativeExecutionState { // A vector operation keeps the labels of the series it reads, or answers under // label sets of its own, so it keeps no ledger of what it handed out. } - NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { + NativeExpr::Scalar(_) => { // A scalar holds nothing: what its consumer reads is the constant the node was - // bound to, which is on the tree rather than in this state. The aggregate is - // rejected above, where the pass has the position to name. + // bound to, which is on the tree rather than in this state. } } // The inputs this node reads, in the order the tree declares them: a unary node's one @@ -3147,12 +3109,11 @@ impl NativeExecutionState { // pushed back to front, so the node the spec names is the one the pass reaches next and // nothing but its own subtree lies between the two. let reads_input = match &node.expr { - NativeExpr::Input(_) | NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { - false - } + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } + | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => { pending.push((input.as_ref(), path.input())); true @@ -3222,6 +3183,14 @@ impl NativeExecutionState { holds_shared[index] = input_holds_shared; } NodeTag::Map => holds_shared[index] = true, + NodeTag::Aggregate => { + // An aggregate hands over one series per group, each under the label set its + // group was keyed by and each closed on its own chunk: two groups are two + // label sets, so nothing above has to join pieces of a series. Its own input + // is read whole -- the fold holds what it accumulated, not the chunks -- so + // nothing about its input's flag reaches it either. + holds_shared[index] = false; + } NodeTag::Scalar => { // A scalar produces one value and no labels, so it is not a stream of series // and nothing above it can read two of them under one label set. The flag @@ -3780,9 +3749,24 @@ impl NativeExecutor { // A scalar produces its value without reading anything, so it is settled from the // start: there is no chunk it owes its consumer. NodeTag::Scalar => true, + // These report as they are handed chunks, so a child with nothing left is all it takes + // for them to have nothing left either. NodeTag::Map | NodeTag::Window | NodeTag::Instant => self.state.nodes[index] .input .is_none_or(|child| self.settled(child)), + // An aggregate answers one series per group and only once its input ended, so its own + // fold has to have run as well: a child that has handed everything up is not enough, + // because every one of those chunks was folded and none of them answered yet. This is + // what keeps it out of the line above, which asks only about the child. + NodeTag::Aggregate => { + runtime + .aggregate + .as_ref() + .is_some_and(|aggregate| aggregate.computed) + && self.state.nodes[index] + .input + .is_none_or(|child| self.settled(child)) + } NodeTag::Vector => { runtime .vector @@ -3980,6 +3964,7 @@ impl NativeExecutor { // does, so its values are counted rather than read as its positions. runtime.pop_queued(evaluated_values) } + NodeTag::Aggregate => self.aggregate_node(index), NodeTag::Vector => { let child = input.ok_or_else(|| missing_input(tag))?; // Which samples are kept at a point is not decided until the whole input has @@ -3999,6 +3984,35 @@ impl NativeExecutor { } } + /// Folds everything the aggregate at `index` can read and hands over one group's series. + /// + /// What it reads is taken to the end of what its child has, because nothing the fold is handed + /// is retained: a chunk is folded as it arrives and dropped, so reading the child dry costs the + /// fold's states and not the input's bytes. Nothing is answered until the input has ended -- + /// before that, a group's value at a point is still missing the members that have not arrived. + /// + /// That is the difference from a binary node. A binary node retains both operands and so can ask + /// whether they hold anything still open; an aggregate keeps none of them and asks the driver + /// instead: [`Self::inputs_finished`] says the caller will hand over nothing more, and + /// [`NativeExecutionState::settled`] says the child has none left to produce. + fn aggregate_node(&mut self, index: usize) -> Result> { + let child = self.state.nodes[index] + .input + .ok_or_else(|| missing_input(NodeTag::Aggregate))?; + while let Some(chunk) = self.pull(child)? { + let path = &self.state.paths[index]; + self.state.runtimes[index].retain_aggregate(path, chunk)?; + } + if !self.inputs_finished || !self.settled(child) { + return Ok(None); + } + let path = &self.state.paths[index]; + self.state.runtimes[index].compute_aggregate(path)?; + // A group's series holds one value per field and point it decided, so its values are counted + // rather than read as its positions, as a binary result's are. + self.state.runtimes[index].pop_queued(evaluated_values) + } + /// Pulls the vector side of the broadcast at `index` until it hands over a chunk, and returns /// that chunk transformed. /// @@ -4226,11 +4240,7 @@ fn valid_values(chunk: &EvaluatedChunk) -> usize { .sum() } -/// The values a chunk a binary node hands over holds. -/// -/// A binary result is an evaluated chunk, so this is [`valid_values`]: one value per field sample -/// that holds one. A raw chunk cannot reach a binary node's queue, and is answered as its positions -/// so an impossible queue cannot under-count. +/// The values an evaluated chunk hands over, or the timestamps a raw source hands over. fn evaluated_values(chunk: &NativeChunk) -> usize { match chunk { NativeChunk::Evaluated(chunk) => valid_values(chunk), @@ -4238,18 +4248,6 @@ fn evaluated_values(chunk: &NativeChunk) -> usize { } } -/// The rejection of a node kind this runtime does not execute yet. -/// -/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call -/// the same node, so a reader can find it in the tree. An aggregate is the one kind left here: a -/// scalar executes as a value node, and a binary node executes both as a vector match and as a -/// constant broadcast. -fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { - DataFusionError::NotImplemented(format!( - "the native runtime does not execute {kind} nodes yet: {kind} at {path}" - )) -} - /// The error for a window whose evaluator this execution already closed. fn closed(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the window at {path} is already closed")) @@ -4331,11 +4329,27 @@ fn truncated_input(path: &NodePath) -> DataFusionError { )) } +/// The rejection of a node kind this runtime does not execute yet. +/// +/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call +/// the same node, so a reader can find it in the tree. +fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { + DataFusionError::NotImplemented(format!( + "the native runtime does not execute {kind} nodes yet: {kind} at {path}" + )) +} + /// The error for a node asked to compute without the runtime its kind holds. fn no_binary(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no binary operation")) } +/// The error for an aggregate node whose fold is missing. +fn no_aggregate(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!("the node at {path} holds no aggregate")) +} + +/// The error for a node asked to run a vector operation without one. fn no_vector(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no vector operation")) } diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index dff37b0c9d8..1efa3234a00 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -45,8 +45,9 @@ use crate::instant_eval::select_instants; use crate::map::MathOp; use crate::native::metrics::{NodeMetric, NodeMetrics, NodePath}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, BoundInstant, + BoundMap, BoundScalar, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, + NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::select::SortOrder; @@ -485,6 +486,7 @@ fn source_runtime(pool: &Arc) -> NodeRuntime { window: None, map: None, binary: None, + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -507,6 +509,7 @@ fn binary_runtime(pool: &Arc, schema: &Arc) -> Nod Arc::new(binary_matching(BinaryOp::Add)), [Arc::clone(schema), Arc::clone(schema)], )), + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -1083,30 +1086,17 @@ fn a_downstream_window_emits_before_its_input_series_ends() { } } -/// Every node kind the runtime does not execute is rejected, naming the kind and its position. +/// Every binary shape the runtime does not execute is rejected, naming the kind and its position. /// /// A scalar is not one of them any more: it is a value node, and a binary node over a constant is -/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes, -/// and the aggregate is still a kind this runtime does not run at all. +/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes; +/// the aggregate kinds and stages this runtime does not execute are refused by +/// `an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named`. #[test] fn an_unimplemented_node_is_reported_not_panicked() { // One grid shared by the nodes that validate by allocation identity. let statement = || grid(0, 120_000, 2); - let aggregate_tree = { - let out = statement(); - NativeNode { - expr: NativeExpr::Aggregate { - input: Box::new(window(input(), Arc::clone(&out))), - op: BoundAggregate { - kind: AggregateKind::Sum, - group: Grouping::By(Arc::new(["job".into()])), - stage: AggregateStage::Complete, - }, - }, - output: NativeType::Samples(evaluated_schema(out)), - } - }; // Two constants make a scalar, which is a legal node and a shape no binary here broadcasts // over: there is no vector to pair either of them with. let two_scalars_tree = { @@ -1138,12 +1128,6 @@ fn an_unimplemented_node_is_reported_not_panicked() { }; let cases = [ - ( - "Aggregate", - "root", - "does not execute Aggregate", - aggregate_tree, - ), ("Binary", "root", "two constants", two_scalars_tree), ("Binary", "root", "matching modifiers", matched_scalar_tree), ]; @@ -6558,6 +6542,158 @@ fn a_constant_on_the_right_multiplies_every_value() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// A `Complete` aggregate over `inner`, carrying `inner`'s fields and evaluated on its grid. +/// +/// A `Complete` aggregate neither adds nor drops a field and lands on its input's grid, so the node's +/// own output is the layout its input produces. +fn aggregate(inner: NativeNode, op: BoundAggregate) -> NativeNode { + let fields = inner + .output + .samples() + .expect("an aggregate reads samples") + .fields + .clone(); + let grid = inner + .output + .grid() + .expect("a complete aggregate reads an evaluated vector") + .clone(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields, + eval_grid: Some(grid), + })); + NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(inner), + op, + }, + output, + } +} + +/// A `Complete` `sum` grouped as `group`. +fn sum_grouping(group: Grouping) -> BoundAggregate { + BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Complete, + } +} + +/// `name{job=...,instance=...}`, the two labels these tests group by. +fn job_instance_metric(name: &str, job: &str, instance: &str) -> MetricName { + MetricName::try_new( + Some(name.into()), + vec![ + Label { + name: "job".into(), + value: job.into(), + }, + Label { + name: "instance".into(), + value: instance.into(), + }, + ], + ) + .unwrap() +} + +/// What one aggregate result reports: its name, its tags and the value at each point it decided. +type GroupResult = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// The groups a run produced, each with the absolute positions it decided, sorted by identity. +/// +/// Sorting is the test's, not the runtime's: a group's place in the answer is the order its first +/// input appeared, which the tests that check it assert on the stream rather than here. +fn groups_of(chunks: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = chunks + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + ( + chunk.series.metric.name().map(str::to_string), + tags, + decided(std::slice::from_ref(chunk)), + ) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// One group's answer under the name and tags `tags` give it, ordered by label name. +fn grouped(tags: &[(&str, &str)], points: Vec<(usize, Option)>) -> GroupResult { + let mut tags: Vec<(String, String)> = tags + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + tags.sort(); + (None, tags, points) +} + +/// One group's answer when every label of its input was dropped. +fn unlabelled(points: Vec<(usize, Option)>) -> GroupResult { + grouped(&[], points) +} + +/// `sum` with no grouping drops every label, name and tags alike, and adds the members per point. +#[test] +fn an_aggregate_drops_every_label_and_adds_its_group_per_point() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let set = ExecutionPlanMetricsSet::new(); + let mut execution = executor(&tree, pool(), &set); + for (metric, values) in [ + ( + job_instance_metric("requests", "api", "a"), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + job_instance_metric("requests", "db", "b"), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + execution + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the two series were not folded into one unlabelled group" + ); + // One group is one series, so it closes exactly once and its chunk says so. + assert!(output.len() == 1 && output[0].is_last); + assert_eq!( + count(&set, &NodePath::root(), NodeMetric::COMPLETED_SERIES), + 1 + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// The side a constant is on decides which operand the vector's value is. /// /// Division is not commutative, so `x / 100` and `100 / x` answering different values is what says @@ -7031,6 +7167,208 @@ fn a_window_above_a_broadcast_merges_the_series_it_collapses() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// Every supported `Complete` kind reaches its matching float kernel through the native runtime. +#[test] +fn complete_aggregate_kinds_fold_two_values() { + let cases = [ + (AggregateKind::Sum, 4.0), + (AggregateKind::Count, 2.0), + (AggregateKind::Avg, 2.0), + (AggregateKind::Min, 1.0), + (AggregateKind::Max, 3.0), + (AggregateKind::Group, 1.0), + (AggregateKind::Stddev, 1.0), + (AggregateKind::Stdvar, 1.0), + ]; + + for (kind, expected) in cases { + let grid = grid(0, 60_000, 1); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + BoundAggregate { + kind, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for value in [1.0, 3.0] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![(0, Some(expected))])], + "{kind:?} did not compute its native complete result" + ); + } +} + +/// The same series, delivered in two chunks, fold to what the whole series folds to. +#[test] +fn an_aggregate_folds_a_series_split_across_chunks_like_a_whole_one() { + let bound = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&bound)); + // A fresh tree per execution: the node's state is the execution's, not the tree's. + let tree = || { + NativeTree::try_new(aggregate( + evaluated_input(grid(0, 60_000, 4)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap() + }; + let (first, second) = ( + job_instance_metric("requests", "api", "a"), + job_instance_metric("requests", "api", "b"), + ); + + // One chunk per series. + let mut whole = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + for (metric, values) in [ + ( + first.clone(), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + second.clone(), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + whole + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + let whole = whole.finish().unwrap(); + + // The first series split in two, the second whole: the group is the same group, folded from + // positions the chunks name absolutely. + let mut split = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first.clone(), + 0, + vec![Some(1.0), Some(2.0)], + false, + )) + .unwrap(); + assert!( + split.next_chunk().unwrap().is_none(), + "a chunk that does not close its series answered for the group" + ); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first, + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + split + .push_evaluated(evaluated_chunk_of( + &schema, + second, + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + + let split = split.finish().unwrap(); + assert_eq!(groups_of(&split), groups_of(&whole)); + assert_eq!( + groups_of(&split), + vec![grouped( + &[("job", "api")], + vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ], + )], + "the split series was folded as two groups or at the wrong positions" + ); +} + +/// A point no member of a group has a value at is not a point of the group's result. +/// +/// Sparse input is what says the fold indexes its states by absolute position: a chunk holding +/// positions 0 and 1 and another holding 1 and 2 leave position 3 to nothing, and the hole between +/// the rows is never filled with a value. +#[test] +fn an_aggregate_reports_only_the_points_a_group_holds_a_value_at() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + // Positions 0 and 1 of the first series, 1 and 2 of the second, and nothing at 3. + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "api", "a"), + &[(0, 2)], + vec![1.0, 2.0], + true, + )) + .unwrap(); + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "db", "b"), + &[(1, 2)], + vec![20.0, 300.0], + true, + )) + .unwrap(); + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![ + // Position 1 is the only one both members hold, and it is their sum. + unlabelled(vec![(0, Some(1.0)), (1, Some(22.0)), (2, Some(300.0))]), + ], + "the group was answered at a point no member holds, or summed the wrong members" + ); +} + +/// An aggregate handed no series produces nothing at all: no group, so no empty series either. +#[test] +fn an_aggregate_over_no_series_produces_nothing() { + let grid = grid(0, 60_000, 4); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let output = execution.finish().unwrap(); + assert!( + output.is_empty(), + "an aggregate with no input answered {} series", + output.len() + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// A broadcast hands one chunk over for every chunk it reads, and closes the series once. /// /// The metrics are what says the chunk-per-chunk behaviour held: one input batch per chunk read, one @@ -7246,6 +7584,411 @@ fn a_broadcast_refuses_two_series_that_share_a_label_set_and_releases_the_reserv ); } +/// A series that ended is one member of a group, not the group: nothing is answered until the input +/// ends. +/// +/// The observable state between the two is what this pins: the fold has read the closing chunk of +/// the first series and answered nothing, and `computed` still says so. +#[test] +fn an_aggregate_answers_nothing_until_its_whole_input_ends() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let computed = |execution: &NativeExecutor| { + execution.state().runtimes[0] + .aggregate + .as_ref() + .expect("the root holds an aggregate") + .computed + }; + + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "api", "a"), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + assert!(!computed(&execution), "the fold finished on one series"); + assert!( + execution.next_chunk().unwrap().is_none(), + "the closing chunk of one series answered for its group" + ); + // The second series closed too, and the input still has not ended. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "db", "b"), + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + assert!( + !computed(&execution), + "the fold finished before the input ended" + ); + assert!( + execution.next_chunk().unwrap().is_none(), + "every series closed, and the group was answered before the input ended" + ); + + execution.end_of_input().unwrap(); + assert!( + !computed(&execution), + "ending the input finished the fold before anything pulled for it" + ); + let mut output = Vec::new(); + while let Some(chunk) = execution.next_chunk().unwrap() { + output.push(chunk); + } + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the group was not answered once the input ended" + ); +} + +/// A chunk of another layout is refused, naming the aggregate and its position, and the execution +/// stops with nothing reserved. +/// +/// The fold indexes its states by absolute position and by field, so a chunk of another field count +/// would be folded into states that mean something else. It is refused rather than folded, and what +/// the refusal leaves behind is a stopped execution whose reservation went back to the pool. +#[test] +fn an_aggregate_refuses_a_chunk_of_another_layout_and_returns_its_reservation() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + // Two fields where the tree bound one: the source node accepts it, the fold refuses it. + let wider = Arc::new(SeriesSchema { + fields: field_count(2), + eval_grid: Some(Arc::clone(&grid)), + }); + execution + .push_evaluated(evaluated_columns( + &wider, + metric(), + 0, + vec![vec![Some(1.0); 4], vec![Some(2.0); 4]], + true, + )) + .unwrap(); + + let error = execution + .next_chunk() + .expect_err("a chunk of another field count was folded"); + assert!( + matches!(error, DataFusionError::Execution(_)), + "the refusal is not an execution error: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("aggregate at root"), + "the refusal does not name the aggregate: {message}" + ); + assert!( + message.contains("field count"), + "the refusal does not say what disagreed: {message}" + ); + assert_eq!(execution.state().reserved_bytes(), 0); + // A stopped execution is not resumed: the fold must not be asked to continue after this. + assert!(execution.next_chunk().is_err()); +} + +/// A pool rejection while the first aggregate block is admitted releases the source and fold. +#[test] +fn aggregate_budget_rejects_first_block_and_releases_everything() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(2_000); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().is_err()); + assert_eq!(pool.reserved(), 0); + assert_eq!(execution.state().reserved_bytes(), 0); +} + +/// The aggregate runtime charges its output queue, releases each output as it is pulled, and keeps +/// queue backing until the runtime itself is released. +#[test] +fn aggregate_runtime_queue_keeps_only_backing_after_pops_until_release() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + for (job, value) in [("api", 1.0), ("db", 2.0)] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", job, "a"), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + execution.end_of_input().unwrap(); + let first = execution.next_chunk().unwrap().expect("first group output"); + assert_eq!(first.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + let entry = std::mem::size_of::<(NativeChunk, usize)>(); + assert_eq!(runtime.queued.len(), 1); + let remaining = runtime.queued[0].1; + let expected = runtime.queued.capacity() * entry + remaining; + assert_eq!(runtime.buffer.reservation.size(), expected); + + let second = execution + .next_chunk() + .unwrap() + .expect("second group output"); + assert_eq!(second.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + assert!(runtime.queued.is_empty()); + let backing = runtime.queued.capacity() * entry; + assert!(backing > 0); + assert_eq!(runtime.buffer.reservation.size(), backing); + assert!(execution.next_chunk().unwrap().is_none()); + // Pulling the terminal None releases the queue backing; only explicit close checks the + // runtime's final cleanup below. + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + + execution.state.close().unwrap(); + assert_eq!(execution.state().runtimes[0].queued.capacity(), 0); + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate whose sparse output has a few points on a large grid can finish in a small pool. +#[test] +fn aggregate_sparse_output_budget_uses_rows_not_grid_length() { + let grid = grid(0, 1, 100_000); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(64 * 1024); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 10, + vec![Some(2.0)], + true, + )) + .unwrap(); + let output = execution.finish().unwrap(); + assert_eq!(groups_of(&output), vec![unlabelled(vec![(10, Some(2.0))])]); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate retains its fold before finish and release returns its reservation. +#[test] +fn aggregate_fold_state_is_released_after_input() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + let runtime = &execution.state.runtimes[0]; + let aggregate = runtime.aggregate.as_ref().unwrap(); + assert_eq!(aggregate.groups(), 1); + assert!(!aggregate.computed); + assert!(pool.reserved() > 0); + execution.state.close().unwrap(); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate kind or stage this runtime does not compute is refused, naming what was found. +#[test] +fn an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named() { + let grid = || grid(0, 60_000, 4); + let kind_tree = NativeTree::try_new(aggregate( + evaluated_input(grid()), + BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let partial_tree = { + let out = grid(); + let group = Grouping::By(Arc::new(["job".into()])); + let state = Arc::new(AggregateStateSchema { + group: group.clone(), + accumulators: fields(), + grid: Arc::clone(&out), + }); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(evaluated_input(Arc::clone(&out))), + op: BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(state), + }) + .unwrap() + }; + + for (found, tree) in [("Quantile", kind_tree), ("Partial", partial_tree)] { + let error = match NativeExecutionState::try_new( + &tree, + pool(), + &ExecutionPlanMetricsSet::new(), + 0, + ) { + Ok(_) => panic!("the runtime accepted a {found} aggregate"), + Err(error) => error, + }; + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "{found}: {error}" + ); + let message = error.to_string(); + assert!(message.contains(found), "{found}: {message}"); + // A `Partial` rejection has to say why it is not computed: the stages it would need. + if found == "Partial" { + assert!(message.contains("Final"), "{message}"); + } + } +} + +/// An aggregate above a window that collects and merges its input waits for that merge. +/// +/// The multiplication drops the metric name, so the two input series land on one label set and the +/// `last_over_time` above them has to collect and merge before it evaluates. The aggregate above the +/// window must wait for that merge: its guard is `inputs_finished && settled(child)`, and with only +/// the first half it finalizes on the pull that follows the caller's last chunk, before the window +/// has answered. +#[test] +fn an_aggregate_above_a_collecting_window_waits_for_the_merge() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = || { + NativeTree::try_new(aggregate( + window_with( + broadcast_right( + evaluated_input(Arc::clone(&grid)), + 2.0, + BinaryOp::Mul, + false, + ), + RangeFunction::LastOverTime, + 60_000, + Some(Arc::clone(&grid)), + Arc::clone(&grid), + ), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap() + }; + + let mut execution = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + // Two series that only differ by name, each holding half of the grid the other leaves empty: + // after the name is dropped they are one series only once merged. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests", "x"), + 0, + vec![Some(1.0), Some(2.0)], + true, + )) + .unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the aggregate answered before the window above it had merged" + ); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests_total", "x"), + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + let out = run(&mut execution); + assert_eq!( + decided(&out), + vec![ + (0, Some(2.0)), + (1, Some(4.0)), + (2, Some(6.0)), + (3, Some(8.0)), + ], + "the aggregate did not fold the merged window result" + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + fn vector_op(inner: NativeNode, op: BoundVectorOp) -> NativeNode { let output = inner.output.clone(); NativeNode { diff --git a/src/promql/src/native/mod.rs b/src/promql/src/native/mod.rs index c9545a9fe69..fd04dcfda32 100644 --- a/src/promql/src/native/mod.rs +++ b/src/promql/src/native/mod.rs @@ -27,8 +27,10 @@ //! - A series closes exactly once, on the chunk that says so, so a consumer can tell "nothing more //! here" from "no more input". +mod aggregate; pub mod batch; mod display; pub mod execute; pub mod metrics; pub mod plan; +pub(crate) mod runs; diff --git a/src/promql/src/native/runs.rs b/src/promql/src/native/runs.rs new file mode 100644 index 00000000000..8de325a5386 --- /dev/null +++ b/src/promql/src/native/runs.rs @@ -0,0 +1,146 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The grid positions a field occupies, as sorted, disjoint runs. +//! +//! A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name and +//! the empty-valued labels — can turn two input series into one label set, and every layer above +//! then reads the two chunks as one series. What decides whether that pair is legal is which +//! positions of each field already hold a value: two values of one field at one position cannot be +//! merged away, while a chunk that fills positions the set left empty merges cleanly, as +//! [`crate::grid_series::merge_same_labelset`] does. +//! +//! The state is kept per label set, per field, as runs, so a chunk of many positions costs one run +//! rather than one entry per position. Keeping it in one place is what keeps the two layers that +//! check it — the native runtime's own [`crate::native::execute::Emitted`] table and the +//! cross-partition [`crate::extension_plan::fused_series::LabelSets`] — from drifting apart: both +//! ask the same functions for the same answers, so the fast path one of them earned is the fast path +//! the other gets. + +/// One run of grid positions, half-open, that hold a value in one field. +#[derive(Clone, Copy, Debug)] +pub(crate) struct Run { + pub(crate) start: usize, + pub(crate) end: usize, +} + +impl Run { + /// Whether two runs share a position. + pub(crate) fn overlaps(self, other: Self) -> bool { + self.start < other.end && other.start < self.end + } +} + +// How many times `absorb` rebuilt the stored runs instead of appending to them. +// +// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows +// only in what it costs, and a test that reads its own thread's figure is not disturbed by another +// test's chunks. +#[cfg(test)] +thread_local! { + static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// How many times this thread's [`absorb`] rebuilt the stored runs. +#[cfg(test)] +pub(crate) fn absorb_rebuilds() -> usize { + ABSORB_REBUILDS.with(std::cell::Cell::get) +} + +/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that +/// touch. +/// +/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and +/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square +/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values +/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves +/// for a single series. +/// +/// One pass per chunk is still a pass over everything the series has stored so far, so it is not +/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, +/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision +/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run +/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends +/// before anything this chunk holds — and it is appended instead, which is amortised constant per +/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk +/// that starts before the last stored run can split a run or reach back past it, so it still costs +/// the full merge. +/// +/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what a +/// look-up by [`Run::overlaps`] on the runs stored here relies on. +pub(crate) fn absorb(stored: &mut Vec, incoming: &[Run]) { + if incoming.is_empty() { + return; + } + if stored + .last() + .is_none_or(|last| incoming[0].start >= last.end) + { + append_runs(stored, incoming); + } else { + merge_runs(stored, incoming); + } +} + +/// Appends `incoming` to `stored`, all of which lies past the last stored run. +/// +/// Nothing already stored has to be placed again: the two sides are each in position order and +/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and +/// the vector's own growth is the only move — amortised over the runs appended. Where its first run +/// touches the end of the last stored run the two are one run, as they are after a merge, so it is +/// joined rather than left beside it. +fn append_runs(stored: &mut Vec, incoming: &[Run]) { + let mut incoming = incoming; + if let Some(last) = stored.last_mut() + && let Some(first) = incoming.first() + && first.start <= last.end + { + last.end = last.end.max(first.end); + incoming = &incoming[1..]; + } + stored.extend_from_slice(incoming); +} + +/// Merges `incoming` into `stored`, where the two can interleave. +/// +/// One pass is enough because each side is sorted and disjoint: whichever run starts first is +/// placed, and it joins the run before it where the two touch. +fn merge_runs(stored: &mut Vec, incoming: &[Run]) { + #[cfg(test)] + ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); + let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); + let (mut a, mut b) = (0, 0); + while a < stored.len() || b < incoming.len() { + let larger = match (stored.get(a), incoming.get(b)) { + (Some(left), Some(right)) if right.start < left.start => { + b += 1; + *right + } + (Some(left), _) => { + a += 1; + *left + } + (None, Some(right)) => { + b += 1; + *right + } + (None, None) => break, + }; + match merged.last_mut() { + Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), + _ => merged.push(larger), + } + } + *stored = merged; +} diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 0a44374f971..7c07f68468d 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -13,17 +13,17 @@ // limitations under the License. use std::sync::Arc; -use std::time::{Duration, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use datafusion::error::DataFusionError; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::MetricName; use datatypes::series::SeriesOrdering; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; use promql::evaluated::EvaluatedChunk; use promql::map::MathOp; use promql::native::execute::{NativeExecutionState, NativeExecutor}; @@ -3149,11 +3149,25 @@ fn large_subquery_child_grids_compile_deferred() { assert_eq!(grid.2, range_ms); } +/// The largest millisecond endpoint the platform can represent as a `SystemTime`. +/// +/// The tests below need an endpoint close enough to `i64::MAX` milliseconds that a one-second +/// rounding or a one-millisecond offset crosses the representable range. Unix-like platforms can +/// hold that instant; Windows represents `SystemTime` as a 64-bit count of 100-nanosecond +/// intervals since 1601, so the same instant does not fit and `UNIX_EPOCH + duration` panics. The +/// behaviour under test is about arithmetic on the millisecond value, not about the platform's +/// clock range, so a platform that cannot represent the instant skips instead of failing. +fn extreme_millisecond_endpoint() -> Option { + UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64)) +} + /// A sampled child endpoint can fit even when the unrounded read-window endpoint does not. #[test] fn a_subquery_accepts_an_in_range_sampled_endpoint() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[1s:1s] offset -1ms)"); - let endpoint = UNIX_EPOCH + Duration::from_millis(i64::MAX as u64); stmt.start = endpoint; stmt.end = endpoint; stmt.interval = Duration::from_millis(1); @@ -3172,6 +3186,9 @@ fn a_subquery_accepts_an_in_range_sampled_endpoint() { /// An anchored fold offset that cannot be represented as i64 is a plan error. #[test] fn a_subquery_with_an_overflowing_fold_offset_is_refused() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[5m:1m] @ 0)"); let PromExpr::Call(call) = &mut stmt.expr else { unreachable!("query parses to a call"); @@ -3179,9 +3196,7 @@ fn a_subquery_with_an_overflowing_fold_offset_is_refused() { let PromExpr::Subquery(subquery) = &mut *call.args.args[0] else { unreachable!("call argument parses to a subquery"); }; - subquery.at = Some(AtModifier::At( - UNIX_EPOCH + Duration::from_millis(i64::MAX as u64), - )); + subquery.at = Some(AtModifier::At(endpoint)); subquery.offset = Some(Offset::Neg(Duration::from_millis(1))); let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime) @@ -3734,3 +3749,526 @@ fn whole_vector_functions_bind_a_vector_operation() { assert_eq!(root, format!(" root Vector {op}"), "`{query}`"); } } + +/// One raw chunk of `metric`, keyed by `key`, closing the series. +/// +/// `timestamps` are milliseconds and `values` one per timestamp, which is what a scan of the series +/// hands the tree. +fn agg_raw_chunk( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], +) -> RawSeriesChunk { + agg_raw_chunk_with_last(schema, key, metric, timestamps, values, true) +} + +/// One raw chunk of `metric`, keyed by `key`, with an explicit series-ending flag. +fn agg_raw_chunk_with_last( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], + is_last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(timestamps.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))), + is_last, + ) + .unwrap() +} + +/// `name` with the labels in `tags`, in the order given. +fn agg_metric(name: &str, tags: &[(&str, &str)]) -> MetricName { + MetricName::try_new( + Some(name.into()), + tags.iter() + .map(|(name, value)| Label { + name: (*name).into(), + value: (*value).into(), + }) + .collect(), + ) + .unwrap() +} + +/// What one result series of a native execution reports. +type AggSeries = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// Reads the outputs a native execution produced. +fn aggregate_results(output: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = output + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + let points = chunk + .series + .time + .positions() + .map(|(position, row)| { + let sample = chunk.series.fields.as_slice()[0] + .iter() + .nth(row) + .expect("the axis and the field are aligned"); + ( + position, + match sample { + SampleRef::Float(value) => Some(value), + _ => None, + }, + ) + }) + .collect(); + (chunk.series.metric.name().map(str::to_string), tags, points) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// Runs `tree` over `chunks`, which enter at the tree's only input, and reads the answer. +/// +/// The answer is read the way a caller reads it: the whole input is pushed, the input is ended, and +/// what the root reports is collected. Each series is reported by its identity and the value it +/// decided at each absolute grid position. +fn run_aggregate_tree(tree: &NativeTree, chunks: Vec) -> Vec { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution + .push_raw(chunk) + .unwrap_or_else(|error| panic!("a source chunk was refused: {error}")); + } + let output = execution + .finish() + .unwrap_or_else(|error| panic!("the tree did not answer: {error}")); + aggregate_results(&output) +} + +/// The 21 grid points of a statement over 100 seconds in 5-second steps. +fn agg_points_0_to_100() -> Vec { + (0..21).collect() +} + +/// The timestamps of the same 21 points, in milliseconds. +fn agg_timestamps_0_to_100() -> Vec { + (0..21).map(|index| index * 5_000).collect() +} + +/// `sum(x)` over two series: one unlabelled group, and every label of both inputs dropped. +#[test] +fn a_compiled_sum_drops_every_label_and_adds_its_group() { + let tree = compile( + &eval_stmt("sum(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &(1..=21).map(f64::from).collect::>(), + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "db"), ("instance", "b")]), + ×tamps, + &(1..=21).map(|value| (2 * value) as f64).collect::>(), + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `3` at the first point, three more at each of the twenty after it: the two members summed. + let values: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(3.0 * (position as f64 + 1.0)))) + .collect(); + assert_eq!(output, vec![(None, Vec::new(), values)]); +} + +#[test] +fn a_compiled_sum_aggregates_series_chunks_after_input_ends() { + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let values: Vec = (1..=21).map(f64::from).collect(); + let cases = [ + ("sum(x)", None, Vec::new()), + ( + "sum by(job)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ( + "sum without(instance)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ]; + + for (query, name, tags) in cases { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = vec![ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..10], + &values[..10], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[10..], + &values[10..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "`{query}` produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + let results = aggregate_results(&output); + let expected_points = values + .iter() + .enumerate() + .map(|(position, value)| (position, Some(*value))) + .collect(); + assert_eq!( + results, + vec![(name.map(str::to_string), tags, expected_points)], + "`{query}`" + ); + } +} + +/// Empty source input is not a member of an aggregate's implicit group. +#[test] +fn a_compiled_aggregate_over_empty_input_produces_no_groups() { + for query in ["sum(x)", "sum by(job)(x)", "sum without(instance)(x)"] { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!( + run_aggregate_tree(&tree, Vec::new()), + Vec::new(), + "`{query}` invented a group for empty input" + ); + } +} + +/// `sum by(job)(x)` over five series: two groups, each the sum of its own members. +/// +/// The members cover different points, so each group is a sum over a subset that changes from point +/// to point: one series holds the whole grid, one holds a single sample between two grid points, and +/// one stops a whole lookback before the grid does. No last sample sits exactly one lookback behind +/// a grid point, so which points each member covers does not depend on how the boundary is read. +#[test] +fn a_compiled_sum_by_groups_its_members_and_sums_each_group() { + let tree = compile( + &eval_stmt("sum by(job)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let all = agg_timestamps_0_to_100(); + let chunks = vec![ + // The whole grid, holding one. + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + &all, + &[1.0; 21], + ), + // One sample halfway between grid points 0 and 1, so only grid point 1 selects it. + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + &[2_500], + &[2.0], + ), + // Four samples between grid points, covering grid points 1 to 4. + agg_raw_chunk( + &schema, + b"e", + agg_metric("x", &[("job", "api"), ("instance", "e")]), + &[2_400, 7_400, 12_400, 17_400], + &[5.0; 4], + ), + // The whole grid, holding ten. + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + &all, + &[10.0; 21], + ), + // One sample halfway between grid points 0 and 1, holding twenty. + agg_raw_chunk( + &schema, + b"d", + agg_metric("x", &[("job", "db"), ("instance", "d")]), + &[2_500], + &[20.0], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `api` is `a` alone at points 0 and 5 to 20, and every member at the points they cover: + // point 1 is `1 + 2 + 5`, points 2 to 4 are `1 + 5`. + let api: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| { + let value = match position { + 1 => 8.0, + 2..=4 => 6.0, + _ => 1.0, + }; + (position, Some(value)) + }) + .collect(); + // `db` is the same shape over its own members: point 1 is `10 + 20`, the rest is `c` alone. + let db: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(if position == 1 { 30.0 } else { 10.0 }))) + .collect(); + assert_eq!( + output, + vec![ + (None, vec![("job".to_string(), "api".to_string())], api), + (None, vec![("job".to_string(), "db".to_string())], db), + ] + ); +} + +/// `sum without(instance)(x)`: `instance` goes, `job` and `zone` stay, and so does `__name__` going. +#[test] +fn a_compiled_sum_without_drops_the_named_label_and_the_name() { + let tree = compile( + &eval_stmt("sum without(instance)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a"), ("zone", "west")]), + ×tamps, + &[1.0; 21], + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b"), ("zone", "east")]), + ×tamps, + &[2.0; 21], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let values_of = |value: f64| -> Vec<(usize, Option)> { + agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(value))) + .collect() + }; + // Two groups, because `zone` is not the dropped label: the two series differ by more than the + // label the grouping removes. + assert_eq!( + output, + vec![ + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "east".to_string()), + ], + values_of(2.0), + ), + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "west".to_string()), + ], + values_of(1.0), + ), + ] + ); +} + +/// `sum by(job)(rate(x[5m]))`: an aggregate over a window, mid-chain. +/// +/// Every series is a counter rising by one per five seconds, and every grid point's window ends on a +/// sample and begins five seconds before one, so the extrapolated rate is the same at every point: +/// `59 / 295 = 0.2` per second, which the double arithmetic lands on the literals below. +#[test] +fn a_compiled_sum_over_a_rate_groups_the_window_it_reads() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + // The window reads 300 seconds, so the samples start a full window before the first grid point. + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + ×tamps, + &values, + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let at_every_point = |value: f64| -> Vec<(usize, Option)> { + (0..21).map(|position| (position, Some(value))).collect() + }; + // One `rate` is `0.19999999999999998`, and the pair of them add to `0.39999999999999997`. + assert_eq!( + output, + vec![ + ( + None, + vec![("job".to_string(), "api".to_string())], + at_every_point(0.39999999999999997), + ), + ( + None, + vec![("job".to_string(), "db".to_string())], + at_every_point(0.19999999999999998), + ), + ] + ); +} + +/// Splitting one counter's samples does not change the rate aggregate or make it answer early. +#[test] +fn a_compiled_sum_over_a_rate_aggregates_series_chunks_after_input_ends() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = [ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..41], + &values[..41], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[41..], + &values[41..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the rate aggregate produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + assert_eq!( + aggregate_results(&output), + vec![( + None, + vec![("job".to_string(), "api".to_string())], + (0..21) + .map(|position| (position, Some(0.19999999999999998))) + .collect(), + )] + ); +} + +/// Empty source input is not a member of this aggregate's implicit group, with or without a window. +#[test] +fn a_compiled_sum_over_a_rate_with_empty_input_produces_no_groups() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!(run_aggregate_tree(&tree, Vec::new()), Vec::new()); +}