diff --git a/src/promql/src/range_eval.rs b/src/promql/src/range_eval.rs index b809fe2284d..070ac2d7a4b 100644 --- a/src/promql/src/range_eval.rs +++ b/src/promql/src/range_eval.rs @@ -17,12 +17,13 @@ use std::sync::Arc; use common_query::prometheus::is_prometheus_stale_nan; +use datafusion::arrow::array::Array; use datafusion::arrow::buffer::Buffer; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; use datatypes::series::histogram::CounterResetHint; use datatypes::series::raw::RawSeriesChunk; -use datatypes::series::samples::{SampleRef, Series, SeriesFields, SeriesSchema}; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Series, SeriesFields, SeriesSchema}; use datatypes::series::time::{EvalGrid, GridSpec, TimeAxis}; use crate::grid_series::{GridSample, GridSeriesBuilder, METADATA_LABELS}; @@ -90,6 +91,10 @@ pub struct RangeEvaluator { drops_metric_name: bool, /// Raised by the first non-empty window, where Prometheus validates function arguments. window_error: Option, + /// Chunks appended by the bulk path, so a test can prove that path is reached rather than + /// only that the output is right, which the per-sample path would also produce. + #[cfg(test)] + bulk_ingests: usize, } struct SeriesState { @@ -181,6 +186,32 @@ impl FieldSamples { self.values.push(value); } + /// Appends samples that are all present, in bulk. + /// + /// Same result as calling [`push`](Self::push) once per sample, including the reset scan, + /// which stays per-sample for a field that maintains the index; the two appends become two + /// `memcpy`s. + fn push_bulk(&mut self, timestamps: &[i64], values: &[f64]) { + debug_assert_eq!(timestamps.len(), values.len()); + if self.index_resets { + let start = self.values.len(); + let mut previous = self.values.last().copied(); + for (offset, &value) in values.iter().enumerate() { + if let Some(previous_value) = previous + && value < previous_value + { + self.resets.push(CounterReset { + position: self.base + (start + offset) as u64, + previous_value, + }); + } + previous = Some(value); + } + } + self.timestamps.extend_from_slice(timestamps); + self.values.extend_from_slice(values); + } + /// Drops the samples before `head` and the resets whose predecessor is one of them. fn advance(&mut self) { let head = std::mem::take(&mut self.head); @@ -299,9 +330,19 @@ impl RangeEvaluator { row: Vec::with_capacity(input.fields.len()), drops_metric_name: false, window_error: None, + #[cfg(test)] + bulk_ingests: 0, }) } + /// Chunks appended without the per-sample loop. Test-only: both paths produce the same + /// output by construction, so a test that only compared output could not tell whether the + /// bulk path ran at all. + #[cfg(test)] + pub(crate) fn bulk_ingest_count(&self) -> usize { + self.bulk_ingests + } + /// Schema of the evaluated series: the input fields on the evaluation grid. pub fn output_schema(&self) -> &Arc { &self.output_schema @@ -431,9 +472,65 @@ impl RangeEvaluator { Ok(()) } + /// Appends the chunk in bulk when every field is an entirely present, unshifted float column. + /// + /// Returns `false` when any condition of the per-sample path applies, so the caller can run + /// it. A chunk carrying a stale marker takes that path too: the marker drops the whole row, + /// and reproducing that here would cost the same scan the loop already makes. + /// + /// A field mixing floats and native histograms never qualifies: `dense()` reports + /// [`DenseSamplesRef::Mixed`] for it, keeping its windows on the histogram path where the + /// mixed-window rules decide the result. + fn ingest_in_bulk( + &mut self, + timestamps: &[i64], + fields: &SeriesFields, + drop_stale: bool, + ) -> bool { + // Lengths are checked by `check_chunk` for raw input and by `push_series` for evaluated + // input, so only the shape of the columns is decided here. + if self.window.offset_ms != 0 { + return false; + } + let mut columns = Vec::with_capacity(fields.len()); + for field in fields.iter() { + let Some(DenseSamplesRef::Float(values)) = field.dense() else { + return false; + }; + if values.null_count() != 0 { + return false; + } + columns.push(values); + } + // A kept sample's shifted time must fit `i64`; with no offset it is the source timestamp, + // which `RawTimestamps` already validated as increasing and which the frame widened into + // this array without narrowing, so there is nothing left to check here. + if drop_stale + && columns.iter().any(|values| { + values + .values() + .iter() + .any(|value| is_prometheus_stale_nan(*value)) + }) + { + return false; + } + for (field, values) in self.fields.iter_mut().zip(columns) { + field.push_bulk(timestamps, values.values()); + } + #[cfg(test)] + { + self.bulk_ingests += 1; + } + true + } + /// Retains the samples of `fields` at `timestamps`. From `raw` input a row holding a stale /// marker in any field is dropped; evaluated input keeps it as a value. fn ingest(&mut self, timestamps: &[i64], fields: &SeriesFields, raw: bool) -> Result<()> { + if self.ingest_in_bulk(timestamps, fields, raw) { + return Ok(()); + } let mut samples: Vec<_> = fields.iter().map(|field| field.iter()).collect(); let mut row = Vec::with_capacity(samples.len()); for ×tamp in timestamps { diff --git a/src/promql/src/range_eval/tests.rs b/src/promql/src/range_eval/tests.rs index 66f38734f98..cfa707c0d8c 100644 --- a/src/promql/src/range_eval/tests.rs +++ b/src/promql/src/range_eval/tests.rs @@ -506,6 +506,71 @@ fn functions_match_range_udfs() { } } +/// A chunk of entirely present float columns is appended in bulk, and that is the path the test +/// asserts: the per-sample loop produces the same output, so comparing values alone could not +/// tell whether the bulk path ran. +#[test] +fn an_ordinary_float_chunk_is_appended_in_bulk() { + let rows: Vec = (0..5) + .map(|index| (index as i64 * 10_000, vec![Some(index as f64), Some(1.0)])) + .collect(); + let mut evaluator = probe_evaluator( + 2, + grid(0, 10_000, 5), + RangeWindow { + range_ms: 20_000, + offset_ms: 0, + anchor_ms: None, + }, + ); + assert!(evaluator.push(&chunk(&rows, 2, 0, true)).unwrap().is_some()); + assert_eq!(evaluator.bulk_ingest_count(), 1); +} + +/// A field holding native histogram samples keeps its windows on the histogram path, where the +/// mixed-window rules decide the result. Letting such a chunk take the bulk path would drop the +/// histograms and read the floats alone, so the chunk has to fall back to the per-sample loop. +#[test] +fn a_chunk_with_histograms_is_not_appended_in_bulk() { + let histogram = "0/0/0/4014000000000000/4014000000000000/1/0_1/4014000000000000///"; + let mut builder = HistogramBatchBuilder::new(); + let index = builder + .push(OwnedHistogram::parse(histogram).view()) + .unwrap(); + let histograms = HistogramSamples::try_new( + UInt32Array::from(vec![None, Some(index), None]), + builder.finish().unwrap(), + ) + .unwrap(); + let samples = Samples::try_from_mixed( + Float64Array::from(vec![Some(1.0), None, Some(2.0)]), + histograms, + ) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &raw_schema(1), + Buffer::from(0u64.to_le_bytes().to_vec()), + Some(metric(0)), + RawTimestamps::try_new(ScalarBuffer::from(vec![0, 10_000, 20_000])).unwrap(), + SeriesFields::try_from_vec(vec![samples]).unwrap(), + true, + ) + .unwrap(); + let mut evaluator = RangeEvaluator::try_new( + &raw_schema(1), + grid(0, 10_000, 3), + RangeWindow { + range_ms: 20_000, + offset_ms: 0, + anchor_ms: None, + }, + RangeFunction::PresentOverTime, + ) + .unwrap(); + evaluator.push(&chunk).unwrap(); + assert_eq!(evaluator.bulk_ingest_count(), 0); +} + #[test] fn empty_input_and_series_without_results() { let window = RangeWindow {