From 66f055e186a79913848eceea125800922e74737e Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 21 Oct 2026 11:39:40 +0800 Subject: [PATCH 01/13] feat(promql): merge native histogram series that share a label set Signed-off-by: Dennis Zhuang --- src/promql/src/grid_series.rs | 44 +++++--------- src/promql/src/grid_series/tests.rs | 93 ++++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 30 deletions(-) diff --git a/src/promql/src/grid_series.rs b/src/promql/src/grid_series.rs index ce24cae6af4..91d04b4c4b6 100644 --- a/src/promql/src/grid_series.rs +++ b/src/promql/src/grid_series.rs @@ -320,7 +320,8 @@ pub(crate) fn label_set(metric: &MetricName) -> LabelSet { } /// Merges series that share a label set into one, as Prometheus does after a function that -/// rewrites labels or drops `__name__`; two of them holding one field at one point fail. +/// rewrites labels or drops `__name__`; two of them holding one field at one point fail, whether +/// the samples are floats or native histograms. pub fn merge_same_labelset(schema: &Arc, series: Vec) -> Result> { let mut groups: HashMap> = HashMap::with_capacity(series.len()); let mut collided = false; @@ -346,8 +347,7 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R // slot indexes `members`, so the merge below picks the member's cursor directly instead of // searching for the member per row. let mut rows: Vec<(usize, usize, usize)> = Vec::new(); - let mut values = vec![None; schema.fields.len()]; - let mut merged = vec![None; schema.fields.len()]; + let mut merged: Vec> = Vec::with_capacity(schema.fields.len()); for (labels, members) in groups { if let [member] = members[..] { output.extend(series[member].take()); @@ -372,7 +372,12 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R .map(|field| field.cursor()) .collect(); for (point, row) in member_series.time.positions() { - if read_row(&mut cursors, row, &mut values)? { + if cursors.iter_mut().any(|cursor| { + !matches!( + cursor.at(row).expect("the axis and the fields are aligned"), + SampleRef::Missing + ) + }) { rows.push((point, member_slot, row)); } } @@ -400,11 +405,12 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R let point = rows[start].0; let end = start + rows[start..].partition_point(|&(p, _, _)| p == point); // Rows of one point merge field by field; only a field present twice collides. - merged.iter_mut().for_each(|value| *value = None); + merged.clear(); + merged.resize_with(schema.fields.len(), || None); for &(_, member_slot, row) in &rows[start..end] { - read_row(&mut cursors[member_slot], row, &mut values)?; - for (slot, value) in merged.iter_mut().zip(&values) { - if value.is_some() { + for (slot, cursor) in merged.iter_mut().zip(&mut cursors[member_slot]) { + let sample = cursor.at(row).expect("the axis and the fields are aligned"); + if !matches!(sample, SampleRef::Missing) { if slot.is_some() { let labels: Vec<_> = labels .iter() @@ -415,11 +421,11 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R labels.join(", ") ))); } - *slot = *value; + *slot = owned(sample)?; } } } - builder.push(point, &merged)?; + builder.push_samples(point, &merged)?; start = end; } let metric = series[members[0]].as_ref().unwrap().metric.clone(); @@ -428,24 +434,6 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R Ok(output) } -/// Reads the float fields of `row` through the member's cursors; returns whether any is present. -fn read_row(cursors: &mut [RowCursor<'_>], row: usize, values: &mut [Option]) -> Result { - let mut present = false; - for (value, cursor) in values.iter_mut().zip(cursors) { - *value = match cursor.at(row).expect("the axis and the fields are aligned") { - SampleRef::Missing => None, - SampleRef::Float(value) => Some(value), - SampleRef::Histogram(_) => { - return Err(DataFusionError::NotImplemented( - "merging histogram series with the same label set".into(), - )); - } - }; - present |= value.is_some(); - } - Ok(present) -} - pub(crate) fn external(error: datatypes::error::Error) -> DataFusionError { DataFusionError::External(Box::new(error)) } diff --git a/src/promql/src/grid_series/tests.rs b/src/promql/src/grid_series/tests.rs index 59e833e3fd3..a3c7d0b3547 100644 --- a/src/promql/src/grid_series/tests.rs +++ b/src/promql/src/grid_series/tests.rs @@ -15,7 +15,7 @@ use std::mem::size_of; use std::sync::Arc; -use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, SeriesSchema}; +use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, Series, SeriesSchema}; use datatypes::series::time::{EvalGrid, GridSpec}; use datatypes::series::{Label, MetricName}; use rand::rngs::StdRng; @@ -25,7 +25,9 @@ use super::reference::{ Labels, Points, assert_pieces_eq, assert_result_eq, points_of, random_dense, run_ranges, series_of, stitch, }; -use super::{GridSeriesBuilder, merge_same_labelset}; +use super::{GridSample, GridSeriesBuilder, merge_same_labelset}; +use crate::grid_series::reference::OwnedHistogram; +use crate::histogram_eval::float_histogram::FloatHistogram; fn grid_schema(fields: usize, grid_len: u64) -> Arc { Arc::new(SeriesSchema { @@ -206,6 +208,93 @@ fn merge_combines_disjoint_fields_at_one_point() { ); } +/// Native histograms merge like floats: members of one label set combine by field and position, +/// histograms and floats keep their type, and a field held twice at one point collides whatever +/// the sample types. +#[test] +fn merge_keeps_histograms_and_rejects_any_doubled_field() { + #[derive(Clone, Copy)] + enum Value { + Float(f64), + Histogram(usize), + } + + let schema = grid_schema(2, 3); + let histograms = [ + OwnedHistogram::parse("0/0/0/3ff0000000000000/3ff0000000000000/0/1_1/3ff0000000000000///"), + OwnedHistogram::parse("0/0/0/4000000000000000/4008000000000000/1/0_1/4000000000000000///"), + ]; + let histogram = + |index: usize| FloatHistogram::try_from_view(&histograms[index].view()).unwrap(); + let sample = |value: Option| { + value.map(|value| match value { + Value::Float(float) => GridSample::Float(float), + Value::Histogram(index) => GridSample::Histogram(Box::new(histogram(index))), + }) + }; + let series = |rows: &[(usize, [Option; 2])]| { + let mut builder = GridSeriesBuilder::try_new(schema.clone()).unwrap(); + for &(point, row) in rows { + builder + .push_samples(point, &[sample(row[0]), sample(row[1])]) + .unwrap(); + } + builder.finish(metric(0)).unwrap().unwrap() + }; + let describe = |series: &Series, field: usize| -> Vec<(usize, String)> { + let samples = &series.fields.as_slice()[field]; + series + .time + .positions() + .filter_map(|(point, row)| match samples.sample(row) { + SampleRef::Missing => None, + SampleRef::Float(value) => Some((point, format!("{value}"))), + SampleRef::Histogram(view) => Some(( + point, + format!("{:?}", FloatHistogram::try_from_view(&view).unwrap()), + )), + }) + .collect() + }; + + let merged = merge_same_labelset( + &schema, + vec![ + series(&[ + (0, [Some(Value::Histogram(0)), None]), + (2, [Some(Value::Float(1.0)), None]), + ]), + series(&[ + (0, [None, Some(Value::Float(2.0))]), + (1, [Some(Value::Histogram(1)), None]), + ]), + ], + ) + .unwrap(); + assert_eq!(merged.len(), 1); + assert_eq!( + describe(&merged[0], 0), + [ + (0, format!("{:?}", histogram(0))), + (1, format!("{:?}", histogram(1))), + (2, "1".to_string()), + ] + ); + assert_eq!(describe(&merged[0], 1), [(0, "2".to_string())]); + + for second in [Value::Float(3.0), Value::Histogram(1)] { + let error = merge_same_labelset( + &schema, + vec![ + series(&[(1, [Some(Value::Histogram(0)), None])]), + series(&[(1, [Some(second), None])]), + ], + ) + .unwrap_err(); + assert!(error.to_string().contains("same labelset"), "{error}"); + } +} + /// Two values a million positions apart are one series with two rows, addressed by two runs: the /// positions between them cost neither a row nor a value buffer sized by the span. #[test] From b16cffd5c618f0a0aa2eb900109e4e6183154227 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 21 Oct 2026 12:08:42 +0800 Subject: [PATCH 02/13] feat(promql): add max_of, min_of and query context functions to native kernels Signed-off-by: Dennis Zhuang --- src/promql/src/map.rs | 5 +- src/promql/src/map/broadcast.rs | 52 +++++ src/promql/src/map/max_min_cases.txt | 65 ++++++ src/promql/src/map/tests.rs | 56 +++++ src/promql/src/map_time.rs | 31 +++ .../src/map_time/duration_seconds_cases.txt | 211 ++++++++++++++++++ src/promql/src/map_time/tests.rs | 32 ++- 7 files changed, 450 insertions(+), 2 deletions(-) create mode 100644 src/promql/src/map/max_min_cases.txt create mode 100644 src/promql/src/map_time/duration_seconds_cases.txt diff --git a/src/promql/src/map.rs b/src/promql/src/map.rs index ffc7383f37e..2d2c4b80a7b 100644 --- a/src/promql/src/map.rs +++ b/src/promql/src/map.rs @@ -23,7 +23,10 @@ mod go_math; #[cfg(test)] mod tests; -pub use broadcast::{ScalarSeriesOperand, map_arithmetic_series, map_compare_series}; +pub use broadcast::{ + ScalarFunction, ScalarSeriesOperand, map_arithmetic_series, map_compare_series, + map_scalar_function, +}; pub use comparison::{Comparison, ComparisonOp, map_compare, scalar_arithmetic, scalar_compare}; use std::f64::consts::PI; diff --git a/src/promql/src/map/broadcast.rs b/src/promql/src/map/broadcast.rs index b94605a7c28..b6b1a4c161b 100644 --- a/src/promql/src/map/broadcast.rs +++ b/src/promql/src/map/broadcast.rs @@ -74,6 +74,58 @@ pub fn map_compare_series( }) } +/// `max_of(a, b)` and `min_of(a, b)`, which Prometheus evaluates with Go's `math.Max` and +/// `math.Min`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ScalarFunction { + MaxOf, + MinOf, +} + +impl ScalarFunction { + /// Go's `math.Max` / `math.Min`: an infinity of the function's own sign wins over NaN, NaN + /// over any other value, and `+0` is greater than `-0`. + pub fn apply(self, a: f64, b: f64) -> f64 { + let (wins, signed_zero_wins) = match self { + Self::MaxOf => (f64::INFINITY, false), + Self::MinOf => (f64::NEG_INFINITY, true), + }; + if a == wins || b == wins { + wins + } else if a.is_nan() || b.is_nan() { + f64::NAN + } else if a == 0.0 && a == b { + if a.is_sign_negative() == signed_zero_wins { + a + } else { + b + } + } else if (a > b) == (self == Self::MaxOf) { + a + } else { + b + } + } +} + +/// Applies `function` to two time-varying scalars, aligned by absolute grid index; a point where +/// either has no value has no result. Both operands must be unlabelled single-field series. +pub fn map_scalar_function( + schema: &Arc, + function: ScalarFunction, + left: &Series, + right: &Series, +) -> Result> { + validate_scalar(schema, left)?; + map_grid_scalar( + schema, + left, + ScalarSeriesOperand::Right(right), + left.metric.clone(), + |left, right, _| Some(function.apply(left, right)), + ) +} + /// Calls `calculate(left, right, vector_sample)` at each index where both operands exist. fn map_grid_scalar( schema: &Arc, diff --git a/src/promql/src/map/max_min_cases.txt b/src/promql/src/map/max_min_cases.txt new file mode 100644 index 00000000000..ecd6c0dc738 --- /dev/null +++ b/src/promql/src/map/max_min_cases.txt @@ -0,0 +1,65 @@ +# Generated by promql-compat cmd/narthex-golden-scalarfn built for amd64: `a b math.Max(a,b) math.Min(a,b)` as hex bits. +0 0 0 0 +0 8000000000000000 0 8000000000000000 +0 3ff0000000000000 3ff0000000000000 0 +0 bff0000000000000 0 bff0000000000000 +0 4004000000000000 4004000000000000 0 +0 7ff0000000000000 7ff0000000000000 0 +0 fff0000000000000 0 fff0000000000000 +0 7ff8000000000001 7ff8000000000001 7ff8000000000001 +8000000000000000 0 0 8000000000000000 +8000000000000000 8000000000000000 8000000000000000 8000000000000000 +8000000000000000 3ff0000000000000 3ff0000000000000 8000000000000000 +8000000000000000 bff0000000000000 8000000000000000 bff0000000000000 +8000000000000000 4004000000000000 4004000000000000 8000000000000000 +8000000000000000 7ff0000000000000 7ff0000000000000 8000000000000000 +8000000000000000 fff0000000000000 8000000000000000 fff0000000000000 +8000000000000000 7ff8000000000001 7ff8000000000001 7ff8000000000001 +3ff0000000000000 0 3ff0000000000000 0 +3ff0000000000000 8000000000000000 3ff0000000000000 8000000000000000 +3ff0000000000000 3ff0000000000000 3ff0000000000000 3ff0000000000000 +3ff0000000000000 bff0000000000000 3ff0000000000000 bff0000000000000 +3ff0000000000000 4004000000000000 4004000000000000 3ff0000000000000 +3ff0000000000000 7ff0000000000000 7ff0000000000000 3ff0000000000000 +3ff0000000000000 fff0000000000000 3ff0000000000000 fff0000000000000 +3ff0000000000000 7ff8000000000001 7ff8000000000001 7ff8000000000001 +bff0000000000000 0 0 bff0000000000000 +bff0000000000000 8000000000000000 8000000000000000 bff0000000000000 +bff0000000000000 3ff0000000000000 3ff0000000000000 bff0000000000000 +bff0000000000000 bff0000000000000 bff0000000000000 bff0000000000000 +bff0000000000000 4004000000000000 4004000000000000 bff0000000000000 +bff0000000000000 7ff0000000000000 7ff0000000000000 bff0000000000000 +bff0000000000000 fff0000000000000 bff0000000000000 fff0000000000000 +bff0000000000000 7ff8000000000001 7ff8000000000001 7ff8000000000001 +4004000000000000 0 4004000000000000 0 +4004000000000000 8000000000000000 4004000000000000 8000000000000000 +4004000000000000 3ff0000000000000 4004000000000000 3ff0000000000000 +4004000000000000 bff0000000000000 4004000000000000 bff0000000000000 +4004000000000000 4004000000000000 4004000000000000 4004000000000000 +4004000000000000 7ff0000000000000 7ff0000000000000 4004000000000000 +4004000000000000 fff0000000000000 4004000000000000 fff0000000000000 +4004000000000000 7ff8000000000001 7ff8000000000001 7ff8000000000001 +7ff0000000000000 0 7ff0000000000000 0 +7ff0000000000000 8000000000000000 7ff0000000000000 8000000000000000 +7ff0000000000000 3ff0000000000000 7ff0000000000000 3ff0000000000000 +7ff0000000000000 bff0000000000000 7ff0000000000000 bff0000000000000 +7ff0000000000000 4004000000000000 7ff0000000000000 4004000000000000 +7ff0000000000000 7ff0000000000000 7ff0000000000000 7ff0000000000000 +7ff0000000000000 fff0000000000000 7ff0000000000000 fff0000000000000 +7ff0000000000000 7ff8000000000001 7ff0000000000000 7ff8000000000001 +fff0000000000000 0 0 fff0000000000000 +fff0000000000000 8000000000000000 8000000000000000 fff0000000000000 +fff0000000000000 3ff0000000000000 3ff0000000000000 fff0000000000000 +fff0000000000000 bff0000000000000 bff0000000000000 fff0000000000000 +fff0000000000000 4004000000000000 4004000000000000 fff0000000000000 +fff0000000000000 7ff0000000000000 7ff0000000000000 fff0000000000000 +fff0000000000000 fff0000000000000 fff0000000000000 fff0000000000000 +fff0000000000000 7ff8000000000001 7ff8000000000001 fff0000000000000 +7ff8000000000001 0 7ff8000000000001 7ff8000000000001 +7ff8000000000001 8000000000000000 7ff8000000000001 7ff8000000000001 +7ff8000000000001 3ff0000000000000 7ff8000000000001 7ff8000000000001 +7ff8000000000001 bff0000000000000 7ff8000000000001 7ff8000000000001 +7ff8000000000001 4004000000000000 7ff8000000000001 7ff8000000000001 +7ff8000000000001 7ff0000000000000 7ff0000000000000 7ff8000000000001 +7ff8000000000001 fff0000000000000 7ff8000000000001 fff0000000000000 +7ff8000000000001 7ff8000000000001 7ff8000000000001 7ff8000000000001 diff --git a/src/promql/src/map/tests.rs b/src/promql/src/map/tests.rs index e0eea771af1..1cc936287fd 100644 --- a/src/promql/src/map/tests.rs +++ b/src/promql/src/map/tests.rs @@ -551,3 +551,59 @@ fn math_drops_reserved_metadata_and_empty_labels_with_the_name() { assert_eq!(output.metric.name(), None); assert_eq!(names, ["__other__", "host"]); } + +/// `max_of` / `min_of` against Go's `math.Max` / `math.Min`, whose signed-zero, NaN and infinity +/// rules differ from `f64::max` / `f64::min`. +#[test] +fn max_of_and_min_of_follow_go() { + let bits = |text: &str| f64::from_bits(u64::from_str_radix(text, 16).unwrap()); + let same = |a: f64, b: f64| a.to_bits() == b.to_bits() || (a.is_nan() && b.is_nan()); + let mut checked = 0; + for line in include_str!("max_min_cases.txt") + .lines() + .filter(|line| !line.starts_with('#')) + { + let values: Vec = line.split_whitespace().map(bits).collect(); + let (a, b) = (values[0], values[1]); + assert!( + same(ScalarFunction::MaxOf.apply(a, b), values[2]), + "max_of({a}, {b})" + ); + assert!( + same(ScalarFunction::MinOf.apply(a, b), values[3]), + "min_of({a}, {b})" + ); + checked += 1; + } + assert_eq!(checked, 64); +} + +/// Time-varying scalars pair by absolute position; a point either one lacks has no result. +#[test] +fn scalar_functions_pair_scalars_by_position() { + let schema = Arc::new(schema(&["scalar"])); + let scalar = |start: usize, values: &[Option]| { + Series::try_new( + &schema, + MetricName::try_new(None, Vec::new()).unwrap(), + TimeAxis::Grid( + GridRange::try_new(schema.eval_grid.clone().unwrap(), start, values.len()).unwrap(), + ), + SeriesFields::try_from_vec(vec![floats(values)]).unwrap(), + ) + .unwrap() + }; + let left = scalar(2, &[Some(1.0), Some(5.0), None, Some(-0.0)]); + let right = scalar(3, &[Some(2.0), Some(7.0), Some(0.0), Some(9.0)]); + let result = map_scalar_function(&schema, ScalarFunction::MaxOf, &left, &right) + .unwrap() + .unwrap(); + assert_eq!( + run_floats(&result, 0), + vec![((3, 1), vec![Some(5.0)]), ((5, 1), vec![Some(0.0)])] + ); + assert!( + run_floats(&result, 0)[1].1[0].unwrap().is_sign_positive(), + "max_of(-0, +0) is +0" + ); +} diff --git a/src/promql/src/map_time.rs b/src/promql/src/map_time.rs index 5101cb7fbfa..afb9fd419aa 100644 --- a/src/promql/src/map_time.rs +++ b/src/promql/src/map_time.rs @@ -43,6 +43,37 @@ pub enum DateFunction { Year, } +/// `start()`, `end()`, `range()` and `step()`: constants of a query, which Prometheus folds into +/// number literals before evaluating it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueryContextFunction { + Start, + End, + Range, + Step, +} + +impl QueryContextFunction { + /// The value for a query from `start_ms` to `end_ms` every `step_ms`. `step()` is 0 when the + /// start and end are equal, as for an instant query. + pub fn value(self, start_ms: i64, end_ms: i64, step_ms: i64) -> f64 { + match self { + Self::Start => start_ms as f64 / 1000.0, + Self::End => end_ms as f64 / 1000.0, + Self::Range => duration_seconds(end_ms - start_ms), + Self::Step if start_ms == end_ms => 0.0, + Self::Step => duration_seconds(step_ms), + } + } +} + +/// Go's `time.Duration.Seconds` of `ms` milliseconds: whole seconds plus the remainder over 1e9, +/// two roundings rather than one division. The nanosecond count saturates as `time.Time.Sub` does. +fn duration_seconds(ms: i64) -> f64 { + let nanos = ms.saturating_mul(1_000_000); + (nanos / 1_000_000_000) as f64 + (nanos % 1_000_000_000) as f64 / 1e9 +} + /// `time()`: the evaluation time in seconds at every point of the schema's grid, unlabelled. pub fn time(schema: &SeriesSchema) -> Result { eval_time_series(schema, |ms| ms as f64 / 1000.0) diff --git a/src/promql/src/map_time/duration_seconds_cases.txt b/src/promql/src/map_time/duration_seconds_cases.txt new file mode 100644 index 00000000000..851c93c2196 --- /dev/null +++ b/src/promql/src/map_time/duration_seconds_cases.txt @@ -0,0 +1,211 @@ +# Generated by promql-compat cmd/narthex-golden-scalarfn built for amd64: `ms time.Duration.Seconds()` for a millisecond difference, as hex bits. +0 0 +1 3f50624dd2f1a9fc +999 3feff7ced916872b +1000 3ff0000000000000 +1001 3ff004189374bc6a +4097 40106353f7ced917 +86400000 40f5180000000000 +9007199254740993 42012e0be826d695 +9223372036854 42012e0be826d4fe +9223372036855 42012e0be826d695 +2340891196754 41e170e5479820c5 +28653324801615 42012e0be826d695 +15150070723101 42012e0be826d695 +5203274071043 41f3623a9570b021 +6052846691025 41f68c71c6306666 +19267443321048 42012e0be826d695 +12341472526366 42012e0be826d695 +20530526999188 42012e0be826d695 +34395057096608 42012e0be826d695 +8401369033625 41ff4c2b7c9a0000 +31640814836996 42012e0be826d695 +29341436692007 42012e0be826d695 +34783771064630 42012e0be826d695 +35122158268884 42012e0be826d695 +1479560899851 41d60c13b0f676c9 +23204167308761 42012e0be826d695 +24788341754728 42012e0be826d695 +31429837266203 42012e0be826d695 +34244401504646 42012e0be826d695 +13068139990387 42012e0be826d695 +7761385230625 41fce9d570ea0000 +9735534642756 42012e0be826d695 +20240948685375 42012e0be826d695 +33614386955765 42012e0be826d695 +4944260193162 41f26b36c612978d +14479808172613 42012e0be826d695 +17797368697448 42012e0be826d695 +14261095223522 42012e0be826d695 +13428055133044 42012e0be826d695 +26352584022491 42012e0be826d695 +12325408233432 42012e0be826d695 +17010777739835 42012e0be826d695 +28430422903295 42012e0be826d695 +29753883088875 42012e0be826d695 +25998949474507 42012e0be826d695 +10226187168843 42012e0be826d695 +17769199632283 42012e0be826d695 +4363496320571 41f0415ab80922d1 +3287173645812 41e87dc941b9fbe7 +12725827928836 42012e0be826d695 +24795910633675 42012e0be826d695 +2828521278844 41e512f9e7db020c +4170359431716 41ef125450f6e979 +3589276087583 41eabe0036f2a7f0 +28524985356220 42012e0be826d695 +30320656727859 42012e0be826d695 +4256148642671 41efb5f5945578d5 +29566978976822 42012e0be826d695 +1648135601414 41d88f236c5a7efa +4780677690357 41f1cf35a3a5b646 +3390646811377 41e94325436c1062 +6912080800569 41f9bfdfba091aa0 +19744540834685 42012e0be826d695 +5048451373608 41f2ce9412d9ba5e +31511077099905 42012e0be826d695 +33060246075211 42012e0be826d695 +29446920826861 42012e0be826d695 +20868689814545 42012e0be826d695 +25732868269994 42012e0be826d695 +26023816056118 42012e0be826d695 +3312873794438 41e8aece284e0419 +8928671130029 4200a185acd03b64 +27201078145192 42012e0be826d695 +28161007422010 42012e0be826d695 +23548861302960 42012e0be826d695 +8403898590605 41ff4e950de9ae14 +2101319756957 41df4fe6133d3f7d +18153789976583 42012e0be826d695 +16078846218429 42012e0be826d695 +14439764018086 42012e0be826d695 +15717223852073 42012e0be826d695 +7594418227955 41fc4a9a033f47ae +1007547111184 41ce06f973978d50 +7653748532646 41fc832ef34a5604 +1542530612607 41d6fc498d26d917 +8959567520889 4200b04135071cac +1963454100123 41dd41fba507df3b +15423611227189 42012e0be826d695 +26921592371531 42012e0be826d695 +31560677483706 42012e0be826d695 +21024346859206 42012e0be826d695 +21411506917513 42012e0be826d695 +14299718455620 42012e0be826d695 +25966354044973 42012e0be826d695 +21480240207871 42012e0be826d695 +13383017194856 42012e0be826d695 +19694757267735 42012e0be826d695 +8329763775671 41ff07e1bbfabc6a +24204352256265 42012e0be826d695 +5166480575653 41f33f23cbfa72b0 +17311660727623 42012e0be826d695 +26768379971433 42012e0be826d695 +21452255664239 42012e0be826d695 +32930407231978 42012e0be826d695 +15551692227912 42012e0be826d695 +23123035545994 42012e0be826d695 +5875644708440 41f5e37392470a3d +18292969058316 42012e0be826d695 +1707825217332 41d972d610553f7d +23633727413655 42012e0be826d695 +24472694155897 42012e0be826d695 +32056439689975 42012e0be826d695 +25133237617897 42012e0be826d695 +10984619967605 42012e0be826d695 +12399994332458 42012e0be826d695 +17017307690557 42012e0be826d695 +24405803592380 42012e0be826d695 +34039839184523 42012e0be826d695 +12086900862154 42012e0be826d695 +4880216990563 41f22e2339e9020c +8291016285084 41fee2ede5d15810 +34263874447617 42012e0be826d695 +5219901505275 41f3721604146666 +3325751404254 41e8c75e0d8820c5 +10069265674283 42012e0be826d695 +2205986208125 41e06f95b4040000 +17331608106692 42012e0be826d695 +13164910887275 42012e0be826d695 +20165189440141 42012e0be826d695 +18100958213119 42012e0be826d695 +33252881487365 42012e0be826d695 +28493023658516 42012e0be826d695 +8853845580432 42007dd7b26374bc +34838873892001 42012e0be826d695 +20214329624262 42012e0be826d695 +31705600540751 42012e0be826d695 +31189001046116 42012e0be826d695 +29476027927210 42012e0be826d695 +26781011366842 42012e0be826d695 +22699855557780 42012e0be826d695 +17749079048291 42012e0be826d695 +30750407110207 42012e0be826d695 +24417907650311 42012e0be826d695 +10509573177644 42012e0be826d695 +1079378703273 41d01580c3d178d5 +34862833947828 42012e0be826d695 +24980588397686 42012e0be826d695 +11209991956240 42012e0be826d695 +8937655194575 4200a5ce5cd4999a +23978885618615 42012e0be826d695 +30065528340854 42012e0be826d695 +2628030736846 41e39492221b126f +2663938618407 41e3d90f474d0625 +34895264552379 42012e0be826d695 +15670995231869 42012e0be826d695 +4191047303785 41ef39c9d0f91eb8 +8437263482680 41ff6e66c7aae148 +3960463773594 41ed81fc33b3020c +4309293785322 41f00da9ad9526e9 +26649392599963 42012e0be826d695 +4157860955829 41eefa7d8b7a872b +22296599249089 42012e0be826d695 +32695131508279 42012e0be826d695 +30635397919632 42012e0be826d695 +15331538117365 42012e0be826d695 +11063348740606 42012e0be826d695 +32371941731295 42012e0be826d695 +13719859399284 42012e0be826d695 +7293417131828 41fb2b8b6abd3f7d +33407534042909 42012e0be826d695 +31297857606586 42012e0be826d695 +1097102363540 41d0591d06e28f5c +25332039849587 42012e0be826d695 +20041642732253 42012e0be826d695 +18752081044817 42012e0be826d695 +29442119239956 42012e0be826d695 +28855454932055 42012e0be826d695 +31958604071351 42012e0be826d695 +18688407702704 42012e0be826d695 +17871169582207 42012e0be826d695 +16003600945447 42012e0be826d695 +5284509633417 41f3afb37c16ac08 +11575618223404 42012e0be826d695 +23263846003261 42012e0be826d695 +14260376149051 42012e0be826d695 +20109100243330 42012e0be826d695 +16359587412236 42012e0be826d695 +34836697130263 42012e0be826d695 +23618221322728 42012e0be826d695 +16214535784284 42012e0be826d695 +8272436198602 41fed135be69a1cb +5401941921460 41f41fb17a175c29 +7235498080071 41faf44f060122d1 +14773466706669 42012e0be826d695 +34843449378732 42012e0be826d695 +23707829047074 42012e0be826d695 +4081628494438 41ee6916a9ce0419 +13573087199008 42012e0be826d695 +21542290482225 42012e0be826d695 +16883004474069 42012e0be826d695 +16172951872856 42012e0be826d695 +5849767460850 41f5cac5e24d999a +12771076879764 42012e0be826d695 +7080148168638 41fa6027cc8a353f +13469889734230 42012e0be826d695 +9296342191848 42012e0be826d695 +753418709293 41c67420eaa58106 +23411833150514 42012e0be826d695 +29448477422232 42012e0be826d695 +29413433948272 42012e0be826d695 diff --git a/src/promql/src/map_time/tests.rs b/src/promql/src/map_time/tests.rs index 05128523a65..3dba9472832 100644 --- a/src/promql/src/map_time/tests.rs +++ b/src/promql/src/map_time/tests.rs @@ -24,7 +24,7 @@ use datatypes::series::samples::{ }; use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis}; -use super::{DateFunction, date_of_eval_time, map_date, time}; +use super::{DateFunction, QueryContextFunction, date_of_eval_time, map_date, time}; use crate::grid_series::merge_same_labelset; use crate::grid_series::reference::{ Labels, Points, canonical_bits, golden_cases, labels_of, metric_of, points_by_labels, @@ -289,3 +289,33 @@ fn map_date_reads_a_sparse_axis_by_its_positions() { ]) ); } + +/// `range()` and `step()` are Go `Duration.Seconds` of a millisecond count, which rounds twice +/// and so differs from `ms / 1000` in the last bit for most values. +#[test] +fn query_context_durations_match_go() { + let mut checked = 0; + for line in include_str!("duration_seconds_cases.txt") + .lines() + .filter(|line| !line.starts_with('#')) + { + let (ms, bits) = line.split_once(' ').unwrap(); + let ms: i64 = ms.parse().unwrap(); + let expected = u64::from_str_radix(bits, 16).unwrap(); + assert_eq!( + QueryContextFunction::Range.value(0, ms, 1).to_bits(), + expected, + "range() over {ms} ms" + ); + assert_eq!( + QueryContextFunction::Step.value(0, 1, ms).to_bits(), + expected, + "step() of {ms} ms" + ); + checked += 1; + } + assert!(checked > 200, "{checked} durations checked"); + assert_eq!(QueryContextFunction::Step.value(5_000, 5_000, 1_000), 0.0); + assert_eq!(QueryContextFunction::Start.value(-1_500, 0, 1), -1.5); + assert_eq!(QueryContextFunction::End.value(0, 2_250, 1), 2.25); +} From 9158f6ccba3ce00b61f25c200cc7b2c3a2e826ed Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Wed, 21 Oct 2026 15:43:20 +0800 Subject: [PATCH 03/13] fix(promql): saturate range() when the endpoint difference leaves i64 Signed-off-by: Dennis Zhuang --- src/promql/src/map_time.rs | 2 +- src/promql/src/map_time/tests.rs | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/promql/src/map_time.rs b/src/promql/src/map_time.rs index afb9fd419aa..c0a2c5b94c3 100644 --- a/src/promql/src/map_time.rs +++ b/src/promql/src/map_time.rs @@ -60,7 +60,7 @@ impl QueryContextFunction { match self { Self::Start => start_ms as f64 / 1000.0, Self::End => end_ms as f64 / 1000.0, - Self::Range => duration_seconds(end_ms - start_ms), + Self::Range => duration_seconds(end_ms.saturating_sub(start_ms)), Self::Step if start_ms == end_ms => 0.0, Self::Step => duration_seconds(step_ms), } diff --git a/src/promql/src/map_time/tests.rs b/src/promql/src/map_time/tests.rs index 3dba9472832..0124b420129 100644 --- a/src/promql/src/map_time/tests.rs +++ b/src/promql/src/map_time/tests.rs @@ -315,6 +315,11 @@ fn query_context_durations_match_go() { checked += 1; } assert!(checked > 200, "{checked} durations checked"); + // Two valid endpoints whose difference leaves `i64` saturate, as `time.Time.Sub` does. + assert_eq!( + QueryContextFunction::Range.value(i64::MIN / 2, i64::MAX / 2 + 1, 1), + QueryContextFunction::Range.value(0, i64::MAX, 1) + ); assert_eq!(QueryContextFunction::Step.value(5_000, 5_000, 1_000), 0.0); assert_eq!(QueryContextFunction::Start.value(-1_500, 0, 1), -1.5); assert_eq!(QueryContextFunction::End.value(0, 2_250, 1), 2.25); From 8543cbf401db9dc4fa7e1ca9aa347f254c6cb83d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 21 Oct 2026 16:38:37 +0800 Subject: [PATCH 04/13] feat(mito): scaffold series format writers and artifact lifecycle (#9788) * feat(mito): add pluggable series artifact reconciliation Signed-off-by: evenyag * feat(mito): prepare independent series reads from SST artifacts Signed-off-by: evenyag * feat(mito): prepare mixed sidecar SST and memtable series reads Signed-off-by: evenyag * refactor(mito): make prepared series sources replaceable Signed-off-by: evenyag * fix(mito): narrow series formats to artifact lifecycle integration Signed-off-by: evenyag * test(mito): populate artifact regression buckets above build threshold Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/AGENTS.md | 3 +- src/mito2/src/engine/scan_test.rs | 1 + src/mito2/src/read/series_candidate.rs | 47 ++ src/mito2/src/series_index.rs | 4 + src/mito2/src/series_index/artifact.rs | 93 ++++ src/mito2/src/series_index/bucket.rs | 1 + src/mito2/src/series_index/builder.rs | 114 ++++- src/mito2/src/series_index/catalog.rs | 76 ++- src/mito2/src/series_index/format.rs | 72 +++ src/mito2/src/series_index/format_tests.rs | 442 ++++++++++++++++++ src/mito2/src/series_index/maintenance.rs | 151 +++++- src/mito2/src/series_index/purger.rs | 49 +- src/mito2/src/series_index/searcher.rs | 1 + src/mito2/src/series_index/tests.rs | 5 +- src/mito2/src/series_index/version.rs | 27 +- .../fixtures/series_index/README.md | 10 + .../fixtures/series_index/legacy_catalog.json | 27 ++ 17 files changed, 1082 insertions(+), 41 deletions(-) create mode 100644 src/mito2/src/series_index/artifact.rs create mode 100644 src/mito2/src/series_index/format.rs create mode 100644 src/mito2/src/series_index/format_tests.rs create mode 100644 tests/compatibility/fixtures/series_index/README.md create mode 100644 tests/compatibility/fixtures/series_index/legacy_catalog.json diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index f70aa622a27..95b9eddae07 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned reconciliation (`maintenance.rs`), shared approximate usage (`task.rs`), catalogs, and index lifecycle | +| `series_index` | `src/mito2/src/series_index/` | Format interfaces (`format.rs`), index writer/searcher, per-SST samples artifacts (`artifact.rs`), bucket planning, catalogs, reconciliation (`maintenance.rs`), and deferred purge | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | @@ -66,6 +66,7 @@ filtered `RecordBatch` stream. and defer new builds when the shared estimate is full. Cleanup still prunes obsolete range entries and expires series indexes to reclaim published usage. Concurrent builds may overshoot; closed-region files are not counted. Series handles and the SST purger own file deletion. +- **Series format PoC**: samples are rebuildable local SST sidecars. Optional writer interfaces (`format.rs`) integrate with reconciliation, catalog publication, and deferred purge. Defaults retain the original Parquet writer and disable samples. Prepared sources, read APIs, and concrete new formats are deferred; ordinary scans only use Parquet index artifacts. - **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. - **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`. diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index c0b155f7311..68469e85324 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1549,6 +1549,7 @@ async fn check_two_phase_series_scan( let store = region.series_index_store.clone().unwrap(); let sequence = region.flushed_sequence(); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 3ba5512a7bd..82210e0ffa7 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -243,6 +243,10 @@ impl SeriesIndexCoverage { .version .series_indexes .values() + // Unimplemented formats cannot replace authoritative SST candidates. + .filter(|index| { + index.entry().format() == crate::series_index::artifact::ArtifactFormat::Parquet + }) .map(|index| { let files: HashSet<_> = uncovered .iter() @@ -778,6 +782,7 @@ mod tests { )); let store = env.access_layer.object_store().clone(); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), @@ -964,6 +969,48 @@ mod tests { } } + #[tokio::test] + async fn index_merge_releases_permits_before_yielding() { + let (_env, scanner, _pruner) = indexed_scanner().await; + let semaphore = Arc::new(Semaphore::new(1)); + let context = scanner.stream_ctx.input.series_index.clone().unwrap(); + let streams = (0..3) + .map(|_| { + index_primary_key_stream( + scanner.stream_ctx.clone(), + context.clone(), + scanner.coverage.indexes[0].clone(), + semaphore.clone(), + ) + }) + .collect(); + let merged = merge_primary_key_streams( + streams, + scanner.memory_pool.clone(), + &scanner.metrics_set, + 0, + "permit-regression", + ) + .unwrap(); + let series = + decode_metric_series(merged, scanner.stream_ctx.input.region_metadata().clone()) + .unwrap(); + let batches = tokio::time::timeout( + std::time::Duration::from_secs(10), + series.try_collect::>(), + ) + .await + .expect("index streams must not hold permits across yield") + .unwrap(); + assert_eq!( + (0..1001) + .map(|tsid| MetricSeriesId { table_id: 1, tsid }) + .collect::>(), + batches.into_iter().flatten().collect::>() + ); + assert_eq!(1, semaphore.available_permits()); + } + #[tokio::test] async fn index_read_failure_after_output_is_propagated() { let (_env, scanner, _pruner) = indexed_scanner().await; diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 3aa70256709..585120628e8 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -16,9 +16,13 @@ //! //! Under development. Index files are currently stored on the local filesystem. +pub(crate) mod artifact; mod bucket; mod builder; mod catalog; +pub(crate) mod format; +#[cfg(test)] +mod format_tests; mod maintenance; mod purger; mod searcher; diff --git a/src/mito2/src/series_index/artifact.rs b/src/mito2/src/series_index/artifact.rs new file mode 100644 index 00000000000..e1412920f6f --- /dev/null +++ b/src/mito2/src/series_index/artifact.rs @@ -0,0 +1,93 @@ +// 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. + +//! Immutable, potentially multi-object format outputs and reader-owned lifetimes. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use serde::{Deserialize, Serialize}; +use store_api::storage::FileId; + +use crate::series_index::purger::IndexFilePurger; +use crate::sst::file::RegionFileId; + +/// Persisted artifact identifiers. Concrete codecs are added with their implementations. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) enum ArtifactFormat { + #[serde(rename = "parquet")] + Parquet, + #[cfg(test)] + #[serde(rename = "fake-samples")] + FakeSamples, + #[cfg(test)] + #[serde(rename = "fake-index")] + FakeIndex, +} + +/// Identifies the exact bytes produced by one backend build. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ArtifactMetadata { + pub(crate) format: ArtifactFormat, + pub(crate) generation: FileId, + /// Every object owned by this build, including metadata and auxiliary files. + pub(crate) objects: Vec, + pub(crate) file_size: u64, +} + +/// One samples artifact derived from an immutable Parquet SST. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SamplesEntry { + pub(crate) file_id: FileId, + pub(crate) artifact: ArtifactMetadata, +} + +#[derive(Debug, Clone)] +pub(crate) struct SamplesFileHandle(Arc); + +#[derive(Debug)] +struct SamplesFileInner { + file_id: RegionFileId, + entry: SamplesEntry, + retired: AtomicBool, + purger: IndexFilePurger, +} + +impl SamplesFileHandle { + pub(crate) fn new(file_id: RegionFileId, entry: SamplesEntry, purger: IndexFilePurger) -> Self { + Self(Arc::new(SamplesFileInner { + file_id, + entry, + retired: AtomicBool::new(false), + purger, + })) + } + + pub(crate) fn entry(&self) -> &SamplesEntry { + &self.0.entry + } + + pub(crate) fn mark_deleted(&self) { + self.0.retired.store(true, Ordering::Release); + } +} + +impl Drop for SamplesFileInner { + fn drop(&mut self) { + if self.retired.load(Ordering::Acquire) { + self.purger + .purge_objects(self.file_id, self.entry.artifact.objects.clone()); + } + } +} diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs index 446948838f6..470caeadd3b 100644 --- a/src/mito2/src/series_index/bucket.rs +++ b/src/mito2/src/series_index/bucket.rs @@ -334,6 +334,7 @@ impl SeriesBucket { .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) .min()?; Some(SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: self.start, diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index 63a51badaeb..2a9db9503e3 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -35,6 +35,7 @@ use crate::series_index::bucket::SeriesBucket; use crate::series_index::catalog::{ RangeIndexEntry, SeriesIndexEntry, range_index_path, series_index_path, series_metadata, }; +use crate::series_index::format::SeriesIndexFormat; use crate::series_index::purger::{IndexFilePurger, IndexFileType, file_operation}; use crate::series_index::version::SeriesIndexFileHandle; use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; @@ -130,6 +131,7 @@ pub(crate) async fn build_series_index( version: &VersionRef, bucket: &SeriesBucket, entry: &SeriesIndexEntry, + backend: Option<&dyn SeriesIndexFormat>, purger: &IndexFilePurger, ) -> Result { let mut sources = Vec::::new(); @@ -180,6 +182,35 @@ pub(crate) async fn build_series_index( ) }; // TODO(yingwen): Deduplicate update-mode rows before series indexes are used by queries. + if let Some(backend) = backend { + let mut writer = backend + .writer(version.metadata.clone(), store.clone(), entry) + .await?; + let result: Result<()> = async { + while let Some(batch) = visible.try_next().await? { + writer.write(&batch).await?; + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup_error) = writer.abort().await { + warn!(cleanup_error; "Failed to abort series-index artifact build"); + } + return Err(error); + } + let artifact = writer.finish().await?; + file_operation(IndexFileType::Series, "build", "success"); + return Ok(SeriesIndexFileHandle::new( + region.region_id, + SeriesIndexEntry { + file_size: artifact.file_size, + artifact: Some(artifact), + ..entry.clone() + }, + purger.clone(), + )); + } let path = series_index_path(region.region_id, entry.index_uuid); let mut writer = SeriesIndexWriter::try_new( version.metadata.clone(), @@ -214,6 +245,79 @@ pub(crate) async fn build_series_index( )) } +/// Converts one immutable SST into an independently rebuildable samples artifact. +pub(crate) async fn build_samples( + store: &ObjectStore, + region: &MitoRegionRef, + version: &VersionRef, + file: FileHandle, + backend: &dyn crate::series_index::format::SamplesFormat, + purger: &IndexFilePurger, +) -> Result> { + use crate::series_index::artifact::{SamplesEntry, SamplesFileHandle}; + let file_id = file.file_id(); + let Some((context, mut selection)) = region + .access_layer + .read_sst(file) + .build_reader_input(&mut ReaderMetrics::default()) + .await? + else { + return Ok(None); + }; + let context = Arc::new(context); + let mapper = FlatProjectionMapper::new( + &version.metadata, + 0..version.metadata.column_metadatas.len(), + )?; + let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; + let generation = store_api::storage::FileId::random(); + let prefix = format!( + "{}/samples/{}/{generation}/", + region.region_id.as_u64(), + file_id.file_id() + ); + let mut writer = backend + .writer(version.metadata.clone(), store.clone(), &prefix, generation) + .await?; + let result: Result<()> = async { + while let Some((row_group, selection)) = selection.pop_first() { + let reader = context + .reader_builder() + .build(context.build_context(row_group, Some(selection), None)) + .await?; + let mut reader = FlatPruneReader::new_with_row_group_reader( + context.clone(), + FlatRowGroupReader::new(context.clone(), reader), + context.pre_filter_mode().skip_fields(), + ); + while let Some(batch) = reader.next_batch().await? { + let batch = match &compat { + Some(compat) => compat.compat(batch)?, + None => batch, + }; + writer.write(&batch).await?; + } + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup) = writer.abort().await { + warn!(cleanup; "Failed to abort samples build"); + } + return Err(error); + } + let artifact = writer.finish().await?; + Ok(Some(SamplesFileHandle::new( + file_id, + SamplesEntry { + file_id: file_id.file_id(), + artifact, + }, + purger.clone(), + ))) +} + #[cfg(test)] mod tests { use std::collections::{BTreeMap, HashMap}; @@ -228,7 +332,8 @@ mod tests { use super::*; use crate::series_index::bucket::{group_files_into_series_buckets, plan_series_indexes}; use crate::series_index::catalog::{ - SeriesIndexCatalog, load_version_control, series_catalog_path, store_catalog, + SeriesIndexCatalog, load_version_control, series_catalog_path, series_index_path, + store_catalog, }; use crate::series_index::purger::series_index_channel; use crate::series_index::tests::prepare_region; @@ -321,6 +426,7 @@ mod tests { &version, &bucket, &entry, + None, &purger, ) .await; @@ -471,7 +577,7 @@ mod tests { .unwrap() .unwrap(); } - let _handle = build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + let _handle = build_series_index(&store, ®ion, &version, &bucket, &entry, None, &purger) .await .unwrap(); let bytes = store @@ -540,7 +646,7 @@ mod tests { }; completed.push(entry.file_id); } - build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + build_series_index(&store, ®ion, &version, &bucket, &entry, None, &purger) .await .map(Some) } @@ -746,7 +852,7 @@ mod tests { .map(|_| ()) } else { let (purger, _receiver) = series_index_channel(store.clone()); - build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + build_series_index(&store, ®ion, &version, &bucket, &entry, None, &purger) .await .map(|_| ()) }; diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index 95cf87f83d3..66355ace359 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -33,6 +33,7 @@ use crate::series_index::version::{ pub(crate) use crate::sst::range_index::range_index_path; const SERIES_DIR: &str = "series"; +const SAMPLES_CATALOG: &str = "samples.json"; const RANGE_CATALOG: &str = "range-index.json"; const SERIES_CATALOG: &str = "series-index.json"; const SERIES_METADATA_KEY: &str = "greptime.series_index"; @@ -60,6 +61,8 @@ pub(crate) struct WindowSequence { /// relies on this complete-coverage contract, not on `source_file_ids`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) artifact: Option, /// Completed size in the catalog; zero in the footer written before completion. pub(crate) file_size: u64, pub(crate) index_uuid: FileId, @@ -83,6 +86,14 @@ pub(crate) struct SeriesIndexEntry { } impl SeriesIndexEntry { + /// Entries written before pluggable formats describe the original Parquet layout. + pub(crate) fn format(&self) -> crate::series_index::artifact::ArtifactFormat { + self.artifact.as_ref().map_or( + crate::series_index::artifact::ArtifactFormat::Parquet, + |artifact| artifact.format, + ) + } + /// Whether this index completely covers an SST in the query's sequence domain. pub(crate) fn covers_file( &self, @@ -116,6 +127,18 @@ pub(crate) struct RangeIndexCatalog { pub(crate) indexes: Vec, } +#[derive(Debug, Default, Serialize, Deserialize)] +pub(crate) struct SamplesCatalog { + pub(crate) entries: Vec, + /// Once samples are enabled this single catalog commits all artifact references together. + pub(crate) series: Vec, + pub(crate) ranges: Vec, +} + +pub(crate) fn samples_catalog_path(region_id: RegionId) -> String { + format!("{}/{SAMPLES_CATALOG}", region_id.as_u64()) +} + pub(crate) fn range_catalog_path(region_id: RegionId) -> String { format!("{}/{RANGE_CATALOG}", region_id.as_u64()) } @@ -172,6 +195,7 @@ where pub(crate) async fn delete_catalogs(store: &ObjectStore, region_id: RegionId) { for path in [ series_catalog_path(region_id), + samples_catalog_path(region_id), range_catalog_path(region_id), ] { if let Err(error) = store.delete(&path).await @@ -194,14 +218,17 @@ pub(crate) async fn load_version_control( let series = load_catalog::(store, &series_catalog_path(region_id)) .await .unwrap_or_default(); - let version = SeriesIndexVersion::new( - range - .indexes + let samples = load_catalog::(store, &samples_catalog_path(region_id)).await; + let (ranges, indexes, sample_entries) = match samples { + Some(catalog) => (catalog.ranges, catalog.series, catalog.entries), + None => (range.indexes, series.indexes, Vec::new()), + }; + let mut version = SeriesIndexVersion::new( + ranges .into_iter() .map(|entry| (entry.file_id, entry)) .collect(), - series - .indexes + indexes .into_iter() .map(|entry| { ( @@ -211,6 +238,17 @@ pub(crate) async fn load_version_control( }) .collect(), ); + for entry in sample_entries { + let id = entry.file_id; + version.samples.insert( + id, + crate::series_index::artifact::SamplesFileHandle::new( + crate::sst::file::RegionFileId::new(region_id, id), + entry, + purger.clone(), + ), + ); + } let control = SeriesIndexVersionControl::default(); control.publish(std::sync::Arc::new(version)); control @@ -258,10 +296,37 @@ mod tests { } } + #[test] + fn legacy_series_catalog_fixture() { + use crate::series_index::artifact::ArtifactFormat; + let json = include_str!( + "../../../../tests/compatibility/fixtures/series_index/legacy_catalog.json" + ); + let catalog: SeriesIndexCatalog = serde_json::from_str(json).unwrap(); + assert_eq!(1, catalog.indexes.len()); + let entry = &catalog.indexes[0]; + assert!(entry.artifact.is_none()); + assert_eq!(ArtifactFormat::Parquet, entry.format()); + assert_eq!( + serde_json::from_str::(json).unwrap(), + serde_json::to_value(&catalog).unwrap() + ); + for (format, name) in [ + (ArtifactFormat::Parquet, "parquet"), + (ArtifactFormat::FakeSamples, "fake-samples"), + (ArtifactFormat::FakeIndex, "fake-index"), + ] { + let encoded = serde_json::to_value(format).unwrap(); + assert_eq!(name, encoded.as_str().unwrap()); + assert_eq!(format, serde_json::from_value(encoded).unwrap()); + } + } + #[test] fn coverage_uses_exclusive_time_end_and_inclusive_file_sequences() { let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(1), @@ -359,6 +424,7 @@ mod tests { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), diff --git a/src/mito2/src/series_index/format.rs b/src/mito2/src/series_index/format.rs new file mode 100644 index 00000000000..b59933f719a --- /dev/null +++ b/src/mito2/src/series_index/format.rs @@ -0,0 +1,72 @@ +// 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. + +//! Storage-format boundaries. Encodings stay outside reconciliation and source scheduling. + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use datatypes::arrow::record_batch::RecordBatch; +use object_store::ObjectStore; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::FileId; + +use crate::error::Result; +use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata}; +use crate::series_index::catalog::SeriesIndexEntry; + +/// Writers own partial outputs until finish succeeds; abort removes every partial object. +/// A failed finish must clean up its output before returning; implementations also own +/// cleanup on cancellation while a writer is live. +#[async_trait] +pub(crate) trait FormatWriter: Send { + async fn write(&mut self, batch: &RecordBatch) -> Result<()>; + async fn finish(self: Box) -> Result; + async fn abort(self: Box) -> Result<()>; +} + +/// Builds series-index artifacts; read interfaces are deferred. +#[async_trait] +pub(crate) trait SeriesIndexFormat: Debug + Send + Sync { + fn format(&self) -> ArtifactFormat; + async fn writer( + &self, + metadata: RegionMetadataRef, + store: ObjectStore, + entry: &SeriesIndexEntry, + ) -> Result>; +} + +/// Builds samples artifacts; read interfaces are deferred. +#[async_trait] +pub(crate) trait SamplesFormat: Debug + Send + Sync { + fn format(&self) -> ArtifactFormat; + /// The caller supplies a unique prefix. Finish reports all objects owned below it. + async fn writer( + &self, + metadata: RegionMetadataRef, + store: ObjectStore, + prefix: &str, + generation: FileId, + ) -> Result>; +} + +/// Optional writers, captured once per reconciliation pass. +/// Without an index backend, builds use the existing Parquet writer. +#[derive(Debug, Clone, Default)] +pub(crate) struct SeriesFormats { + pub(crate) index: Option>, + pub(crate) samples: Option>, +} diff --git a/src/mito2/src/series_index/format_tests.rs b/src/mito2/src/series_index/format_tests.rs new file mode 100644 index 00000000000..a7e0c86069d --- /dev/null +++ b/src/mito2/src/series_index/format_tests.rs @@ -0,0 +1,442 @@ +// 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. + +//! Artifact lifecycle tests; fake formats have no codec. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use datatypes::arrow::record_batch::RecordBatch; +use futures::TryStreamExt; +use object_store::ObjectStore; +use object_store::layers::mock::{self, MockLayerBuilder, oio}; +use snafu::ResultExt; +use store_api::metadata::RegionMetadataRef; +use store_api::region_engine::RegionEngine; +use store_api::storage::FileId; + +use crate::config::MitoConfig; +use crate::error::{OpenDalSnafu, Result, UnexpectedSnafu}; +use crate::series_index::artifact::{ArtifactFormat, ArtifactMetadata}; +use crate::series_index::catalog::{SeriesIndexEntry, load_version_control}; +use crate::series_index::format::{FormatWriter, SamplesFormat, SeriesFormats, SeriesIndexFormat}; +use crate::series_index::maintenance::reconcile_series_indexes; +use crate::series_index::purger::series_index_channel; +use crate::series_index::tests::prepare_region_with_config; +use crate::test_util::TestEnv; + +#[derive(Debug, Default)] +struct FakeSamples { + fail: Arc, +} + +struct FakeWriter { + store: ObjectStore, + artifact: ArtifactMetadata, + fail: Arc, +} + +#[async_trait] +impl FormatWriter for FakeWriter { + async fn write(&mut self, _batch: &RecordBatch) -> Result<()> { + if self.fail.load(Ordering::Relaxed) { + return UnexpectedSnafu { + reason: "injected samples write failure", + } + .fail(); + } + Ok(()) + } + async fn finish(self: Box) -> Result { + Ok(self.artifact) + } + async fn abort(self: Box) -> Result<()> { + for path in &self.artifact.objects { + self.store.delete(path).await.context(OpenDalSnafu)?; + } + Ok(()) + } +} + +#[async_trait] +impl SamplesFormat for FakeSamples { + fn format(&self) -> ArtifactFormat { + ArtifactFormat::FakeSamples + } + async fn writer( + &self, + _metadata: RegionMetadataRef, + store: ObjectStore, + prefix: &str, + generation: FileId, + ) -> Result> { + let objects = vec![format!("{prefix}data"), format!("{prefix}directory")]; + for path in &objects { + store.write(path, "fake").await.context(OpenDalSnafu)?; + } + Ok(Box::new(FakeWriter { + store, + artifact: ArtifactMetadata { + format: self.format(), + generation, + objects, + file_size: 8, + }, + fail: self.fail.clone(), + })) + } +} + +#[tokio::test] +async fn format_skeleton_build_reopen_and_retire() { + let mut env = TestEnv::with_prefix("format-skeleton").await; + let (engine, region) = prepare_region_with_config( + &mut env, + &[1000, 2000], + MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + ) + .await; + let fake = Arc::new(FakeSamples::default()); + *region.series_index_version_control.formats.write().unwrap() = Arc::new(SeriesFormats { + samples: Some(fake.clone()), + ..Default::default() + }); + let store = region.series_index_store.clone().unwrap(); + let (purger, mut purges) = series_index_channel(store.clone()); + let stats = reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + .await + .unwrap(); + assert_eq!(2, stats.built_samples); + let snapshot = region.series_index_version(); + assert_eq!(2, snapshot.samples.len()); + let reopened = load_version_control(&store, region.region_id, &purger).await; + assert_eq!(2, reopened.current().samples.len()); + assert!(snapshot.disk_usage() >= 16); + let stats = reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + .await + .unwrap(); + assert_eq!(0, stats.built_samples); + + // Rebuilding a missing sidecar replaces its generation but retains old readers' files. + let old = snapshot.samples.values().next().unwrap().clone(); + store + .delete(&old.entry().artifact.objects[0]) + .await + .unwrap(); + let stats = reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + .await + .unwrap(); + assert_eq!(1, stats.built_samples); + assert!(purges.try_recv().is_err()); + drop(snapshot); + assert!(purges.try_recv().is_err()); + drop(old); + let request = purges.try_recv().unwrap(); + assert_eq!(2, request.objects.unwrap().len()); + engine.stop().await.unwrap(); +} + +#[tokio::test] +async fn format_skeleton_abort_removes_all_objects() { + let mut env = TestEnv::with_prefix("format-abort").await; + let (engine, region) = crate::series_index::tests::prepare_region(&mut env).await; + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let (purger, _receiver) = series_index_channel(store.clone()); + let version = region.version(); + let file = version + .ssts + .levels() + .iter() + .flat_map(|l| l.files()) + .next() + .unwrap() + .clone(); + let fake = FakeSamples::default(); + fake.fail.store(true, Ordering::Relaxed); + assert!( + crate::series_index::builder::build_samples( + &store, ®ion, &version, file, &fake, &purger + ) + .await + .is_err() + ); + assert!( + store + .list_with(&format!("{}/samples/", region.region_id.as_u64())) + .recursive(true) + .await + .unwrap() + .iter() + .all(|entry| entry.metadata().is_dir()) + ); + assert!(region.series_index_version().samples.is_empty()); + engine.stop().await.unwrap(); +} + +#[derive(Debug, Default)] +struct FakeIndex; + +#[async_trait] +impl SeriesIndexFormat for FakeIndex { + fn format(&self) -> ArtifactFormat { + ArtifactFormat::FakeIndex + } + + async fn writer( + &self, + metadata: RegionMetadataRef, + store: ObjectStore, + entry: &SeriesIndexEntry, + ) -> Result> { + let path = format!( + "{}/series/{}.fake", + metadata.region_id.as_u64(), + entry.index_uuid + ); + store.write(&path, "fake").await.context(OpenDalSnafu)?; + Ok(Box::new(FakeWriter { + store, + artifact: ArtifactMetadata { + format: self.format(), + generation: entry.index_uuid, + objects: vec![path], + file_size: 4, + }, + fail: Arc::new(AtomicBool::new(false)), + })) + } +} + +#[tokio::test] +async fn format_index_rebuild_is_scoped_and_scans_fall_back_to_ssts() { + let mut env = TestEnv::with_prefix("format-index").await; + // Index builds require four SSTs per bucket. + let timestamps = [1000, 2000, 3000, 4000, 101000, 102000, 103000, 104000]; + let (engine, region) = prepare_region_with_config( + &mut env, + ×tamps, + MitoConfig { + experimental_enable_series_index: true, + experimental_series_scan_v2: true, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }, + ) + .await; + *region.series_index_version_control.formats.write().unwrap() = Arc::new(SeriesFormats { + index: Some(Arc::new(FakeIndex)), + ..Default::default() + }); + let store = region.series_index_store.clone().unwrap(); + let (purger, mut purges) = series_index_channel(store.clone()); + let reconcile = || { + reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + }; + assert_eq!(2, reconcile().await.unwrap().built_series); + let before = region.series_index_version(); + assert!( + before + .series_indexes + .values() + .all(|h| h.entry().format() == ArtifactFormat::FakeIndex) + ); + assert_eq!(0, reconcile().await.unwrap().built_series); + let scanner = engine + .scanner( + region.region_id, + store_api::storage::ScanRequest { + distribution: Some(store_api::storage::TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + let batches = scanner + .scan() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + timestamps.len(), + batches.iter().map(|batch| batch.num_rows()).sum::() + ); + drop(scanner); + + let missing = before.series_indexes.values().next().unwrap(); + store + .delete(&missing.entry().artifact.as_ref().unwrap().objects[0]) + .await + .unwrap(); + assert_eq!(1, reconcile().await.unwrap().built_series); + let after = region.series_index_version(); + assert_eq!( + 1, + after + .series_indexes + .keys() + .filter(|id| before.series_indexes.contains_key(id)) + .count() + ); + assert!(purges.try_recv().is_err()); + drop(before); + assert!(purges.try_recv().is_ok()); + + // Disabling the injected writer restores ordinary Parquet artifacts without a read adapter. + *region.series_index_version_control.formats.write().unwrap() = + Arc::new(SeriesFormats::default()); + assert_eq!(2, reconcile().await.unwrap().built_series); + assert!( + region + .series_index_version() + .series_indexes + .values() + .all(|h| h.entry().artifact.is_none()) + ); + assert_eq!(0, reconcile().await.unwrap().built_series); + engine.stop().await.unwrap(); +} + +struct FailedCatalogWriter(oio::Writer); + +impl mock::Write for FailedCatalogWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.0.write(buffer).await + } + async fn close(&mut self) -> mock::Result { + Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected artifact catalog failure", + )) + } + async fn abort(&mut self) -> mock::Result<()> { + self.0.abort().await + } +} + +#[tokio::test] +async fn format_catalog_failure_preserves_published_artifacts() { + let mut env = TestEnv::with_prefix("format-catalog-failure").await; + let (engine, region) = crate::series_index::tests::prepare_region(&mut env).await; + *region.series_index_version_control.formats.write().unwrap() = Arc::new(SeriesFormats { + index: Some(Arc::new(FakeIndex)), + samples: Some(Arc::new(FakeSamples::default())), + }); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let (purger, mut purges) = series_index_channel(store.clone()); + reconcile_series_indexes( + 0, + store.clone(), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + .await + .unwrap(); + let before = region.series_index_version(); + let catalog_path = crate::series_index::catalog::samples_catalog_path(region.region_id); + let catalog = store.read(&catalog_path).await.unwrap().to_bytes(); + let old = before.samples.values().next().unwrap(); + store + .delete(&old.entry().artifact.objects[0]) + .await + .unwrap(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(|path, _, inner| { + if path.ends_with("samples.json") { + Box::new(FailedCatalogWriter(inner)) + } else { + inner + } + })) + .build() + .unwrap(); + let error = reconcile_series_indexes( + 0, + store.clone().layer(layer), + region.clone(), + Duration::from_secs(100), + 0, + purger.clone(), + false, + true, + ) + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("injected artifact catalog failure")); + assert!(Arc::ptr_eq(&before, ®ion.series_index_version())); + assert_eq!(catalog, store.read(&catalog_path).await.unwrap().to_bytes()); + let cleanup = purges.try_recv().unwrap().objects.unwrap(); + assert_eq!(2, cleanup.len()); + assert!( + cleanup + .iter() + .all(|path| !old.entry().artifact.objects.contains(path)) + ); + assert!(purges.try_recv().is_err()); + let reopened = load_version_control(&store, region.region_id, &purger).await; + assert_eq!( + old.entry(), + reopened.current().samples[&old.entry().file_id].entry() + ); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index f5221bb1ebe..232e17ac266 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -20,6 +20,7 @@ use std::time::{Duration, Instant}; use common_telemetry::{debug, info}; use object_store::ObjectStore; +use snafu::ResultExt; use store_api::storage::{FileId, RegionId}; use crate::error::Result; @@ -40,16 +41,23 @@ use crate::series_index::version::{SeriesIndexFileHandle, SeriesIndexVersion}; /// Retires newly completed series files unless their snapshot is published. #[derive(Default)] -struct UnpublishedSeriesFiles(Vec); +struct UnpublishedSeriesFiles( + Vec, + Vec, +); impl UnpublishedSeriesFiles { fn disarm(&mut self) { self.0.clear(); + self.1.clear(); } } impl Drop for UnpublishedSeriesFiles { fn drop(&mut self) { + for handle in &self.1 { + handle.mark_deleted(); + } for handle in &self.0 { handle.mark_deleted(); } @@ -58,6 +66,8 @@ impl Drop for UnpublishedSeriesFiles { #[derive(Debug, Default)] pub(crate) struct ReconcileStats { + pub(crate) built_samples: usize, + pub(crate) removed_samples: usize, pub(crate) source_files: usize, pub(crate) built_range: usize, pub(crate) built_series: usize, @@ -69,7 +79,13 @@ pub(crate) struct ReconcileStats { impl ReconcileStats { fn changed(&self) -> bool { - self.built_range + self.built_series + self.removed_range + self.removed_series > 0 + self.built_samples + + self.removed_samples + + self.built_range + + self.built_series + + self.removed_range + + self.removed_series + > 0 } } @@ -201,12 +217,43 @@ async fn build_index_version( } (false, _) => Vec::new(), }; - let plan = plan_series_indexes( - buckets, - current.index_buckets.clone(), - version.options.ttl, - now_ms, + let formats = region + .series_index_version_control + .formats + .read() + .unwrap() + .clone(); + let selected_format = formats.index.as_ref().map_or( + crate::series_index::artifact::ArtifactFormat::Parquet, + |backend| backend.format(), ); + let mut unusable = HashSet::new(); + for handle in current.series_indexes.values() { + let entry = handle.entry(); + let mut reusable = entry.format() == selected_format; + // Legacy Parquet indexes retain their existing maintenance behavior. + if reusable && let Some(artifact) = &entry.artifact { + for object in &artifact.objects { + reusable &= store + .exists(object) + .await + .context(crate::error::OpenDalSnafu)?; + } + } + if !reusable { + unusable.insert(entry.index_uuid); + } + } + let mut index_buckets = current.index_buckets.clone(); + for bucket in index_buckets.values_mut() { + if bucket.index_ids.iter().any(|id| unusable.contains(id)) { + // Force replacement only for the affected bucket. Retain its geometry and + // handles until a replacement is built and its catalog is published. + bucket.compaction_window_secs = 0; + bucket.window_sequences.clear(); + } + } + let plan = plan_series_indexes(buckets, index_buckets, version.options.ttl, now_ms); stats.computed_buckets = plan.computed_buckets; stats.skipped_buckets = plan.skipped_buckets; let mut next = prune_index_version(¤t, &visible, &plan.expired_index_ids, &mut stats); @@ -218,6 +265,7 @@ async fn build_index_version( } if plan.builds.is_empty() && next.is_none() + && formats.samples.is_none() && (!enable_range_index || current.range_indexes.len() == visible.len()) { return Ok((None, stats)); @@ -225,13 +273,50 @@ async fn build_index_version( let SeriesIndexVersion { mut range_indexes, mut series_indexes, + mut samples, .. } = next.unwrap_or_else(|| SeriesIndexVersion { + samples: current.samples.clone(), range_indexes: current.range_indexes.clone(), series_indexes: current.series_indexes.clone(), index_buckets: Default::default(), }); let index_buckets = plan.index_buckets; + if let Some(backend) = &formats.samples { + for file in &files { + let id = file.file_id().file_id(); + let reusable = if let Some(handle) = samples.get(&id) { + let artifact = &handle.entry().artifact; + let mut exists = artifact.format == backend.format(); + for path in &artifact.objects { + exists &= store + .exists(path) + .await + .context(crate::error::OpenDalSnafu)?; + } + exists + } else { + false + }; + if reusable { + continue; + } + if let Some(handle) = crate::series_index::builder::build_samples( + store, + region, + version, + file.clone(), + backend.as_ref(), + purger, + ) + .await? + { + unpublished.1.push(handle.clone()); + samples.insert(id, handle); + stats.built_samples += 1; + } + } + } for id in &plan.superseded_index_ids { series_indexes.remove(id); } @@ -246,8 +331,16 @@ async fn build_index_version( range_indexes.insert(entry.file_id, entry); } } - let series_handle = - build_series_index(store, region, version, &bucket, &expected, purger).await?; + let series_handle = build_series_index( + store, + region, + version, + &bucket, + &expected, + formats.index.as_deref(), + purger, + ) + .await?; unpublished.0.push(series_handle.clone()); stats.built_series += 1; series_indexes.insert(expected.index_uuid, series_handle); @@ -276,6 +369,7 @@ async fn build_index_version( return Ok((None, stats)); } let next = SeriesIndexVersion { + samples, range_indexes, series_indexes, index_buckets, @@ -293,7 +387,8 @@ fn prune_index_version( expired_index_ids: &[FileId], stats: &mut ReconcileStats, ) -> Option { - if current.range_indexes.keys().all(|id| visible.contains(id)) + if current.samples.keys().all(|id| visible.contains(id)) + && current.range_indexes.keys().all(|id| visible.contains(id)) && expired_index_ids .iter() .all(|id| !current.series_indexes.contains_key(id)) @@ -301,10 +396,13 @@ fn prune_index_version( return None; } let mut next = SeriesIndexVersion { + samples: current.samples.clone(), range_indexes: current.range_indexes.clone(), series_indexes: current.series_indexes.clone(), index_buckets: current.index_buckets.clone(), }; + next.samples.retain(|id, _| visible.contains(id)); + stats.removed_samples = current.samples.len() - next.samples.len(); next.range_indexes .retain(|file_id, _| visible.contains(file_id)); for id in expired_index_ids { @@ -322,6 +420,31 @@ async fn persist_index_catalogs( next: &SeriesIndexVersion, stats: &ReconcileStats, ) -> Result<()> { + let samples_path = crate::series_index::catalog::samples_catalog_path(region_id); + if !next.samples.is_empty() + || stats.removed_samples > 0 + || store + .exists(&samples_path) + .await + .context(crate::error::OpenDalSnafu)? + { + // A single object write commits all references. A failed write leaves the prior + // catalog usable while the unpublished-output guard retires the new artifacts. + return store_catalog( + store, + &samples_path, + &crate::series_index::catalog::SamplesCatalog { + entries: next.samples.values().map(|h| h.entry().clone()).collect(), + series: next + .series_indexes + .values() + .map(|h| h.entry().clone()) + .collect(), + ranges: next.range_indexes.values().copied().collect(), + }, + ) + .await; + } if stats.built_range + stats.removed_range > 0 { let mut range_entries = next.range_indexes.values().copied().collect::>(); range_entries @@ -364,6 +487,13 @@ async fn persist_index_catalogs( /// Publishes the snapshot and retires series files absent from the new version. fn publish_index_version(region: &MitoRegionRef, next: Arc) { let previous = region.series_index_version_control.publish(next.clone()); + for (id, handle) in &previous.samples { + if next.samples.get(id).is_none_or(|next| { + next.entry().artifact.generation != handle.entry().artifact.generation + }) { + handle.mark_deleted(); + } + } for (id, handle) in &previous.series_indexes { if !next.series_indexes.contains_key(id) { // Purge only after readers release their retained handles. @@ -405,6 +535,7 @@ mod tests { SeriesIndexFileHandle::new( RegionId::new(1, 1), SeriesIndexEntry { + artifact: None, index_uuid: series_id, file_size: 20, bucket_start: Timestamp::new_second(0), diff --git a/src/mito2/src/series_index/purger.rs b/src/mito2/src/series_index/purger.rs index a3eb6d89e38..66eda84c327 100644 --- a/src/mito2/src/series_index/purger.rs +++ b/src/mito2/src/series_index/purger.rs @@ -39,9 +39,10 @@ impl IndexFileType { } } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone)] pub(crate) struct PurgeRequest { pub(crate) file_id: RegionFileId, + pub(crate) objects: Option>, } #[derive(Clone)] @@ -57,6 +58,13 @@ impl Debug for IndexFilePurger { } impl IndexFilePurger { + pub(crate) fn purge_objects(&self, file_id: RegionFileId, objects: Vec) { + self.purge(PurgeRequest { + file_id, + objects: Some(objects), + }); + } + pub(crate) fn purge(&self, request: PurgeRequest) { if let Err(error) = self.sender.send(request) { let store = self.store.clone(); @@ -87,17 +95,24 @@ pub(crate) async fn run_index_purge_task( } async fn purge_file(store: &ObjectStore, request: PurgeRequest) { - let path = series_index_path(request.file_id.region_id(), request.file_id.file_id()); - match store.delete(&path).await { - Ok(()) => { - file_operation(IndexFileType::Series, "delete", "success"); - } - Err(error) if error.kind() == ErrorKind::NotFound => { - file_operation(IndexFileType::Series, "delete", "success"); - } - Err(error) => { - file_operation(IndexFileType::Series, "delete", "failure"); - warn!(error; "Failed to delete series index, index_type: {}, path: {}, phase: deletion", IndexFileType::Series.as_str(), path); + let objects = request.objects.unwrap_or_else(|| { + vec![series_index_path( + request.file_id.region_id(), + request.file_id.file_id(), + )] + }); + for path in objects { + match store.delete(&path).await { + Ok(()) => { + file_operation(IndexFileType::Series, "delete", "success"); + } + Err(error) if error.kind() == ErrorKind::NotFound => { + file_operation(IndexFileType::Series, "delete", "success"); + } + Err(error) => { + file_operation(IndexFileType::Series, "delete", "failure"); + warn!(error; "Failed to delete series index, index_type: {}, path: {}, phase: deletion", IndexFileType::Series.as_str(), path); + } } } } @@ -129,7 +144,10 @@ mod tests { let path = series_index_path(file_id.region_id(), file_id.file_id()); store.write(&path, "index").await.unwrap(); paths.push(path); - purger.purge(PurgeRequest { file_id }); + purger.purge(PurgeRequest { + file_id, + objects: None, + }); } drop(purger); tokio::time::timeout( @@ -151,7 +169,10 @@ mod tests { let file_id = RegionFileId::new(RegionId::new(1, 1), FileId::random()); let path = series_index_path(file_id.region_id(), file_id.file_id()); store.write(&path, "index").await.unwrap(); - purger.purge(PurgeRequest { file_id }); + purger.purge(PurgeRequest { + file_id, + objects: None, + }); tokio::time::timeout(Duration::from_secs(10), async { while store.exists(&path).await.unwrap() { tokio::task::yield_now().await; diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs index 0b9befe4b8a..a2589b3bda9 100644 --- a/src/mito2/src/series_index/searcher.rs +++ b/src/mito2/src/series_index/searcher.rs @@ -312,6 +312,7 @@ mod tests { ) -> (SeriesIndexFileHandle, UnboundedReceiver) { let (purger, receiver) = series_index_channel(store.clone()); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: common_time::Timestamp::new_second(0), diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index 631bdc75b87..145c2088cce 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -63,7 +63,7 @@ pub(crate) async fn prepare_region_with_timestamps( .await } -async fn prepare_region_with_config( +pub(crate) async fn prepare_region_with_config( env: &mut TestEnv, timestamps: &[i64], config: MitoConfig, @@ -726,7 +726,7 @@ async fn test_failed_reconcile_retires_only_unpublished_series() { .unwrap(); ranges.insert(entry.file_id, entry); } - let handle = build_series_index(store, ®ion, &version, bucket, entry, purger) + let handle = build_series_index(store, ®ion, &version, bucket, entry, None, purger) .await .unwrap(); let reused_path = series_index_path(region.region_id, entry.index_uuid); @@ -891,6 +891,7 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + artifact: None, file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index fa2115fc1f2..7d819a0df55 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -82,9 +82,15 @@ struct SeriesIndexFileHandleInner { impl Drop for SeriesIndexFileHandleInner { fn drop(&mut self) { if self.deleted.load(Ordering::Acquire) { - self.purger.purge(PurgeRequest { - file_id: self.file_id, - }); + if let Some(artifact) = &self.entry.artifact { + self.purger + .purge_objects(self.file_id, artifact.objects.clone()); + } else { + self.purger.purge(PurgeRequest { + file_id: self.file_id, + objects: None, + }); + } } } } @@ -94,6 +100,7 @@ impl Drop for SeriesIndexFileHandleInner { pub(crate) struct SeriesIndexVersion { /// Range indexes for visible SSTs; reconciliation removes IDs absent from its SST snapshot. /// Physical deletion is independently handled by the SST file purger. + pub(crate) samples: HashMap, pub(crate) range_indexes: HashMap, pub(crate) series_indexes: HashMap, pub(crate) index_buckets: BTreeMap, @@ -110,6 +117,7 @@ impl SeriesIndexVersion { IndexBucket::from_entry(handle.entry()).insert_into(&mut index_buckets); } Self { + samples: Default::default(), range_indexes, series_indexes, index_buckets, @@ -118,10 +126,15 @@ impl SeriesIndexVersion { /// Approximate installed usage; old snapshots and unpublished outputs are excluded. pub(crate) fn disk_usage(&self) -> u64 { - self.range_indexes + self.samples .values() - .map(|entry| entry.file_size) + .map(|handle| handle.entry().artifact.file_size) .sum::() + + self + .range_indexes + .values() + .map(|entry| entry.file_size) + .sum::() + self .series_indexes .values() @@ -130,6 +143,9 @@ impl SeriesIndexVersion { } fn mark_all_deleted(&self) { + self.samples + .values() + .for_each(|handle| handle.mark_deleted()); self.series_indexes .values() .for_each(SeriesIndexFileHandle::mark_deleted); @@ -140,6 +156,7 @@ impl SeriesIndexVersion { #[derive(Debug, Default)] pub(crate) struct SeriesIndexVersionControl { current: RwLock>, + pub(crate) formats: RwLock>, } impl SeriesIndexVersionControl { diff --git a/tests/compatibility/fixtures/series_index/README.md b/tests/compatibility/fixtures/series_index/README.md new file mode 100644 index 00000000000..29c3000f838 --- /dev/null +++ b/tests/compatibility/fixtures/series_index/README.md @@ -0,0 +1,10 @@ +# Series-index catalog compatibility + +`legacy_catalog.json` is a synthetic fixture matching `SeriesIndexEntry` at +`3ad5f963285e51166fd8789920e4b3bf2cf60423`, before optional artifact metadata. +It is not captured from a released binary. Its entry must remain a reusable +Parquet index and serialize without adding an artifact field. + +Run `cargo nextest run -p mito2 -E 'test(legacy_series_catalog_fixture)'`. +The owning test also checks the persisted artifact format names. Ordinary +index-reader tests exercise Parquet files without a full-primary-key column. diff --git a/tests/compatibility/fixtures/series_index/legacy_catalog.json b/tests/compatibility/fixtures/series_index/legacy_catalog.json new file mode 100644 index 00000000000..e119e602a8d --- /dev/null +++ b/tests/compatibility/fixtures/series_index/legacy_catalog.json @@ -0,0 +1,27 @@ +{ + "indexes": [ + { + "file_size": 123, + "index_uuid": "00000000-0000-0000-0000-000000000001", + "bucket_start": { + "value": 0, + "unit": "Second" + }, + "bucket_end": { + "value": 100, + "unit": "Second" + }, + "source_file_ids": [], + "min_file_sequence": 1, + "max_file_sequence": 2, + "compaction_window_secs": 100, + "window_sequences": { + "0": { + "start": 0, + "end": 100, + "max_sequence": 2 + } + } + } + ] +} From ee4c06d49815bfc60dbe67404c50cd0217ccce45 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 16:57:22 +0800 Subject: [PATCH 05/13] test(promql): hold the batch path to the per-window one on random input (#9739) * test(promql): hold the batch path to the per-window one on random input Two comments claimed the opposite of what the code relies on: a batch's spans were described as absolute indices that survive a release, but a span is recorded against the retained arrays as they stand, and `advance` drains and shifts them. Nothing enforces the difference -- the batch is correct only because `evaluate` flushes it before advancing, so a later change that let a batch outlive one call would return wrong counter values with no error. Both now say that, and `PendingBatch::spans` carries the invariant where a reader looks for it. The batch path had no randomized comparison at all: the differential tests drive a probe kernel and the function tests cover no counter function, so only fixed cases exercised it. A new differential test runs rate and increase over 200 seeds of one to three fields, three series through one evaluator, random chunk splits and mixed histogram samples, comparing the batch path against the same evaluator with its batch form removed. Two assertions that could only have failed on a broken helper are gone, and the counter-reset flag is documented as governing both readers of the index, not just the per-window one. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): assert the batch is flushed before the sample drain `PendingBatch.spans` index the retained sample arrays, which `FieldSamples::advance` drains and shifts. The batch is correct only because `evaluate` flushes it just above that drain, and nothing enforced the ordering: a future change that let a batch outlive the call would compute counter values against shifted samples with no error. The assertion holds on every tested batch path. It is skipped when no batch kernel is installed, since `pending` is only ever filled on the kernel branch. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): state the flush-before-advance rule once, and assert it plainly The batch assert could not fail on its `batch.is_none()` disjunct: `accumulate` is reached only from inside `if let Some(kernel) = self.batch.as_mut()`, and `batch` is never cleared, so a non-empty `pending` implies a kernel. Drop the disjunct and name who guarantees the order in the message instead. The rule that a batch must flush before `advance` was restated in six places. Keep it on `PendingBatch::spans`, where a reader asks why a span dies with the batch, and in the assertion message. The other four now state only the site-specific fact each needs. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): drop the last restatement of flush-before-advance The rule now lives only on PendingBatch::spans and in the debug_assert message, as the review asked; the comment above the final flush restated it a third time. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/range_eval.rs | 24 +-- src/promql/src/range_eval/function.rs | 5 + src/promql/src/range_eval/tests.rs | 210 ++++++++++++++++++++++---- 3 files changed, 201 insertions(+), 38 deletions(-) diff --git a/src/promql/src/range_eval.rs b/src/promql/src/range_eval.rs index 428f9f0f9af..1564fec0d0a 100644 --- a/src/promql/src/range_eval.rs +++ b/src/promql/src/range_eval.rs @@ -80,7 +80,8 @@ pub(crate) struct FloatWindows<'a> { pub resets: &'a [CounterReset], /// Absolute position of `timestamps[0]`: a reset position is `base` plus a sample index. pub base: u64, - /// One range per window, ascending and left-open at the start. + /// One range per window, ascending. Each is a half-open index range `[start, end)`; the time + /// window it was derived from is left-open, `(T - range_ms, T]`. pub windows: &'a [Range], /// The evaluation timestamp of each window. pub eval_ts: &'a [i64], @@ -90,8 +91,9 @@ pub(crate) struct FloatWindows<'a> { impl FloatWindows<'_> { /// The resets inside window `index`. /// - /// `windows` are absolute indices into `timestamps`, so a window whose start has already - /// been released still maps to the right reset. + /// `windows` index the retained arrays as they stand in this batch, not absolutely: each + /// span is relative to `timestamps` and `values`, which is why `base` plus a span start is + /// the absolute position of the sample the span starts at. pub(crate) fn resets_of(&self, index: usize) -> &[CounterReset] { let span = &self.windows[index]; resets_between( @@ -167,6 +169,10 @@ struct PendingBatch { /// Evaluation timestamp of each accumulated point. eval_ts: Vec, /// The retained-sample range of each point's window, one list per field. + /// + /// A span is relative to the retained arrays at the time it was recorded, not an absolute + /// index, so [`FieldSamples::advance`] invalidates it: that is why `evaluate` flushes the + /// batch before advancing. spans: Vec>>, /// One output column per field, written by the batch kernel. columns: Vec>>, @@ -202,9 +208,6 @@ impl PendingBatch { /// Records one point for the next batch call: its evaluation timestamp and, per field, the range /// of retained samples its window covers. -/// -/// The ranges are relative to the retained samples, so samples dropped before the batch flushes -/// would shift them; [`flush_batch`] therefore runs before any `advance`. fn accumulate(pending: &mut PendingBatch, index: usize, eval_ts: i64, fields: &[FieldSamples]) { debug_assert_eq!(pending.spans.len(), fields.len()); pending.indices.push(index); @@ -236,8 +239,6 @@ fn flush_batch( *batch_calls += 1; } for (index, field) in fields.iter().enumerate() { - // `pending.spans` index the retained arrays absolutely, so no adjustment is needed for - // samples `advance` has released: those sit before every span start. let batch = FloatWindows { timestamps: &field.timestamps, values: &field.values, @@ -1043,8 +1044,6 @@ impl RangeEvaluator { } state.cursor += 1; } - // The accumulated spans index the retained arrays, which `advance` below rewrites, so the - // batch has to flush first. if let Some(kernel) = self.batch.as_mut() { flush_batch( kernel, @@ -1073,6 +1072,11 @@ impl RangeEvaluator { } state.cursor = decided; } + debug_assert!( + self.pending.is_empty(), + "`evaluate` flushes the batch before `advance` rewrites the retained arrays a \ + pending span indexes" + ); for field in &mut self.fields { // After the last window, samples are still ingested so that the overflow and // histogram checks do not depend on how the series is chunked, but none are kept. diff --git a/src/promql/src/range_eval/function.rs b/src/promql/src/range_eval/function.rs index f2765d1a543..88888bb15f8 100644 --- a/src/promql/src/range_eval/function.rs +++ b/src/promql/src/range_eval/function.rs @@ -73,6 +73,11 @@ impl RangeFunction { /// the same as saying `IS_COUNTER` in `extrapolated_rate`: `rate` and `increase`. A kernel /// that corrects counters without the index being maintained would silently drop every /// correction, so a new counter function has to be added here as well. + /// + /// It governs both readers of the index: the per-window kernel above, and the batch kernel, + /// which reads [`FloatWindows::resets_of`]. When it returns `false`, `resets` stays empty and + /// `resets_of` yields nothing, so a batch kernel for a counter function would skip every + /// correction and report `delta`-like values. pub(crate) fn indexes_counter_resets(self) -> bool { matches!(self, Self::Rate | Self::Increase) } diff --git a/src/promql/src/range_eval/tests.rs b/src/promql/src/range_eval/tests.rs index c11af5e3c02..f3ea5d3d952 100644 --- a/src/promql/src/range_eval/tests.rs +++ b/src/promql/src/range_eval/tests.rs @@ -721,6 +721,185 @@ fn functions_match_range_udfs() { } } +/// Replaces the float sample of `field` at every row index in `rows` with a native histogram +/// sample, so the windows holding one take the mixed-window path. +fn mix_in_histograms(chunks: &mut [RawSeriesChunk], field: usize, rows: &BTreeSet) { + let histogram = "0/0/0/4014000000000000/4014000000000000/1/0_1/4014000000000000///"; + let view = OwnedHistogram::parse(histogram); + let mut start = 0; + for piece in chunks { + let len = piece.time.len(); + let local: Vec = rows + .range(start..start + len) + .map(|&row| row - start) + .collect(); + start += len; + if local.is_empty() { + continue; + } + let mut builder = HistogramBatchBuilder::new(); + let indices: UInt32Array = (0..len) + .map(|row| { + local + .binary_search(&row) + .is_ok() + .then(|| builder.push(view.view()).unwrap()) + }) + .collect(); + // A position is a float or a histogram, never both, so the float sample of a mixed-in row + // goes. + let values: Float64Array = piece.fields.as_slice()[field] + .iter() + .enumerate() + .map(|(row, sample)| match sample { + _ if local.binary_search(&row).is_ok() => None, + SampleRef::Float(value) => Some(value), + SampleRef::Missing => None, + SampleRef::Histogram(_) => unreachable!("histograms are mixed in once"), + }) + .collect(); + let histograms = HistogramSamples::try_new(indices, builder.finish().unwrap()).unwrap(); + let mixed = Samples::try_from_mixed(values, histograms).unwrap(); + let fields = piece + .fields + .as_slice() + .iter() + .enumerate() + .map(|(index, samples)| { + if index == field { + mixed.clone() + } else { + samples.clone() + } + }) + .collect(); + piece.fields = SeriesFields::try_from_vec(fields).unwrap(); + } +} + +/// Asserts two evaluator results are identical, including native histogram samples, which the +/// dense form [`expand`] rejects. +fn assert_same_series(actual: Option<&Series>, expected: Option<&Series>, context: &str) { + match (actual, expected) { + (None, None) => {} + (Some(actual), Some(expected)) => { + assert_eq!(actual.metric, expected.metric, "{context}: metric"); + let (TimeAxis::Grid(actual_range), TimeAxis::Grid(expected_range)) = + (&actual.time, &expected.time) + else { + panic!( + "{context}: expected grid series, got {:?} and {:?}", + actual.time, expected.time + ); + }; + assert_eq!(actual_range, expected_range, "{context}: grid range"); + for (field, (got, want)) in actual + .fields + .as_slice() + .iter() + .zip(expected.fields.as_slice()) + .enumerate() + { + for (row, (got, want)) in got.iter().zip(want.iter()).enumerate() { + let context = format!("{context}: field {field} row {row}"); + match (got, want) { + (SampleRef::Missing, SampleRef::Missing) => {} + (SampleRef::Float(got), SampleRef::Float(want)) => { + assert_eq!(got.to_bits(), want.to_bits(), "{context}: {got} vs {want}") + } + (SampleRef::Histogram(got), SampleRef::Histogram(want)) => { + let got = FloatHistogram::try_from_view(&got).unwrap(); + let want = FloatHistogram::try_from_view(&want).unwrap(); + // Both paths run the same histogram code over the same samples, so + // this is an exact comparison: an epsilon of zero. + assert_histogram_close(&got, &want, 0.0, &context); + } + (SampleRef::Missing, _) => { + panic!("{context}: only the per-window path produced a sample") + } + (_, SampleRef::Missing) => { + panic!("{context}: only the batch path produced a sample") + } + _ => panic!("{context}: a float on one path and a histogram on the other"), + } + } + } + } + (actual, expected) => panic!( + "{context}: one path produced a series, the other none: {actual:?} vs {expected:?}" + ), + } +} + +/// `rate` and `increase` over random input, computed twice: by the batch kernel and by the +/// per-window path. +/// +/// Both evaluators are built by `try_new` and differ in nothing but the dispatch shape, and both +/// kernels reach `extrapolated_rate_of`, so the results have to be bit-identical. A difference +/// means the batch read a different set of samples, evaluated at a different time, or corrected a +/// different reset: a window bound that disagrees with the per-window path. +/// +/// The rows mix gaps, stale markers, counter resets and native histogram samples, the last so +/// that some grid points take the mixed-window fallback, where a pending batch has to flush +/// before the point is decided on the per-window path. The random chunk splits are what make +/// `advance` drop samples mid-series, so spans that outlived one would read shifted samples. +#[test] +fn batch_counter_functions_match_the_per_window_path_on_random_input() { + let mut batch_calls = 0; + for seed in 0..200 { + let mut rng = StdRng::seed_from_u64(seed); + let fields = rng.random_range(1..=3); + let (grid, mut window) = random_window(&mut rng); + // An anchored selector reads one window and installs no batch kernel, so both paths would + // run the same code; every seed stays on the grid-driven batch path. + window.anchor_ms = None; + let function = if rng.random_bool(0.5) { + RangeFunction::Rate + } else { + RangeFunction::Increase + }; + let mut batched = + RangeEvaluator::try_new(&raw_schema(fields), grid.clone(), window, function).unwrap(); + let mut per_window = + RangeEvaluator::try_new(&raw_schema(fields), grid.clone(), window, function).unwrap(); + // The per-window path is the same evaluator without the batch kernel. `with_kernel`, + // which the other differential tests use, rejects native histogram samples, and these + // series carry some, so removing the batch form is the only switch that leaves the rest + // of the evaluation identical. + per_window.batch = None; + // One evaluator serves several series, so state must not leak between them. + for host in 0..3 { + let rows = random_rows(&mut rng, fields); + let histogram_field = rng.random_range(0..fields); + let histogram_rows: BTreeSet = + (0..rows.len()).filter(|_| rng.random_bool(0.1)).collect(); + let mut chunks = random_chunks(&mut rng, &rows, fields, host); + mix_in_histograms(&mut chunks, histogram_field, &histogram_rows); + // A push hands out the positions it decided, not only on the last chunk, so what each + // path produced is every chunk it handed out, folded back into one series. + let batched_chunks = collect_chunks(&mut batched, &chunks); + let per_window_chunks = collect_chunks(&mut per_window, &chunks); + + let batched_result = fold_chunks(&batched_chunks, batched.output_schema()); + let per_window_result = fold_chunks(&per_window_chunks, per_window.output_schema()); + + let context = format!("seed {seed} host {host}"); + assert_same_series( + batched_result.as_ref(), + per_window_result.as_ref(), + &context, + ); + } + batch_calls += batched.batch_call_count(); + } + // The comparison is worth nothing unless the batch kernel ran at all: counting the calls is + // what tells a batch-path comparison from one that silently fell back to per-window. + assert!( + batch_calls > 300, + "only {batch_calls} batch kernel calls: the batch path was barely reached" + ); +} + /// 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. @@ -786,31 +965,6 @@ fn a_chunk_with_histograms_is_not_appended_in_bulk() { assert_eq!(evaluator.bulk_ingest_count(), 0); } -/// `rate` over a float-only series takes the batch path. The per-window path produces the same -/// output, so a test comparing values alone could not tell which one ran. -#[test] -fn rate_over_floats_is_evaluated_in_batches() { - let rows: Vec = (0..20) - .map(|index| (index as i64 * 10_000, vec![Some(index as f64)])) - .collect(); - let mut evaluator = RangeEvaluator::try_new( - &raw_schema(1), - grid(0, 10_000, 20), - RangeWindow { - range_ms: 30_000, - offset_ms: 0, - anchor_ms: None, - }, - RangeFunction::Rate, - ) - .unwrap(); - evaluator.push(&chunk(&rows, 1, 0, true)).unwrap(); - assert!( - evaluator.batch_call_count() > 0, - "rate over floats must reach the batch kernel" - ); -} - /// A batch spanning more grid points than `EVALUATION_BATCH` flushes in the middle of one /// evaluation, and each field's windows come from that field's own samples. /// @@ -881,7 +1035,6 @@ fn a_batch_spanning_fields_and_the_flush_threshold_matches_the_per_window_path() .push(&chunk(&rows, fields, 0, true)) .unwrap() .expect("the series completes"); - assert_eq!(per_window.batch_call_count(), 0); // Fold the chunk back into one series first: a chunk's range reports the positions its push // decided, which can include leading and trailing positions holding no value, while the @@ -899,8 +1052,9 @@ fn a_batch_spanning_fields_and_the_flush_threshold_matches_the_per_window_path() /// A window holding a native histogram sample keeps its grid point on the per-window path, since /// the batch kernel reads the floats alone and cannot reproduce the mixed-window rules. /// -/// The series is laid out so every grid point's window contains the histogram: one point fewer -/// than the first test, so no point can fall before it. +/// The series is laid out so every grid point's window contains the histogram: three samples +/// 10 s apart with the histogram on the first, and a window reaching 30 s back, so a window +/// covering any grid point reaches timestamp 0. #[test] fn a_histogram_window_keeps_rate_on_the_per_window_path() { let histogram = "0/0/0/4014000000000000/4014000000000000/1/0_1/4014000000000000///"; From 77a7d099e6a7dfcea7706c28798dde267a2b3a11 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 17:02:18 +0800 Subject: [PATCH 06/13] docs(promql): add a map for src/promql (#9796) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs(promql): add a map for src/promql `src/promql` has no guide of its own, and the two things that cost the most to learn the hard way are not visible from any one file: which of the two evaluation paths is live, and what has to line up for a new extension node to reach a datanode at all. The map records the module layout, that coupling (substrait payload branch, physical planner branch, `Commutativity` arm), the native model's reachability, and the gotchas that fail silently: a missing commutativity arm, a label set that closes twice under one label, a raw frame decoded with a binding other than the one it was written with. It describes mechanisms and where they live, not which branch or PR adds what, so it stays true as the surrounding work lands. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): correct what the frame schema does and does not say Two claims in the map were wrong about the code it describes. The guide said a raw frame's schema metadata explains how to decode the primary key. It does not: the metadata carries the format keys only — layout, content, labels, time — and the decoder compares those and the field list, never a label identity. `unpack_raw_series` still requires the caller's `RawLabelSource`, so the guide now says to decode with the binding the stream was produced under, which is the mistake it exists to prevent. A future decode descriptor can be described here once it is in the base and this text can describe the API that actually exists. The raw-scan gap named one way past the partition problem as though it were the only one. `RegionScanExec::try_to_series_scan` documents the other — finish partition preparation before converting — so both are stated now, with the scanner-side override kept as the actual requirement. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): say the storage producer exists and name what is left The "cannot be planned" bullet listed the scanner-side override as missing, but mito2 implements `RegionScanner::raw_series_schema` — #9669 landed it, so the trait's `Ok(None)` default only refuses engines that have not implemented it. Listing an implemented override as a gap misreads how far the raw path actually reaches. Separate the two things: the storage producer is available, and what is missing is planner wiring plus the partition-preparation order. Merge current narthex in so the paragraph describes the branch it will land on rather than a stale base. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- AGENTS.md | 1 + src/promql/AGENTS.md | 110 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 src/promql/AGENTS.md diff --git a/AGENTS.md b/AGENTS.md index 20b2689a401..4ff990a5542 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -62,6 +62,7 @@ a module map, read/write paths, change-coupling points, and gotchas: - [`src/servers/AGENTS.md`](src/servers/AGENTS.md) - [`src/operator/AGENTS.md`](src/operator/AGENTS.md) - [`src/mito2/AGENTS.md`](src/mito2/AGENTS.md) +- [`src/promql/AGENTS.md`](src/promql/AGENTS.md) - [`src/metric-engine/AGENTS.md`](src/metric-engine/AGENTS.md) - [`src/flow/AGENTS.md`](src/flow/AGENTS.md) - [`src/frontend/AGENTS.md`](src/frontend/AGENTS.md) diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md new file mode 100644 index 00000000000..dc90c5a0eec --- /dev/null +++ b/src/promql/AGENTS.md @@ -0,0 +1,110 @@ +# promql — Agent & Contributor Guide + +Navigation map for `src/promql`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`promql` owns PromQL evaluation: the series/evaluated data model, range and +instant evaluation, function kernels, vector matching, and the PromQL extension +operators. Statement-to-plan translation lives in `src/query/src/promql/`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Data model | `src/promql/src/grid_series.rs`, `src/promql/src/evaluated.rs` | Grid series, evaluated chunks | +| Evaluation | `src/promql/src/range_eval.rs`, `src/promql/src/instant_eval.rs` | Range and instant evaluation | +| Kernels | `src/promql/src/functions.rs`, `src/promql/src/functions/` | PromQL functions | +| Matching | `src/promql/src/vector_match.rs` | Binary vector matching | +| Extension operators | `src/promql/src/extension_plan.rs`, `src/promql/src/extension_plan/` | Logical nodes and execution plans | +| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics | + +## Two evaluation paths, and their status + +The easiest thing to get wrong here. Both exist; only one runs. + +- **Extension operators are live.** `PromPlanner` builds `SeriesDivide`, + `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, `EmptyMetric` and + `UnionDistinctOn`; `PromExtensionPlanner` + (`src/promql/src/extension_plan/planner.rs`) maps each logical node to its + execution plan, one branch per node, `None` for anything else. +- **Nothing plans into the native model yet.** `src/query/src/promql/native_compile.rs` + compiles a statement into a `NativeTree`, but no planner calls it, so a tree + built here reaches no execution. `src/promql/src/native/` holds the tree, the + chunk protocol, display and metrics, and — once a runtime lands — the driver + that runs a tree node by node. Compiling and running are separate concerns + owned by separate files; do not assume a tree is executable because it + compiles. + +## Change coupling + +- **An extension node reaches a datanode only if three things line up**: a payload + branch in `src/common/substrait/src/extension_serializer.rs`, a physical branch + in `PromExtensionPlanner`, and a `Commutativity::Commutative` arm in + `check_extension_plan` (`src/query/src/dist_plan/commutativity.rs`). Missing the + third is silent: the node falls through to `Unsupported`, the distributed + boundary is placed above it, and it never reaches a region — with no error. +- A new logical node also needs its physical planner reachable from + `DfQueryPlanner::create_physical_plan`, assembled in + `src/query/src/query_engine/state.rs`. +- User-visible changes need planner coverage plus a sqlness case under + `tests/cases/standalone/common/promql/`. + +## Gotchas + +- **One operation, two implementations; build on the live one.** `InstantManipulate` + routes to `InstantManipulateExec`. The native alternative + `NativeInstantManipulateExec` is crate-private and + `InstantManipulate::try_to_native_execution_plan` has no caller on any branch — + its doc is written for a caller that does not exist. +- **A native window requires non-interleaved input.** Only `Contiguous` and + `KeyThenTime` are accepted; `Interleaved` is refused when compiled. +- **A series closes exactly once**, on the chunk that says so — per input series, + not per label set. A node that rewrites labels can hand two input series over + under one label set, and each closes on its own chunk, so a label set may be + closed more than once; a merge above them must not read the second close as a + defect. That closing chunk is the only way a consumer distinguishes "nothing more + here" from "no more input": dropping it leaves a series open, and a child that + merges regions must recompute finality globally rather than pass a region-local + flag through. +- **Raw-frame labels come from the binding, not the schema.** A raw series frame + carries an opaque encoded primary key while its schema metadata describes only + the format: `greptime.series.layout`, `.content`, `.labels` and `.time`, plus + the producer ordering the adapter adds. Nothing in that metadata binds the key + to a decode — the frame decoder compares those format keys and the field list, + and never a label identity. `unpack_raw_series` therefore requires a + caller-supplied `RawLabelSource`: decode with the source that was actually bound + when the stream was produced, because the frame schema does not replace that + binding. The metadata describes format, not label identity. +- **A wrong label binding is silently wrong, not an error.** Storage validates the + binding it is given (metadata identity, tag-ness, name matches, no duplicates) + but cannot tell whether that set is the one the query meant, and the metric name + is not part of that check. + +## Testing + +```bash +cargo nextest run -p promql +cargo nextest run -p query # planner and native-compile changes +cargo sqlness bare -t # PromQL cases live under tests/cases/standalone/common/promql/ +``` + +## Not wired up yet + +The storage side of the raw path exists; what is missing is planner wiring. Both +gaps below are structural, not bugs. + +- **The native tree has no planner entry point.** `native_compile::compile` has no + caller, and a `NativeTree` becomes a physical plan only once something builds it + while it still holds the AST and the selector binding — the point where the + matchers, `__field__` and the lookback are still available. Compiling later, from + a plan, loses them. +- **No planner reaches a raw scan.** mito2 implements + `RegionScanner::raw_series_schema` (`mito2/src/read/series_scan.rs`), so + `RegionScanExec::try_to_series_scan` can build `SeriesScanExec` from a prepared + scanner; the trait's `Ok(None)` default only refuses engines that have not + implemented it. What is missing is the planner side: nothing selects that + conversion, and `SeriesScanExec::children()` is empty, so wrapping a + `RegionScanExec` before partition preparation leaves `ParallelizeScan` unable to + assign partitions. Either give the partitioner a child it can see, or convert to + the native leaf only after partition preparation — the second is the contract + `RegionScanExec::try_to_series_scan` documents. From 1156cfbe0cc54cbce6d4550ae999dd28a33d91f8 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 17:02:37 +0800 Subject: [PATCH 07/13] chore(query): drop the prepared raw-series RPC skeleton (#9797) The raw-series preparation RPC was merged into `narthex` for a plan that never landed: no server or frontend dispatches `greptime.prepare_raw_series.v1`, and the three `Datanode` methods, `bind_raw_series_query` and the `QueryEngine` trio have no call sites and no implementors. `PreparedNamedRawSeries` had no user outside them. The wire envelope never shipped -- no tag contains the commit that added it -- so the compatibility fixtures go with it. The harness drives `wire/` fixtures from the Rust tests that own their context conversion rather than from SQLness discovery, and has no retirement mechanism to record instead. Left in place: `RawSeriesRequest` and its extension key, `raw_series_frame_ordering`, `unpack_raw_series` and `RawLabelSource`, all of which the surviving series-frame path still uses. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/meta/src/node_manager.rs | 120 +----- src/common/query/src/request.rs | 5 - .../query/src/request/named_raw_series.rs | 387 ------------------ src/query/src/query_engine.rs | 40 -- src/query/src/region_query.rs | 16 +- .../wire/named_raw_series_prepare/README.md | 32 -- .../v1_supported.json | 1 - .../v1_unsupported.json | 1 - 8 files changed, 3 insertions(+), 599 deletions(-) delete mode 100644 src/common/query/src/request/named_raw_series.rs delete mode 100644 tests/compatibility/wire/named_raw_series_prepare/README.md delete mode 100644 tests/compatibility/wire/named_raw_series_prepare/v1_supported.json delete mode 100644 tests/compatibility/wire/named_raw_series_prepare/v1_unsupported.json diff --git a/src/common/meta/src/node_manager.rs b/src/common/meta/src/node_manager.rs index 512952c9c99..d18e2cc8471 100644 --- a/src/common/meta/src/node_manager.rs +++ b/src/common/meta/src/node_manager.rs @@ -18,21 +18,10 @@ use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse}; use api::v1::region::{InsertRequests, RegionRequest}; pub use common_base::AffectedRows; -use common_error::ext::{BoxedError, PlainError}; -use common_error::status_code::StatusCode; -use common_query::error::Error as CommonQueryError; -use common_query::request::{ - PreparedNamedRawSeries, QueryRequest, RAW_SERIES_REQUEST_EXTENSION_KEY, RawSeriesRequest, -}; -use common_query::series_frame::{raw_series_frame_ordering, unpack_raw_series}; +use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; -use common_recordbatch::adapter::DfRecordBatchStreamAdapter; -use datatypes::series::SeriesOrdering; -use datatypes::series::samples::SeriesSchema; -use snafu::{ResultExt, location}; -use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; -use crate::error::{ExternalSnafu, Result}; +use crate::error::Result; use crate::peer::Peer; /// The trait for handling requests to datanode. @@ -43,111 +32,6 @@ pub trait Datanode: Send + Sync { /// Handles query requests async fn handle_query(&self, request: QueryRequest) -> Result; - - /// Prepares a bound raw-series query without executing it. - /// - /// Implementations that do not support preflight return `Ok(None)` without execution or a - /// validation guarantee; failures other than unsupported preparation are returned as errors. - async fn prepare_raw_series_query( - &self, - _request: QueryRequest, - _binding: RawSeriesRequest, - ) -> Result> { - Ok(None) - } - - /// Non-consumingly checks whether a named raw-series preflight is supported. - /// - /// Returns the approved binding and ordering on support, or `None` when the request is - /// explicitly unsupported; failures that are not unsupported preparation are returned as - /// errors. This point-in-time check does not reserve a plan, and execution revalidates - /// support. - async fn prepare_named_raw_series_query( - &self, - _request: QueryRequest, - ) -> Result> { - Ok(None) - } - - /// Handles a bound raw-series query and decodes its series frames. - /// - /// The binding is an expectation supplied by the already-authorized internal planner; it is - /// not evidence of provenance or an independent authorization grant. - async fn handle_raw_series_query( - &self, - request: QueryRequest, - binding: RawSeriesRequest, - ) -> Result { - let (request, expected_schema, label_source) = bind_raw_series_query(request, &binding)?; - let stream = self.handle_query(request).await?; - let ordering = raw_series_frame_ordering(stream.schema().arrow_schema().as_ref()); - let stream = Box::pin(DfRecordBatchStreamAdapter::new(stream)); - ordering - .and_then(|ordering| unpack_raw_series(stream, expected_schema, label_source, ordering)) - .map_err(|error| { - BoxedError::new(CommonQueryError::GeneralDataFusion { - error, - location: location!(), - }) - }) - .context(ExternalSnafu) - } -} - -/// Binds a raw-series request to the query context without executing the query. -pub fn bind_raw_series_query( - mut request: QueryRequest, - binding: &RawSeriesRequest, -) -> Result<(QueryRequest, Arc, RawLabelSource)> { - if request.region_id != binding.physical_metadata.region_id { - return Err(BoxedError::new(PlainError::new( - "raw-series request region does not match binding".to_string(), - StatusCode::InvalidArguments, - ))) - .context(ExternalSnafu); - } - let (expected_schema, label_source) = binding - .bind(Arc::new(binding.physical_metadata.clone())) - .map_err(|error| { - BoxedError::new(CommonQueryError::GeneralDataFusion { - error, - location: location!(), - }) - }) - .context(ExternalSnafu)?; - if request - .header - .as_ref() - .and_then(|header| header.query_context.as_ref()) - .is_some_and(|context| { - context - .extensions - .contains_key(RAW_SERIES_REQUEST_EXTENSION_KEY) - }) - { - return Err(BoxedError::new(PlainError::new( - "raw-series request marker already exists".to_string(), - StatusCode::InvalidArguments, - ))) - .context(ExternalSnafu); - } - let extension = binding - .to_extension_value() - .map_err(|error| { - BoxedError::new(CommonQueryError::GeneralDataFusion { - error, - location: location!(), - }) - }) - .context(ExternalSnafu)?; - request - .header - .get_or_insert_default() - .query_context - .get_or_insert_default() - .extensions - .insert(RAW_SERIES_REQUEST_EXTENSION_KEY.to_string(), extension); - Ok((request, expected_schema, label_source)) } impl dyn Datanode + '_ { diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 829fa40b22d..c61ff41da7e 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -14,13 +14,8 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; -mod named_raw_series; mod raw_series_request; -pub use self::named_raw_series::{ - PREPARE_RAW_SERIES_ACTION, PreparedNamedRawSeries, decode_prepare_raw_series_response, - encode_prepare_raw_series_response, -}; pub use self::raw_series_request::{ RAW_SERIES_REQUEST_EXTENSION_KEY, RawSeriesRequest, RawSeriesTag, take_raw_series_request, }; diff --git a/src/common/query/src/request/named_raw_series.rs b/src/common/query/src/request/named_raw_series.rs deleted file mode 100644 index 1a6cf8f8521..00000000000 --- a/src/common/query/src/request/named_raw_series.rs +++ /dev/null @@ -1,387 +0,0 @@ -// 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. - -//! Named raw-series preflight contract. -//! -//! The contract is the JSON body of the [`PREPARE_RAW_SERIES_ACTION`] action. -//! It carries the already-bound raw-series request and the ordering the -//! approving node can guarantee for it, or an explicit unsupported marker. - -use datafusion_common::{DataFusionError, Result as DataFusionResult}; -use datatypes::series::SeriesOrdering; -use serde::{Deserialize, Deserializer, Serialize}; - -use crate::request::RawSeriesRequest; - -/// Action name of the named raw-series preflight contract. -pub const PREPARE_RAW_SERIES_ACTION: &str = "greptime.prepare_raw_series.v1"; - -/// A named raw-series request approved for preflight. -#[derive(Clone, Debug)] -pub struct PreparedNamedRawSeries { - /// The bound raw-series request approved by the preparing node. - pub binding: RawSeriesRequest, - /// The ordering guaranteed by the preparing node for the binding. - pub ordering: SeriesOrdering, -} - -/// Encodes a preflight response body for [`PREPARE_RAW_SERIES_ACTION`]. -/// -/// `Some` approves the request after revalidating its binding through the -/// existing [`RawSeriesRequest`] validation, while `None` encodes the frozen -/// unsupported shape. Not every [`SeriesOrdering`] can be reported: interleaved -/// orderings are rejected. -pub fn encode_prepare_raw_series_response( - supported: Option<&PreparedNamedRawSeries>, -) -> DataFusionResult> { - let supported = match supported { - Some(prepared) => { - prepared.binding.to_extension_value()?; - Some(Supported { - binding: prepared.binding.clone(), - ordering: WireOrdering::from_series_ordering(prepared.ordering)?, - }) - } - None => None, - }; - let response = PrepareRawSeriesResponse { - version: PREPARE_RAW_SERIES_VERSION, - supported, - }; - serde_json::to_vec(&response).map_err(|_| invalid_response_error()) -} - -/// Decodes a preflight response body for [`PREPARE_RAW_SERIES_ACTION`]. -/// -/// Returns `Some` for an approving body and `None` for the frozen unsupported -/// shape. The body is decoded directly as a typed JSON envelope, so duplicate -/// keys, unknown fields, unknown versions, and malformed bindings are rejected. -pub fn decode_prepare_raw_series_response( - body: &[u8], -) -> DataFusionResult> { - let response: PrepareRawSeriesResponse = - serde_json::from_slice(body).map_err(|_| invalid_response_error())?; - if response.version != PREPARE_RAW_SERIES_VERSION { - return Err(unsupported_version_error()); - } - let Some(supported) = response.supported else { - return Ok(None); - }; - supported.binding.to_extension_value()?; - Ok(Some(PreparedNamedRawSeries { - binding: supported.binding, - ordering: supported.ordering.into_series_ordering(), - })) -} - -/// Wire-format version of the named raw-series preflight contract. -const PREPARE_RAW_SERIES_VERSION: u32 = 1; - -/// Private JSON envelope of [`PREPARE_RAW_SERIES_ACTION`]. -/// -/// `supported` is required but nullable: a missing field is rejected instead of -/// defaulting, while an explicit `null` is the unsupported encoding. -#[derive(Debug, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -struct PrepareRawSeriesResponse { - version: u32, - #[serde(deserialize_with = "deserialize_required_supported")] - supported: Option, -} - -/// Private shape of an approving response. -#[derive(Debug, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -struct Supported { - binding: RawSeriesRequest, - ordering: WireOrdering, -} - -/// Private mirror of the orderings this contract can report. -#[derive(Clone, Copy, Debug, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -enum WireOrdering { - Contiguous, - ReadKeyThenTime, -} - -impl WireOrdering { - /// Rejects orderings the contract cannot report. - fn from_series_ordering(ordering: SeriesOrdering) -> DataFusionResult { - match ordering { - SeriesOrdering::Contiguous => Ok(Self::Contiguous), - SeriesOrdering::KeyThenTime => Ok(Self::ReadKeyThenTime), - SeriesOrdering::Interleaved => Err(unsupported_ordering_error()), - } - } - - fn into_series_ordering(self) -> SeriesOrdering { - match self { - Self::Contiguous => SeriesOrdering::Contiguous, - Self::ReadKeyThenTime => SeriesOrdering::KeyThenTime, - } - } -} - -/// Requires the `supported` field while still accepting an explicit `null`. -fn deserialize_required_supported<'de, D>(deserializer: D) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - Option::::deserialize(deserializer) -} - -fn invalid_response_error() -> DataFusionError { - DataFusionError::Plan("Invalid prepare raw-series response".to_string()) -} - -fn unsupported_version_error() -> DataFusionError { - DataFusionError::Plan("Unsupported prepare raw-series response version".to_string()) -} - -fn unsupported_ordering_error() -> DataFusionError { - DataFusionError::Plan("Unsupported prepare raw-series response ordering".to_string()) -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use store_api::storage::RegionId; - - use super::*; - - const SUPPORTED_FIXTURE: &str = include_str!( - "../../../../../tests/compatibility/wire/named_raw_series_prepare/v1_supported.json" - ); - const UNSUPPORTED_FIXTURE: &str = include_str!( - "../../../../../tests/compatibility/wire/named_raw_series_prepare/v1_unsupported.json" - ); - const BINDING_FIXTURE: &str = - include_str!("../../../../../tests/compatibility/wire/raw_series_request/v1_request.json"); - - fn fixture_binding() -> RawSeriesRequest { - RawSeriesRequest::from_extension_value(BINDING_FIXTURE).unwrap() - } - - fn fixture_body() -> String { - format!( - r#"{{"version":1,"supported":{{"binding":{},"ordering":"contiguous"}}}}"#, - BINDING_FIXTURE.trim() - ) - } - - #[test] - fn frozen_fixtures_decode_supported_and_unsupported() { - let prepared = decode_prepare_raw_series_response(SUPPORTED_FIXTURE.as_bytes()) - .unwrap() - .expect("supported fixture must decode as approved"); - assert_eq!(prepared.ordering, SeriesOrdering::Contiguous); - assert_eq!( - prepared.binding.physical_metadata.region_id, - RegionId::new(1, 1) - ); - assert_eq!( - prepared.binding.to_extension_value().unwrap(), - fixture_binding().to_extension_value().unwrap() - ); - - let (schema, _) = prepared - .binding - .bind(Arc::new(prepared.binding.physical_metadata.clone())) - .unwrap(); - assert_eq!( - schema - .fields - .iter() - .map(|field| field.name.as_ref()) - .collect::>(), - ["aux", "value"] - ); - - assert!( - decode_prepare_raw_series_response(UNSUPPORTED_FIXTURE.as_bytes()) - .unwrap() - .is_none() - ); - } - - #[test] - fn prepare_response_round_trips_both_orderings() { - for (ordering, wire_name) in [ - (SeriesOrdering::Contiguous, "contiguous"), - (SeriesOrdering::KeyThenTime, "read_key_then_time"), - ] { - let prepared = PreparedNamedRawSeries { - binding: fixture_binding(), - ordering, - }; - let encoded = encode_prepare_raw_series_response(Some(&prepared)).unwrap(); - let json: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); - assert_eq!(json["version"].as_u64(), Some(1)); - assert_eq!(json["supported"]["ordering"].as_str(), Some(wire_name)); - assert!(json["supported"]["binding"].is_object()); - - let decoded = decode_prepare_raw_series_response(&encoded) - .unwrap() - .expect("approved response must decode as supported"); - assert_eq!(decoded.ordering, ordering); - assert_eq!( - decoded.binding.to_extension_value().unwrap(), - prepared.binding.to_extension_value().unwrap() - ); - } - } - - #[test] - fn unsupported_prepare_response_has_frozen_shape() { - let encoded = encode_prepare_raw_series_response(None).unwrap(); - assert_eq!( - std::str::from_utf8(&encoded).unwrap(), - r#"{"version":1,"supported":null}"# - ); - assert!( - decode_prepare_raw_series_response(&encoded) - .unwrap() - .is_none() - ); - } - - #[test] - fn prepare_response_encoding_rejects_interleaved_and_malformed_binding() { - let interleaved = PreparedNamedRawSeries { - binding: fixture_binding(), - ordering: SeriesOrdering::Interleaved, - }; - assert!(encode_prepare_raw_series_response(Some(&interleaved)).is_err()); - - let mut malformed = fixture_binding(); - malformed.field_ids.push(3); - let malformed = PreparedNamedRawSeries { - binding: malformed, - ordering: SeriesOrdering::Contiguous, - }; - assert!(encode_prepare_raw_series_response(Some(&malformed)).is_err()); - } - - #[test] - fn prepare_response_rejects_malformed_bodies() { - let body = fixture_body(); - let cases = [ - ("empty body", String::new()), - ("not json", "{".to_string()), - ("top-level array", "[]".to_string()), - ("missing version", r#"{"supported":null}"#.to_string()), - ( - "zero version", - r#"{"version":0,"supported":null}"#.to_string(), - ), - ( - "unknown version", - r#"{"version":2,"supported":null}"#.to_string(), - ), - ( - "unknown envelope field", - r#"{"version":1,"supported":null,"extra":true}"#.to_string(), - ), - ("missing supported", r#"{"version":1}"#.to_string()), - ( - "duplicate version", - r#"{"version":1,"version":1,"supported":null}"#.to_string(), - ), - ( - "duplicate supported", - r#"{"version":1,"supported":null,"supported":null}"#.to_string(), - ), - ( - "supported missing binding", - r#"{"version":1,"supported":{"ordering":"contiguous"}}"#.to_string(), - ), - ( - "supported missing ordering", - body.replace(r#","ordering":"contiguous""#, ""), - ), - ( - "supported unknown field", - body.replace( - r#""ordering":"contiguous""#, - r#""ordering":"contiguous","extra":true"#, - ), - ), - ( - "interleaved ordering", - body.replace(r#""ordering":"contiguous""#, r#""ordering":"interleaved""#), - ), - ( - "unknown ordering", - body.replace(r#""ordering":"contiguous""#, r#""ordering":"future""#), - ), - ( - "non-string ordering", - body.replace(r#""ordering":"contiguous""#, r#""ordering":1"#), - ), - ( - "unknown nested binding field", - body.replace( - r#""metric_name":"cpu""#, - r#""metric_name":"cpu","extra":true"#, - ), - ), - ( - "duplicate nested binding field", - body.replace( - r#""metric_name":"cpu""#, - r#""metric_name":"cpu","metric_name":"cpu""#, - ), - ), - ( - "unknown nested tag field", - body.replace( - r#"{"column_id":1,"name":"host"}"#, - r#"{"column_id":1,"name":"host","extra":true}"#, - ), - ), - ( - "duplicate nested tag field", - body.replace( - r#"{"column_id":1,"name":"host"}"#, - r#"{"column_id":1,"name":"host","name":"host"}"#, - ), - ), - ( - "malformed binding version", - body.replace( - r#"{"version":1,"physical_metadata""#, - r#"{"version":2,"physical_metadata""#, - ), - ), - ( - "duplicate binding field id", - body.replace(r#""field_ids":[3,2]"#, r#""field_ids":[3,3]"#), - ), - ( - "non-field binding id", - body.replace(r#""field_ids":[3,2]"#, r#""field_ids":[1]"#), - ), - ]; - - for (name, candidate) in cases { - assert!( - decode_prepare_raw_series_response(candidate.as_bytes()).is_err(), - "must reject {name}" - ); - } - assert!(decode_prepare_raw_series_response(body.as_bytes()).is_ok()); - } -} diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index 51f7d7ce477..c512b9c1e5e 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -29,11 +29,9 @@ use common_function::handlers::{ FlowServiceHandlerRef, ProcedureServiceHandlerRef, TableMutationHandlerRef, }; use common_query::Output; -use common_query::request::{PreparedNamedRawSeries, RawSeriesRequest}; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; use datafusion_expr::{AggregateUDF, LogicalPlan, WindowUDF}; -use datatypes::series::SeriesOrdering; pub use default_serializer::{DefaultPlanDecoder, DefaultSerializer}; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; @@ -79,44 +77,6 @@ pub trait QueryEngine: Send + Sync { /// Execute the given [`LogicalPlan`]. async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result; - /// Non-consumingly checks whether a native single-region raw-series scan is supported. - /// Returns its ordering on support, `None` when explicitly unsupported, and propagates errors. - /// This point-in-time check does not reserve a plan; execution revalidates support. - async fn prepare_raw_series( - &self, - _plan: LogicalPlan, - _query_ctx: QueryContextRef, - _binding: RawSeriesRequest, - ) -> Result> { - Ok(None) - } - - /// Non-consumingly checks whether a named raw-series preflight is supported. - /// - /// Returns the approved binding and ordering on support, or `None` when the plan is - /// explicitly unsupported. Errors propagate unchanged. This point-in-time check does not - /// reserve a plan; execution revalidates support. - async fn prepare_named_raw_series( - &self, - _plan: LogicalPlan, - _query_ctx: QueryContextRef, - ) -> Result> { - Ok(None) - } - - /// Execute a native single-region raw-series scan. - async fn execute_raw_series( - &self, - _plan: LogicalPlan, - _query_ctx: QueryContextRef, - _binding: RawSeriesRequest, - ) -> Result { - crate::error::UnimplementedSnafu { - operation: "native raw-series query", - } - .fail() - } - /// Register an aggregate function. /// /// # Panics diff --git a/src/query/src/region_query.rs b/src/query/src/region_query.rs index 96c843893bb..4d2dfb68d60 100644 --- a/src/query/src/region_query.rs +++ b/src/query/src/region_query.rs @@ -18,7 +18,7 @@ use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; use async_trait::async_trait; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; -use common_query::request::{PreparedNamedRawSeries, QueryRequest}; +use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use partition::manager::PartitionRuleManagerRef; use session::ReadPreference; @@ -68,20 +68,6 @@ pub trait RegionQueryHandler: Send + Sync { request: QueryRequest, ) -> Result; - /// Non-consumingly checks whether a named raw-series preflight is supported. - /// - /// Returns the approved binding and ordering on support, or `None` when the request is - /// explicitly unsupported; failures that are not unsupported preparation are returned as - /// errors. This point-in-time check does not reserve a plan, and execution revalidates - /// support. - async fn prepare_named_raw_series_query( - &self, - _target: &RegionQueryTarget, - _request: QueryRequest, - ) -> Result> { - Ok(None) - } - async fn handle_remote_dyn_filter_update( &self, target: &RegionQueryTarget, diff --git a/tests/compatibility/wire/named_raw_series_prepare/README.md b/tests/compatibility/wire/named_raw_series_prepare/README.md deleted file mode 100644 index 0175bdbeb3d..00000000000 --- a/tests/compatibility/wire/named_raw_series_prepare/README.md +++ /dev/null @@ -1,32 +0,0 @@ -# Named raw-series preflight contract v1 - -`v1_supported.json` and `v1_unsupported.json` are hand-authored, frozen JSON -bodies of the named raw-series preflight action -`greptime.prepare_raw_series.v1` (`PREPARE_RAW_SERIES_ACTION`). - -- `v1_supported.json` approves the request with the `contiguous` ordering. Its - nested `binding` value is the same version-1 request frozen in - [`../raw_series_request/v1_request.json`](../raw_series_request/v1_request.json) - (identical JSON text; the envelope embeds it without the source file's - trailing newline). -- `v1_unsupported.json` is the frozen unsupported shape, exactly - `{"version":1,"supported":null}`: the version-1 envelope with an explicit - `null` `supported` field. - -The envelope requires `version` (only `1` is supported) and a present-but-nullable -`supported` field. An approving `supported` object rejects unknown fields and -carries exactly the existing raw-series request plus an ordering of `contiguous` -or `read_key_then_time`; `interleaved`, unknown or missing orderings, duplicate -keys, unknown fields, and malformed bindings are rejected. The nested request -keeps its existing deserializer and validation rules; the nested physical -metadata is not checked for unknown fields by this contract. - -Run the focused check: - -```shell -cargo nextest run -p common-query named_raw_series -``` - -The fixture proves the version-1 envelope decodes as written and that the -unsupported shape is frozen. It does not prove interoperability with an old -binary; it records the contract this revision writes and reads. diff --git a/tests/compatibility/wire/named_raw_series_prepare/v1_supported.json b/tests/compatibility/wire/named_raw_series_prepare/v1_supported.json deleted file mode 100644 index e201cebe3b0..00000000000 --- a/tests/compatibility/wire/named_raw_series_prepare/v1_supported.json +++ /dev/null @@ -1 +0,0 @@ -{"version":1,"supported":{"binding":{"version":1,"physical_metadata":{"column_metadatas":[{"column_schema":{"name":"host","data_type":{"Int64":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Tag","column_id":1},{"column_schema":{"name":"value","data_type":{"Int32":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Field","column_id":2},{"column_schema":{"name":"aux","data_type":{"Int64":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Field","column_id":3},{"column_schema":{"name":"ts","data_type":{"Timestamp":{"Millisecond":null}},"is_nullable":false,"is_time_index":true,"default_constraint":null,"metadata":{}},"semantic_type":"Timestamp","column_id":4}],"primary_key":[1],"region_id":4294967297,"schema_version":0,"primary_key_encoding":"dense"},"metric_name":"cpu","logical_tags":[{"column_id":1,"name":"host"}],"field_ids":[3,2]},"ordering":"contiguous"}} diff --git a/tests/compatibility/wire/named_raw_series_prepare/v1_unsupported.json b/tests/compatibility/wire/named_raw_series_prepare/v1_unsupported.json deleted file mode 100644 index cd2aeccb8a6..00000000000 --- a/tests/compatibility/wire/named_raw_series_prepare/v1_unsupported.json +++ /dev/null @@ -1 +0,0 @@ -{"version":1,"supported":null} From ee54ded60fce10cb9295383735aedb63a72a1cfc Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 17:24:07 +0800 Subject: [PATCH 08/13] docs(promql): state the real reason an instant selection needs no ordering check (#9785) * docs(promql): state the real reason an instant selection needs no ordering check `select_instants` reads one series' complete axis, so the executor must collect a series to its last chunk before evaluating it: a single buffer closed by `is_last` and keyed by nothing else. That needs each series' chunks contiguous, which `Contiguous` and `KeyThenTime` both guarantee, and `Interleaved` never reaches it because `unpack_raw_series` rejects that ordering at the raw-frame boundary. The comment said the opposite: that interleaved input was enough because `select_instants` holds no cross-series state. The conclusion stood, but the stated reason named a property the executor does not have. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): stop deriving the instant arm's ordering check from another path The comment justified skipping an ordering check here by pointing at `unpack_raw_series`, which rejects `Interleaved` at the raw-frame boundary. This validator never runs through that call, so the rejection was never a premise it could use, and the conclusion did not follow. The requirement it was reaching for does exist, one layer down: `NativeInstantManipulateExec` collects a series to its `is_last` chunk into one buffer keyed by nothing else, so it needs each series' chunks contiguous. This validator does not check that. `Instant` has no implementation in the native tree yet and `try_to_native_execution_plan` has no callers, so the requirement is unenforced rather than met in advance; it is settled when `Instant` lands. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): refuse an interleaved scan at the compiler entry, and in the validator The `Instant` validation arm replaced an ordering check with a comment saying no check was needed. The argument was that `unpack_raw_series` rejects `Interleaved`, which holds only for the legacy `InstantManipulate` path and says nothing about this tree. Meanwhile the refusal in `compile` sat after the bare-selector early return, so `compile(eval_stmt("x"), schema, Interleaved)` returned `Ok` against `compile`'s own rustdoc, and the same hole was open to every nested shape. The requirement is real for every shape: a native evaluation reads one series to its last chunk before the next. The refusal now runs before `compile` dispatches on the shape, so the bare selector, the range function and its parameterised form all inherit it, and a shape added later inherits it too. `NativeNode::validate` keeps a check of its own, for a tree built some other way: `Instant` and `Window` now share `ensure_not_interleaved` instead of one branch having the check and the other a comment about it. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): link the validator by the node that has it `NativeTree` has no `validate`; the method lives on `NativeNode`. Both doc comments linked the tree, so rustdoc reports an unresolved link and the published API docs lose the reference. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): assert the control records the scan's ordering The contiguous control in the interleaved-refusal test only checked that the bare selector compiled, while its comment said the ordering is recorded. Assert it on the tree's input instead. Drop the plain range case from the loop, which the ordering table already compiles under `Interleaved`. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/native/plan.rs | 38 ++++++++++------ src/promql/src/native/plan/tests.rs | 48 ++++++++++++++++---- src/query/src/promql/native_compile.rs | 39 +++++++++------- src/query/src/promql/native_compile/tests.rs | 48 ++++++++++++++++++-- 4 files changed, 129 insertions(+), 44 deletions(-) diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 6cca42134f3..f5715b21f72 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -122,8 +122,8 @@ pub struct AggregateStateSchema { /// A scan or exchange boundary bound as a node input. /// /// The binding records what the input was created for: the output a consumer type-checks against, -/// and the ordering it guarantees. A window reads the ordering, rejecting an interleaved input; no -/// other consumer reads it yet. +/// and the ordering it guarantees. An instant selection and a window both read the ordering, +/// rejecting an interleaved input; no other consumer reads it yet. #[derive(Clone, Debug)] pub struct BoundInput { /// What the input produces: raw samples from a scan, or state from an exchange boundary. @@ -329,14 +329,13 @@ impl NativeNode { ))); } } + ensure_not_interleaved(input)?; ensure_samples(&self.output, "an instant selection")?; // A selection is one source row per grid point: `select_instants` picks row // indices and every field is taken at them, so it carries one value per input // field. ensure_same_fields(&input.output, &self.output, "an instant selection")?; ensure_grid(&self.output, &op.output_grid, "an instant selection")?; - // `select_instants` reads one series at a time and holds no cross-series state, so - // an interleaved input is enough for it: no ordering check here. } NativeExpr::Window { input, op } => { input.validate()?; @@ -354,17 +353,7 @@ impl NativeNode { "a window reads an empty range: its range must be positive", )); } - // A range evaluator consumes one series to its last chunk before it starts the - // next, so interleaved chunks of different series cannot be fed to it. Only an - // input node carries an ordering, so only there is there one to check. - if let NativeExpr::Input(binding) = &input.expr - && binding.ordering == SeriesOrdering::Interleaved - { - return Err(invalid( - "a window reads one series to its last chunk before the next, so \ - interleaved chunks cannot be fed to it", - )); - } + ensure_not_interleaved(input)?; // What is checked here and just below is that the declared input grid equals the // input's grid, and that the output grid matches the declared one. The grid is // not checked to cover the window's read range @@ -624,6 +613,25 @@ fn ensure_grid(ty: &NativeType, expected: &Arc, what: &str) -> Result< } } +/// Checks that `input`, when it is a [`NativeExpr::Input`], does not declare interleaved series. +/// +/// A native evaluation reads one series to its last chunk before the next: a range evaluator folds +/// a series completely before it starts another, and an instant selection's consumer collects one +/// series' chunks into one buffer. Neither can be handed chunks of different series that +/// interleave, which is the one rule both consumers state. Only an input node carries an ordering, +/// so only there is there one to check. +fn ensure_not_interleaved(input: &NativeNode) -> Result<()> { + if let NativeExpr::Input(binding) = &input.expr + && binding.ordering == SeriesOrdering::Interleaved + { + return Err(invalid( + "a native evaluation reads one series to its last chunk before the next, so \ + interleaved chunks cannot be fed to it", + )); + } + Ok(()) +} + fn invalid(reason: &str) -> DataFusionError { // A rejected tree is a compiler defect, not a query error: the compiler only builds a node // after deciding that this combination is supported. `Internal` keeps that meaning on the wire diff --git a/src/promql/src/native/plan/tests.rs b/src/promql/src/native/plan/tests.rs index 1c42c3b32af..ba01eb53cdd 100644 --- a/src/promql/src/native/plan/tests.rs +++ b/src/promql/src/native/plan/tests.rs @@ -96,6 +96,18 @@ fn input_with_fields(count: usize) -> NativeNode { } } +/// A scan of raw samples whose chunks of different series may interleave. +fn interleaved_input() -> NativeNode { + let output = raw(); + NativeNode { + expr: NativeExpr::Input(BoundInput { + output: output.clone(), + ordering: SeriesOrdering::Interleaved, + }), + output, + } +} + /// A constant scalar, which is a scalar type evaluated on the statement's grid. fn scalar_on(value: f64, grid: Arc) -> NativeNode { NativeNode { @@ -457,15 +469,35 @@ fn a_window_checks_its_declared_input_grid() { /// chunks of different series may not interleave. #[test] fn a_window_cannot_read_interleaved_input() { - let output = raw(); - let interleaved = NativeNode { - expr: NativeExpr::Input(BoundInput { - output: output.clone(), - ordering: SeriesOrdering::Interleaved, - }), - output, + let error = window_over(interleaved_input(), None) + .validate() + .unwrap_err(); + assert!( + error + .to_string() + .contains("interleaved chunks cannot be fed to it"), + "{error}" + ); +} + +/// An instant selection's consumer collects one series' chunks into one buffer, so it reads a +/// series to its last chunk before the next: it refuses an interleaved input for the reason a +/// window does. +#[test] +fn an_instant_selection_cannot_read_interleaved_input() { + let out = grid(0, 60_000, 4); + let node = NativeNode { + expr: NativeExpr::Instant { + input: Box::new(interleaved_input()), + op: BoundInstant { + lookback_ms: 300_000, + offset_ms: 0, + output_grid: Arc::clone(&out), + }, + }, + output: evaluated(Arc::new([]), out), }; - let error = window_over(interleaved, None).validate().unwrap_err(); + let error = node.validate().unwrap_err(); assert!( error .to_string() diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 88e5023aa42..1be936d7980 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -34,7 +34,9 @@ //! The statement's grid, the selector's lookback, window, offset and anchor are the only numbers //! derived here; the legality of the tree is not checked here. [`NativeTree::try_new`] is the //! single authority on that, so a node added later cannot be legal in one place and illegal in -//! another. +//! another. What is refused before a tree exists is the scan's shape: no shape below can consume +//! chunks of different series that interleave, so an `Interleaved` ordering is refused once, at the +//! entry, and [`NativeNode::validate`] refuses the same ordering for a tree built some other way. use std::sync::Arc; use std::time::Duration; @@ -72,12 +74,12 @@ const SUPPORTED: &str = "a bare selector, a range function over a selector — \ /// one. /// /// `ordering` is what the caller's scan guarantees, and the tree records it rather than assuming -/// one. Only `Contiguous` and `KeyThenTime` are accepted: a native range window consumes one series -/// to its last chunk before the next, so chunks of different series may not interleave, and an -/// `Interleaved` declaration is refused here — which is where the refusal belongs, since a tree that -/// declares it under a window cannot be built by this compiler at all. [`NativeTree::validate`]'s -/// window branch checks the same thing, so if both land this refusal comes first and the tree check -/// is a backstop for a tree built some other way. +/// one. Only `Contiguous` and `KeyThenTime` are accepted: every shape this compiler accepts reads +/// one series to its last chunk before the next, so an `Interleaved` declaration is refused at the +/// entry, before any shape is built — the bare selector included, which returns before a window +/// exists — and a shape added later inherits the refusal instead of restating it. +/// [`NativeNode::validate`] refuses the same ordering in its `Instant` and `Window` branches, as a +/// backstop for a tree built some other way. /// /// The tree carries no metric name and no matchers, so `rate(x{job="a"}[5m])` and `rate(y[5m])` /// compile to identical trees, and an `EXPLAIN` built on this display cannot tell them apart. That @@ -88,6 +90,17 @@ pub fn compile( input: Arc, ordering: SeriesOrdering, ) -> Result { + // Every shape this compiler accepts reads one series to its last chunk before the next, so raw + // samples whose chunks of different series interleave cannot be fed to any of them. The refusal + // is here rather than in a shape's branch so that a shape added later inherits it, and it comes + // before the early return below for a bare selector. + if ordering == SeriesOrdering::Interleaved { + return Err(not_implemented( + "raw samples whose series interleave, which this compiler cannot consume: a native \ + evaluation reads one series to its last chunk before the next, so chunks of different \ + series may not interleave", + )); + } // A bare selector is the other shape: an instant selection over the same raw input. if let PromExpr::VectorSelector(selector) = unparen(&stmt.expr) { return instant_selection(stmt, selector, input, ordering); @@ -126,16 +139,8 @@ pub fn compile( func.name, argument ))); }; - // A native range window reads one series to its last chunk before it reads the next, so raw - // samples whose chunks of different series interleave cannot be folded by it. The declaration - // is refused here rather than recorded onto an input the tree is built over and fails on. - if ordering == SeriesOrdering::Interleaved { - return Err(not_implemented( - "raw samples whose series interleave, which this compiler cannot consume: its range \ - window reads one series to its last chunk before the next, so chunks of different \ - series may not interleave", - )); - } + // The ordering was refused at the entry above, so every shape below is built over a scan whose + // series' chunks do not interleave. window_selection(stmt, function, vs, range, input, ordering) } diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 5e375858111..05b4d07ead8 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -631,11 +631,11 @@ fn the_numeric_edges_of_a_statement_are_bound_as_given() { SeriesOrdering::KeyThenTime, Expected::Rejected(Rejection::Plan), ), - // A range window reads one series to its last chunk before the next, so a scan whose - // chunks of different series interleave is refused instead of recorded onto an input the - // window would fail on. + // Every shape reads one series to its last chunk before the next, so a scan whose chunks + // of different series interleave is refused instead of recorded onto an input that shape + // would fail on. ( - "an interleaved ordering, which a window cannot consume to its last chunk", + "an interleaved ordering, which no shape can consume to its last chunk", eval_stmt("rate(x[5m])"), SeriesOrdering::Interleaved, Expected::Rejected(Rejection::NotImplemented), @@ -707,6 +707,46 @@ fn the_numeric_edges_of_a_statement_are_bound_as_given() { } } +/// An interleaved scan is refused whatever shape the statement has. The refusal sits at the entry +/// of the compiler rather than in one shape's branch, because a bare selector returns before any +/// window exists and would otherwise be accepted with an ordering it cannot be evaluated on. +#[test] +fn an_interleaved_scan_is_refused_whatever_the_shape() { + // The plain range function is the ordering table's case, which compiles `rate(x[5m])` under + // `Interleaved` and expects the same refusal; these two cover the shapes that table cannot, + // because one returns before a window exists and the other takes a parameter. + for query in ["x", "quantile_over_time(0.9, x[5m])"] { + let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::Interleaved) + .expect_err(&format!("`{query}` compiled over an interleaved scan")); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "`{query}` did not report NotImplemented: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("reads one series to its last chunk before the next"), + "`{query}` did not say why the ordering is refused: {message}" + ); + } + // The same scan compiles a bare selector when its chunks are contiguous, and the ordering it + // was compiled under is recorded on the tree's input rather than dropped: that is what makes + // the refusals above about the ordering rather than about the shape. + let tree = compile(&eval_stmt("x"), raw_schema(), SeriesOrdering::Contiguous) + .expect("a contiguous scan compiles a bare selector"); + let (input, _) = root_instant(&tree); + let NativeExpr::Input(bound) = &input.expr else { + panic!( + "an instant selection's input is not an input: {:?}", + input.expr + ); + }; + assert_eq!( + bound.ordering, + SeriesOrdering::Contiguous, + "the instant selection dropped the scan's ordering" + ); +} + /// A shape outside the subset is refused, and the error says what was found so the query can be /// rewritten rather than guessed at. #[test] From e31b4884ade7b85c4fa837da0adff46eb0606a6a Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 18:47:18 +0800 Subject: [PATCH 09/13] fix(promql): bound an instant selection by the memory pool, not a point count (#9801) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(promql): bound an instant selection by the memory pool, not a point count `MAX_INSTANT_MANIPULATE_OUTPUT_POINTS` capped an instant selection at a million grid points per series. It arrived in a commit whose message was only "cap max points": no doc, no derivation, no test of the bound itself. It bounded neither the output (series count is unbounded) nor the payload (a histogram row is kilobytes), and because the comparison was `>`, exactly a million points passed. `InstantManipulateExec` — the operator this path actually runs — had no memory accounting at all, so deleting the count alone would have turned a refusal into an unbounded allocation. It now registers one reservation per partition and grows it by the workspace the selection is about to allocate, before it allocates: the two vectors `select_instants` returns, the copy of the indices `take` reads with, the output time column, and one row per point per column — where a histogram column is charged the widest row it holds, because one source row can be selected by every point of the grid and is copied for each of them. The reservation is released once the batch is built, since from then on it belongs to whoever reads it. The native operators keep their own reservations; `NativeInstantManipulateExec`'s existing per-point term already covered the same three arrays, and this only says so where it is computed. `select_instants` no longer refuses anything: it documents that it allocates 16 bytes per selected point and reserves none of that, which is now every caller's job. A query over a longer range therefore succeeds when the pool can hold it and is refused with `ResourcesExhausted` when it cannot, instead of being refused at a count that says nothing about either. One bound is gone with the count and is not replaced here: a resolved `query.memory_pool_size` of zero leaves DataFusion's unbounded pool attached, so an operator that leaves that option unset has nothing refusing on this path. The operator doc says so. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): reserve an instant selection for the points it clips to The workspace estimate charged a series for the whole planned grid, while select_instants clips the range to the series' first and last sample and its lookback. A week-long query at a one-minute step therefore reserved ~413 KB for a series that produces five points, and a small pool refused it before selecting anything. Reserve for the clipped candidate count instead, derived with the same alignment select_instants uses; a series that does span the grid is still charged for it. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): keep one clipping, shared by the selection and its reservation The reservation and `manipulate` carried the same clipping arithmetic twice, with a comment asking the two copies to stay in line. They cannot drift now: `clipped_range` is the only copy, and both call it. No behaviour change. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/extension_plan/instant_manipulate.rs | 561 ++++++++++++++++-- .../instant_manipulate/native.rs | 11 +- .../instant_manipulate/native_exec.rs | 7 + .../instant_manipulate/typed_tests.rs | 8 + src/promql/src/instant_eval.rs | 13 +- src/promql/src/instant_eval/tests.rs | 15 +- 6 files changed, 530 insertions(+), 85 deletions(-) diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 85b08f5ca41..255d51be2c4 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -20,7 +20,10 @@ use std::sync::Arc; use std::task::{Context, Poll}; use common_query::prelude::{greptime_native_histogram, greptime_value}; -use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray, UInt64Array}; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, LargeListArray, LargeStringArray, ListArray, StringArray, + StructArray, TimestampMillisecondArray, UInt64Array, +}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -28,6 +31,7 @@ use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DFSchema, DFSchemaRef, ScalarValue}; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::logical_expr::{ EmptyRelation, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, }; @@ -58,7 +62,7 @@ use crate::extension_plan::{ nanoseconds_per_native_tick, prometheus_stale_sample_column, resolve_column_name, serialize_column_index, timestamp_unit, }; -use crate::instant_eval::{MAX_INSTANT_MANIPULATE_OUTPUT_POINTS, select_instants}; +use crate::instant_eval::select_instants; use crate::metrics::PROMQL_SERIES_COUNT; fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] { @@ -75,6 +79,15 @@ fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] { /// This plan will try to align the input time series, for every timestamp between /// `start` and `end` with step `interval`. Find in the `lookback` range if data /// is missing at the given timestamp. +/// +/// The memory pool bounds one instant query, not a point count: the execution reserves the +/// workspace its selection allocates before allocating it, and a pool that cannot meet the +/// reservation refuses the query with `ResourcesExhausted` (see +/// [`InstantManipulateStream::manipulate_reserved`]). That bound exists only where the pool does: +/// a resolved `query.memory_pool_size` of zero leaves DataFusion's unbounded pool attached, whose +/// growth always succeeds, so an operator that leaves the option unset has no bound on this path +/// at all. What one query can cost still grows with the number of series, which no reservation of +/// a single series accounts for. #[derive(Debug, PartialEq, Eq, Hash)] pub struct InstantManipulate { start: Millisecond, @@ -530,6 +543,10 @@ impl ExecutionPlan for InstantManipulateExec { count: num_series.clone(), }); + // The reservation is the operator's own, registered before the child is polled so that a + // pool which cannot hold one series' workspace refuses it rather than the child's memory. + let reservation = + MemoryConsumer::new("InstantManipulateStream").register(context.memory_pool()); let input = self.input.execute(partition, context)?; let schema = input.schema(); let time_index = schema @@ -559,6 +576,7 @@ impl ExecutionPlan for InstantManipulateExec { input, metric: baseline_metric, num_series, + reservation, })) } @@ -638,6 +656,8 @@ pub struct InstantManipulateStream { metric: BaselineMetrics, /// Number of series processed. num_series: Count, + /// What the workspace of the selection being built holds, reserved before it is allocated. + reservation: MemoryReservation, } impl RecordBatchStream for InstantManipulateStream { @@ -656,7 +676,7 @@ impl Stream for InstantManipulateStream { if batch.num_rows() != 0 { self.num_series.add(1); } - let result = Ok(batch).and_then(|batch| self.manipulate(batch)); + let result = self.manipulate_reserved(batch); self.metric.elapsed_compute().add_elapsed(timer); Poll::Ready(Some(result)) } @@ -671,6 +691,24 @@ impl Stream for InstantManipulateStream { } impl InstantManipulateStream { + /// Manipulates one complete series with the workspace the selection allocates reserved first. + /// + /// One grid point of the selection is decided per selected sample position, so the workspace is + /// proportional to the points of the clipped grid rather than to the samples of the series: the + /// two arrays [`select_instants`] returns, the take indices the batch is built with, and one + /// value per point for every column of the batch. That is what the point-count limit this + /// operator used to carry stood for, and [`Self::workspace_bytes`] derives it term by term. + /// + /// The reservation goes back as soon as the batch is built, because what it covers is this + /// operator's workspace and not its output: the batch belongs to whoever reads it from here and + /// is charged to the pool there, as the output of any operator is. + fn manipulate_reserved(&mut self, input: RecordBatch) -> DataFusionResult { + self.reservation.try_grow(self.workspace_bytes(&input))?; + let result = self.manipulate(input); + self.reservation.free(); + result + } + /// Manipulates one complete series sorted by timestamp. The planner enforces /// this input contract with a sort followed by [`SeriesDivide`]. /// @@ -722,7 +760,6 @@ impl InstantManipulateStream { return self.take_record_batch_optional(input, indices, evaluation_timestamps); } } - let scale = nanoseconds_per_native_tick(self.time_unit); let stale_sample_columns = self.field_indices.map(|index| { index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref())) }); @@ -738,48 +775,8 @@ impl InstantManipulateStream { let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let len = timestamps.len(); - // Shift the native-tick timeline in i128 before comparing samples. Doing - // this in the Arrow storage unit can overflow even when the shifted - // PromQL millisecond evaluation time is representable. - let to_nanoseconds = - |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000; - let first_ns = to_nanoseconds(timestamps[0]); - let last_ns = to_nanoseconds(timestamps[len - 1]); - // An exact sample remains useful with zero lookback. Otherwise the lower - // boundary is exclusive, so subtract one nanosecond from its final window. - let last_useful = if self.lookback_delta == 0 { - last_ns - } else { - last_ns + (self.lookback_delta as i128) * 1_000_000 - 1 - }; - let first_ms = (first_ns + 999_999).div_euclid(1_000_000); - let last_ms = last_useful.div_euclid(1_000_000); - let query_start = self.start as i128; - let query_end = self.end as i128; - let interval = self.interval as i128; - let max_start = first_ms.max(query_start); - let min_end = last_ms.min(query_end); - let (aligned_start, aligned_end) = if max_start > min_end { - (1, 0) - } else { - ( - query_start + (max_start - query_start) / interval * interval, - query_end - (query_end - min_end) / interval * interval, - ) - }; - let estimated_points = if aligned_end >= aligned_start { - (aligned_end - aligned_start) / interval + 1 - } else { - 0 - }; - if estimated_points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 { - return Err(DataFusionError::Execution(format!( - "InstantManipulate output points exceed limit: {estimated_points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}" - ))); - } - let estimated_points = estimated_points as usize; - let aligned_start = aligned_start as i64; - let aligned_end = aligned_end as i64; + let (aligned_start, aligned_end) = self.clipped_range(timestamps); + let estimated_points = self.range_points(aligned_start, aligned_end); // Narrow `i64` nanoseconds whenever the batch fits them, exact `i128` otherwise. let narrow = (timestamp_null_count == 0) .then(|| { @@ -794,8 +791,19 @@ impl InstantManipulateStream { }) .flatten(); if let Some(scan) = &narrow { - debug_assert_eq!(i64::try_from(first_ns), Ok(scan.shift(timestamps[0]))); - debug_assert_eq!(i64::try_from(last_ns), Ok(scan.shift(timestamps[len - 1]))); + // The clip converted the series bounds to milliseconds; this asserts the scan the + // selection walks agrees with them on those two samples. + let scale = nanoseconds_per_native_tick(self.time_unit); + let to_nanoseconds = + |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000; + debug_assert_eq!( + i64::try_from(to_nanoseconds(timestamps[0])), + Ok(scan.shift(timestamps[0])) + ); + debug_assert_eq!( + i64::try_from(to_nanoseconds(timestamps[len - 1])), + Ok(scan.shift(timestamps[len - 1])) + ); } let (take_indices, aligned_ts) = match narrow { Some(scan) => self.scan_instants::( @@ -895,6 +903,215 @@ impl InstantManipulateStream { .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; Ok(result) } + + /// What one selected series' workspace holds, in bytes, before it is allocated. + /// + /// The figure follows the reasoning of `InstantRuntime::gather_bytes` in `native::execute` and + /// of `NativeInstantStream::finish_series` in `native_exec`: a selection decides one point of + /// its grid per sample position it reads, so what it allocates grows with the points of the + /// clipped grid and not with the samples of the series. Each term below is one allocation, and + /// the ones this path shares with those are named as they are there. + /// + /// Per selected point, [`Self::clipped_points`] points being the most it can decide: + /// + /// - A `u64` source row index and an `i64` evaluation timestamp: the two vectors + /// [`select_instants`] returns, and the scan of [`Self::scan_instants`] returns the same pair. + /// This is that node's "two arrays [`select_instants`] returns" term. + /// - The same indices again, for the `UInt64Array` [`Self::take_record_batch_optional`] takes + /// the column values with: `take_indices.clone()` is a second copy of them. + /// - The timestamps again, as an allowance for the output's time column: they are moved into it + /// when they fit and copied when they do not, and a move costs nothing to bound. + /// + /// Per column, one row per point: `take` allocates a row of the column for every point, so a + /// column costs its widest row per point -- the same term as that node's "one `f64` per field" + /// -- and one byte per point for the validity buffer it allocates, charged a byte for a bit so + /// that the figure stays an upper bound. A reused TSID column is sliced instead of taken + /// whenever the output is no longer than it and materialized a row per point otherwise, so the + /// same row bounds both. [`row_bytes`] gives the width, and a column of a type the series + /// frames of this path do not carry is charged for the bytes it already occupies once rather + /// than as a row per point. + /// + /// Per native histogram column, the widest row is not one value: a histogram row is a struct of + /// scalar columns whose values are copied and bucket lists whose values are copied for the row + /// the point selects, so a point costs the scalars, the list offsets and the widest payload the + /// column holds. That node charges its histogram payload the same way, "at the largest one the + /// field holds and per *point* rather than per row", because one source row can be selected by + /// every point of the grid and still be copied for each of them. + /// + /// Not charged, deliberately: the headers this selection builds per batch rather than per point + /// (the sample iterators and `Samples` clones of the fast path are constant per field), the + /// `Arc` bumps of a shared column, and the input batch, which belongs to the child. No term + /// doubles: every vector here is created with its exact capacity and `take` sizes its output by + /// the point count, unlike the builders that node charges twice. + fn workspace_bytes(&self, input: &RecordBatch) -> usize { + let points = self.clipped_points(input); + let mut bytes = points.saturating_mul(4 * std::mem::size_of::()); + for (index, column) in input.columns().iter().enumerate() { + // The time column of the output is built from the timestamps counted above. + if index == self.time_index { + continue; + } + bytes = bytes.saturating_add(match row_bytes(column.as_ref()) { + Some(width) => points.saturating_mul(width.saturating_add(1)), + None => column.get_buffer_memory_size(), + }); + } + // The `Arc<[ArrayRef]>` this batch is assembled from, one pointer per column. + bytes.saturating_add( + input + .num_columns() + .saturating_mul(std::mem::size_of::()), + ) + } + + /// The points [`select_instants`] decides for the series `input` holds. + /// + /// That selection clips this operator's `start..=end` to the series it reads before it walks + /// the grid, so the points it allocates for are the ones of the clipped range and not the ones + /// of the whole planned grid: a series whose samples end long before the query does reaches no + /// later point, yet stepping the whole planned range by the interval reserved the workspace of + /// every point of it, and a pool that could hold everything such a series selects refused the + /// query before selecting anything. + /// + /// The bounds, the clip and the alignment below are [`select_instants`]'s, kept the same line + /// for line because the reservation must cover what that function allocates: a point it misses + /// is memory no reservation bounds, and a point it invents is memory reserved for nothing. + /// + /// This batch holds one complete series -- [`SeriesDivide`] upstream guarantees it -- so the + /// first and last timestamp of its time column are the series bounds and reading them is + /// constant work, no matter how many samples the series holds. + fn clipped_points(&self, input: &RecordBatch) -> usize { + if self.end < self.start || self.interval <= 0 { + return 0; + } + let time = input.column(self.time_index); + if time.is_empty() { + return 0; + } + let Some((timestamps, _)) = timestamp_array_to_primitive(time) else { + return 0; + }; + let (start, end) = self.clipped_range(timestamps.values()); + self.range_points(start, end) + } + + /// This operator's grid clipped to the instants one series is selected at. + /// + /// `ticks` are that series' sample times in the operator's native tick unit, sorted and + /// non-empty, and the interval must be positive. The lower bound moves to the first instant at + /// or after the first sample and the upper bound to the last instant at or before the last + /// sample widened by the lookback, both rounded onto the grid. An instant outside the returned + /// range has no sample to select, so [`Self::scan_instants`] never walks it and nothing is + /// allocated for it. Returns `(1, 0)` -- an empty range -- when nothing is left. + /// + /// [`Self::manipulate`] and [`Self::clipped_points`] both call this. They used to carry the + /// same arithmetic twice, on the reasoning that the reservation must mirror the selection; one + /// copy is the stronger guarantee, because the two cannot drift apart at all. + fn clipped_range(&self, ticks: &[i64]) -> (Millisecond, Millisecond) { + // Shift the native-tick timeline in i128 before comparing samples. Doing this in the Arrow + // storage unit can overflow even when the shifted PromQL millisecond evaluation time is + // representable. + let scale = nanoseconds_per_native_tick(self.time_unit); + let to_nanoseconds = + |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000; + let first_ns = to_nanoseconds(ticks[0]); + let last_ns = to_nanoseconds(ticks[ticks.len() - 1]); + // An exact sample remains useful with zero lookback. Otherwise the lower boundary is + // exclusive, so subtract one nanosecond from its final window. + let last_useful = if self.lookback_delta == 0 { + last_ns + } else { + last_ns + (self.lookback_delta as i128) * 1_000_000 - 1 + }; + let first_ms = (first_ns + 999_999).div_euclid(1_000_000); + let last_ms = last_useful.div_euclid(1_000_000); + let query_start = self.start as i128; + let query_end = self.end as i128; + let interval = self.interval as i128; + let max_start = first_ms.max(query_start); + let min_end = last_ms.min(query_end); + if max_start > min_end { + return (1, 0); + } + ( + (query_start + (max_start - query_start) / interval * interval) as i64, + (query_end - (query_end - min_end) / interval * interval) as i64, + ) + } + + /// The number of instants in `start..=end`, stepping by the interval. + fn range_points(&self, start: Millisecond, end: Millisecond) -> usize { + if end < start || self.interval <= 0 { + return 0; + } + usize::try_from((end as i128 - start as i128) / self.interval as i128 + 1) + .unwrap_or(usize::MAX) + } +} + +/// The bytes one selected point copies out of `column`, `None` for a column type this path does not +/// carry. +/// +/// `take` allocates one row of the column per point, so one point costs the widest row the column +/// holds: the width of its value for a fixed-width column, the longest value for a byte column, and +/// the whole row for a struct such as a native histogram -- its scalar columns, the offsets of its +/// list columns and the values those lists hold for the widest row of the column. +fn row_bytes(column: &dyn Array) -> Option { + if let Some(width) = column.data_type().primitive_width() { + return Some(width); + } + // `primitive_width` reports a `Boolean` as unknown, and its values are one byte wide. + if let DataType::Boolean = column.data_type() { + return Some(1); + } + if let Some(values) = column.as_any().downcast_ref::() { + return Some(values.offsets().lengths().max().unwrap_or(0)); + } + if let Some(values) = column.as_any().downcast_ref::() { + return Some(values.offsets().lengths().max().unwrap_or(0)); + } + let values = column.as_any().downcast_ref::()?; + Some( + values + .columns() + .iter() + .map(|child| match child.data_type() { + // One offset per row of a list column, and the values the widest row holds. + DataType::List(_) => { + std::mem::size_of::().saturating_add(widest_list_row(child)) + } + DataType::LargeList(_) => { + std::mem::size_of::().saturating_add(widest_list_row(child)) + } + // The scalars of a row are fixed width; a nested column counted at the width of the + // widest value this crate stores keeps the sum an upper bound. + _ => child + .data_type() + .primitive_width() + .unwrap_or(std::mem::size_of::()), + }) + .fold(0usize, usize::saturating_add), + ) +} + +/// The values the widest row of a list column holds, in bytes. +/// +/// A list column of a histogram row is one bucket side or one bucket list of it, so one row of it +/// is one list's values: the widest row times the width of a value bounds what a point copies. +fn widest_list_row(column: &dyn Array) -> usize { + let (widest, values) = if let Some(list) = column.as_any().downcast_ref::() { + (list.offsets().lengths().max().unwrap_or(0), list.values()) + } else if let Some(list) = column.as_any().downcast_ref::() { + (list.offsets().lengths().max().unwrap_or(0), list.values()) + } else { + return 0; + }; + widest.saturating_mul( + values + .data_type() + .primitive_width() + .unwrap_or(std::mem::size_of::()), + ) } fn reuse_constant_column(array: &Arc, len: usize) -> DataFusionResult> { @@ -925,6 +1142,8 @@ mod test { use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::logical_expr::{ EmptyRelation, Extension, LogicalPlan, Projection, UserDefinedLogicalNodeCore, }; @@ -937,6 +1156,21 @@ mod test { TIME_INDEX_COLUMN, native_histogram, prepare_test_data, prepare_test_data_with_stale_marker, }; + /// An unbounded reservation, for the streams built outside an operator's `execute`. + fn unbounded() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("InstantManipulateStreamTest").register(&pool) + } + + /// A task context whose pool holds `bytes`, so a reservation larger than that is refused. + fn bounded_task_context(bytes: usize) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(bytes))) + .build() + .unwrap(); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) + } + async fn do_normalize_test( start: Millisecond, end: Millisecond, @@ -1553,8 +1787,15 @@ mod test { ); } - #[tokio::test] - async fn manipulate_should_reject_too_many_output_points() { + /// One series of one `value` column holding `rows` of `(timestamp, value)` decided over + /// `start..=end` by `interval` with `lookback`. + fn one_series( + rows: Vec<(Millisecond, f64)>, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + lookback: Millisecond, + ) -> Arc { let schema = Arc::new(Schema::new(vec![ Field::new( TIME_INDEX_COLUMN, @@ -1563,24 +1804,24 @@ mod test { ), Field::new("value", DataType::Float64, true), ])); + let (timestamps, values): (Vec, Vec) = rows.into_iter().unzip(); let batch = RecordBatch::try_new( schema.clone(), vec![ - Arc::new(TimestampMillisecondArray::from(vec![0])), - Arc::new(Float64Array::from(vec![1.0])), + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(Float64Array::from(values)), ], ) .unwrap(); let input = Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); - let too_many_points = MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as Millisecond + 1; - let normalize_exec = Arc::new(InstantManipulateExec { + Arc::new(InstantManipulateExec { offset: 0, - start: 0, - end: too_many_points, - lookback_delta: too_many_points + 1, - interval: 1, + start, + end, + lookback_delta: lookback, + interval, time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, @@ -1588,15 +1829,200 @@ mod test { properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), - }); - let session_context = SessionContext::default(); - let err = datafusion::physical_plan::collect(normalize_exec, session_context.task_ctx()) - .await - .unwrap_err(); + }) + } + + /// A single row of one series selected at every point of a grid `points` long, with `lookback` + /// wide enough that the row answers all of them. + fn one_row_selected_everywhere(points: Millisecond) -> Arc { + one_series(vec![(0, 1.0)], 0, points, 1, points + 1) + } + + /// A grid of a million and one points used to be refused by a point count of the kernel's own. + /// What refuses it now is the pool, which cannot hold the workspace of the selection. + #[tokio::test] + async fn manipulate_reports_the_pool_that_cannot_hold_the_workspace() { + let points = 1_000_001; + let err = datafusion::physical_plan::collect( + one_row_selected_everywhere(points), + bounded_task_context(1024), + ) + .await + .unwrap_err(); assert!( - err.to_string() - .contains("InstantManipulate output points exceed limit") + matches!(err, DataFusionError::ResourcesExhausted(_)), + "the pool did not refuse the workspace: {err}" + ); + assert!( + !err.to_string().contains("exceed limit"), + "a point limit answered in the pool's place: {err}" + ); + } + + /// The same grid, past what the removed limit allowed, is decided against a pool that holds it. + #[tokio::test] + async fn manipulate_decides_a_grid_past_the_removed_point_limit() { + let points = 1_000_001; + let batches = datafusion::physical_plan::collect( + one_row_selected_everywhere(points), + // This series' one sample answers every point of the grid under that wide lookback, so + // its clipped grid is the whole grid. The estimate of `manipulate_reserved` for this + // schema is 41 bytes per point (four `u64`, one `f64` and a validity bit), so this + // holds the million-odd points it reserves for with room to spare. + bounded_task_context(64 * 1024 * 1024), + ) + .await + .unwrap(); + + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!( + rows, + points as usize + 1, + "the grid the removed limit refused was not decided" + ); + } + + /// A series that ends at its first sample is charged for the points it selects, not for the + /// grid of the whole query range. + /// + /// Seven days at a 60 s step is 10 081 grid points, and the whole-grid figure of this schema is + /// 41 bytes per point: 413 337 bytes, which a 256 KiB pool refuses. The series holds one sample + /// at `t=0` and its lookback is 5 min, so `select_instants` clips it to the six-point grid + /// `0..=300 s`, from which the five points `0, 60, ..., 240 s` are selected; the reservation + /// covers those six points -- 262 bytes -- and the pool holds it. Before the reservation + /// followed the clipped range, the query was refused before anything was selected. + #[tokio::test] + async fn series_that_ends_at_its_first_sample_is_not_charged_for_the_grid() { + let batches = datafusion::physical_plan::collect( + one_series( + vec![(0, 1.0)], + 0, + 7 * 24 * 60 * 60 * 1_000, + 60_000, + 5 * 60 * 1_000, + ), + bounded_task_context(256 * 1024), + ) + .await + .unwrap(); + + let timestamps: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("the output time column is in milliseconds") + .values() + .to_vec() + }) + .collect(); + assert_eq!(timestamps, vec![0, 60_000, 120_000, 180_000, 240_000]); + } + + /// A series whose samples span the whole query range is still charged for the whole grid, so + /// the clipped figure is a bound on what is selected and not a reservation of nothing. + #[tokio::test] + async fn series_that_spans_the_whole_grid_is_still_charged_for_it() { + let end = 7 * 24 * 60 * 60 * 1_000; + let spanning = || one_series(vec![(0, 1.0), (end, 2.0)], 0, end, 60_000, 5 * 60 * 1_000); + + let err = datafusion::physical_plan::collect(spanning(), bounded_task_context(256 * 1024)) + .await + .unwrap_err(); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "a pool below the clipped grid's workspace did not refuse it: {err}" + ); + + let batches = + datafusion::physical_plan::collect(spanning(), bounded_task_context(1024 * 1024)) + .await + .unwrap(); + // A sample answers only the points of its own lookback window, so the two samples decide + // six rows: the five of the first 5 min and the point of the second sample. The grid they + // reach, and the one the reservation covers, is the whole 10 081 points. + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!( + rows, 6, + "the series spanning the whole grid was not decided" + ); + } + + /// The figure `manipulate_reserved` grows is the clipped one: 262 bytes for a series that ends + /// at its first sample, 413 337 bytes for one that spans the same seven-day grid of 10 081 + /// points at a 60 s step. Per point it is the four `u64` of the two index vectors and the + /// timestamp allowance, the `f64` and validity bit `take` allocates for `value`, and one + /// `ArrayRef` per column. + #[test] + fn workspace_bytes_follows_the_clipped_grid_not_the_planned_one() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let end = 7 * 24 * 60 * 60 * 1_000; + let stream = InstantManipulateStream { + offset: 0, + start: 0, + end, + lookback_delta: 5 * 60 * 1_000, + interval: 60_000, + time_index: 0, + time_unit: TimeUnit::Millisecond, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: schema.clone(), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new( + vec![], + schema.clone(), + None, + ) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + reservation: unbounded(), + }; + let series = |timestamps: Vec| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.clone())), + Arc::new(Float64Array::from(vec![1.0; timestamps.len()])), + ], + ) + .unwrap() + }; + + // One sample at `t=0` under a 5 min lookback clips to the six points `0, 60, ..., 300 s`; + // the exclusive lower bound leaves five of them selected. + let short = series(vec![0]); + assert_eq!(stream.clipped_points(&short), 6); + assert_eq!( + stream.workspace_bytes(&short), + 6 * 32 + 6 * 9 + 2 * std::mem::size_of::() + ); + // Samples at both ends reach every point of the planned grid, and are charged for it. + let long = series(vec![0, end]); + assert_eq!(stream.clipped_points(&long), 10_081); + assert_eq!( + stream.workspace_bytes(&long), + 10_081 * 32 + 10_081 * 9 + 2 * std::mem::size_of::() + ); + // An input without samples selects nothing and reserves nothing beyond the batch's columns. + let empty = series(vec![]); + assert_eq!(stream.clipped_points(&empty), 0); + assert_eq!( + stream.workspace_bytes(&empty), + 2 * std::mem::size_of::() ); } @@ -1921,6 +2347,7 @@ mod test { ), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), }; let output = stream.manipulate(input).unwrap(); @@ -1988,6 +2415,7 @@ mod test { ), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), }; let output = stream.manipulate(input).unwrap(); @@ -2053,6 +2481,7 @@ mod test { input: Box::pin(input), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), }; let waker = noop_waker_ref(); let mut cx = Context::from_waker(waker); @@ -2136,6 +2565,7 @@ mod test { ), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), }; let output = stream.manipulate(input).unwrap(); @@ -2208,6 +2638,7 @@ mod test { ), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), }; let output = stream.manipulate(input).unwrap(); assert_eq!(output.num_rows(), 1); diff --git a/src/promql/src/extension_plan/instant_manipulate/native.rs b/src/promql/src/extension_plan/instant_manipulate/native.rs index f7df2a5feeb..ea2a28c2054 100644 --- a/src/promql/src/extension_plan/instant_manipulate/native.rs +++ b/src/promql/src/extension_plan/instant_manipulate/native.rs @@ -27,7 +27,7 @@ use datatypes::series::samples::{SampleRef, Samples, SeriesFields, SeriesSchema} mod tests; use crate::extension_plan::instant_manipulate::InstantManipulate; -use crate::instant_eval::{MAX_INSTANT_MANIPULATE_OUTPUT_POINTS, select_instants}; +use crate::instant_eval::select_instants; #[derive(Clone)] enum Column { @@ -160,6 +160,10 @@ impl NativeInstantMapping { } /// Estimates the clipped evaluation grid for complete source timestamp bounds. + /// + /// This count is what the caller reserves the workspace of a selection for before it selects, + /// so nothing but the reservation bounds the grid: a grid too long to index is refused as + /// exhausted rather than by a point limit of its own. pub(crate) fn estimated_points(&self, first: i64, last: i64) -> DataFusionResult { let first = first as i128; let last = last as i128; @@ -185,11 +189,6 @@ impl NativeInstantMapping { } else { 0 }; - if points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 { - return Err(DataFusionError::Execution(format!( - "InstantManipulate output points exceed limit: {points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}" - ))); - } usize::try_from(points).map_err(|_| { DataFusionError::ResourcesExhausted("native instant point count overflow".into()) }) diff --git a/src/promql/src/extension_plan/instant_manipulate/native_exec.rs b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs index 383d341b259..e2b41cb8aa4 100644 --- a/src/promql/src/extension_plan/instant_manipulate/native_exec.rs +++ b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs @@ -311,6 +311,13 @@ impl NativeInstantStream { n.checked_add(field_count.checked_mul(std::mem::size_of::>())?) }) .and_then(|n| { + // Per point of the clipped grid, the most `select_instants` and the mapping of its + // selection can allocate: 24 bytes per output column, the labels of the row, and + // 48 for the two vectors the selection returns and the time column it is mapped + // into -- a `u64` source row index, an `i64` evaluation timestamp and the `i64` + // of the output column, at 8 bytes each. The point count is what stands in for + // the point limit `estimated_points` used to carry, so the selection arrays are + // reserved for here rather than refused there. points .checked_mul( self.schema diff --git a/src/promql/src/extension_plan/instant_manipulate/typed_tests.rs b/src/promql/src/extension_plan/instant_manipulate/typed_tests.rs index f2949eb3ec1..97f97a095d4 100644 --- a/src/promql/src/extension_plan/instant_manipulate/typed_tests.rs +++ b/src/promql/src/extension_plan/instant_manipulate/typed_tests.rs @@ -22,6 +22,7 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::SessionContext; +use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::physical_plan::memory::MemoryStream; use datatypes::arrow::buffer::ScalarBuffer; use datatypes::series::raw::RawTimestamps; @@ -32,6 +33,12 @@ use crate::extension_plan::range_manipulate::ScanNanoseconds; use crate::extension_plan::test_util::TIME_INDEX_COLUMN; use crate::instant_eval::select_instants; +/// An unbounded reservation, for the streams built outside an operator's `execute`. +fn unbounded() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("InstantManipulateStreamTest").register(&pool) +} + fn stream( batch: &RecordBatch, start: i64, @@ -56,6 +63,7 @@ fn stream( input: Box::pin(MemoryStream::try_new(vec![], schema, None).unwrap()), metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), num_series: Count::new(), + reservation: unbounded(), } } diff --git a/src/promql/src/instant_eval.rs b/src/promql/src/instant_eval.rs index 956bfe93bd9..7c15d26ed69 100644 --- a/src/promql/src/instant_eval.rs +++ b/src/promql/src/instant_eval.rs @@ -22,13 +22,17 @@ use datatypes::series::samples::{SampleRef, SeriesFields}; #[cfg(test)] mod tests; -pub(crate) const MAX_INSTANT_MANIPULATE_OUTPUT_POINTS: usize = 1_000_000; - /// Selects source rows for an aligned PromQL instant-evaluation range. /// /// Returns original logical row indices and evaluation timestamps from a complete source axis /// of strictly increasing milliseconds. Any provided field's stale marker suppresses selection; /// missing samples are not stale and remain selectable. +/// +/// The two vectors it returns hold one `u64` source row index and one `i64` evaluation timestamp +/// per selected point, so what it allocates is 16 bytes per point of the clipped grid rather than +/// anything proportional to the samples of the axis. It reserves none of that: every caller grows +/// a memory reservation for those two vectors before it calls, which is what bounds an instant +/// selection now that the kernel carries no point limit of its own. pub fn select_instants( time: &RawTimestamps, fields: &SeriesFields, @@ -99,11 +103,6 @@ pub fn select_instants( } else { 0 }; - if estimated_points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 { - return Err(DataFusionError::Execution(format!( - "InstantManipulate output points exceed limit: {estimated_points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}" - ))); - } let mut indices = Vec::with_capacity(estimated_points as usize); let mut evaluation_timestamps = Vec::with_capacity(estimated_points as usize); let mut sample_iters: Vec<_> = fields.iter().map(|field| field.iter()).collect(); diff --git a/src/promql/src/instant_eval/tests.rs b/src/promql/src/instant_eval/tests.rs index c13b5fe21b4..9df63e91f52 100644 --- a/src/promql/src/instant_eval/tests.rs +++ b/src/promql/src/instant_eval/tests.rs @@ -190,7 +190,7 @@ fn stale_sample_in_any_field_suppresses_without_falling_back() { } #[test] -fn offset_extreme_timestamps_and_point_limit_are_checked_without_large_fixture() { +fn offset_extreme_timestamps_are_checked_without_large_fixture() { let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![Some(7.0)]))); assert_eq!( select(&[i64::MAX], &fields, i64::MAX..=i64::MAX, 1, 0, 0).unwrap(), @@ -209,18 +209,19 @@ fn offset_extreme_timestamps_and_point_limit_are_checked_without_large_fixture() (vec![0], vec![110]) ); + // A grid of a million and one points used to be refused by a point limit of the kernel's own. + // The kernel decides it now, and the reservation its caller grows before the call is what + // stands in for that limit; the stale marker still suppresses the points row zero would answer, + // so only the one sample after it is selected. let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ Some(f64::from_bits( common_query::prometheus::PROMETHEUS_STALE_NAN_BITS, )), Some(1.0), ]))); - let result = select(&[0, 1_000_000], &fields, 0..=1_000_000, 1, 0, 0); - let error = result.unwrap_err().to_string(); - assert!( - error.contains("output points exceed limit: 1000001 > 1000000"), - "{error}" - ); + let (rows, times) = select(&[0, 1_000_000], &fields, 0..=1_000_000, 1, 0, 0).unwrap(); + assert_eq!(rows, [1]); + assert_eq!(times, [1_000_000]); } #[test] From c87956467169e38e7509f9ba80105a1eb2886b4f Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 18:53:57 +0800 Subject: [PATCH 10/13] chore(common-query): drop the raw-series request extension transport (#9815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(query): drop the prepared raw-series RPC skeleton The raw-series preparation RPC was merged into `narthex` for a plan that never landed: no server or frontend dispatches `greptime.prepare_raw_series.v1`, and the three `Datanode` methods, `bind_raw_series_query` and the `QueryEngine` trio have no call sites and no implementors. `PreparedNamedRawSeries` had no user outside them. The wire envelope never shipped -- no tag contains the commit that added it -- so the compatibility fixtures go with it. The harness drives `wire/` fixtures from the Rust tests that own their context conversion rather than from SQLness discovery, and has no retirement mechanism to record instead. Left in place: `RawSeriesRequest` and its extension key, `raw_series_frame_ordering`, `unpack_raw_series` and `RawLabelSource`, all of which the surviving series-frame path still uses. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(common-query): drop the raw-series request extension transport The prepared raw-series RPC carried a `RawSeriesRequest` through a query-context extension, and with that RPC gone (`607de914464`) nothing in `src/` writes or reads the extension any more. What remains is the binding the table scan actually uses: `RawSeriesRequest`, `RawSeriesTag`, `bind`, `validate` and `same_encoding_identity`, unchanged. Deleted: the extension key, `take_raw_series_request`, and the two serialization methods that only existed to encode and decode that transport, plus the wire fixture whose only reader was the test for them. The fixture never shipped — the commit that introduced it is in no tag. The tests that covered `validate` are kept, with their rejections now asserted directly instead of through a JSON round-trip: an unsupported version, an empty field list, a duplicate field, and a non-field column id are all still refused. The binding test is untouched, including its order, identity and `Arc` assertions. `pub` items in a library crate never trip `dead_code`, so the compiler could not have found this; the greps above are the evidence. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/src/request.rs | 4 +- .../query/src/request/raw_series_request.rs | 129 ++++++++---------- .../wire/raw_series_request/README.md | 3 - .../wire/raw_series_request/v1_request.json | 1 - 4 files changed, 58 insertions(+), 79 deletions(-) delete mode 100644 tests/compatibility/wire/raw_series_request/README.md delete mode 100644 tests/compatibility/wire/raw_series_request/v1_request.json diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index c61ff41da7e..77971440447 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -16,9 +16,7 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; mod raw_series_request; -pub use self::raw_series_request::{ - RAW_SERIES_REQUEST_EXTENSION_KEY, RawSeriesRequest, RawSeriesTag, take_raw_series_request, -}; +pub use self::raw_series_request::{RawSeriesRequest, RawSeriesTag}; use std::cell::Cell; use std::sync::Arc; diff --git a/src/common/query/src/request/raw_series_request.rs b/src/common/query/src/request/raw_series_request.rs index 2631703088b..0dd02b5290d 100644 --- a/src/common/query/src/request/raw_series_request.rs +++ b/src/common/query/src/request/raw_series_request.rs @@ -16,7 +16,6 @@ use std::collections::HashSet; use std::sync::Arc; use api::v1::SemanticType; -use api::v1::region::RegionRequestHeader; use datafusion_common::{DataFusionError, Result as DataFusionResult}; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use serde::{Deserialize, Serialize}; @@ -24,9 +23,6 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_engine::{BoundTag, RawLabelSource}; use store_api::storage::ColumnId; -/// Reserved query-context extension containing an explicit internal raw-series request. -pub const RAW_SERIES_REQUEST_EXTENSION_KEY: &str = "greptime.raw_series_request"; - /// A logical label mapped to its physical region column. #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -54,21 +50,6 @@ pub struct RawSeriesRequest { } impl RawSeriesRequest { - /// Serializes the control after validating its version and field binding. - pub fn to_extension_value(&self) -> DataFusionResult { - self.validate()?; - serde_json::to_string(self) - .map_err(|_| DataFusionError::Plan("Invalid raw-series request".into())) - } - - /// Parses and validates a strict versioned control value. - pub fn from_extension_value(value: &str) -> DataFusionResult { - let request: Self = serde_json::from_str(value) - .map_err(|_| DataFusionError::Plan("Invalid raw-series request".into()))?; - request.validate()?; - Ok(request) - } - /// Binds expected output fields and labels to the actual opened scanner metadata. pub fn bind( &self, @@ -138,83 +119,87 @@ fn same_encoding_identity(left: &RegionMetadata, right: &RegionMetadata) -> bool && left.time_index_type() == right.time_index_type() } -/// Removes and parses the optional raw control from a mutable request header. -pub fn take_raw_series_request( - header: &mut Option, -) -> DataFusionResult> { - let Some(value) = header - .as_mut() - .and_then(|header| header.query_context.as_mut()) - .and_then(|context| context.extensions.remove(RAW_SERIES_REQUEST_EXTENSION_KEY)) - else { - return Ok(None); - }; - RawSeriesRequest::from_extension_value(&value).map(Some) -} - #[cfg(test)] mod tests { use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; use super::*; - #[test] - fn strict_versioned_extension_fixture_is_parseable() { - let fixture = include_str!( - "../../../../../tests/compatibility/wire/raw_series_request/v1_request.json" - ); - let request = RawSeriesRequest::from_extension_value(fixture).unwrap(); - assert_eq!(request.version, 1); - assert_eq!(request.metric_name, "cpu"); - assert_eq!(request.logical_tags[0].name, "host"); - assert_eq!(request.logical_tags[0].column_id, 1); - assert_eq!(request.field_ids, [3, 2]); - assert_eq!(request.physical_metadata.region_id, RegionId::new(1, 1)); - } - - #[test] - fn absent_marker_remains_an_ordinary_request() { - let mut header = None; - assert!(take_raw_series_request(&mut header).unwrap().is_none()); + /// Builds the physical metadata of a one-tag, two-field region. + fn fixture_metadata() -> RegionMetadata { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("aux", ConcreteDataType::int64_datatype(), false), + semantic_type: SemanticType::Field, + column_id: 3, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 4, + }) + .primary_key(vec![1]); + builder.build().unwrap() } + /// Builds a request asking for fields in descending column order. fn fixture_request() -> RawSeriesRequest { - let fixture = include_str!( - "../../../../../tests/compatibility/wire/raw_series_request/v1_request.json" - ); - RawSeriesRequest::from_extension_value(fixture).unwrap() + RawSeriesRequest { + version: 1, + physical_metadata: fixture_metadata(), + metric_name: "cpu".to_string(), + logical_tags: vec![RawSeriesTag { + column_id: 1, + name: "host".to_string(), + }], + field_ids: vec![3, 2], + } } #[test] - fn rejects_unknown_fields_and_versions() { - let mut version: serde_json::Value = serde_json::from_str(include_str!( - "../../../../../tests/compatibility/wire/raw_series_request/v1_request.json" - )) - .unwrap(); - version["version"] = serde_json::json!(2); - assert!(RawSeriesRequest::from_extension_value(&version.to_string()).is_err()); - - let mut unknown: serde_json::Value = serde_json::from_str(include_str!( - "../../../../../tests/compatibility/wire/raw_series_request/v1_request.json" - )) - .unwrap(); - unknown["logical_tags"][0]["extra"] = serde_json::json!(true); - assert!(RawSeriesRequest::from_extension_value(&unknown.to_string()).is_err()); + fn rejects_unsupported_version() { + let mut request = fixture_request(); + request.version = 2; + assert!(request.validate().is_err()); } #[test] fn rejects_empty_duplicate_and_non_field_ids() { let mut request = fixture_request(); request.field_ids.clear(); - assert!(request.to_extension_value().is_err()); + assert!(request.validate().is_err()); let mut request = fixture_request(); request.field_ids.push(3); - assert!(request.to_extension_value().is_err()); + assert!(request.validate().is_err()); let mut request = fixture_request(); request.field_ids = vec![1]; - assert!(request.to_extension_value().is_err()); + assert!(request.validate().is_err()); } #[test] diff --git a/tests/compatibility/wire/raw_series_request/README.md b/tests/compatibility/wire/raw_series_request/README.md deleted file mode 100644 index 30d3631a36d..00000000000 --- a/tests/compatibility/wire/raw_series_request/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Raw-series request control v1 - -`v1_request.json` is a hand-authored frozen JSON example consumed by the common-query request tests. The version-1 envelope and logical-tag objects reject unknown fields, and unsupported control versions are rejected. Physical metadata keeps its existing deserializer and compatibility rules. A legacy request without the reserved marker remains an ordinary request. A new client talking to an old server may receive an ordinary-schema response, which the raw unpacker rejects; this fixture does not claim cross-version raw execution support. diff --git a/tests/compatibility/wire/raw_series_request/v1_request.json b/tests/compatibility/wire/raw_series_request/v1_request.json deleted file mode 100644 index 1ff6446db3d..00000000000 --- a/tests/compatibility/wire/raw_series_request/v1_request.json +++ /dev/null @@ -1 +0,0 @@ -{"version":1,"physical_metadata":{"column_metadatas":[{"column_schema":{"name":"host","data_type":{"Int64":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Tag","column_id":1},{"column_schema":{"name":"value","data_type":{"Int32":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Field","column_id":2},{"column_schema":{"name":"aux","data_type":{"Int64":{}},"is_nullable":false,"is_time_index":false,"default_constraint":null,"metadata":{}},"semantic_type":"Field","column_id":3},{"column_schema":{"name":"ts","data_type":{"Timestamp":{"Millisecond":null}},"is_nullable":false,"is_time_index":true,"default_constraint":null,"metadata":{}},"semantic_type":"Timestamp","column_id":4}],"primary_key":[1],"region_id":4294967297,"schema_version":0,"primary_key_encoding":"dense"},"metric_name":"cpu","logical_tags":[{"column_id":1,"name":"host"}],"field_ids":[3,2]} From f235434b7f0f97af74ee46534ffe5b5618bacdeb Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 19:09:00 +0800 Subject: [PATCH 11/13] feat(promql): compile nested expressions recursively (#9782) * feat(promql): compile nested expressions instead of fixed query shapes The native compiler matched three shapes: a bare selector, `range_fn(x[r])`, and `range_fn(n, x[r])`. Every composition of those kernels was refused, so `abs(rate(x[5m]))` and `clamp_min(rate(x[5m]), 0)` had no path through it even though the kernels, the `Map` node and the `Sum` node already existed. `compile` is now a thin entry over one recursive rule, `compile_expr`: a parenthesis compiles its inner expression, a vector selector becomes the instant selection it already was, a call dispatches on its signature, and anything else is refused naming what was found. Composition follows from that rule rather than from a list of accepted query shapes -- `abs(rate(x[5m]))` is a map over whatever its argument compiles to, and `clamp_min(abs(rate(x[5m])), 0)` is two maps over one window because the inner call compiles once. Math functions lower onto the node that already computes them: one `Map` carrying one `MathOp`, over the node its vector argument compiles to. The 28 names and their parameters are one explicit table, so supporting another is a decision made in one place; a name outside it is refused as having no native kernel rather than silently treated as math. `round` without an argument binds the kernel's documented default. A scalar parameter is folded from the statement: number literals, parentheses, unary negation, and arithmetic whose two sides both fold. The operator goes through `ArithmeticOp::calculate`, the same function the kernels use, so `%`, `^` and `atan2` keep Prometheus' semantics instead of being redefined in the compiler -- `(-7) % 3` is `-1`, signed as Go's `math.Mod` is. A parameter the statement cannot decide, `time()` being the real case, is refused rather than bound to a value that would be wrong at every point but one. `ArithmeticOp::calculate` widens from `pub(crate)` to `pub` for the fold; its body is unchanged. Nothing else is added: the plan nodes, the kernels and the tree validation are the ones that were there, and validation still happens once per statement, at the root. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): match the arithmetic token by its id, not by its text `arithmetic_op` read the operator through the token's `Display` output. That is a second source of truth for what an operator is: the id is stable and is what the rest of the planner matches on (`prom_token_to_binary_expr_builder`), while the rendered text is a presentation detail. A formatting change would silently stop `%`, `^` and `atan2` from folding -- not wrong, but quietly less than it says it does. Also corrects the doc comment that claimed the parser keeps its token ids to itself. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): say what the validator compares, and make the fold test bite The map's doc claimed `NativeTree::validate` compares the map's output grid with its input's by identity. It compares them by `EvalGrid::spec`, and only an input node's output is compared by identity, so a map that allocated an equal grid of its own would validate the same way. The comment says that now, and the two assertions in the nested-map test compare spec and field names instead of `Arc::ptr_eq`, which failed a map that allocated an equal grid for no behavioural reason. The constant-fold test could not catch what its doc said it guarded: `go_math`'s modulo is `x % y` and its power agrees with `f64::powf` on `2 ^ 10`, so replacing the kernel arithmetic with a hand-written match still passed. The power case is now `1.0000006850274261 ^ 9.80019037e8`, where Go's repeated squaring accumulates error the standard library's `pow` does not, and where the two disagree above the last bit. Replacing the fold with `%`/`powf` fails it. `pi()` and `time()` have no vector to be pointwise over, so the "no arguments" branch explained them wrongly; without it they fall to "has no native kernel", which is accurate, and a listed math name called with no argument at all is covered by the signature's arity check. `math_op` now reads the name first and returns the vector it found, so both rejections come from one place. `argument_position` indexes a fixed array: its three callers pass literals, so there is no fourth position to default to. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): refuse an interleaved scan at the compiler entry Every accepted shape reads one series to its last chunk before the next, so the ordering has to be refused before the statement's shape is dispatched on; the check sat in the range function branch, which let the bare selector and a math function over it compile under Interleaved. The window branch of NativeTree::validate stays as the backstop for a tree built another way. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/map.rs | 2 +- src/query/src/promql/native_compile.rs | 477 +++++++++++++++---- src/query/src/promql/native_compile/tests.rs | 369 +++++++++++++- 3 files changed, 735 insertions(+), 113 deletions(-) diff --git a/src/promql/src/map.rs b/src/promql/src/map.rs index 2d2c4b80a7b..cb3c3cc6f80 100644 --- a/src/promql/src/map.rs +++ b/src/promql/src/map.rs @@ -318,7 +318,7 @@ fn go_min(x: f64, y: f64) -> f64 { impl ArithmeticOp { /// Prometheus `vectorElemBinop` for two floats. - pub(crate) fn calculate(self, left: f64, right: f64) -> f64 { + pub fn calculate(self, left: f64, right: f64) -> f64 { match self { Self::Add => left + right, Self::Subtract => left - right, diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 1be936d7980..56ea7445b76 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -16,20 +16,31 @@ //! //! This is the front half of the native evaluation model: it reads the parsed statement and writes //! down what it means as bound nodes. It runs nothing, and nothing calls it yet, so the subset it -//! accepts is deliberately three shapes: +//! accepts is one recursive rule: //! //! ```text -//! e.g. x -//! ([]) e.g. rate(x[5m]) -//! (, []) e.g. quantile_over_time(0.9, x[5m]) +//! e.g. x +//! ([]) e.g. rate(x[5m]) +//! (, []) e.g. quantile_over_time(0.9, x[5m]) +//! ([, ]*) e.g. abs(rate(x[5m])), clamp_min(x, 0) //! ``` //! -//! which become `Input` (raw samples) → `Instant` (the last sample within the statement's lookback) -//! and `Input` (raw samples) → `Window` (the range function over the statement's grid). A selector -//! may carry `offset` in either shape and, in the window shape, an `@` anchor. Everything else is a -//! [`DataFusionError::NotImplemented`] that names what was found and quotes the shapes above, -//! because rejecting a query the compiler does not understand is the only honest alternative to -//! compiling it into something else. +//! which become `Input` (raw samples) → `Instant` (the last sample within the statement's lookback), +//! `Input` (raw samples) → `Window` (the range function over the statement's grid) and +//! `` → `Map` (a pointwise math function over what its argument evaluates to). A +//! selector may carry `offset` in either selection shape and, in the window shape, an `@` anchor. +//! Everything else is a [`DataFusionError::NotImplemented`] that names what was found and quotes the +//! shapes above, because rejecting a query the compiler does not understand is the only honest +//! alternative to compiling it into something else. +//! +//! Two decisions keep the recursion from growing a branch per query. A constant argument — a +//! quantile, a rounding step, a clamp bound — is folded out of the statement by [`constant_value`], +//! so `quantile_over_time(1 - 0.1, x[5m])` binds the number it means rather than the expression it +//! was written as. A math function is bound by [`math_op`] from a list of names, so +//! `clamp_min(abs(rate(x[5m])), 0)` is two maps over one window because the arguments nest, not +//! because the compiler knows the query. Arithmetic inside a fold is [`ArithmeticOp::calculate`]'s +//! alone, so the number this compiler binds and the number the kernel computes cannot disagree +//! about `%`, `^` or `atan2`. //! //! The statement's grid, the selector's lookback, window, offset and anchor are the only numbers //! derived here; the legality of the tree is not checked here. [`NativeTree::try_new`] is the @@ -45,25 +56,28 @@ use datafusion::error::{DataFusionError, Result}; use datatypes::series::SeriesOrdering; use datatypes::series::samples::SeriesSchema; use datatypes::series::time::{EvalGrid, GridSpec}; +use promql::map::{ArithmeticOp, MathOp}; use promql::native::plan::{ - BoundInput, BoundInstant, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, + BoundInput, BoundInstant, BoundMap, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, }; use promql::range_eval::{RangeFunction, RangeWindow}; +use promql_parser::parser::token::TokenType; +use promql_parser::parser::value::ValueType; use promql_parser::parser::{ AtModifier, Call, EvalStmt, Expr as PromExpr, MatrixSelector, NumberLiteral, Offset, - VectorSelector, + VectorSelector, token, }; use crate::promql::label_values::signed_millis_since_epoch; /// The shapes this compiler accepts, quoted in every rejection. -const SUPPORTED: &str = "a bare selector, a range function over a selector — \ - `([])` — or a parameterised one, \ - `(, [])`, for example `x`, \ - `rate(x[5m])` or `quantile_over_time(0.9, x[5m])`"; +const SUPPORTED: &str = "a selector, a range function over a range selector — \ + `([])`, with its constants in front of \ + the selector as in `quantile_over_time(0.9, x[5m])` — or a math function \ + over any of them, as in `abs(rate(x[5m]))`: for example `x`, \ + `rate(x[5m])`, `quantile_over_time(0.9, x[5m])` or `abs(rate(x[5m]))`"; -/// Compiles `stmt` into a tree that reads `input` and either selects from it or folds it with the -/// statement's one window. +/// Compiles `stmt` into a tree that reads `input` and evaluates the statement's expression over it. /// /// `input` is the raw series schema the caller bound from the table, and the caller has to apply the /// selector's matchers before binding it: `__field__`, `__schema__` and `__database__` select rows @@ -92,8 +106,7 @@ pub fn compile( ) -> Result { // Every shape this compiler accepts reads one series to its last chunk before the next, so raw // samples whose chunks of different series interleave cannot be fed to any of them. The refusal - // is here rather than in a shape's branch so that a shape added later inherits it, and it comes - // before the early return below for a bare selector. + // is here rather than in a shape's branch so that a shape added later inherits it. if ordering == SeriesOrdering::Interleaved { return Err(not_implemented( "raw samples whose series interleave, which this compiler cannot consume: a native \ @@ -101,47 +114,137 @@ pub fn compile( series may not interleave", )); } - // A bare selector is the other shape: an instant selection over the same raw input. - if let PromExpr::VectorSelector(selector) = unparen(&stmt.expr) { - return instant_selection(stmt, selector, input, ordering); - } - let PromExpr::Call(Call { func, args }) = unparen(&stmt.expr) else { - return Err(not_implemented(&format!("the expression `{}`", stmt.expr))); + let context = Context { + stmt, + input, + ordering, }; - let Some(signature) = range_signature(func.name) else { + // The whole tree is validated once, here: a node compiled as another node's input is checked as + // part of the tree rather than on its own, so the same node cannot be legal in one place and + // illegal in another. + NativeTree::try_new(compile_expr(&context, &stmt.expr)?) +} + +/// What every node of one compilation derives from the statement rather than from the expression. +/// +/// The three are the same for the whole tree: one statement decides the grid, the lookback and the +/// anchors, one caller binds the raw input, and one scan guarantees the ordering. Passing them as +/// one value is what lets [`compile_expr`] recurse without restating what it was handed. +struct Context<'a> { + /// The statement being compiled. + stmt: &'a EvalStmt, + /// The raw series the caller bound; shared by every input node rather than copied per node. + input: Arc, + /// The ordering the caller's scan guarantees. + ordering: SeriesOrdering, +} + +/// Compiles `expr` into the node it denotes, recursing into whatever it is made of. +/// +/// One rule per kind of expression and nothing per query: a math function is compiled by compiling +/// its argument, so a query nests as deeply as it was written. A kind this compiler has no node for +/// is refused with what was found, which is why the rejection is useful without a coverage list. +fn compile_expr(context: &Context, expr: &PromExpr) -> Result { + match expr { + // Parentheses group, so they change how a statement was typed rather than what it asks for: + // the expression inside them is the one that compiles. + PromExpr::Paren(paren) => compile_expr(context, &paren.expr), + PromExpr::Call(call) => compile_call(context, call), + PromExpr::VectorSelector(selector) => instant_selection(context, selector), + found => Err(not_implemented(&format!("the expression `{found}`"))), + } +} + +/// Compiles a call: a range function opens a window over the selector it is given, and a math +/// function maps the value of its first argument. +/// +/// The two name lists are disjoint, and a range function is decided first because it binds its own +/// window rather than its argument's value. A name in neither list is refused here, which is where +/// the difference between "this compiler cannot fold this function" and "this name means nothing" +/// is still visible. +fn compile_call(context: &Context, call: &Call) -> Result { + let Call { func, args } = call; + if let Some(signature) = range_signature(func.name) { + return window_call(context, func.name, signature, &args.args); + } + // The name is looked up before the call's arguments are read, so a name this compiler has no + // kernel for is refused as such whichever arguments it was given: `pi()` and `time()` have no + // vector to map, and the rejection names the name rather than describing a vector they never + // wrote. + let Some((op, vector)) = math_op(func.name, &args.args)? else { return Err(not_implemented(&format!( - "a call to `{}`, which has no native range kernel", + "a call to `{}`, which has no native kernel", func.name ))); }; + map_call(context, func.name, op, vector) +} + +/// Binds a range function over the range selector it is given, and the window that selector opens. +/// +/// The parameters come first and the range selector is last, for both signatures: a function without +/// parameters has no `parameters` to read. +fn window_call( + context: &Context, + name: &str, + signature: RangeSignature, + args: &[Box], +) -> Result { // The count is the signature's, so the rejection names what this function takes rather than // what one argument happens to be. let expected = signature.parameters() + 1; - if args.args.len() != expected { + if args.len() != expected { return Err(not_implemented(&format!( - "a call to `{}` with {} argument(s), where it takes {expected}: {}", - func.name, - args.args.len(), + "a call to `{name}` with {} argument(s), where it takes {expected}: {}", + args.len(), signature.arguments() ))); } - // The parameters come first and the range selector is last, for both signatures: a function - // without parameters has no `parameters` to read. - let (parameters, selector) = args.args.split_at(signature.parameters()); + let (parameters, selector) = args.split_at(signature.parameters()); let function = match signature { RangeSignature::Selector(function) => function, - RangeSignature::Scalar(bind) => bind(scalar_parameter(func.name, ¶meters[0])?), + RangeSignature::Scalar(bind) => bind(constant_argument(name, ¶meters[0], 0)?), }; let argument = &selector[0]; let PromExpr::MatrixSelector(MatrixSelector { vs, range }) = unparen(argument) else { return Err(not_implemented(&format!( - "a call to `{}` whose argument is `{}`", - func.name, argument + "a call to `{name}` whose argument is `{argument}`" ))); }; - // The ordering was refused at the entry above, so every shape below is built over a scan whose - // series' chunks do not interleave. - window_selection(stmt, function, vs, range, input, ordering) + window_selection(context, function, vs, range) +} + +/// Wraps `op` over the value of `vector`, which is the first argument of a math function. +/// +/// The output is the input's own type: a pointwise function moves no value to another grid point +/// and reports one value per field it was given, so the map's schema is the input's allocation +/// rather than a copy that merely agrees with it. [`NativeTree::validate`] does not require that +/// sharing: it compares the map's grid with its input's by [`EvalGrid::spec`] and its field count +/// with its input's, so a map that allocated an equal grid of its own would validate the same way. +/// The only comparison by identity is an input node's, whose output has to be the very binding it +/// was bound to. +fn map_call(context: &Context, name: &str, op: MathOp, vector: &PromExpr) -> Result { + // A math function is pointwise over a vector, so its first argument has to be one. The parser + // checks the same thing and the tree's validation checks the compiled input's type again, but + // neither can name the argument the client wrote: this is the check that can. + if vector.value_type() != ValueType::Vector { + return Err(not_implemented(&format!( + "a call to `{name}` whose first argument is `{vector}`, which is a {}, where the \ + function takes a vector", + vector.value_type() + ))); + } + let input = compile_expr(context, vector)?; + let output = input.output.clone(); + Ok(NativeNode { + expr: NativeExpr::Map { + input: Box::new(input), + op: BoundMap { + ops: Arc::from([op]), + }, + }, + output, + }) } /// Binds `function` over the range selector `vs[range]` into a window over the raw input. @@ -150,39 +253,37 @@ pub fn compile( /// whichever function reads them, so a parameterised function only decides *what* folds the window, /// never *which* window is folded. fn window_selection( - stmt: &EvalStmt, + context: &Context, function: RangeFunction, vs: &VectorSelector, range: &Duration, - input: Arc, - ordering: SeriesOrdering, -) -> Result { +) -> Result { // An `@` anchor is step-invariant: every step folds the same samples, which the window's // `anchor_ms` tells `RangeEvaluator` to replay over the grid. let range_ms = range_millis(range.as_millis())?; let offset_ms = selector_offset_ms(&vs.offset)?; - let anchor_ms = selector_anchor_ms(&vs.at, stmt)?; + let anchor_ms = selector_anchor_ms(&vs.at, context.stmt)?; // The grid is the last thing built: a query refused above, or a window that cannot exist, // never pays for a grid nothing would read. - let output_grid = statement_grid(stmt)?; + let output_grid = statement_grid(context.stmt)?; // A range function reports one value per input field, so the layout is the input's. let output = NativeType::Samples(Arc::new(SeriesSchema { - fields: Arc::clone(&input.fields), + fields: Arc::clone(&context.input.fields), eval_grid: Some(Arc::clone(&output_grid)), })); // One allocation shared by the binding and the node's output: `NativeTree::try_new` compares // them by identity, because a compiler that meant one schema has no reason to allocate it // twice. - let raw = NativeType::Samples(input); + let raw = NativeType::Samples(Arc::clone(&context.input)); - NativeTree::try_new(NativeNode { + Ok(NativeNode { expr: NativeExpr::Window { input: Box::new(NativeNode { expr: NativeExpr::Input(BoundInput { output: raw.clone(), - ordering, + ordering: context.ordering, }), output: raw, }), @@ -206,12 +307,7 @@ fn window_selection( /// /// The lookback, the offset and the grid are the whole of an instant selection: it reports the last /// sample at or before each grid point, within the statement's lookback. -fn instant_selection( - stmt: &EvalStmt, - selector: &VectorSelector, - input: Arc, - ordering: SeriesOrdering, -) -> Result { +fn instant_selection(context: &Context, selector: &VectorSelector) -> Result { // An `@` anchor is bound on a range window, which carries `anchor_ms`; an instant selection has // no such field, so the anchor has nowhere to go. Refusing it is the only alternative to // dropping it, which would silently select around every grid point instead of one anchor. @@ -222,25 +318,25 @@ fn instant_selection( anchoring is bound on a range window, as in `rate(x[5m] @ 1234)`" ))); } - let lookback_ms = lookback_millis(stmt.lookback_delta.as_millis())?; + let lookback_ms = lookback_millis(context.stmt.lookback_delta.as_millis())?; let offset_ms = selector_offset_ms(&selector.offset)?; - let output_grid = statement_grid(stmt)?; + let output_grid = statement_grid(context.stmt)?; // An instant selection reports one value per input field, so the layout is the input's. let output = NativeType::Samples(Arc::new(SeriesSchema { - fields: Arc::clone(&input.fields), + fields: Arc::clone(&context.input.fields), eval_grid: Some(Arc::clone(&output_grid)), })); // One allocation shared by the binding and the node's output, as in the window shape. - let raw = NativeType::Samples(input); + let raw = NativeType::Samples(Arc::clone(&context.input)); - NativeTree::try_new(NativeNode { + Ok(NativeNode { expr: NativeExpr::Instant { input: Box::new(NativeNode { expr: NativeExpr::Input(BoundInput { output: raw.clone(), - ordering, + ordering: context.ordering, }), output: raw, }), @@ -356,10 +452,10 @@ fn grid_error(error: datatypes::error::Error) -> DataFusionError { enum RangeSignature { /// A function that reads only its range selector, as in `rate(x[5m])`. Selector(RangeFunction), - /// A function that reads one number literal before its range selector, and the function that - /// number binds, as in `quantile_over_time(0.9, x[5m])`. + /// A function that reads one constant before its range selector, and the function that constant + /// binds, as in `quantile_over_time(0.9, x[5m])`. /// - /// The number is kept out of the variant name because only the call site knows it. + /// The constant is kept out of the variant name because only the call site knows it. Scalar(fn(f64) -> RangeFunction), } @@ -381,7 +477,7 @@ impl RangeSignature { } } -/// The signature the PromQL name calls, or `None` when this compiler has no kernel for it. +/// The signature the PromQL name calls, or `None` when this compiler has no range kernel for it. /// /// The map is explicit on purpose, as in [`range_function`]: a name that is not listed is rejected, /// so supporting one more function is a decision made in one place rather than a consequence of a @@ -394,29 +490,235 @@ fn range_signature(name: &str) -> Option { } } -/// The number a parameterised range function was given, as the literal it was written as. +/// What a call to a math function has to be: the constants it reads after its vector, and how to +/// bind them. /// -/// A number literal is the only parameter this compiler binds. A parameter whose value is only -/// known at run time, `time()` being the real case, would make the kernel's constant a query-time -/// value, which the tree has no field for. +/// The shapes are separated by how many constants the function reads, because that is the only +/// thing the variants differ in: a name maps to one kernel, and the constants are bound from the +/// call. `round` is the one function whose constant may be left out, so it is its own shape: the +/// default is a decision about the function rather than about the constant. +#[derive(Clone, Copy)] +enum MathSignature { + /// A function of its vector alone, as in `abs`. + Unary(MathOp), + /// A function of its vector and one constant, as in `clamp_min`. + OneConstant(fn(f64) -> MathOp), + /// A function of its vector and one constant it may be called without, as in `round`. + DefaultedConstant { + default: f64, + bind: fn(f64) -> MathOp, + }, + /// A function of its vector and two constants, as in `clamp`. + TwoConstants(fn(f64, f64) -> MathOp), +} + +impl MathSignature { + /// How many constants the function reads when every one of them is written. + fn parameters(self) -> usize { + match self { + Self::Unary(_) => 0, + Self::OneConstant(_) | Self::DefaultedConstant { .. } => 1, + Self::TwoConstants(_) => 2, + } + } + + /// How many constants the function has to be given. + fn required(self) -> usize { + match self { + Self::DefaultedConstant { .. } => 0, + _ => self.parameters(), + } + } + + /// What the call's arguments are, for the arity rejection to name them. + fn arguments(self) -> &'static str { + match self { + Self::Unary(_) => "its vector", + Self::OneConstant(_) => "its vector and one constant", + Self::DefaultedConstant { .. } => "its vector and one optional constant", + Self::TwoConstants(_) => "its vector and two constants", + } + } +} + +/// The signature the PromQL name calls, or `None` when this compiler has no math kernel for it. /// -/// A parameter written as an expression the existing planner folds to a literal is refused here -/// and is a gap this stage leaves open on purpose: `quantile_over_time(1 - 0.1, x[5m])` and -/// `quantile_over_time(-(0.9), x[5m])` plan today through `try_build_literal_expr`, which folds -/// `Paren`, `Unary` and literal-only `Binary`, and reach this function still unfolded. They are -/// refused rather than answered from the wrong number, and folding them here would be the work -/// to close the gap. +/// The map is explicit on purpose, as in [`range_function`]: a name that is not listed is refused +/// rather than treated as math because it looks mathematical, so adding one more function is a +/// decision made in one place. A name in this list is a promise that its kernel exists, which is why +/// arity is a property of the variant rather than of the caller. +fn math_signature(name: &str) -> Option { + match name { + "abs" => Some(MathSignature::Unary(MathOp::Abs)), + "ceil" => Some(MathSignature::Unary(MathOp::Ceil)), + "floor" => Some(MathSignature::Unary(MathOp::Floor)), + "sqrt" => Some(MathSignature::Unary(MathOp::Sqrt)), + "exp" => Some(MathSignature::Unary(MathOp::Exp)), + "ln" => Some(MathSignature::Unary(MathOp::Ln)), + "log2" => Some(MathSignature::Unary(MathOp::Log2)), + "log10" => Some(MathSignature::Unary(MathOp::Log10)), + "sin" => Some(MathSignature::Unary(MathOp::Sin)), + "cos" => Some(MathSignature::Unary(MathOp::Cos)), + "tan" => Some(MathSignature::Unary(MathOp::Tan)), + "asin" => Some(MathSignature::Unary(MathOp::Asin)), + "acos" => Some(MathSignature::Unary(MathOp::Acos)), + "atan" => Some(MathSignature::Unary(MathOp::Atan)), + "sinh" => Some(MathSignature::Unary(MathOp::Sinh)), + "cosh" => Some(MathSignature::Unary(MathOp::Cosh)), + "tanh" => Some(MathSignature::Unary(MathOp::Tanh)), + "asinh" => Some(MathSignature::Unary(MathOp::Asinh)), + "acosh" => Some(MathSignature::Unary(MathOp::Acosh)), + "atanh" => Some(MathSignature::Unary(MathOp::Atanh)), + "deg" => Some(MathSignature::Unary(MathOp::Deg)), + "rad" => Some(MathSignature::Unary(MathOp::Rad)), + "sgn" => Some(MathSignature::Unary(MathOp::Sgn)), + // `round`'s nearest defaults to one, so `round(x)` is `round(x, 1)` and its variant always + // carries the number it rounds to. + "round" => Some(MathSignature::DefaultedConstant { + default: 1.0, + bind: |nearest| MathOp::Round { nearest }, + }), + "clamp" => Some(MathSignature::TwoConstants(|min, max| MathOp::Clamp { + min, + max, + })), + "clamp_min" => Some(MathSignature::OneConstant(|min| MathOp::ClampMin { min })), + "clamp_max" => Some(MathSignature::OneConstant(|max| MathOp::ClampMax { max })), + _ => None, + } +} + +/// The pointwise kernel `name` denotes over its vector, and that vector, or `None` when this +/// compiler has no math kernel for that name. +/// +/// A name that is not in [`math_signature`] is not math: the caller reports that it has no kernel +/// instead of compiling it into a kernel that would compute something else. `arguments` is the +/// call's whole argument list, because the vector is its first argument: the arity is the +/// signature's and counts the vector in, so a call with no argument at all is refused as a call +/// that wrote none rather than as a shape this compiler cannot bind. The arguments after the vector +/// are the call's constants, in the order they were written, and each is folded to the number it +/// means by [`constant_value`], because the tree binds a kernel's constants and has no field for an +/// expression. +fn math_op<'a>( + name: &str, + arguments: &'a [Box], +) -> Result> { + let Some(signature) = math_signature(name) else { + return Ok(None); + }; + // The vector is the function's first argument, so a call that wrote no argument at all is + // refused by the same arity check as one that wrote too few constants. + let Some((vector, parameters)) = arguments.split_first() else { + return Err(arity_error(name, signature, 0)); + }; + let given = parameters.len() + 1; + if given < signature.required() + 1 || given > signature.parameters() + 1 { + return Err(arity_error(name, signature, given)); + } + // The vector is the first argument, so the constants are read from the second onwards. + let op = match signature { + MathSignature::Unary(op) => op, + MathSignature::OneConstant(bind) => bind(constant_argument(name, ¶meters[0], 1)?), + MathSignature::DefaultedConstant { default, bind } => bind(match parameters.first() { + Some(parameter) => constant_argument(name, parameter, 1)?, + None => default, + }), + MathSignature::TwoConstants(bind) => bind( + constant_argument(name, ¶meters[0], 1)?, + constant_argument(name, ¶meters[1], 2)?, + ), + }; + Ok(Some((op, vector))) +} + +/// The rejection of a math call that wrote `given` arguments where `signature` takes another count. +/// +/// The count is the signature's, so the message says what this function takes rather than what one +/// argument happens to be. The vector is counted in: it is the function's first argument, and a +/// reader counting the call's arguments should not have to add it themselves. +fn arity_error(name: &str, signature: MathSignature, given: usize) -> DataFusionError { + not_implemented(&format!( + "a call to `{name}` with {given} argument(s), where it takes {}", + signature.arguments() + )) +} + +/// The constant `argument` folds to, or a rejection naming the argument and where it was read. +/// +/// A parameter is a constant of the statement: a quantile, a rounding step or a clamp bound is a +/// field of the tree, so a parameter whose value is only known at run time — `time()` being the real +/// case — has nowhere to go. It is refused rather than answered from the wrong number, and the +/// rejection names the expression, because that is what the query has to change. /// /// A quantile outside `[0, 1]` is not refused here: Prometheus binds it and the kernel reports -/// `-Infinity` below zero and `+Infinity` above one, so refusing it would be this compiler -/// inventing a rule the query engine does not have. -fn scalar_parameter(name: &str, parameter: &PromExpr) -> Result { - match unparen(parameter) { - PromExpr::NumberLiteral(NumberLiteral { val }) => Ok(*val), - found => Err(not_implemented(&format!( - "a call to `{name}` whose first argument is `{found}`, where it has to be a number \ - literal" - ))), +/// `-Infinity` below zero and `+Infinity` above one, so refusing it would be this compiler inventing +/// a rule the query engine does not have. +fn constant_argument(name: &str, argument: &PromExpr, position: usize) -> Result { + constant_value(argument).ok_or_else(|| { + not_implemented(&format!( + "a call to `{name}` whose {} is `{argument}`, where it has to be a constant expression \ + this compiler folds to a number literal", + argument_position(position) + )) + }) +} + +/// How a rejection names the argument a parameter was read from. +/// +/// The three positions are the ones a signature can read: the vector, and the two constants behind +/// it. Every caller passes the literal position its signature reads, so there is no fourth position +/// to name and no default to invent one for. +fn argument_position(position: usize) -> &'static str { + const POSITIONS: [&str; 3] = ["first argument", "second argument", "third argument"]; + POSITIONS[position] +} + +/// The number `expr` folds to, or `None` when its value is not the statement's to know. +/// +/// Only what the statement itself decides is folded: a number literal, a group around one, a negated +/// one, and arithmetic whose two operands both fold. A call, a selector, a range selector, an +/// aggregate and a subquery are `None`, because their value comes from the data or from the +/// evaluation time, and a kernel's constant is a field of the tree rather than something bound while +/// it runs. Returning `None` rather than an error is what lets the caller report where the +/// expression was found, which is the useful half of the refusal. +/// +/// Arithmetic is [`ArithmeticOp::calculate`]'s and nothing else's: that is where Prometheus' +/// semantics for `%`, `^` and `atan2` live, the last two through `go_math`, so defining them again +/// here would be a second answer for the two to disagree over. An operator with no +/// [`ArithmeticOp`] — a comparison, `and`, `or`, `unless` — is not arithmetic and does not fold. +fn constant_value(expr: &PromExpr) -> Option { + match expr { + PromExpr::NumberLiteral(NumberLiteral { val }) => Some(*val), + PromExpr::Paren(paren) => constant_value(&paren.expr), + PromExpr::Unary(unary) => constant_value(&unary.expr).map(|value| -value), + PromExpr::Binary(binary) => { + let op = arithmetic_op(&binary.op)?; + let left = constant_value(&binary.lhs)?; + let right = constant_value(&binary.rhs)?; + Some(op.calculate(left, right)) + } + // Calls (`time()`, `scalar(v)`), selectors, range selectors, aggregates and subqueries all + // read something this compiler cannot know before the query runs. + _ => None, + } +} + +/// The arithmetic operator a PromQL token denotes, or `None` when the token is not arithmetic. +/// +/// The token is matched by its id, which is stable and is what the rest of the planner matches on, +/// rather than by the text it displays as: a symbol this compiler does not recognise is not +/// arithmetic here rather than a wrong operator. A comparison and a set operator are the cases that +/// reach this and fold nothing, which is this batch's boundary rather than an accident of the map. +fn arithmetic_op(op: &TokenType) -> Option { + match op.id() { + token::T_ADD => Some(ArithmeticOp::Add), + token::T_SUB => Some(ArithmeticOp::Subtract), + token::T_MUL => Some(ArithmeticOp::Multiply), + token::T_DIV => Some(ArithmeticOp::Divide), + token::T_MOD => Some(ArithmeticOp::Modulo), + token::T_POW => Some(ArithmeticOp::Power), + token::T_ATAN2 => Some(ArithmeticOp::Atan2), + _ => None, } } @@ -452,7 +754,8 @@ fn range_function(name: &str) -> Option { /// /// Parentheses group, so they say nothing about the expression inside them: `rate((x[5m]))` asks /// for exactly what `rate(x[5m])` asks for. Refusing it would reject a statement on the way it was -/// typed rather than on what it means. +/// typed rather than on what it means. A range selector is not an expression this compiler +/// compiles on its own, so the call that reads one looks through the grouping itself. fn unparen(mut expr: &PromExpr) -> &PromExpr { while let PromExpr::Paren(paren) = expr { expr = &paren.expr; diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 05b4d07ead8..f16bb74b072 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -18,8 +18,9 @@ use std::time::{Duration, UNIX_EPOCH}; use datafusion::error::DataFusionError; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; +use promql::map::MathOp; use promql::native::plan::{ - BoundInstant, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, + BoundInstant, BoundMap, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, }; use promql::range_eval::RangeFunction; use promql_parser::parser::{EvalStmt, Expr as PromExpr, NumberLiteral}; @@ -82,6 +83,18 @@ fn with_call_arguments(query: &str, edit: impl FnOnce(&mut Vec>)) stmt } +/// The range selector out of `rate(x[5m])`. +/// +/// The parser type-checks function arguments, so a range selector can only be handed to a function +/// that does not take one by editing a tree the parser accepted. +fn a_matrix_selector() -> Box { + let mut stmt = eval_stmt("rate(x[5m])"); + let PromExpr::Call(call) = &mut stmt.expr else { + unreachable!("`rate(x[5m])` parses to a call"); + }; + call.args.args.remove(0) +} + /// The window at the root of `tree`, for the range-function shape. fn root_window(tree: &NativeTree) -> (&NativeNode, &BoundWindow) { let NativeExpr::Window { input, op } = &tree.root().expr else { @@ -101,6 +114,32 @@ fn root_instant(tree: &NativeTree) -> (&NativeNode, &BoundInstant) { (input, op) } +/// The input and the operation of a `Map` node. +fn map_parts(node: &NativeNode) -> (&NativeNode, &BoundMap) { + let NativeExpr::Map { input, op } = &node.expr else { + panic!("the node is not a map: {:?}", node.expr); + }; + (input, op) +} + +/// The one operation a `Map` node applies. +fn map_op(node: &NativeNode) -> MathOp { + let (_, op) = map_parts(node); + let [only] = &op.ops[..] else { + panic!("the map holds {} operation(s)", op.ops.len()); + }; + *only +} + +/// The schema a node's evaluated output carries. +fn evaluated(node: &NativeNode) -> &Arc { + let Some(schema) = node.output.samples() else { + panic!("the node does not produce samples: {:?}", node.output); + }; + assert!(schema.eval_grid.is_some(), "the node is not evaluated"); + schema +} + /// `rate(x[5m])` becomes a window over the raw input, on the statement's grid. #[test] fn a_rate_over_a_selector_becomes_a_window_over_raw_input() { @@ -277,27 +316,29 @@ fn a_parenthesised_selector_binds_as_written() { } } -/// A parameter the planner folds to a literal is still refused, and that gap is deliberate: the -/// kernel's constant is a field of the tree, so binding `1 - 0.1` means folding it here rather than -/// carrying an expression. A refusal is the agreed behaviour for this stage; an answer computed -/// from the wrong number would not be. +/// A parameter written as an expression binds the number it folds to, because the kernel's constant +/// is a field of the tree: `1 - 0.1` is the quantile `0.9` written two ways, and carrying the +/// expression would bind nothing at all. The value is asserted exactly, so a fold that produced a +/// neighbouring number rather than the arithmetic's own result would fail here. #[test] -fn a_parameter_folded_from_an_expression_is_refused() { - for query in [ - "quantile_over_time(1 - 0.1, x[5m])", - "quantile_over_time(-(0.9), x[5m])", - ] { - let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) - .expect_err(&format!("`{query}` compiled, and it should not have")); - assert!( - matches!(error, DataFusionError::NotImplemented(_)), - "`{query}` did not report NotImplemented: {error}" - ); - let message = error.to_string(); - assert!( - message.contains("number literal"), - "`{query}` did not say what the argument has to be: {message}" +fn a_parameter_folded_from_an_expression_binds_its_value() { + // `1 - 0.1` is `0.9` bit for bit in IEEE754, and `-(0.9)` is the negative of the same number: + // negation is exact, and the fold does not decide whether a quantile is inside `[0, 1]`. + let cases = [ + ("quantile_over_time(1 - 0.1, x[5m])", 0.9), + ("quantile_over_time(-(0.9), x[5m])", -0.9), + ]; + + for (query, quantile) in cases { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) + .unwrap_or_else(|error| panic!("`{query}` was refused: {error}")); + let (_, op) = root_window(&tree); + assert_eq!( + op.function, + RangeFunction::QuantileOverTime(quantile), + "`{query}` did not bind the quantile it folds to" ); + assert_eq!(op.window.range_ms, 300_000, "`{query}` lost its window"); } } @@ -369,6 +410,23 @@ fn a_call_with_the_wrong_number_of_arguments_names_the_count_it_takes() { "rate", "where it takes 1", ), + // A math function's constants are counted the same way: `clamp` reads two of them after its + // vector, so one is one too few. + ( + with_call_arguments("clamp(x, 0, 1)", |args| { + args.pop(); + }), + "clamp", + "where it takes its vector and two constants", + ), + // A listed math name called with no argument at all: the vector is counted as the first + // argument, so a call that wrote none is one argument short of `abs` rather than a shape + // the compiler has no words for. + ( + with_call_arguments("abs(x)", |args| args.clear()), + "abs", + "with 0 argument(s), where it takes its vector", + ), ]; for (stmt, name, expected) in cases { @@ -713,9 +771,9 @@ fn the_numeric_edges_of_a_statement_are_bound_as_given() { #[test] fn an_interleaved_scan_is_refused_whatever_the_shape() { // The plain range function is the ordering table's case, which compiles `rate(x[5m])` under - // `Interleaved` and expects the same refusal; these two cover the shapes that table cannot, - // because one returns before a window exists and the other takes a parameter. - for query in ["x", "quantile_over_time(0.9, x[5m])"] { + // `Interleaved` and expects the same refusal; these cover the shapes that table cannot: the + // bare selector, a math function over one, and a parameterised range function. + for query in ["x", "abs(x)", "quantile_over_time(0.9, x[5m])"] { let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::Interleaved) .expect_err(&format!("`{query}` compiled over an interleaved scan")); assert!( @@ -765,12 +823,16 @@ fn an_unsupported_shape_is_rejected_and_names_what_it_found() { // A function with no kernel, whether or not it takes a constant argument. ( "predict_linear(x[5m], 300)", - "a call to `predict_linear`, which has no native range kernel", + "a call to `predict_linear`, which has no native kernel", ), ( "absent(x)", - "a call to `absent`, which has no native range kernel", + "a call to `absent`, which has no native kernel", ), + // The two functions PromQL calls with no arguments at all: neither has a kernel, and the + // rejection names the name rather than describing a vector they were never given. + ("pi()", "a call to `pi`, which has no native kernel"), + ("time()", "a call to `time`, which has no native kernel"), ]; for (query, named) in cases { @@ -840,3 +902,260 @@ fn parentheses_around_the_accepted_shape_are_transparent() { ); } } +/// A math function compiles by compiling the expression it is applied to, so its node sits over the +/// node that expression produces: `abs(rate(x[5m]))` is a map over the window `rate(x[5m])` is, not +/// a fourth shape the compiler knows by name. The same rule reads `abs(x)` over an instant +/// selection, which is why the two are asserted together: nothing below the map is specific to it. +#[test] +fn a_math_function_maps_over_the_node_its_argument_compiles_to() { + let tree = compile( + &eval_stmt("abs(rate(x[5m]))"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`abs(rate(x[5m]))` was refused: {error}")); + + assert!( + matches!(map_op(tree.root()), MathOp::Abs), + "the root is not an abs: {:?}", + map_parts(tree.root()).1.ops + ); + let (window, _) = map_parts(tree.root()); + let NativeExpr::Window { input, op } = &window.expr else { + panic!("the map's input is not a window: {:?}", window.expr); + }; + assert_eq!(op.function, RangeFunction::Rate); + assert_eq!(op.window.range_ms, 300_000); + let NativeExpr::Input(bound) = &input.expr else { + panic!("the window's input is not an input: {:?}", input.expr); + }; + assert!( + matches!(&bound.output, NativeType::Samples(schema) if schema.eval_grid.is_none()), + "the window's input is not raw samples: {:?}", + bound.output + ); + assert_eq!(bound.ordering, SeriesOrdering::KeyThenTime); + + // The same rule over the other selection shape: a map over an instant selection. + let tree = compile( + &eval_stmt("abs(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`abs(x)` was refused: {error}")); + assert!(matches!(map_op(tree.root()), MathOp::Abs)); + let (instant, _) = map_parts(tree.root()); + assert!( + matches!(instant.expr, NativeExpr::Instant { .. }), + "`abs(x)`'s map is not over an instant selection: {:?}", + instant.expr + ); +} + +/// Nested math functions are two maps over one window, and every map reports that window's layout. +/// +/// The maps are separate nodes because the arguments nest, not because the compiler knows +/// `clamp_min(abs(...))`: a compiler that bound the pair to one node would fail the tree assertions +/// here. The grids are compared by their spec, which is what [`NativeTree::validate`] compares: a +/// map that allocated a second, equal grid would still be the window's layout, while one that +/// reported a different grid or different fields would not. +#[test] +fn nested_math_functions_are_two_maps_over_one_window() { + let tree = compile( + &eval_stmt("clamp_min(abs(rate(x[5m])), 0)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`clamp_min(abs(rate(x[5m])), 0)` was refused: {error}")); + + let MathOp::ClampMin { min } = map_op(tree.root()) else { + panic!( + "the root is not a clamp_min: {:?}", + map_parts(tree.root()).1.ops + ); + }; + assert_eq!( + min, 0.0, + "the clamp bound is not the number the query wrote" + ); + + let (inner, _) = map_parts(tree.root()); + assert!( + matches!(map_op(inner), MathOp::Abs), + "the inner node is not an abs: {:?}", + map_parts(inner).1.ops + ); + let (window, op) = map_parts(inner); + let NativeExpr::Window { op: bound, .. } = &window.expr else { + panic!("the inner map's input is not a window: {:?}", window.expr); + }; + assert_eq!(bound.function, RangeFunction::Rate); + assert_eq!(op.ops.len(), 1, "a map applies one operation"); + + let window_schema = evaluated(window); + let window_grid = window_schema + .eval_grid + .as_ref() + .expect("the window is evaluated"); + for (what, node) in [("the clamp_min", tree.root()), ("the abs", inner)] { + let schema = evaluated(node); + assert_eq!( + schema + .eval_grid + .as_ref() + .expect("the map is evaluated") + .spec(), + window_grid.spec(), + "`{what}` reports a grid the window does not" + ); + assert!( + schema + .fields + .iter() + .map(|field| &field.name) + .eq(window_schema.fields.iter().map(|field| &field.name)), + "`{what}` reports fields the window does not" + ); + } +} + +/// Every listed math name binds its own kernel, with the constants the call wrote. The display names +/// the operation and its constants, so a name that fell through to a neighbour's variant — or a +/// parameter that was bound from the wrong argument — shows up in a line someone would read. +#[test] +fn the_supported_math_functions_map_to_their_kernels() { + let cases = [ + ("abs(x)", "abs"), + ("ceil(x)", "ceil"), + ("floor(x)", "floor"), + ("sqrt(x)", "sqrt"), + ("exp(x)", "exp"), + ("ln(x)", "ln"), + ("log2(x)", "log2"), + ("log10(x)", "log10"), + ("sin(x)", "sin"), + ("cos(x)", "cos"), + ("tan(x)", "tan"), + ("asin(x)", "asin"), + ("acos(x)", "acos"), + ("atan(x)", "atan"), + ("sinh(x)", "sinh"), + ("cosh(x)", "cosh"), + ("tanh(x)", "tanh"), + ("asinh(x)", "asinh"), + ("acosh(x)", "acosh"), + ("atanh(x)", "atanh"), + ("deg(x)", "deg"), + ("rad(x)", "rad"), + ("sgn(x)", "sgn"), + // The three shapes that carry constants, including the default `round` binds one for. + ("round(x)", "round 1"), + ("round(x, 0.5)", "round 0.5"), + ("clamp(x, 0, 1)", "clamp 0 1"), + ("clamp_min(x, 0)", "clamp_min 0"), + ("clamp_max(x, 1)", "clamp_max 1"), + ]; + + for (query, op) in cases { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) + .unwrap_or_else(|error| panic!("`{query}` was refused: {error}")); + let display = tree.to_string(); + let root = display.lines().nth(1).expect("a compiled tree has a root"); + assert_eq!( + root, + format!(" root Map {op}"), + "`{query}` bound the wrong kernel" + ); + } +} + +/// The constants a function folds are folded by the kernel's own arithmetic, not by a second copy of +/// it: `%` keeps the sign of its left operand and `^` is `go_math`'s power. Both numbers below are +/// ones a hand-written copy would get wrong: a Euclidean modulo folds `-7 % 3` to 2, `*` folds the +/// power to about 9.8e8, and Rust's own `f64::powf` folds it to 3.623449445683154e291 where the +/// kernel gives 3.6234491766672406e291, because `go_math::pow` accumulates error by repeated +/// squaring as Go's `math.Pow` does. That is the point: the compiler must not disagree with the +/// kernel about what an operator means. +#[test] +fn the_constant_fold_is_the_kernels_arithmetic() { + // Go's `math.Mod` gives -1 for `-7 % 3`; a Euclidean modulo would give 2. + let tree = compile( + &eval_stmt("clamp_min(x, (-7) % 3)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`clamp_min(x, (-7) % 3)` was refused: {error}")); + let MathOp::ClampMin { min } = map_op(tree.root()) else { + panic!( + "the root is not a clamp_min: {:?}", + map_parts(tree.root()).1.ops + ); + }; + assert_eq!(min, -1.0, "`-7 % 3` did not fold with `math.Mod`'s sign"); + + // `1.0000006850274261 ^ 9.80019037e8` is a power over a base barely above one: as multiplication + // it would be about 9.8e8, and it is neither a comparison nor a set operator, which fold nothing + // at all. It is also the case `go_math` keeps its own `pow` for: Go's repeated squaring + // accumulates error the standard library's `pow` does not, so the two disagree well above the + // last bit and a fold that hand-wrote `powf` binds a different number here. + let tree = compile( + &eval_stmt("round(x, 1.0000006850274261 ^ 9.80019037e8)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| { + panic!("`round(x, 1.0000006850274261 ^ 9.80019037e8)` was refused: {error}") + }); + let MathOp::Round { nearest } = map_op(tree.root()) else { + panic!( + "the root is not a round: {:?}", + map_parts(tree.root()).1.ops + ); + }; + assert_eq!( + nearest, 3.6234491766672406e291, + "`1.0000006850274261 ^ 9.80019037e8` did not fold with `go_math`'s power" + ); +} + +/// A math function is pointwise over a vector, so an argument that is not one is refused and the +/// rejection names the argument rather than the function alone. The parser refuses `abs(1)` too, so +/// the AST is edited the way the arity test edits one: a shape the parser would have caught is still +/// a shape the compiler must not compile into something else. +#[test] +fn a_math_function_over_a_non_vector_argument_is_refused_and_names_it() { + let cases = [ + // A scalar, which a pointwise function has no positions to be pointwise over. + ( + with_call_arguments("abs(x)", |args| { + *args[0] = PromExpr::NumberLiteral(NumberLiteral::new(1.0)); + }), + "a call to `abs` whose first argument is `1`, which is a scalar", + ), + // A range selector, which belongs in a window rather than under a map. + ( + with_call_arguments("abs(x)", |args| { + args[0] = a_matrix_selector(); + }), + "a call to `abs` whose first argument is `x[5m]`, which is a matrix", + ), + ]; + + for (stmt, named) in cases { + let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime) + .expect_err("a math function over a non-vector compiled"); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "the refusal is not NotImplemented: {error}" + ); + let message = error.to_string(); + assert!( + message.contains(named), + "the refusal did not name the argument it found: {message}" + ); + assert!( + message.contains("where the function takes a vector"), + "the refusal did not say what the argument has to be: {message}" + ); + } +} From 7309800aef64f29548268fc43aff3a0eb50adfc7 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 19:11:24 +0800 Subject: [PATCH 12/13] feat(datatypes): describe label decoding in the raw series frame schema (#9786) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(datatypes): describe label decoding in the raw series frame schema A raw series frame written with key-encoded labels carries the primary key as opaque bytes; only the binding held by the scanner says how to turn those bytes back into a metric name and labels. A consumer that receives the frame without that binding cannot decode it, so the schema now describes its own label layout. The description is compact JSON under `greptime.series.label_decode`, carrying a version, the metric name, the primary-key encoding, the ordered key columns with their physical types, and the logical tags to expose -- exactly what decoding reads. It is derived from the binding, so there is one writer, and a reader that does not implement the version it finds rejects the frame instead of guessing. `unpack_raw_series` keeps its binding argument and now requires it to agree with the description the schema declares, so the two cannot drift. `PrimaryKeyLabelDecoder` reduces both of its inputs to one internal key layout, so the metadata-backed and description-backed paths share their validation instead of growing a second decoder. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(datatypes): compare label descriptions by what their order means The frame's declared label-decode description was compared to the binding's with plain equality, but the description has one part whose order carries no meaning: the decoder sorts logical tags by name when it binds them, so the same tags listed in another order denote the same labels and were being refused. `same_label_layout` canonicalises exactly that part and leaves `primary_key` in place, because for a dense key the array position is the encoding position: the same columns in another order are a different key, not the same key reordered. Both fields now say which order they mean. Three checks go with it. The binding stays for a stated reason -- it establishes identity, the physical region key layout a decoded identity came from -- rather than because a count of call sites kept it. The post-pack schema comparison in the series scan could not fire: the leaf already re-derives that schema and compares it, and every `SeriesSchema` field is covered by the checks between them. And two test groups re-asserted what the tracked v1_metadata fixture pins case by case. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): decode a raw frame from the frame alone `unpack_raw_series` took the label mode, the ordering and the caller's binding along with the stream, so a consumer that received a frame from a peer still needed a local copy of the binding to read it — the thing the frame's own description exists to avoid, and the reason the description was freezing a wire format for nobody. The stream is now the only input: the field schema, the label mode, the label decode description and the producer ordering all come from the frame's schema, and every frame is validated whether or not a caller asked for it. `expected` stays as an optional check on the schema that is declared — field names and their order, and no evaluation grid — and supplies nothing to the decoder, so a caller cannot smuggle in a layout the frame does not declare. Three call sites compared the caller's binding against the frame and all three were load-bearing: a peer's response is not trusted to be the region that was planned, and a native plan's child is an arbitrary execution plan. Those became `ensure_frame_matches_label_source`, which states what it compares: the label layout, not the region. Two regions with the same tag ids and names compare equal, and the doc now says so instead of claiming to establish identity. A stream with no scan-bound identity reports `NO_BOUND_LABEL_SOURCE` and answers `label_decode_spec()` instead; every consumer that decides whether to decode reads the description, so the placeholder never stands in for "these labels are complete". The merge contract is the description, not the binding: inputs are compared by the layout they declare, and a binding that could not be expressed on the wire — a `RegionMetadata` snapshot — no longer makes the merge stricter than the frames it merges. `logical_tags` gets one rule rather than two: the writer canonicalizes a copy, because the caller's binding keeps its own order and the table scan compares producer tag lists positionally, and readers compare by equality. The text embeds `ConcreteDataType`'s serde, which `LABEL_DECODE_VERSION` has to cover. Rather than mirror all 25 datatype variants into a type this module owns, a frozen fixture records the text of every variant and an exhaustive match makes a new variant fail to compile until it is recorded. The version policy is written next to the constant: exact-text comparison across nodes, unknown versions rejected, any same-version change breaking — with the planning-time fallback stated as intended and marked not implemented. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(store-api): make every key-encoded raw stream declare its layout A raw series stream whose chunks carry key-encoded labels had to be decoded against a scan-bound `RawLabelSource`, and two consumers compensated for producers that stayed silent by deriving a description from that binding: `pack_raw_series` and the merge's `declared_layout`. `materialize_raw_labels` did not compensate, so an unframed producer stream reached it with no description and a chunk carrying no labels, and failed. The stream itself is the authority for its chunks. A key-encoded producer knows its own region metadata and can state the layout it writes, exactly as a frame states the layout it carries, so the split contract — "a producer declares its identity, a frame declares its layout" — is replaced by one rule: a stream whose chunks carry key-encoded labels declares the description, and a stream whose chunks carry complete labels answers `None`. mito2's `RawStream` caches the derivation at construction, so the declaration exists before the first poll and for an empty stream. The trait documentation states the rule and no longer exempts unframed producers. Both fallbacks are gone: `pack_raw_series` reads the stream's declaration only, and rejects a key-encoded stream that declares none rather than packing it as an explicit-label frame, which would describe chunks carrying no labels as if they carried all of them. The merge reads the declaration only. Test fixtures that model key-encoded producers now declare their layout too; fixtures that carry complete labels keep answering `None`. The new mito2 test decodes a real produced key through the declared description for both dense and sparse encodings, and checks the declaration is present before polling even for an empty stream. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(common-query): let the declaration decide the frame mode alone `pack_raw_series` rejected a key-encoded stream that declared no description outright, so a stream that stays silent could not be packed at all. That check was redundant: the declaration already decides the frame's label mode, and the first chunk whose labels disagree with it is refused by `RawSeriesChunkCursor` with "raw chunk label mode does not match frame mode", so no such stream can pack data either way. What the check added was an early error for a producer that already breaks the contract, in exchange for reading the stream's scan binding here — in the one function whose rule is that the declaration is the only decode input. Drop it, so pack keeps deciding the mode from the declaration and nothing else. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(store-api): assert the fixture's declared label decode description The mock prepared source derives a `LabelDecodeSpec` from its label source and reports it through `RawSeriesStream::label_decode_spec`, but no assertion read it, so the derivation could drift from the chunks the fixture actually builds. Check it against the description the fixture's own key columns imply: an explicit source declares nothing, a primary-key source declares the metric, the dense encoding, the one key column and its tag. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): let the frame's declaration be the only label authority Two things still treated a caller-supplied binding as authoritative after the declaration became the decode input: - A merged stream forwarded its first input's scan-bound identity, though a merge is not a region scan and decodes from the declaration alone. It now reports no bound identity. - `KeyColumn.tag` repeated an invariant every key column is a tag. The check moves to the one place a decoder is bound, so a snapshot built without validation is still rejected instead of silently decoded. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/src/series_frame.rs | 249 +++++- .../src/series_frame/record_batches/tests.rs | 8 +- src/common/query/src/series_frame/tests.rs | 846 ++++++++++++++---- .../query/src/series_frame/tests/lifecycle.rs | 16 +- src/datatypes/src/series/frame.rs | 330 ++++++- src/datatypes/src/series/frame/tests.rs | 527 +++++++++++ src/mito-codec/src/error.rs | 10 + src/mito-codec/src/labels.rs | 199 +++- src/mito-codec/src/labels/tests.rs | 113 ++- src/mito2/src/engine/scan_test.rs | 251 +++++- src/mito2/src/read/series_scan/raw.rs | 16 + .../instant_manipulate/native_exec.rs | 28 +- src/promql/src/source.rs | 42 +- src/promql/src/source/merge.rs | 127 +-- src/promql/src/source/merge/tests.rs | 456 ++++++---- src/promql/src/source/merge/tests/payload.rs | 35 +- src/promql/src/source/tests.rs | 83 ++ src/query/src/dist_plan/merge_scan.rs | 93 ++ src/store-api/src/region_engine.rs | 38 +- .../region_engine/prepared_series_tests.rs | 73 ++ src/table/src/table/scan.rs | 21 +- src/table/src/table/scan/series_tests.rs | 18 +- .../wire/raw_series_label_decode/README.md | 75 ++ .../raw_series_label_decode/v1_metadata.json | 111 +++ .../v1_physical_types.json | 102 +++ 25 files changed, 3265 insertions(+), 602 deletions(-) create mode 100644 tests/compatibility/wire/raw_series_label_decode/README.md create mode 100644 tests/compatibility/wire/raw_series_label_decode/v1_metadata.json create mode 100644 tests/compatibility/wire/raw_series_label_decode/v1_physical_types.json diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs index e202e78c4ca..d787106c634 100644 --- a/src/common/query/src/series_frame.rs +++ b/src/common/query/src/series_frame.rs @@ -15,10 +15,18 @@ //! Streaming boundary between finite raw-series chunks and DataFusion record batches. //! //! Schema fields identify the ordered logical payload fields and codec layout; additional -//! application metadata is permitted when stable across the stream. The caller-provided -//! [`RawLabelSource`] is the label binding, not authenticated by the frame schema or field -//! provenance. This adapter preserves scan candidates only; it is not a residual filter, -//! authorization check, remote transport, or compaction boundary. Only +//! application metadata is permitted when stable across the stream. A key-encoded frame also +//! declares its own label decode description, and the frame is the sole source of decode +//! information: [`unpack_raw_series`] reads the fields, the label mode, the description, and the +//! producer ordering from the received schema, keeps the description on the returned stream, and +//! therefore needs no [`RawLabelSource`]. A caller that holds a scan-bound identity — the physical +//! region key layout a decoded identity came from — compares that identity's declared label +//! decode layout with [`ensure_frame_matches_label_source`] instead of passing it as a decode +//! input; that comparison is equality of the description alone and proves nothing about region +//! identity or provenance, and neither side is authenticated as field provenance. This adapter +//! preserves +//! scan candidates only; it is not a residual filter, authorization check, remote transport, or +//! compaction boundary. Only //! [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] are accepted; //! `Interleaved` is rejected before any poll. @@ -32,14 +40,19 @@ use datafusion_common::DataFusionError; use datatypes::arrow::datatypes::{Schema as ArrowSchema, SchemaRef as ArrowSchemaRef}; use datatypes::error::InvalidSeriesSnafu; use datatypes::series::frame::{ + LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, LabelDecodeSpec, LabelDecodeTag, RawLabelMode, RawSeriesChunkCursor, RawSeriesFrameBuilder, RawSeriesFrameDecoder, - raw_series_frame_schema, + raw_series_frame_label_decode, raw_series_frame_schema_with_label_decode, + raw_series_frame_series_schema, }; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SeriesSchema; use datatypes::series::{MetricName, SeriesOrdering}; use futures_util::Stream; -use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::{ + NO_BOUND_LABEL_SOURCE, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, +}; mod record_batches; @@ -50,13 +63,35 @@ mod tests; const ORDERING_KEY: &str = "greptime.series.ordering"; -/// Builds the canonical raw frame schema with the producer's guaranteed ordering. +/// Builds the canonical raw frame schema with the producer's label binding and ordering. /// /// Ordering records a producer guarantee; it is not inferred from frame contents. `Interleaved` -/// has no v1 wire representation and is rejected before the schema is assembled. +/// has no v1 wire representation and is rejected before the schema is assembled. The label binding +/// determines both the label mode and, for a primary-key binding, the self-describing label decode +/// description written into the schema, so a frame can never declare a label layout it does not +/// describe. pub fn raw_series_frame_schema_with_ordering( + schema: &SeriesSchema, + label_source: &RawLabelSource, + ordering: SeriesOrdering, +) -> datafusion_common::Result { + frame_schema_with( + schema, + label_mode(label_source), + label_decode_spec(label_source).as_ref(), + ordering, + ) +} + +/// Assembles a declared frame schema from an explicit label mode and decode description. +/// +/// The mode and description pair is checked by +/// [`raw_series_frame_schema_with_label_decode`], so a key mode without a description and an +/// explicit mode with one are both rejected here instead of reaching the wire. +fn frame_schema_with( schema: &SeriesSchema, mode: RawLabelMode, + decode: Option<&LabelDecodeSpec>, ordering: SeriesOrdering, ) -> datafusion_common::Result { let ordering = match ordering { @@ -68,7 +103,8 @@ pub fn raw_series_frame_schema_with_ordering( )); } }; - let schema = raw_series_frame_schema(schema, mode).map_err(codec_error)?; + let schema = + raw_series_frame_schema_with_label_decode(schema, mode, decode).map_err(codec_error)?; let mut metadata = schema.metadata().clone(); metadata.insert(ORDERING_KEY.to_string(), ordering.to_string()); Ok(Arc::new(ArrowSchema::new_with_metadata( @@ -77,6 +113,56 @@ pub fn raw_series_frame_schema_with_ordering( ))) } +/// Derives the label decode description a frame must carry for this binding. +/// +/// A primary-key binding reduces to exactly what decoding reads: the metric name, the primary-key +/// encoding, the ordered key columns with their physical types, and the logical tags. Region +/// metadata guarantees the key list holds every tag column and nothing else, so the description +/// needs no per-column role. Explicit bindings already carry complete labels and declare no +/// description. +/// +/// The result is canonical: its logical tags are in [`LabelDecodeSpec::canonicalize`] order, which +/// is what the writer serializes, so a derived description and a received one for the same labels +/// are the same value and the same text. Derivation cannot fail, because a physical type is copied +/// rather than translated; a description a reader cannot honour is rejected when it is read, and a +/// decoder rejects the tags it cannot bind. +pub fn label_decode_spec(source: &RawLabelSource) -> Option { + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = source + else { + return None; + }; + let primary_key = physical_metadata + .primary_key_columns() + .map(|column| LabelDecodeColumn { + column_id: column.column_id, + name: column.column_schema.name.clone(), + physical_type: column.column_schema.data_type.clone(), + }) + .collect(); + let mut description = LabelDecodeSpec { + version: LABEL_DECODE_VERSION, + metric_name: metric_name.to_string(), + encoding: match physical_metadata.primary_key_encoding { + PrimaryKeyEncoding::Dense => LabelDecodeEncoding::Dense, + PrimaryKeyEncoding::Sparse => LabelDecodeEncoding::Sparse, + }, + primary_key, + logical_tags: logical_tags + .iter() + .map(|tag| LabelDecodeTag { + column_id: tag.column_id, + name: tag.name.to_string(), + }) + .collect(), + }; + description.canonicalize(); + Some(description) +} + /// Reads the declared producer ordering, conservatively treating legacy frames as contiguous. /// /// Only the two orderings representable on the wire can be returned. @@ -99,17 +185,23 @@ pub fn raw_series_frame_ordering( /// work bound, not a resident-memory cap. A ready batch is yielded without polling the source /// again. Producers must keep each series' chunks contiguous. Validation checks adjacent chunk /// conditions but does not certify nonadjacent recurrence for `Contiguous` ordering. +/// +/// The stream's own declaration is the only decode input: a frame unpacked without a scan-bound +/// [`RawLabelSource`] round trips unchanged instead of being re-declared from a binding it does not +/// have, and nothing is derived from a binding here. The declaration also decides the frame's label +/// mode, so a chunk whose labels disagree with it is rejected when that chunk is coded. pub fn pack_raw_series( source: SendableRawSeriesStream, ) -> datafusion_common::Result { let ordering = source.ordering(); ensure_supported_ordering(ordering)?; let schema = source.schema().clone(); - let label_source = source.label_source().clone(); - let mode = label_mode(&label_source); - // The declared frame schema carries the producer ordering; the shared builder's schema - // intentionally omits it, so every ready batch is declared before it is yielded. - let frame_schema = raw_series_frame_schema_with_ordering(&schema, mode, ordering)?; + let decode = source.label_decode_spec().cloned(); + let mode = declared_label_mode(&decode); + // The declared frame schema carries the producer ordering and the label decode description; + // the shared builder's schema intentionally omits both, so every ready batch is declared under + // this schema before it is yielded, including when no row is ever produced. + let frame_schema = frame_schema_with(&schema, mode, decode.as_ref(), ordering)?; let builder = RawSeriesFrameBuilder::try_new( &schema, mode, @@ -129,33 +221,48 @@ pub fn pack_raw_series( })) } -/// Unpacks DataFusion batches using the caller's explicit scan binding and ordering. +/// Unpacks DataFusion batches declared by [`pack_raw_series`]. /// -/// Only [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] streams are -/// accepted; an `Interleaved` declaration is rejected before the stream is returned, including -/// for an empty stream. Producers must keep each series' chunks contiguous. Validation checks -/// adjacent chunk conditions but does not certify nonadjacent recurrence for `Contiguous` -/// ordering. +/// The frame is the sole source of decode information: the payload fields, the label mode, the +/// label decode description, and the producer ordering are all read from the stream's own schema, +/// so unpacking a frame a producer wrote needs no local copy of its binding. Only +/// [`SeriesOrdering::Contiguous`] and [`SeriesOrdering::KeyThenTime`] streams are accepted; an +/// `Interleaved` declaration is rejected before the stream is returned, including for an empty +/// stream. Producers must keep each series' chunks contiguous. Validation checks adjacent chunk +/// conditions but does not certify nonadjacent recurrence for `Contiguous` ordering. /// -/// The supplied binding is retained unchanged and is not validated as field provenance. +/// Every frame is validated here, unconditionally. The declared layout must be this revision's, the +/// label decode description must parse, name a known version, and agree with the declared label +/// mode — a key frame must declare one and an explicit frame must not, since explicit labels are +/// already complete. The declared fields and the empty batch are then decoded once, so a stream that +/// could not be decoded is rejected before it is returned. +/// +/// `expected` is a check, never an input. When it is given it must match the schema the frame +/// declares — the same logical fields in the same order, and no evaluation grid — and a mismatch is +/// an error. The decoder and the returned stream's schema are built from the frame either way, so a +/// caller cannot supply a field layout the frame does not declare. pub fn unpack_raw_series( source: DfSendableRecordBatchStream, - expected: Arc, - label_source: RawLabelSource, - ordering: SeriesOrdering, + expected: Option>, ) -> datafusion_common::Result { - ensure_supported_ordering(ordering)?; - let mode = label_mode(&label_source); let frame_schema = source.schema(); - let decoder = RawSeriesFrameDecoder::try_new(&expected, mode).map_err(codec_error)?; + let declared = raw_series_frame_label_decode(frame_schema.as_ref()).map_err(codec_error)?; + let mode = declared_label_mode(&declared); + let schema = + Arc::new(raw_series_frame_series_schema(frame_schema.as_ref()).map_err(codec_error)?); + if let Some(expected) = &expected { + ensure_declared_schema(&schema, expected)?; + } + let ordering = raw_series_frame_ordering(frame_schema.as_ref())?; + let decoder = RawSeriesFrameDecoder::try_new(&schema, mode).map_err(codec_error)?; // Validate even empty streams' declared schema and layout before returning the stream. let empty = DfRecordBatch::new_empty(frame_schema.clone()); decoder.decode(&empty).map_err(codec_error)?; Ok(Box::pin(UnpackStream { source: Some(source), - schema: expected, + schema, frame_schema, - label_source, + label_decode_spec: declared, ordering, validator: SequenceValidator::new(ordering), pending: Vec::new().into_iter(), @@ -163,6 +270,66 @@ pub fn unpack_raw_series( })) } +/// Rejects an `expected` series schema the frame does not declare. +/// +/// Only the two things a caller can get wrong are checked, because those are the only two things +/// [`SeriesSchema`] holds: the logical fields and their order, and an evaluation grid a raw frame +/// cannot carry. +fn ensure_declared_schema( + declared: &SeriesSchema, + expected: &SeriesSchema, +) -> datafusion_common::Result<()> { + if expected.eval_grid.is_some() { + return Err(invalid( + "expected raw series schema declares an evaluation grid the frame cannot carry", + )); + } + if expected.fields.len() != declared.fields.len() + || expected + .fields + .iter() + .zip(declared.fields.iter()) + .any(|(expected, declared)| expected.name != declared.name) + { + return Err(invalid( + "raw frame fields do not match the expected series schema", + )); + } + Ok(()) +} + +/// Checks that a frame declares the label decode layout a scan-bound binding describes. +/// +/// This compares one thing: the description parsed from the frame's +/// `greptime.series.label_decode` metadata against the description derived from `label_source`, by +/// plain equality. A difference in the version, the metric name, the primary-key encoding, the +/// ordered key columns and their physical types, or the logical tags is a mismatch. Nothing else is +/// compared, so agreeing descriptions do **not** establish that these are the region's chunks: no +/// region id, table id, or provenance is part of either side, and two regions whose tags have the +/// same ids, names, and types produce equal descriptions. The check answers "does this frame describe +/// the label layout this binding expects", never "is this the region I asked for". +/// +/// It belongs to the caller that holds both sides: a frame describes its own labels, so +/// [`unpack_raw_series`] decodes it without a binding, and a caller that received the frame from a +/// peer compares the declaration here instead of decoding under whichever side it trusts. +/// +/// The comparison is plain equality. It used to canonicalize logical tag order on both sides, which +/// gave local comparisons one rule and the byte-exact remote comparison another; the writer now +/// canonicalizes instead (see [`LabelDecodeSpec::canonicalize`]), so a description that reaches a +/// reader already has one spelling and equality says everything there is to say about it. +pub fn ensure_frame_matches_label_source( + frame: &ArrowSchema, + label_source: &RawLabelSource, +) -> datafusion_common::Result<()> { + let declared = raw_series_frame_label_decode(frame).map_err(codec_error)?; + if declared != label_decode_spec(label_source) { + return Err(invalid( + "raw frame label decode description does not match the supplied binding", + )); + } + Ok(()) +} + fn label_mode(source: &RawLabelSource) -> RawLabelMode { match source { RawLabelSource::PrimaryKey { .. } => RawLabelMode::Key, @@ -170,6 +337,15 @@ fn label_mode(source: &RawLabelSource) -> RawLabelMode { } } +/// Returns the label mode a frame declares: a key frame describes its key, an explicit frame +/// carries complete labels and declares no description. +fn declared_label_mode(declared: &Option) -> RawLabelMode { + match declared { + Some(_) => RawLabelMode::Key, + None => RawLabelMode::Explicit, + } +} + /// Rejects orderings this adapter cannot validate; a bad declaration is an error, not a fallback. fn ensure_supported_ordering(ordering: SeriesOrdering) -> datafusion_common::Result<()> { match ordering { @@ -408,7 +584,8 @@ struct UnpackStream { source: Option, schema: Arc, frame_schema: datatypes::arrow::datatypes::SchemaRef, - label_source: RawLabelSource, + /// The description parsed from the frame schema; the only input decoding needs. + label_decode_spec: Option, ordering: SeriesOrdering, validator: SequenceValidator, pending: std::vec::IntoIter, @@ -420,8 +597,18 @@ impl RawSeriesStream for UnpackStream { &self.schema } + /// Returns [`NO_BOUND_LABEL_SOURCE`]: a frame establishes its own label layout and carries no + /// scan-bound identity, so decoding reads [`Self::label_decode_spec`] and never this. There is + /// no physical region key layout to report and none is invented from the description, which + /// deliberately carries no region metadata. fn label_source(&self) -> &RawLabelSource { - &self.label_source + &NO_BOUND_LABEL_SOURCE + } + + /// Returns the label decode description the frame declared, which is what decoding reads; an + /// explicit-label frame declares none. + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + self.label_decode_spec.as_ref() } fn ordering(&self) -> SeriesOrdering { diff --git a/src/common/query/src/series_frame/record_batches/tests.rs b/src/common/query/src/series_frame/record_batches/tests.rs index b24ac5b9431..7d4948e7d18 100644 --- a/src/common/query/src/series_frame/record_batches/tests.rs +++ b/src/common/query/src/series_frame/record_batches/tests.rs @@ -1130,13 +1130,7 @@ async fn exact_int64_filter_sort_adapter_frame_pipeline() { .into(), eval_grid: None, }); - let mut decoded = super::super::unpack_raw_series( - frames, - expected, - store_api::region_engine::RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut decoded = super::super::unpack_raw_series(frames, Some(expected)).unwrap(); let chunk = decoded.next().await.unwrap().unwrap(); assert!(chunk.is_last); assert_eq!(chunk.time.as_buffer().as_ref(), &[1, 2]); diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs index feabbe9d019..e306fa98ddf 100644 --- a/src/common/query/src/series_frame/tests.rs +++ b/src/common/query/src/series_frame/tests.rs @@ -28,13 +28,13 @@ use datafusion::physical_plan::RecordBatchStream; use datatypes::arrow::array::{Array, Float64Array, UInt32Array}; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; use datatypes::arrow::compute::concat_batches; -use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use datatypes::arrow::datatypes::{DataType, Schema, SchemaRef}; use datatypes::arrow::ipc::{reader::StreamReader, writer::StreamWriter}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::series::Label; use datatypes::series::MetricName; use datatypes::series::SeriesOrdering; -use datatypes::series::frame::{RawLabelMode, raw_series_frame_schema}; +use datatypes::series::frame::{LabelDecodeSpec, RawLabelMode, raw_series_frame_schema}; use datatypes::series::histogram::{BucketSideView, CounterResetHint, HistogramBatchBuilder}; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ @@ -45,7 +45,10 @@ use store_api::region_engine::{ BoundTag, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, }; -use super::{pack_raw_series, raw_series_frame_ordering, unpack_raw_series}; +use super::{ + ensure_frame_matches_label_source, label_decode_spec, pack_raw_series, + raw_series_frame_ordering, unpack_raw_series, +}; fn schema() -> Arc { Arc::new(SeriesSchema { @@ -101,6 +104,8 @@ struct RawMock { schema: Arc, items: VecDeque>, binding: RawLabelSource, + /// The description this fixture declares, derived once from `binding`. + label_decode_spec: Option, ordering: SeriesOrdering, polls: Arc, drops: Arc, @@ -144,6 +149,9 @@ impl RawSeriesStream for RawMock { fn label_source(&self) -> &RawLabelSource { &self.binding } + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + self.label_decode_spec.as_ref() + } fn ordering(&self) -> SeriesOrdering { self.ordering } @@ -239,6 +247,7 @@ fn raw_never_polled( schema: schema(), items: VecDeque::new(), binding: RawLabelSource::Explicit, + label_decode_spec: None, ordering, polls: polls.clone(), drops: drops.clone(), @@ -260,11 +269,15 @@ fn raw_items( ) -> (SendableRawSeriesStream, Arc, Arc) { let polls = Arc::new(AtomicUsize::new(0)); let drops = Arc::new(AtomicUsize::new(0)); + // A fixture holding a key binding declares the layout it encodes, like a real producer; a + // fixture with explicit labels already carries complete ones and declares none. + let label_decode_spec = label_decode_spec(&binding); ( Box::pin(RawMock { schema, items: items.into(), binding, + label_decode_spec, ordering, polls: polls.clone(), drops: drops.clone(), @@ -382,13 +395,7 @@ async fn contiguous_adapter_does_not_certify_nonadjacent_recurrence() { assert_eq!(batches[0].num_rows(), 3); let frame_schema = batches[0].schema(); let (source, _, _) = df_source(batches, frame_schema); - let mut decoded = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut decoded = unpack_raw_series(source, Some(schema())).unwrap(); for (key, time, value) in expected { let actual = decoded.next().await.unwrap().unwrap(); @@ -484,13 +491,7 @@ async fn roundtrip_explicit_unordered_contiguous_preserves_complete_chunk_conten panic_on_poll: false, emitted: 0, }); - let mut output = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut output = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &output.next().await.unwrap().unwrap(), b"complete-b", @@ -535,13 +536,7 @@ async fn key_then_time_valid_roundtrip() { panic_on_poll: false, emitted: 0, }); - let mut out = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .unwrap(); + let mut out = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &out.next().await.unwrap().unwrap(), b"a", @@ -667,10 +662,20 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { .unwrap(), ); } - let frame_schema = frames[0].schema(); + // The ordering is a property of the frame, so the frames must declare the case's one for + // unpack to validate the same sequence the producer rejected. + let frame_schema = super::raw_series_frame_schema_with_ordering( + &schema(), + &RawLabelSource::Explicit, + case.order, + ) + .unwrap(); + let frames = frames + .into_iter() + .map(|frame| frame.with_schema(frame_schema.clone()).unwrap()) + .collect(); let (source, _, _) = df_source(frames, frame_schema); - let mut decoded = - unpack_raw_series(source, schema(), RawLabelSource::Explicit, case.order).unwrap(); + let mut decoded = unpack_raw_series(source, Some(schema())).unwrap(); for (key, time, last, host) in &case.prefix { let chunk = decoded.next().await.unwrap().unwrap(); assert_chunk(&chunk, key, &[*time], &[Some(*time as f64)], *last); @@ -693,13 +698,7 @@ async fn invalid_sequence_reports_valid_prefix_then_one_terminal_error() { .unwrap(); let frame_schema = frame[0].schema(); let (source, _, _) = df_source(frame, frame_schema); - let mut decoded = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut decoded = unpack_raw_series(source, Some(schema())).unwrap(); let first = decoded.next().await.unwrap().unwrap(); assert_chunk(&first, b"unfinished", &[7], &[Some(7.0)], false); assert_eq!( @@ -750,8 +749,11 @@ fn primary_binding(metadata: store_api::metadata::RegionMetadataRef) -> RawLabel } } +/// A key frame states its whole label layout, so unpacking keeps exactly that description and no +/// binding: the description is what a consumer decodes from, and the physical region metadata the +/// producer read under never reaches the receiver. #[tokio::test] -async fn primary_key_binding_is_retained_exactly() { +async fn the_declared_description_is_retained_and_no_binding_is_invented() { use api::v1::SemanticType; use datatypes::schema::ColumnSchema; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; @@ -779,6 +781,7 @@ async fn primary_key_binding_is_retained_exactly() { .primary_key(vec![0]); let metadata = Arc::new(builder.build().unwrap()); let binding = primary_binding(metadata.clone()); + let description = label_decode_spec(&binding).expect("a primary-key binding describes its key"); let (input, _, _) = raw_custom( vec![ RawSeriesChunk::try_new( @@ -806,25 +809,25 @@ async fn primary_key_binding_is_retained_exactly() { panic_on_poll: false, emitted: 0, }); - let mut out = unpack_raw_series(source, schema(), binding, SeriesOrdering::Contiguous).unwrap(); + let mut out = unpack_raw_series(source, None).unwrap(); let decoded = out.next().await.unwrap().unwrap(); assert_eq!(decoded.primary_key.as_slice(), b"encoded-complete-key"); assert!(decoded.metric.is_none()); - let retained = out.label_source(); - match retained { - RawLabelSource::PrimaryKey { - metric_name, - physical_metadata, - logical_tags, - } => { - assert!(Arc::ptr_eq(physical_metadata, &metadata)); - assert_eq!(metric_name.as_ref(), "cpu"); - assert_eq!(logical_tags.as_ref().len(), 1); - assert_eq!(logical_tags[0].column_id, 0); - assert_eq!(logical_tags[0].name.as_ref(), "host"); - } - _ => panic!("binding changed"), - } + assert_eq!( + out.label_decode_spec(), + Some(&description), + "the declaration is retained exactly as the frame states it" + ); + assert_eq!(description.metric_name, "cpu"); + assert_eq!(description.primary_key.len(), 1); + assert_eq!(description.primary_key[0].column_id, 0); + assert_eq!(description.primary_key[0].name, "host"); + assert_eq!(description.logical_tags.len(), 1); + assert_eq!(description.logical_tags[0].name, "host"); + assert!( + matches!(out.label_source(), RawLabelSource::Explicit), + "the physical region metadata the producer read under is not carried across the frame" + ); } #[tokio::test] @@ -946,6 +949,7 @@ async fn locally_invalid_public_raw_chunk_is_rejected() { schema: schema(), items: VecDeque::from([Ok(bad)]), binding: RawLabelSource::Explicit, + label_decode_spec: None, ordering: SeriesOrdering::Contiguous, polls, drops: Arc::new(AtomicUsize::new(0)), @@ -977,13 +981,7 @@ async fn empty_streams_keep_self_describing_schema() { // A declared-but-empty stream is still validated against the expected schema before any // poll, and then ends without producing rows. let (source, polls, _) = df_source(vec![], frame_schema.clone()); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); assert_eq!(polls.load(Ordering::SeqCst), 0); assert_eq!(unpacked.schema().fields.len(), 1); assert_eq!(unpacked.schema().fields[0].name.as_ref(), "value"); @@ -1000,6 +998,7 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { schema: schema(), items: VecDeque::from([Err(e)]), binding: RawLabelSource::Explicit, + label_decode_spec: None, ordering: SeriesOrdering::Contiguous, polls: polls.clone(), drops: drops.clone(), @@ -1037,13 +1036,7 @@ async fn upstream_error_status_is_preserved_once_and_releases_source() { panic_on_poll: false, emitted: 0, }); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let error = unpacked.next().await.unwrap().unwrap_err(); assert_eq!(common_error::ext::ErrorExt::status_code(&error), status); assert!(unpacked.next().await.is_none()); @@ -1096,15 +1089,7 @@ async fn dropping_adapter_releases_upstream_and_pending_item_is_resumed() { async fn unpack_rejects_mismatched_batch_schema_before_polling() { let wrong = datatypes::arrow::datatypes::Schema::empty().into(); let (source, polls, drops) = df_source(vec![], wrong); - assert!( - unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous - ) - .is_err() - ); + assert!(unpack_raw_series(source, Some(schema())).is_err()); assert_eq!(polls.load(Ordering::SeqCst), 0); assert_eq!(drops.load(Ordering::SeqCst), 1); } @@ -1123,6 +1108,7 @@ async fn invalid_grid_schema_rejected_for_empty_raw_stream() { schema: Arc::new(invalid), items: VecDeque::new(), binding: RawLabelSource::Explicit, + label_decode_spec: None, ordering: SeriesOrdering::Contiguous, polls: polls.clone(), drops: Arc::new(AtomicUsize::new(0)), @@ -1138,9 +1124,11 @@ async fn invalid_grid_schema_rejected_for_empty_raw_stream() { #[tokio::test] async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { let mut batches = Vec::new(); - let frame_schema = datatypes::series::frame::raw_series_frame_schema( + // `KeyThenTime` is what makes a recurring key an error, so the frame has to declare it. + let frame_schema = super::raw_series_frame_schema_with_ordering( &schema(), - datatypes::series::frame::RawLabelMode::Explicit, + &RawLabelSource::Explicit, + SeriesOrdering::KeyThenTime, ) .unwrap(); for (key, time) in [ @@ -1153,13 +1141,7 @@ async fn unpack_batch_sequence_error_is_atomic_for_all_rows() { } let multi = concat_batches(&frame_schema, &batches).unwrap(); let (source, _, _) = df_source(vec![multi], frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); assert!(unpacked.next().await.unwrap().is_err()); assert!(unpacked.next().await.is_none()); } @@ -1216,13 +1198,7 @@ async fn unpack_rejects_late_invalid_frame_row_before_prefix() { second = RecordBatch::try_new(frame_schema.clone(), columns).unwrap(); let multi = concat_batches(&frame_schema, &[first, second]).unwrap(); let (source, _, _) = df_source(vec![multi], frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let error = unpacked.next().await.unwrap().unwrap_err(); assert_eq!( common_error::ext::ErrorExt::status_code(&error), @@ -1264,13 +1240,7 @@ async fn frame_application_metadata_is_stable_and_schema_change_rejected() { .with_schema(declared.clone()) .unwrap(); let (source, _, _) = df_source(vec![first.clone(), stable_second], declared.clone()); - let mut stable = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut stable = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &stable.next().await.unwrap().unwrap(), b"a", @@ -1305,13 +1275,7 @@ async fn frame_application_metadata_is_stable_and_schema_change_rejected() { .with_schema(changed_schema) .unwrap(); let (source, _, _) = df_source(vec![first.clone(), bad_batch], declared.clone()); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &unpacked.next().await.unwrap().unwrap(), b"a", @@ -1367,13 +1331,7 @@ async fn unpack_bounds_empty_batches_per_poll_and_eventually_yields_row() { let mut batches = vec![RecordBatch::new_empty(frame_schema.clone()); 10_000]; batches.push(row); let (source, polls, _) = df_source(batches, frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); let waker = Waker::from(wake.clone()); @@ -1427,13 +1385,7 @@ fn unpack_yields_at_empty_batch_budget_before_observing_eof() { .unwrap(); let batches = vec![RecordBatch::new_empty(frame_schema.clone()); BUDGET]; let (source, polls, _) = df_source(batches, frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); let waker = Waker::from(wake.clone()); @@ -1462,13 +1414,7 @@ fn unpack_propagates_upstream_pending_without_extra_wake() { frame_schema, 1, ); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); let waker = Waker::from(wake.clone()); @@ -1493,13 +1439,7 @@ fn unpack_reports_bad_batch_at_budget_boundary_before_yielding() { datatypes::arrow::datatypes::Schema::empty().into(), )); let (source, polls, _) = df_source(batches, frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); let waker = Waker::from(wake.clone()); @@ -1541,13 +1481,7 @@ fn unpack_surfaces_upstream_error_after_budget_yield_once() { panic_on_poll: false, emitted: 0, }); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); let wake = Arc::new(CountingWake(AtomicUsize::new(0))); let waker = Waker::from(wake.clone()); @@ -1579,13 +1513,7 @@ async fn unpack_drains_decoded_rows_before_polling_another_batch() { let frame_schema = batches[0].schema(); let multi = concat_batches(&frame_schema, &batches).unwrap(); let (source, polls, _) = df_source(vec![multi], frame_schema); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &unpacked.next().await.unwrap().unwrap(), b"a", @@ -1633,15 +1561,7 @@ async fn unpack_preflights_empty_and_grid_schemas_without_polling() { ) .unwrap(); let (source, polls, drops) = df_source(vec![], frame_schema); - assert!( - unpack_raw_series( - source, - Arc::new(expected), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .is_err() - ); + assert!(unpack_raw_series(source, Some(Arc::new(expected))).is_err()); assert_eq!(polls.load(Ordering::SeqCst), 0); assert_eq!(drops.load(Ordering::SeqCst), 1); } @@ -1683,6 +1603,7 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() schema: expected.clone(), items: VecDeque::from([Ok(input_chunk)]), binding: RawLabelSource::Explicit, + label_decode_spec: None, ordering: SeriesOrdering::KeyThenTime, polls: Arc::new(AtomicUsize::new(0)), drops: Arc::new(AtomicUsize::new(0)), @@ -1716,13 +1637,7 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() ); let (source, _, _) = df_source(ipc_batches.clone(), frame_schema.clone()); - let mut decoded = unpack_raw_series( - source, - expected.clone(), - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .unwrap(); + let mut decoded = unpack_raw_series(source, Some(expected.clone())).unwrap(); let actual = decoded.next().await.unwrap().unwrap(); assert_chunk( &actual, @@ -1750,15 +1665,7 @@ async fn named_frame_fields_roundtrip_and_reject_wrong_bindings_before_polling() let wrong_schema = named_fields_schema(wrong_names); let batches = if empty { vec![] } else { ipc_batches.clone() }; let (source, polls, drops) = df_source(batches, frame_schema.clone()); - assert!( - unpack_raw_series( - source, - wrong_schema, - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .is_err() - ); + assert!(unpack_raw_series(source, Some(wrong_schema)).is_err()); assert_eq!(polls.load(Ordering::SeqCst), 0); assert_eq!(drops.load(Ordering::SeqCst), 1); } @@ -2193,13 +2100,7 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { let frame_schema = reader.schema(); let ipc_batches = reader.map(Result::unwrap).collect::>(); let (source, _, _) = df_source(ipc_batches, frame_schema); - let mut decoded = unpack_raw_series( - source, - schema.clone(), - RawLabelSource::Explicit, - SeriesOrdering::KeyThenTime, - ) - .unwrap(); + let mut decoded = unpack_raw_series(source, Some(schema.clone())).unwrap(); struct Expected { time: i64, @@ -2391,7 +2292,7 @@ async fn interleaved_ordering_is_rejected_before_any_poll() { assert!( super::raw_series_frame_schema_with_ordering( &schema(), - datatypes::series::frame::RawLabelMode::Explicit, + &RawLabelSource::Explicit, SeriesOrdering::Interleaved, ) .is_err(), @@ -2414,20 +2315,27 @@ async fn interleaved_ordering_is_rejected_before_any_poll() { ); assert_eq!(drops.load(Ordering::SeqCst), 1); - let frame_schema = datatypes::series::frame::raw_series_frame_schema( + // An unpack has no ordering argument, so the only way an `Interleaved` stream could reach it is + // a frame that declares one; that spelling has no wire representation and is rejected. + let frame = datatypes::series::frame::raw_series_frame_schema( &schema(), datatypes::series::frame::RawLabelMode::Explicit, ) .unwrap(); - let (source, polls, drops) = df_never_polled(frame_schema); - assert!( - unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Interleaved, - ) - .is_err() + let mut metadata = frame.metadata().clone(); + metadata.insert( + "greptime.series.ordering".to_string(), + "interleaved".to_string(), + ); + let frame = Arc::new(Schema::new_with_metadata(frame.fields().clone(), metadata)); + let (source, polls, drops) = df_never_polled(frame); + let error = match unpack_raw_series(source, Some(schema())) { + Ok(_) => panic!("an interleaved declaration must be rejected"), + Err(error) => error, + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments ); assert_eq!(polls.load(Ordering::SeqCst), 0); assert_eq!(drops.load(Ordering::SeqCst), 1); @@ -2482,13 +2390,7 @@ async fn unpack_rejects_unrelated_metadata_drift_and_empty_schema_changes() { RecordBatch::try_new(changed, first.columns().to_vec()).unwrap() }; let (source, _, _) = df_source(vec![first.clone(), bad], base.clone()); - let mut unpacked = unpack_raw_series( - source, - schema(), - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut unpacked = unpack_raw_series(source, Some(schema())).unwrap(); assert_chunk( &unpacked.next().await.unwrap().unwrap(), b"a", @@ -2609,8 +2511,7 @@ async fn pack_declares_ordering_on_every_batch_and_ipc_roundtrips_it() { assert_eq!(ipc_batches.len(), expected_batch_count); assert!(ipc_batches.iter().all(|batch| batch.schema() == declared)); let (source, _, _) = df_source(ipc_batches, declared); - let mut decoded = - unpack_raw_series(source, schema(), RawLabelSource::Explicit, ordering).unwrap(); + let mut decoded = unpack_raw_series(source, Some(schema())).unwrap(); if is_empty { assert!(decoded.next().await.is_none()); } else { @@ -2628,3 +2529,552 @@ async fn pack_declares_ordering_on_every_batch_and_ipc_roundtrips_it() { } } } + +/// A primary-key binding over a region whose only tag is `host`. +fn key_binding(metric: &str) -> RawLabelSource { + use api::v1::SemanticType; + use datatypes::schema::ColumnSchema; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::storage::RegionId; + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + datatypes::data_type::ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + RawLabelSource::PrimaryKey { + metric_name: Arc::from(metric), + physical_metadata: metadata, + logical_tags: Arc::from([BoundTag { + column_id: 0, + name: Arc::from("host"), + }]), + } +} + +/// A primary-key binding over a region whose two tags are `host` then `zone`. +fn two_tag_binding(metric: &str) -> RawLabelSource { + use api::v1::SemanticType; + use datatypes::schema::ColumnSchema; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + use store_api::storage::RegionId; + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "host", + datatypes::data_type::ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "zone", + datatypes::data_type::ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }) + .primary_key(vec![0, 1]); + let metadata = Arc::new(builder.build().unwrap()); + RawLabelSource::PrimaryKey { + metric_name: Arc::from(metric), + physical_metadata: metadata, + logical_tags: Arc::from([ + BoundTag { + column_id: 0, + name: Arc::from("host"), + }, + BoundTag { + column_id: 1, + name: Arc::from("zone"), + }, + ]), + } +} + +/// A key frame whose metadata carries exactly the supplied description and nothing else. +fn key_frame_with(description: &datatypes::series::frame::LabelDecodeSpec) -> SchemaRef { + let layout = raw_series_frame_schema(&schema(), RawLabelMode::Key).unwrap(); + let mut metadata = layout.metadata().clone(); + metadata.insert( + "greptime.series.label_decode".to_string(), + serde_json::to_string(description).unwrap(), + ); + Arc::new(Schema::new_with_metadata(layout.fields().clone(), metadata)) +} + +/// Logical tag order is fixed by the writer, so the wire form has one spelling and the comparison +/// is plain equality. The ordered key columns are the encoding, so the same columns in another +/// order are a different layout and are rejected. +/// +/// The tag order is canonical rather than meaningful: the decoder sorts tags by name when it binds +/// them (`src/mito-codec/src/labels.rs:235`), so sorting in the writer cannot change a decoded label +/// set. What it does change is that the written text no longer depends on the caller's order, which +/// is what makes one comparison rule enough for both a local check and the byte-exact remote one. +#[tokio::test] +async fn frame_check_compares_the_canonical_description_and_key_order_exactly() { + let binding = two_tag_binding("cpu"); + let description = label_decode_spec(&binding).expect("a primary-key binding describes its key"); + assert_eq!( + description.logical_tags.len(), + 2, + "two tags to canonicalize" + ); + + // A description derived from the binding is already canonical, whichever order the binding + // lists its tags in. + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = &binding + else { + unreachable!("the fixture is a primary-key binding") + }; + let mut reversed_tags = logical_tags.to_vec(); + reversed_tags.reverse(); + let reordered_binding = RawLabelSource::PrimaryKey { + metric_name: metric_name.clone(), + physical_metadata: physical_metadata.clone(), + logical_tags: reversed_tags.into(), + }; + assert_eq!( + label_decode_spec(&reordered_binding).expect("a primary-key binding describes its key"), + description, + "a binding's tag order does not reach the description" + ); + let reordered = + label_decode_spec(&reordered_binding).expect("a primary-key binding describes its key"); + assert_eq!( + key_frame_with(&reordered).metadata(), + key_frame_with(&description).metadata(), + "and so two producers with the same tags write the same bytes" + ); + + // A frame that is not canonical is compared as written instead of being guessed equal. + let mut non_canonical = description.clone(); + non_canonical.logical_tags.reverse(); + let frame = key_frame_with(&non_canonical); + assert!( + ensure_frame_matches_label_source(frame.as_ref(), &binding).is_err(), + "the reader does not reorder what a writer should have ordered" + ); + + // The ordered key columns are the encoding and are never reordered. + let mut reordered_key = description; + reordered_key.primary_key.reverse(); + let frame = key_frame_with(&reordered_key); + assert!( + ensure_frame_matches_label_source(frame.as_ref(), &binding).is_err(), + "primary-key column order is the encoding and is compared as written" + ); +} + +/// The identity check a caller performs on its own compares the frame's declaration with the +/// binding it holds: a different metric, another key layout, or an explicit frame with no +/// declaration at all is a disagreement, and the frame alone still decides what is decoded. +#[tokio::test] +async fn frame_check_rejects_a_binding_the_frame_does_not_describe() { + let description = + label_decode_spec(&key_binding("cpu")).expect("a primary-key binding describes its key"); + let frame = key_frame_with(&description); + assert!( + ensure_frame_matches_label_source(frame.as_ref(), &key_binding("cpu")).is_ok(), + "the same binding is the same layout" + ); + assert!( + ensure_frame_matches_label_source(frame.as_ref(), &key_binding("other_metric")).is_err(), + "another metric is another label layout" + ); + assert!( + ensure_frame_matches_label_source(frame.as_ref(), &RawLabelSource::Explicit).is_err(), + "an explicit identity expects a frame that carries complete labels" + ); + + let explicit = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + assert!( + ensure_frame_matches_label_source(explicit.as_ref(), &key_binding("cpu")).is_err(), + "a frame declaring no layout cannot satisfy a primary-key identity" + ); + assert!( + ensure_frame_matches_label_source(explicit.as_ref(), &RawLabelSource::Explicit).is_ok(), + "an explicit frame is what an explicit identity expects" + ); +} + +/// The description parsed at unpack time is retained on the returned stream, so a consumer that +/// holds only the frame can decode; an explicit frame retains none. +#[tokio::test] +async fn unpack_retains_the_declared_description_without_a_binding() { + let binding = key_binding("cpu"); + let description = label_decode_spec(&binding).expect("a primary-key binding describes its key"); + + let (source, _, _) = df_source(vec![], key_frame_with(&description)); + let unpacked = unpack_raw_series(source, None).unwrap(); + assert_eq!(unpacked.label_decode_spec(), Some(&description)); + assert!( + matches!(unpacked.label_source(), RawLabelSource::Explicit), + "no physical region key layout is invented for a frame that carries none" + ); + + let explicit = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let (source, _, _) = df_source(vec![], explicit); + let unpacked = unpack_raw_series(source, None).unwrap(); + assert_eq!( + unpacked.label_decode_spec(), + None, + "explicit labels are already complete and need no decoder" + ); +} + +/// A key frame whose `greptime.series.label_decode` entry is exactly the supplied text. +fn key_frame_with_encoded(encoded: &str) -> SchemaRef { + let layout = raw_series_frame_schema(&schema(), RawLabelMode::Key).unwrap(); + let mut metadata = layout.metadata().clone(); + metadata.insert( + "greptime.series.label_decode".to_string(), + encoded.to_string(), + ); + Arc::new(Schema::new_with_metadata(layout.fields().clone(), metadata)) +} + +/// The description is validated whether or not a binding is supplied, so trusting the frame never +/// means skipping validation: an unknown version, an unknown field, a malformed entry, a logical +/// tag outside the key, a key frame without a description, and an explicit frame carrying one are +/// all rejected before the stream is returned. +#[tokio::test] +async fn unpack_rejects_invalid_descriptions_without_a_binding() { + let description = + label_decode_spec(&key_binding("cpu")).expect("a primary-key binding describes its key"); + let encoded = |value: &serde_json::Value| serde_json::to_string(value).unwrap(); + let mut unknown_version = serde_json::to_value(&description).unwrap(); + unknown_version["version"] = + serde_json::json!(datatypes::series::frame::LABEL_DECODE_VERSION + 1); + let mut unknown_field = serde_json::to_value(&description).unwrap(); + unknown_field["future_field"] = serde_json::json!(1); + let mut unlisted_tag = serde_json::to_value(&description).unwrap(); + unlisted_tag["logical_tags"] = serde_json::json!([{ "column_id": 7, "name": "absent" }]); + let cases = [ + ("unknown version", encoded(&unknown_version)), + ("unknown field", encoded(&unknown_field)), + ( + "logical tag outside the primary key", + encoded(&unlisted_tag), + ), + ("malformed entry", "{".to_string()), + ]; + for (label, text) in cases { + let (source, polls, _) = df_source(vec![], key_frame_with_encoded(&text)); + let error = match unpack_raw_series(source, None) { + Ok(_) => panic!("{label}: an invalid description must be rejected"), + Err(error) => error, + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments, + "{label}" + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "{label}: a rejected frame is never polled" + ); + } + + let key_without_description = raw_series_frame_schema(&schema(), RawLabelMode::Key).unwrap(); + let (source, _, _) = df_source(vec![], key_without_description); + assert!( + unpack_raw_series(source, None).is_err(), + "a key frame nobody can decode is rejected instead of being trusted" + ); + + let explicit = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let mut metadata = explicit.metadata().clone(); + metadata.insert( + "greptime.series.label_decode".to_string(), + encoded(&serde_json::to_value(&description).unwrap()), + ); + let explicit_with_description = Arc::new(Schema::new_with_metadata( + explicit.fields().clone(), + metadata, + )); + let (source, _, _) = df_source(vec![], explicit_with_description); + assert!( + unpack_raw_series(source, None).is_err(), + "an explicit frame declares no description, so one is a contradiction" + ); +} + +fn key_chunk() -> RawSeriesChunk { + chunk_with(b"encoded-complete-key", &[1], vec![Some(3.0)], None, true) +} + +/// Every packed key frame declares the label decode description derived from its binding, so a +/// consumer needs no scan-side label source. Explicit frames declare none. +#[tokio::test] +async fn pack_declares_the_label_decode_description_of_its_binding() { + let binding = key_binding("cpu"); + let (input, _, _) = raw_custom( + vec![key_chunk()], + SeriesOrdering::Contiguous, + binding.clone(), + ); + let mut packed = pack_raw_series(input).unwrap(); + let declared = packed.schema(); + let canonical = raw_series_frame_schema(&schema(), RawLabelMode::Key).unwrap(); + assert_eq!( + declared.metadata().len(), + canonical.metadata().len() + 2, + "a key frame adds the producer ordering and the label decode description" + ); + assert_eq!( + datatypes::series::frame::raw_series_frame_label_decode(&declared).unwrap(), + label_decode_spec(&binding), + "the declared description is derived from the binding, not from the first row" + ); + let batch = packed.next().await.unwrap().unwrap(); + assert_eq!(batch.schema(), declared); + + let (input, _, _) = raw_custom( + vec![key_chunk()], + SeriesOrdering::Contiguous, + RawLabelSource::Explicit, + ); + let packed = pack_raw_series(input).unwrap(); + let declared = packed.schema(); + assert_eq!( + datatypes::series::frame::raw_series_frame_label_decode(&declared).unwrap(), + None, + "explicit labels are already complete" + ); + assert_eq!(declared.metadata().len(), canonical.metadata().len() + 1); +} + +/// A key frame that omits its description is not decodable, and unpacking says so instead of falling +/// back to a caller-side layout; the description that is present survives the round trip. +#[tokio::test] +async fn unpack_round_trips_a_key_frame_and_rejects_one_without_a_description() { + let binding = key_binding("cpu"); + let (input, _, _) = raw_custom( + vec![key_chunk()], + SeriesOrdering::Contiguous, + binding.clone(), + ); + let mut packed = pack_raw_series(input).unwrap(); + let declared = packed.schema(); + let batch = packed.next().await.unwrap().unwrap(); + + // The frame and the binding agree, so the round trip decodes. + let (source, _, _) = df_source(vec![batch.clone()], declared.clone()); + let mut decoded = unpack_raw_series(source, Some(schema())).unwrap(); + let chunk = decoded.next().await.unwrap().unwrap(); + assert_eq!(chunk.primary_key.as_slice(), b"encoded-complete-key"); + + // A frame that omits the description is not decodable from the schema alone. + let canonical = raw_series_frame_schema(&schema(), RawLabelMode::Key).unwrap(); + let mut metadata = canonical.metadata().clone(); + metadata.insert( + "greptime.series.ordering".to_string(), + "contiguous".to_string(), + ); + let stripped = Arc::new(datatypes::arrow::datatypes::Schema::new_with_metadata( + declared.fields().clone(), + metadata, + )); + let stripped_batch = RecordBatch::try_new(stripped.clone(), batch.columns().to_vec()).unwrap(); + let (source, _, _) = df_source(vec![stripped_batch], stripped); + let error = match unpack_raw_series(source, Some(schema())) { + Ok(_) => panic!("a key frame without a description must be rejected"), + Err(error) => error, + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments + ); +} + +/// The ordering a frame declares is the only ordering a consumer has: unpacking a `KeyThenTime` +/// frame reports `KeyThenTime` although nobody passed it in. +#[tokio::test] +async fn unpack_reports_the_ordering_the_frame_declares() { + let (input, _, _) = raw_custom( + vec![key_chunk()], + SeriesOrdering::KeyThenTime, + key_binding("cpu"), + ); + let packed = pack_raw_series(input).unwrap(); + assert_eq!( + raw_series_frame_ordering(packed.schema().as_ref()).unwrap(), + SeriesOrdering::KeyThenTime, + "the producer declaration reaches the wire" + ); + + let unpacked = unpack_raw_series(packed, None).unwrap(); + assert_eq!( + unpacked.ordering(), + SeriesOrdering::KeyThenTime, + "unpacking derives the ordering from the frame, not from a caller" + ); + + // A frame that declares no ordering is a legacy contiguous frame. + let frame = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let (source, _, _) = df_source(vec![], frame); + assert_eq!( + unpack_raw_series(source, None).unwrap().ordering(), + SeriesOrdering::Contiguous + ); +} + +/// Reserved layout metadata this revision does not write is rejected instead of being read as the +/// current layout, so a frame cannot be decoded under a contract it did not declare. +#[tokio::test] +async fn unpack_rejects_unknown_layout_metadata() { + let cases = [ + ("greptime.series.layout", "v2"), + ("greptime.series.content", "series"), + ("greptime.series.time", "timestamps_v2"), + ("greptime.series.ordering", "hash_then_time"), + ]; + for (key, value) in cases { + let frame = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + let mut metadata = frame.metadata().clone(); + metadata.insert(key.to_string(), value.to_string()); + let frame = Arc::new(Schema::new_with_metadata(frame.fields().clone(), metadata)); + let (source, polls, _) = df_source(vec![], frame); + assert!( + unpack_raw_series(source, None).is_err(), + "{key} = {value} is not a declaration this revision reads" + ); + assert_eq!(polls.load(Ordering::SeqCst), 0, "{key}"); + } +} + +/// `expected` is a check, not an input: a schema the frame does not declare is rejected, and the +/// same frame unpacks without it because the frame is the only source of decode information. +#[tokio::test] +async fn unpack_rejects_an_expected_schema_the_frame_does_not_declare() { + let frame = raw_series_frame_schema(&schema(), RawLabelMode::Explicit).unwrap(); + + let other_field = Arc::new(SeriesSchema { + fields: vec![FieldSpec { + name: Arc::from("other_value"), + }] + .into(), + eval_grid: None, + }); + let (source, polls, _) = df_source(vec![], frame.clone()); + let error = match unpack_raw_series(source, Some(other_field)) { + Ok(_) => panic!("a field layout the frame does not declare must be rejected"), + Err(error) => error, + }; + assert_eq!( + crate::error::datafusion_status_code::(&error, None), + StatusCode::InvalidArguments + ); + assert_eq!(polls.load(Ordering::SeqCst), 0, "nothing is polled"); + + let two_fields = Arc::new(SeriesSchema { + fields: vec![ + FieldSpec { + name: Arc::from("value"), + }, + FieldSpec { + name: Arc::from("second"), + }, + ] + .into(), + eval_grid: None, + }); + let (source, _, _) = df_source(vec![], frame.clone()); + assert!( + unpack_raw_series(source, Some(two_fields)).is_err(), + "the same names in a longer list are another layout" + ); + + let (source, _, _) = df_source(vec![], frame.clone()); + assert!( + unpack_raw_series(source, Some(schema())).is_ok(), + "the schema the frame declares is accepted" + ); + + let (source, _, _) = df_source(vec![], frame); + let unpacked = unpack_raw_series(source, None).unwrap(); + assert_eq!( + unpacked + .schema() + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>(), + vec!["value"], + "the returned schema is the frame's, whether or not an expectation was given" + ); +} + +/// Unpacking a frame and packing it again keeps both declarations, so a relayed stream is the same +/// stream: the description still decodes the same labels and the ordering still holds. +#[tokio::test] +async fn unpack_pack_unpack_preserves_the_description_and_the_ordering() { + let binding = key_binding("cpu"); + let description = label_decode_spec(&binding).expect("a primary-key binding describes its key"); + let (input, _, _) = raw_custom(vec![key_chunk()], SeriesOrdering::KeyThenTime, binding); + let packed = pack_raw_series(input).unwrap(); + assert_eq!( + datatypes::series::frame::raw_series_frame_label_decode(packed.schema().as_ref()).unwrap(), + Some(description.clone()) + ); + + let unpacked = unpack_raw_series(packed, None).unwrap(); + assert_eq!(unpacked.label_decode_spec(), Some(&description)); + assert_eq!(unpacked.ordering(), SeriesOrdering::KeyThenTime); + + let repacked = pack_raw_series(unpacked).unwrap(); + assert_eq!( + datatypes::series::frame::raw_series_frame_label_decode(repacked.schema().as_ref()) + .unwrap(), + Some(description.clone()), + "the repacked frame declares what the unpacked stream declared" + ); + assert_eq!( + raw_series_frame_ordering(repacked.schema().as_ref()).unwrap(), + SeriesOrdering::KeyThenTime, + "the repacked frame keeps the ordering it was unpacked under" + ); + + let mut again = unpack_raw_series(repacked, None).unwrap(); + assert_eq!(again.label_decode_spec(), Some(&description)); + assert_eq!(again.ordering(), SeriesOrdering::KeyThenTime); + let chunk = again.next().await.unwrap().unwrap(); + assert_eq!(chunk.primary_key.as_slice(), b"encoded-complete-key"); + assert!(chunk.metric.is_none(), "key chunks stay undecoded"); + assert!(again.next().await.is_none()); +} diff --git a/src/common/query/src/series_frame/tests/lifecycle.rs b/src/common/query/src/series_frame/tests/lifecycle.rs index b8558c89b03..54f088a5098 100644 --- a/src/common/query/src/series_frame/tests/lifecycle.rs +++ b/src/common/query/src/series_frame/tests/lifecycle.rs @@ -193,13 +193,7 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { chunk: Some(chunk), }); let frames = pack_raw_series(input).unwrap(); - let mut decoded = unpack_raw_series( - frames, - schema, - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut decoded = unpack_raw_series(frames, Some(schema)).unwrap(); let mut output = Vec::new(); while let Some(item) = decoded.next().await { output.push(item.unwrap()); @@ -425,13 +419,7 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() chunk: Some(chunk), }); let frames = pack_raw_series(input).unwrap(); - let mut decoded = unpack_raw_series( - frames, - schema, - RawLabelSource::Explicit, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let mut decoded = unpack_raw_series(frames, Some(schema)).unwrap(); let mut output = Vec::new(); while let Some(item) = decoded.next().await { output.push(item.unwrap()); diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs index da6ff38d1ae..2dcab10523c 100644 --- a/src/datatypes/src/series/frame.rs +++ b/src/datatypes/src/series/frame.rs @@ -18,11 +18,45 @@ //! chunk and [`RawSeriesFrameBuilder`] emits rows directly as Arrow arrays once a row or estimated //! byte bound is reached. Decoding reuses one prepared schema per stream through //! [`RawSeriesFrameDecoder`]. Column layout, metadata, and wire strings stay on the frozen `v1` -//! contract. +//! contract; a key-encoded frame additionally declares a versioned [`LabelDecodeSpec`] so its +//! labels can be decoded from the received schema alone. +//! +//! # Label decode description versioning +//! +//! The description is one `greptime.series.label_decode` metadata entry, and remote plans compare +//! `greptime.series.*` metadata **as exact text, byte for byte**: a frontend and a datanode that +//! write different text for the same labels disagree about the schema, and the query fails instead +//! of falling back. The rules that follow from that are: +//! +//! - A reader rejects a version it does not implement, rather than guessing. This module writes and +//! reads [`LABEL_DECODE_VERSION`] only. +//! - [`LABEL_DECODE_VERSION`] covers the whole text, including each key column's +//! [`LabelDecodeColumn::physical_type`]. That field is a [`ConcreteDataType`], so the text embeds +//! [`ConcreteDataType`]'s serde shape: it is not spelled here, and a change in that type's +//! serialization does alter this text. What keeps such a change from going unnoticed is the +//! frozen fixture `tests/compatibility/wire/raw_series_label_decode/v1_physical_types.json`, +//! which records the text of every datatype variant and fails the crate's tests when any of them +//! is serialized differently. +//! - Any change to the shape of this text — a renamed or added field, a different +//! [`ConcreteDataType`] spelling, a different `logical_tags` order — is a **breaking** change at +//! the same version, because it makes two nodes of different builds disagree. Such a change must +//! bump [`LABEL_DECODE_VERSION`] and update that fixture with it. +//! - The text is canonical: [`LabelDecodeSpec::canonicalize`] fixes the only part whose order +//! carries no meaning, so equal labels always write equal bytes. +//! +//! Deployment rule: during a rolling upgrade a frame may carry a version the receiving node does not +//! implement, and a byte-for-byte schema comparison may fail for that reason alone. The intended +//! behaviour is that such a mismatch disables the native series path at planning time and falls back +//! to the generic path, instead of failing the query at execution. +//! +//! TODO(narthex): implement that planning-time fallback. No planner reads +//! `raw_series_frame_label_decode` to decide support today, so a version or text mismatch currently +//! surfaces as an execution-time error from the raw-frame decode or the remote schema comparison. #[cfg(test)] mod tests; +use std::collections::HashSet; use std::sync::Arc; use arrow::array::{ @@ -33,12 +67,14 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::concat; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use serde::{Deserialize, Serialize}; +use crate::data_type::ConcreteDataType; use crate::error::{InvalidSeriesSnafu, Result}; use crate::series::histogram::{HistogramBatch, fields as histogram_fields}; use crate::series::raw::{RawSeriesChunk, RawTimestamps}; use crate::series::samples::{ - DenseSamplesRef, HistogramSamples, Samples, SeriesFields, SeriesSchema, + DenseSamplesRef, FieldSpec, HistogramSamples, Samples, SeriesFields, SeriesSchema, }; use crate::series::{Label, MetricName, invalid_series}; @@ -46,6 +82,12 @@ const LAYOUT: &str = "greptime.series.layout"; const CONTENT: &str = "greptime.series.content"; const LABELS: &str = "greptime.series.labels"; const TIME: &str = "greptime.series.time"; +const GRID: &str = "greptime.series.grid"; +/// Metadata key of the versioned label decode description carried by key-encoded frames. +const LABEL_DECODE: &str = "greptime.series.label_decode"; + +/// Version of the label decode description this revision writes and reads. +pub const LABEL_DECODE_VERSION: u32 = 1; /// Whether a raw frame carries only the complete key or explicit metric labels as well. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -56,7 +98,166 @@ pub enum RawLabelMode { Explicit, } +/// Primary-key encoding of the labels a key-encoded frame carries. +/// +/// This mirrors the storage contract's `PrimaryKeyEncoding` without making this crate depend on +/// it; Dense and Sparse keys use different decode layouts. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LabelDecodeEncoding { + /// Dense key: every primary-key field is encoded in order, with no ids on the wire. + Dense, + /// Sparse key: `(column id, value)` pairs, so the layout only frames the field stream. + Sparse, +} + +/// One ordered key column needed to decode key-encoded labels. +/// +/// Only tag columns reach this list: region metadata rejects a primary key that holds the time +/// index or a non-`Tag` column, and it requires the primary key to name every tag column. The +/// array's position is therefore the encoding position, with nothing to disambiguate. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LabelDecodeColumn { + /// Physical column id of this key column. + pub column_id: u32, + /// Physical column name; a logical tag of the same id must repeat it. + pub name: String, + /// Physical data type of the encoded value, in [`ConcreteDataType`]'s own serde form. Dense + /// keys carry no types on the wire, so the codec is reconstructed from this exact ordered type + /// list. + pub physical_type: ConcreteDataType, +} + +/// One logical tag name bound to a primary-key column. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LabelDecodeTag { + /// Physical column id this logical tag reads. + pub column_id: u32, + /// Logical tag name exposed to queries. + pub name: String, +} + +/// Versioned, self-contained description of how to decode key-encoded labels. +/// +/// It carries exactly what decoding a primary-key label source reads and nothing else: no region +/// identity, no field columns, no time index. A consumer that receives the frame schema can +/// therefore reconstruct an equivalent decoder without a scan-side binding, while a description +/// that is missing, of an unknown version, or structurally inconsistent is a hard error instead of +/// a silent fallback. +/// +/// It is JSON-encoded in one schema metadata entry, always through +/// [`raw_series_frame_schema_with_label_decode`], so both the producing and the expecting side of a +/// `greptime.series.*` exact metadata comparison see byte-identical text. The text is +/// self-describing but not a general interchange format: readers reject unknown fields here only +/// at this structure's own levels, while [`LabelDecodeColumn::physical_type`] is decoded by +/// [`ConcreteDataType`]'s own representation and is therefore exactly as stable as that type's +/// serde. The frozen text this revision writes is pinned by +/// `tests/compatibility/wire/raw_series_label_decode/v1_physical_types.json`; see the module +/// documentation for what a change to it means. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LabelDecodeSpec { + /// Description version; see [`LABEL_DECODE_VERSION`]. + pub version: u32, + /// Metric name attached to every decoded label set; may be empty. + pub metric_name: String, + /// Primary-key encoding that determines the decode layout. + pub encoding: LabelDecodeEncoding, + /// Every tag column of the primary key, in primary-key encoding order. This is exactly the + /// set of encoded key fields: region metadata forbids a primary key from naming the time + /// index or a field column, and requires it to name every tag column, so no time-index or + /// field entry can appear here. The order is significant and is compared as written: a dense + /// decode reads field *n* at position *n*, so the same columns in another order describe a + /// different key, not the same key in another order. + pub primary_key: Vec, + /// Logical tags, as bound by the producer, in [`LabelDecodeSpec::canonicalize`] order. + /// + /// The order carries no meaning: the decoder sorts the tags by name when it binds them + /// (`src/mito-codec/src/labels.rs:235`), so a decoded label set does not depend on it. It is + /// nevertheless fixed rather than left to the caller, because this description is compared as + /// exact text across nodes: two producers with the same tags in another order must write + /// byte-identical text, or one side rejects the other as a schema mismatch. + /// + /// The semantic checks on these tags (reserved and internal names, name equality with the bound + /// column, membership in `primary_key`) are the decoder's, matching what a typed binding gets. + pub logical_tags: Vec, +} + +impl LabelDecodeSpec { + /// Puts the order-insensitive part of the description in its canonical order. + /// + /// [`raw_series_frame_schema_with_label_decode`] applies this to a copy before serializing, and + /// so does the description a binding derives (in `common_query::series_frame`), so both sides of + /// a text comparison produce the same bytes and a reader can compare descriptions structurally. + /// `primary_key` is deliberately untouched: its order is the encoding order. + /// + /// The writer canonicalizes a **copy** on purpose. A caller's binding keeps its own + /// `logical_tags` order, and that order is meaningful elsewhere: `src/table/src/table/scan.rs` + /// compares two bindings' tag lists positionally (`lt.iter().zip(rt.iter())`), so reordering the + /// caller's list here would make two bindings for the same labels compare unequal. Only the + /// serialized text is canonical; nothing outside it is rewritten. + /// + /// Sorting by `(column_id, name)` is a total order for a validated description, which names + /// each tag column and each tag name at most once, so the canonical form is unique. + pub fn canonicalize(&mut self) { + self.logical_tags + .sort_unstable_by(|a, b| (a.column_id, &a.name).cmp(&(b.column_id, &b.name))); + } + + /// Checks the version and the structural invariants a decoder relies on. + /// + /// Semantic binding checks (reserved names, internal columns, name equality against the bound + /// column) stay in the decoder so both a typed binding and a decoded description fail with the + /// same errors. + pub fn validate(&self) -> Result<()> { + if self.version != LABEL_DECODE_VERSION { + return Err(err(format!( + "unsupported label decode description version {}", + self.version + ))); + } + let mut column_ids = HashSet::with_capacity(self.primary_key.len()); + for column in &self.primary_key { + if !column_ids.insert(column.column_id) { + return Err(err(format!( + "duplicate primary-key column id {} in label decode description", + column.column_id + ))); + } + } + let mut tag_ids = HashSet::with_capacity(self.logical_tags.len()); + let mut tag_names = HashSet::with_capacity(self.logical_tags.len()); + for tag in &self.logical_tags { + if !tag_ids.insert(tag.column_id) { + return Err(err(format!( + "duplicate logical tag column id {} in label decode description", + tag.column_id + ))); + } + if !tag_names.insert(tag.name.as_str()) { + return Err(err(format!( + "duplicate logical tag {} in label decode description", + tag.name + ))); + } + if !column_ids.contains(&tag.column_id) { + return Err(err(format!( + "logical tag {} is not a primary-key column of the label decode description", + tag.name + ))); + } + } + Ok(()) + } +} + /// Builds the typed schema for a raw frame, preserving logical field names and order. +/// +/// This is the column layout shared with the frame builder and decoder; it deliberately carries no +/// label decode description. A key-encoded stream must be declared through +/// [`raw_series_frame_schema_with_label_decode`], which refuses such a layout without one. pub fn raw_series_frame_schema(schema: &SeriesSchema, mode: RawLabelMode) -> Result { if schema.eval_grid.is_some() { return invalid_series("raw series schema cannot have an evaluation grid"); @@ -137,6 +338,129 @@ pub fn raw_series_frame_schema(schema: &SeriesSchema, mode: RawLabelMode) -> Res Ok(Arc::new(Schema::new_with_metadata(columns, metadata))) } +/// Builds the typed schema for a raw frame together with its label decode description. +/// +/// The description is written only for [`RawLabelMode::Key`], where labels stay encoded in +/// `read_key`; explicit-label frames carry complete labels and must not declare one. Both +/// directions are enforced here so a frame cannot claim a label layout it does not describe. +pub fn raw_series_frame_schema_with_label_decode( + schema: &SeriesSchema, + mode: RawLabelMode, + decode: Option<&LabelDecodeSpec>, +) -> Result { + match (mode, decode) { + (RawLabelMode::Key, None) => { + return Err(err( + "key-encoded raw frames must carry a label decode description", + )); + } + (RawLabelMode::Explicit, Some(_)) => { + return Err(err("explicit raw frames carry no label decode description")); + } + (RawLabelMode::Key, Some(_)) | (RawLabelMode::Explicit, None) => {} + } + let layout = raw_series_frame_schema(schema, mode)?; + let Some(decode) = decode else { + return Ok(layout); + }; + // Writing is faithful: the description records what the producer declares, and nothing is + // rewritten or dropped. The one exception is the order of `logical_tags`, which carries no + // meaning to a decoder and which this text nevertheless has to fix so two producers of the same + // labels write the same bytes; a copy is canonicalized and the caller's value is untouched. A + // description this revision cannot honour is rejected when it is read, which is the only + // direction a consumer can rely on. + let mut canonical = decode.clone(); + canonical.canonicalize(); + let encoded = serde_json::to_string(&canonical) + .map_err(|e| err(format!("cannot encode label decode description: {e}")))?; + let mut metadata = layout.metadata().clone(); + metadata.insert(LABEL_DECODE.to_string(), encoded); + Ok(Arc::new(Schema::new_with_metadata( + layout.fields().clone(), + metadata, + ))) +} + +/// Recovers the label decode description a raw frame schema declares. +/// +/// Returns `None` for an explicit-label frame, which must not declare one. A key-encoded frame must +/// declare one, so a missing description is an error rather than a reason to reuse a caller-side +/// binding. A present description is parsed strictly: an unknown version, a missing or unknown +/// field at this structure's own levels, a type mismatch, or a description whose logical tags do +/// not name listed key columns is an error, never a silent default. Only the description's own +/// levels reject unknown fields; [`LabelDecodeColumn::physical_type`] is a +/// [`ConcreteDataType`] and accepts whatever that type's own serde accepts. +pub fn raw_series_frame_label_decode(schema: &Schema) -> Result> { + let mode = match schema.metadata().get(LABELS).map(String::as_str) { + Some("key") => RawLabelMode::Key, + Some("explicit") => RawLabelMode::Explicit, + _ => return Err(err("raw frame schema declares no known label mode")), + }; + let decode = match schema.metadata().get(LABEL_DECODE) { + None => None, + Some(encoded) => { + let decode: LabelDecodeSpec = serde_json::from_str(encoded) + .map_err(|e| err(format!("malformed label decode description: {e}")))?; + decode.validate()?; + Some(decode) + } + }; + match (mode, &decode) { + (RawLabelMode::Key, None) => Err(err( + "key-encoded raw frame declares no label decode description", + )), + (RawLabelMode::Explicit, Some(_)) => Err(err( + "explicit raw frame declares a label decode description", + )), + (RawLabelMode::Key, Some(_)) | (RawLabelMode::Explicit, None) => Ok(decode), + } +} + +/// Reconstructs the series schema a raw frame declares. +/// +/// A raw frame names its payload fields by name only: every member of the `fields` struct carries +/// the same payload type ([`raw_series_frame_schema`]), and [`SeriesSchema`] holds nothing else per +/// field. The reconstruction is therefore exact rather than a projection — the declared names in +/// declared order, and no evaluation grid, which a raw frame cannot carry. +/// +/// The reserved layout, content, and time metadata are checked here together with the absence of an +/// evaluation grid, so a frame this revision cannot read is rejected instead of being silently +/// interpreted as the current layout. +pub fn raw_series_frame_series_schema(schema: &Schema) -> Result { + for (key, expected) in [(LAYOUT, "v1"), (CONTENT, "raw"), (TIME, "timestamps")] { + if schema.metadata().get(key).map(String::as_str) != Some(expected) { + return Err(err(format!("raw frame declares no known {key} value"))); + } + } + if schema + .metadata() + .keys() + .any(|key| key == GRID || key.starts_with("greptime.series.grid.")) + { + return Err(err("raw series schema cannot have an evaluation grid")); + } + let fields = schema + .fields() + .iter() + .find(|field| field.name() == "fields") + .ok_or_else(|| err("raw frame declares no fields column"))?; + let DataType::Struct(members) = fields.data_type() else { + return Err(err("raw frame fields column is not a struct")); + }; + if members.is_empty() { + return Err(err("raw series schema must have fields")); + } + Ok(SeriesSchema { + fields: members + .iter() + .map(|member| FieldSpec { + name: Arc::from(member.name().as_str()), + }) + .collect(), + eval_grid: None, + }) +} + /// Bounded multi-row builder for finite raw frame batches. /// /// Rows accumulate until the row limit or the estimated byte limit is reached; the next @@ -871,7 +1195,7 @@ fn decode_raw_series_frame( if actual .metadata() .keys() - .any(|key| key == "greptime.series.grid" || key.starts_with("greptime.series.grid.")) + .any(|key| key == GRID || key.starts_with("greptime.series.grid.")) || actual.fields() != expected.fields() { return invalid_series("raw frame schema mismatch"); diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs index 976a433ee27..619db9d3e89 100644 --- a/src/datatypes/src/series/frame/tests.rs +++ b/src/datatypes/src/series/frame/tests.rs @@ -3058,3 +3058,530 @@ fn decoder_reuses_cached_schema_and_rejects_schema_evolution() { assert!(explicit_decoder.decode(&first).is_err()); assert!(key_decoder.decode(&nameless).is_err()); } + +/// A dense description carrying the ordered key columns and the logical tags the decoder reads. +fn dense_description() -> LabelDecodeSpec { + LabelDecodeSpec { + version: LABEL_DECODE_VERSION, + metric_name: "cpu_usage".to_string(), + encoding: LabelDecodeEncoding::Dense, + primary_key: vec![ + LabelDecodeColumn { + column_id: 1, + name: "host".to_string(), + physical_type: crate::data_type::ConcreteDataType::string_datatype(), + }, + LabelDecodeColumn { + column_id: 2, + name: "region".to_string(), + physical_type: crate::data_type::ConcreteDataType::string_datatype(), + }, + ], + logical_tags: vec![LabelDecodeTag { + column_id: 1, + name: "host".to_string(), + }], + } +} + +#[test] +fn label_decode_description_roundtrips_through_schema_metadata() { + let description = dense_description(); + let key = raw_series_frame_schema_with_label_decode( + &test_schema(&["value"]), + RawLabelMode::Key, + Some(&description), + ) + .unwrap(); + assert_eq!( + raw_series_frame_label_decode(&key).unwrap(), + Some(description.clone()) + ); + // The description is metadata-only: the column contract is untouched. + assert_eq!( + key.fields(), + raw_series_frame_schema(&test_schema(&["value"]), RawLabelMode::Key) + .unwrap() + .fields() + ); + assert_eq!(key.metadata().get(LABELS).map(String::as_str), Some("key")); +} + +#[test] +fn explicit_frames_declare_no_label_decode_description() { + let schema = test_schema(&["value"]); + let explicit = + raw_series_frame_schema_with_label_decode(&schema, RawLabelMode::Explicit, None).unwrap(); + assert!(explicit.metadata().get(LABEL_DECODE).is_none()); + assert_eq!(raw_series_frame_label_decode(&explicit).unwrap(), None); + // A key frame without a description, or an explicit frame with one, is rejected instead of + // being written as a frame the consumer cannot decode. + assert!( + raw_series_frame_schema_with_label_decode(&schema, RawLabelMode::Key, None).is_err(), + "key-encoded frames must describe their labels" + ); + assert!( + raw_series_frame_schema_with_label_decode( + &schema, + RawLabelMode::Explicit, + Some(&dense_description()) + ) + .is_err(), + "explicit frames carry complete labels" + ); +} + +/// The raw frame's own columns with exactly the supplied metadata map. +fn schema_with_metadata(metadata: Vec<(String, String)>) -> Schema { + let layout = raw_series_frame_schema(&test_schema(&["value"]), RawLabelMode::Key).unwrap(); + Schema::new_with_metadata(layout.fields().clone(), metadata.into_iter().collect()) +} + +/// The four metadata entries every raw frame declares, with the given label mode. +fn base_metadata(labels: &str) -> Vec<(String, String)> { + [ + ("greptime.series.layout", "v1"), + ("greptime.series.content", "raw"), + ("greptime.series.labels", labels), + ("greptime.series.time", "timestamps"), + ] + .into_iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect() +} + +#[test] +fn label_decode_description_rejects_inconsistent_bindings() { + let schema = test_schema(&["value"]); + let mut cases = Vec::new(); + + let mut duplicate_column = dense_description(); + duplicate_column.primary_key[1].column_id = 1; + cases.push(("duplicate primary-key column", duplicate_column)); + + let mut unbound_tag = dense_description(); + unbound_tag.logical_tags[0].column_id = 3; + cases.push(("logical tag outside the primary key", unbound_tag)); + + let mut duplicate_tag = dense_description(); + let tag = duplicate_tag.logical_tags[0].clone(); + duplicate_tag.logical_tags.push(tag); + cases.push(("duplicate logical tag", duplicate_tag)); + + for (label, description) in cases { + // Writing is faithful: the description records what the producer declares. Reading is + // where a description this revision cannot honour is rejected, so a producer cannot turn a + // bad binding into a frame a consumer would decode under silently. + let frame = raw_series_frame_schema_with_label_decode( + &schema, + RawLabelMode::Key, + Some(&description), + ) + .unwrap_or_else(|error| panic!("{label}: {error}")); + // The writer records what the producer declares, so the rejection has to come from the + // reader; one assertion covers it, since the frame it reads is the text just written. + assert!( + raw_series_frame_label_decode(&frame).is_err(), + "{label}: an inconsistent description is rejected when read" + ); + } +} + +/// A frame declaring no label mode has no label layout to describe. The rest of the parse and +/// validation rejections (unknown version, missing/unknown/wrongly typed fields, malformed JSON, +/// and mode and description disagreeing in either direction) are pinned case for case by +/// [`raw_series_label_decode_metadata_fixture`], so they are not re-asserted here. +#[test] +fn label_decode_rejects_a_frame_without_a_label_mode() { + let without_mode = base_metadata("key") + .into_iter() + .filter(|(key, _)| key != LABELS) + .collect(); + assert!( + raw_series_frame_label_decode(&schema_with_metadata(without_mode)).is_err(), + "a frame without a label mode has no label layout to describe" + ); +} + +#[test] +fn raw_series_label_decode_metadata_fixture() { + use std::collections::HashMap; + + #[derive(serde::Deserialize)] + struct Case { + name: String, + metadata: HashMap, + /// The encoding a decoding consumer expects, or `null` when the fixture must be rejected. + encoding: Option, + } + + let cases: Vec = serde_json::from_str(include_str!( + "../../../../../tests/compatibility/wire/raw_series_label_decode/v1_metadata.json" + )) + .unwrap(); + for case in cases { + let layout = raw_series_frame_schema(&test_schema(&["value"]), RawLabelMode::Key).unwrap(); + let metadata = case + .metadata + .iter() + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let frame_schema = Schema::new_with_metadata(layout.fields().clone(), metadata); + match case.encoding.as_deref() { + None => assert!( + raw_series_frame_label_decode(&frame_schema).is_err(), + "{}: must be rejected", + case.name + ), + Some(expected) => { + let decode = raw_series_frame_label_decode(&frame_schema) + .unwrap_or_else(|error| panic!("{}: {error}", case.name)) + .unwrap_or_else(|| panic!("{}: no description", case.name)); + let encoding = match decode.encoding { + LabelDecodeEncoding::Dense => "dense", + LabelDecodeEncoding::Sparse => "sparse", + }; + assert_eq!(encoding, expected, "{}", case.name); + assert_eq!(decode.version, LABEL_DECODE_VERSION, "{}", case.name); + assert_eq!(decode.metric_name, "cpu_usage", "{}", case.name); + assert_eq!( + decode.primary_key, + vec![ + LabelDecodeColumn { + column_id: 1, + name: "host".to_string(), + physical_type: crate::data_type::ConcreteDataType::string_datatype(), + }, + LabelDecodeColumn { + column_id: 2, + name: "shard".to_string(), + physical_type: crate::data_type::ConcreteDataType::uint32_datatype(), + }, + ], + "{}: the ordered key columns are decoded in place", + case.name + ); + assert_eq!( + decode.logical_tags, + vec![ + LabelDecodeTag { + column_id: 1, + name: "host".to_string(), + }, + LabelDecodeTag { + column_id: 2, + name: "shard".to_string(), + }, + ], + "{}", + case.name + ); + if expected == "dense" { + // The fixture text is the encoder's own output, byte for byte: re-encoding the + // description it decodes to must reproduce it exactly. + let roundtrip = raw_series_frame_schema_with_label_decode( + &test_schema(&["value"]), + RawLabelMode::Key, + Some(&decode), + ) + .unwrap(); + assert_eq!( + roundtrip.metadata().get(LABEL_DECODE), + case.metadata.get(LABEL_DECODE), + "{}: re-encoding reproduces the fixture text", + case.name + ); + } + } + } + } +} + +/// A tag column typed as a timestamp is described like any other key column: the description +/// carries its physical type, and the text is exactly what the encoder wrote. +#[test] +fn label_decode_description_carries_a_timestamp_typed_tag_column() { + let mut description = dense_description(); + description.primary_key[1].name = "observed".to_string(); + description.primary_key[1].physical_type = + crate::data_type::ConcreteDataType::timestamp_millisecond_datatype(); + description.logical_tags.push(LabelDecodeTag { + column_id: 2, + name: "observed".to_string(), + }); + let frame = raw_series_frame_schema_with_label_decode( + &test_schema(&["value"]), + RawLabelMode::Key, + Some(&description), + ) + .unwrap(); + let text = frame.metadata().get(LABEL_DECODE).unwrap(); + // The expected text is the encoder's real output for a timestamp-typed tag column. + assert!( + text.contains("\"physical_type\":{\"Timestamp\":{\"Millisecond\":null}}"), + "the encoder wrote a timestamp type: {text}" + ); + let recovered = raw_series_frame_label_decode(&frame).unwrap().unwrap(); + assert_eq!(recovered, description); + assert_eq!( + recovered.primary_key[1].physical_type, + crate::data_type::ConcreteDataType::timestamp_millisecond_datatype() + ); +} + +/// Two producers that list the same logical tags in a different order write the same bytes. +/// +/// The tag order carries no meaning to a decoder, so the writer sorts its own copy before +/// serializing and leaves the caller's value alone. The encoding order of `primary_key` is a +/// different thing entirely: it is the order of the key columns, so it is written as given. +#[test] +fn label_decode_text_is_independent_of_the_callers_tag_order() { + let mut ascending = dense_description(); + ascending.logical_tags = vec![ + LabelDecodeTag { + column_id: 1, + name: "host".to_string(), + }, + LabelDecodeTag { + column_id: 2, + name: "region".to_string(), + }, + ]; + let mut descending = ascending.clone(); + descending.logical_tags.reverse(); + // The encoding order of the key columns is not the tag order: reverse it in the second + // description so a writer that sorted `primary_key` too would be caught below. + descending.primary_key.reverse(); + + let write = |description: &LabelDecodeSpec| { + let frame = raw_series_frame_schema_with_label_decode( + &test_schema(&["value"]), + RawLabelMode::Key, + Some(description), + ) + .unwrap(); + frame.metadata().get(LABEL_DECODE).unwrap().clone() + }; + + let same_tags = { + let mut reordered = ascending.clone(); + reordered.logical_tags.reverse(); + write(&reordered) + }; + assert_eq!( + write(&ascending), + same_tags, + "two tag orders wrote different bytes" + ); + assert_eq!( + ascending.logical_tags, + vec![ + LabelDecodeTag { + column_id: 1, + name: "host".to_string(), + }, + LabelDecodeTag { + column_id: 2, + name: "region".to_string(), + }, + ], + "the writer reordered the caller's description" + ); + // The key columns keep the order they were given, so the two descriptions are different + // frames rather than the same bytes: that is the part canonicalization must not touch. + assert_ne!( + write(&ascending), + write(&descending), + "the writer reordered the encoding columns" + ); + let recovered = raw_series_frame_label_decode( + &raw_series_frame_schema_with_label_decode( + &test_schema(&["value"]), + RawLabelMode::Key, + Some(&descending), + ) + .unwrap(), + ) + .unwrap() + .unwrap(); + assert_eq!( + recovered.primary_key, descending.primary_key, + "the recovered key order is not the encoding order" + ); +} +/// One entry of `v1_physical_types.json`: a [`ConcreteDataType`] variant and the exact text this +/// revision's serde writes for a representative value of it. +#[derive(serde::Deserialize)] +struct FrozenPhysicalType { + /// The [`ConcreteDataType`] variant, as [`concrete_datatype_variants`] spells it. + variant: String, + /// The `serde_json` text this revision writes for that variant's representative value. + text: String, +} + +/// Spells every [`ConcreteDataType`] variant exactly once, together with one representative value of +/// it. +/// +/// Everything that has to notice a new variant is derived from this one list: +/// +/// - [`variant_name`] matches on it with **no wildcard arm**, so adding a variant to +/// [`ConcreteDataType`] is a compile error until the variant is added here. +/// - [`datatype_representatives`] yields one value per listed variant, so the same edit is what +/// supplies the value the tests read the serde text from. +/// - [`VARIANT_NAMES`] is generated from it too, which is what +/// [`the_frozen_text_covers_every_datatype_variant`] requires a fixture entry for. +macro_rules! concrete_datatype_variants { + ($($variant:ident => $value:expr),+ $(,)?) => { + /// The [`ConcreteDataType`] variant a value belongs to, spelled as the fixture names it. + /// + /// Exhaustive on purpose: a new variant of [`ConcreteDataType`] makes this match fail to + /// compile, which is what makes a missing fixture entry impossible to land. + fn variant_name(data_type: &ConcreteDataType) -> &'static str { + match data_type { + $(ConcreteDataType::$variant(_) => stringify!($variant),)+ + } + } + + /// One representative value per variant, in the order listed above. + fn datatype_representatives() -> Vec { + vec![$($value),+] + } + + /// Every [`ConcreteDataType`] variant name, in the order listed above. + const VARIANT_NAMES: &[&str] = &[$(stringify!($variant)),+]; + }; +} + +concrete_datatype_variants!( + Null => ConcreteDataType::null_datatype(), + Boolean => ConcreteDataType::boolean_datatype(), + Int8 => ConcreteDataType::int8_datatype(), + Int16 => ConcreteDataType::int16_datatype(), + Int32 => ConcreteDataType::int32_datatype(), + Int64 => ConcreteDataType::int64_datatype(), + UInt8 => ConcreteDataType::uint8_datatype(), + UInt16 => ConcreteDataType::uint16_datatype(), + UInt32 => ConcreteDataType::uint32_datatype(), + UInt64 => ConcreteDataType::uint64_datatype(), + Float32 => ConcreteDataType::float32_datatype(), + Float64 => ConcreteDataType::float64_datatype(), + Decimal128 => ConcreteDataType::decimal128_datatype(10, 2), + Binary => ConcreteDataType::binary_datatype(), + String => ConcreteDataType::string_datatype(), + Date => ConcreteDataType::date_datatype(), + Timestamp => ConcreteDataType::timestamp_millisecond_datatype(), + Time => ConcreteDataType::time_microsecond_datatype(), + Duration => ConcreteDataType::duration_nanosecond_datatype(), + Interval => ConcreteDataType::interval_month_day_nano_datatype(), + List => ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())), + Dictionary => ConcreteDataType::dictionary_datatype( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ), + Struct => ConcreteDataType::struct_datatype(crate::types::StructType::new(Arc::new(vec![ + crate::types::StructField::new("host", ConcreteDataType::string_datatype(), false), + crate::types::StructField::new("port", ConcreteDataType::uint16_datatype(), true), + ]))), + Json => ConcreteDataType::json_datatype(), + Vector => ConcreteDataType::vector_datatype(3), +); + +/// The frozen text of every [`ConcreteDataType`] variant, as a label decode description embeds it. +const FROZEN_PHYSICAL_TYPES: &str = include_str!( + "../../../../../tests/compatibility/wire/raw_series_label_decode/v1_physical_types.json" +); + +fn frozen_physical_types() -> Vec { + serde_json::from_str(FROZEN_PHYSICAL_TYPES).expect("the frozen physical-type fixture parses") +} + +/// The representative value of one variant, named the way the fixture names it. +fn representative(variant: &str) -> ConcreteDataType { + VARIANT_NAMES + .iter() + .zip(datatype_representatives()) + .find(|(name, _)| **name == variant) + .unwrap_or_else(|| panic!("{variant}: no representative value is listed for this variant")) + .1 +} + +/// The text a description embeds for a key column's type is not this module's to spell: it is +/// [`ConcreteDataType`]'s own serde, so it changes whenever that type's serde does. +/// +/// `v1_physical_types.json` freezes that text for every variant. This test fails when a variant no +/// longer parses back from its frozen text, when it no longer re-serializes to it, or when the +/// parsed value is no longer the type that variant names — which is what a silent serde change +/// looks like. Since the frame text is compared byte for byte across nodes, such a change is a +/// [`LABEL_DECODE_VERSION`] bump, and this test is the alarm that raises it. +#[test] +fn every_datatype_round_trips_through_its_frozen_text() { + use crate::data_type::DataType as _; + + for entry in frozen_physical_types() { + let parsed: ConcreteDataType = serde_json::from_str(&entry.text).unwrap_or_else(|error| { + panic!("{}: {} cannot be read: {error}", entry.variant, entry.text) + }); + assert_eq!( + serde_json::to_string(&parsed).unwrap(), + entry.text, + "{}: re-serializes to the frozen text", + entry.variant + ); + assert_eq!( + parsed.name(), + representative(&entry.variant).name(), + "{}: the frozen text still names the same type", + entry.variant + ); + assert_eq!( + parsed, + representative(&entry.variant), + "{}: the frozen text still describes the same value", + entry.variant + ); + } +} + +/// The fixture above only freezes a variant if it has an entry, so every variant must have one. +/// +/// This test is what makes that hold: [`variant_name`] matches every [`ConcreteDataType`] variant +/// with no wildcard arm, so a new variant does not compile until it is listed in +/// [`concrete_datatype_variants`]; that listing also extends [`VARIANT_NAMES`] and +/// [`datatype_representatives`], and this test then requires a fixture entry for it and rejects an +/// entry naming a variant that does not exist. +#[test] +fn the_frozen_text_covers_every_datatype_variant() { + let frozen = frozen_physical_types(); + + for data_type in datatype_representatives() { + let variant = variant_name(&data_type); + let entry = frozen + .iter() + .find(|entry| entry.variant == variant) + .unwrap_or_else(|| { + panic!("{variant}: v1_physical_types.json has no entry for this datatype variant") + }); + // The entry's text is this variant's text: read it back and name the variant it describes. + let parsed: ConcreteDataType = serde_json::from_str(&entry.text) + .unwrap_or_else(|error| panic!("{variant}: {} is unreadable: {error}", entry.text)); + assert_eq!( + variant_name(&parsed), + variant, + "{variant}: the entry's text describes this variant" + ); + } + + // Nothing in the fixture names a variant that does not exist, and no variant is listed twice. + for entry in &frozen { + assert!( + VARIANT_NAMES.contains(&entry.variant.as_str()), + "{}: the fixture names a variant ConcreteDataType does not have", + entry.variant + ); + } + assert_eq!( + frozen.len(), + VARIANT_NAMES.len(), + "one fixture entry per datatype variant, no duplicates" + ); +} diff --git a/src/mito-codec/src/error.rs b/src/mito-codec/src/error.rs index b266ea9b04f..e88af039d88 100644 --- a/src/mito-codec/src/error.rs +++ b/src/mito-codec/src/error.rs @@ -81,6 +81,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid label decode description"))] + InvalidLabelDecodeDescription { + // Box the source to reduce the size of the error. + #[snafu(source(from(datatypes::error::Error, Box::new)))] + source: Box, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Encode null value"))] IndexEncodeNull { #[snafu(implicit)] @@ -108,6 +117,7 @@ impl ErrorExt for Error { StatusCode::InvalidArguments } NotSupportedField { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, + InvalidLabelDecodeDescription { source, .. } => source.status_code(), InvalidSparsePrimaryKey { .. } | InvalidDensePrimaryKey { .. } => { StatusCode::InvalidArguments } diff --git a/src/mito-codec/src/labels.rs b/src/mito-codec/src/labels.rs index fe4e3bfbde0..92e28574ffa 100644 --- a/src/mito-codec/src/labels.rs +++ b/src/mito-codec/src/labels.rs @@ -17,11 +17,12 @@ use std::sync::Arc; use api::v1::SemanticType; use common_telemetry::debug; +use datatypes::data_type::ConcreteDataType; +use datatypes::series::frame::{LabelDecodeEncoding, LabelDecodeSpec}; use datatypes::series::{Label, MetricName}; use datatypes::value::Value; -use snafu::OptionExt; +use snafu::{OptionExt, ResultExt}; use store_api::codec::PrimaryKeyEncoding; -use store_api::metadata::RegionMetadataRef; use store_api::metric_engine_consts::is_metric_engine_internal_column; use store_api::region_engine::RawLabelSource; use store_api::storage::{ @@ -30,8 +31,10 @@ use store_api::storage::{ }; use crate::error::{ - InvalidDensePrimaryKeySnafu, InvalidSparsePrimaryKeySnafu, Result, UnsupportedOperationSnafu, + InvalidDensePrimaryKeySnafu, InvalidLabelDecodeDescriptionSnafu, InvalidSparsePrimaryKeySnafu, + Result, UnsupportedOperationSnafu, }; +use crate::row_converter::SortField; use crate::row_converter::dense::DensePrimaryKeyCodec; use crate::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView}; @@ -68,10 +71,41 @@ struct BoundLabel { position: usize, } -/// Decodes labels on demand using the metadata snapshot bound to a raw-label source. -pub struct PrimaryKeyLabelDecoder { +/// One ordered key column of the minimal decode layout. +struct KeyColumn { + id: ColumnId, + name: String, + data_type: ConcreteDataType, +} + +/// Minimal primary-key layout labels can be decoded from. +/// +/// It holds only what decoding reads: the metric name, the primary-key encoding, and the ordered +/// primary-key columns with their physical types. A scan-bound [`RawLabelSource::PrimaryKey`] and a +/// frame schema's [`LabelDecodeSpec`] both reduce to this view, so both bind the same decoder. +/// Every key column is a `Tag` column, which is checked when a decoder is bound, so the layout +/// carries no per-column role. +struct KeyLayout { metric_name: Arc, - physical_metadata: RegionMetadataRef, + encoding: PrimaryKeyEncoding, + columns: Vec, +} + +impl KeyLayout { + /// Finds a listed primary-key column. + fn column(&self, id: ColumnId) -> Option<&KeyColumn> { + self.columns.iter().find(|column| column.id == id) + } + + /// Returns the position of a listed primary-key column. + fn position(&self, id: ColumnId) -> Option { + self.columns.iter().position(|column| column.id == id) + } +} + +/// Decodes labels on demand using the minimal primary-key layout bound to a raw-label source. +pub struct PrimaryKeyLabelDecoder { + key: KeyLayout, complete: bool, include_name: bool, labels: Vec, @@ -85,6 +119,9 @@ impl PrimaryKeyLabelDecoder { /// Explicit sources already contain complete labels and are rejected. Dense decoding /// requires the full key layout described by the bound snapshot; compatibility with /// older key layouts must be normalized upstream. + /// + /// Every primary-key column must be a `Tag` column; [`RegionMetadata::validate`] enforces it, so + /// this only rejects a snapshot built without validation. pub fn new(source: &RawLabelSource, projection: LabelProjection) -> Result { let RawLabelSource::PrimaryKey { metric_name, @@ -98,36 +135,113 @@ impl PrimaryKeyLabelDecoder { .fail(); }; let metadata = physical_metadata.as_ref(); - let mut labels = Vec::with_capacity(logical_tags.len()); + if let Some(column) = metadata + .primary_key_columns() + .find(|column| column.semantic_type != SemanticType::Tag) + { + return InvalidSparsePrimaryKeySnafu { + reason: format!( + "primary-key column {} is not a tag column", + column.column_schema.name + ), + } + .fail(); + } + let key = KeyLayout { + metric_name: metric_name.clone(), + encoding: metadata.primary_key_encoding, + columns: metadata + .primary_key_columns() + .map(|column| KeyColumn { + id: column.column_id, + name: column.column_schema.name.clone(), + data_type: column.column_schema.data_type.clone(), + }) + .collect(), + }; + Self::bind( + key, + logical_tags + .iter() + .map(|tag| (tag.column_id, tag.name.as_ref())), + projection, + ) + } + + /// Binds a decoder to a frame schema's label decode description and projection. + /// + /// This is the decoding path for a consumer that received only the raw frame schema: the + /// description carries the metric name, the primary-key encoding, the ordered key columns with + /// their physical types, and the logical tags. The description is validated first, so a version + /// or structure this revision cannot honour is an error instead of a best-effort decode. + /// + /// The description lists exactly the encoded key fields, and region metadata guarantees each + /// of them is a tag column, so no per-column role is carried or re-checked; an array position + /// is an encoding position. The logical tags are still checked here, one by one, exactly as + /// for a typed binding, which is the only place the decoder inspects them (a listed column no + /// tag selects is not examined). + pub fn from_decode_spec(spec: &LabelDecodeSpec, projection: LabelProjection) -> Result { + spec.validate() + .context(InvalidLabelDecodeDescriptionSnafu)?; + let key = KeyLayout { + metric_name: Arc::from(spec.metric_name.as_str()), + encoding: match spec.encoding { + LabelDecodeEncoding::Dense => PrimaryKeyEncoding::Dense, + LabelDecodeEncoding::Sparse => PrimaryKeyEncoding::Sparse, + }, + columns: spec + .primary_key + .iter() + .map(|column| KeyColumn { + id: column.column_id, + name: column.name.clone(), + data_type: column.physical_type.clone(), + }) + .collect(), + }; + Self::bind( + key, + spec.logical_tags + .iter() + .map(|tag| (tag.column_id, tag.name.as_str())), + projection, + ) + } + + /// Binds the ordered logical tags of `key` to one projection and one decode layout. + fn bind<'a, I>(key: KeyLayout, tags: I, projection: LabelProjection) -> Result + where + I: IntoIterator, + { + let mut labels = Vec::new(); let mut names = HashSet::new(); let mut ids = HashSet::new(); - for tag in logical_tags.iter() { - let column = metadata.column_by_id(tag.column_id); - if !names.insert(tag.name.clone()) - || !ids.insert(tag.column_id) - || tag.name.as_ref() == "__name__" - || is_internal_column(&tag.name) - || is_metric_engine_internal_column(&tag.name) - || ReservedColumnId::is_reserved(tag.column_id) - || !column.is_some_and(|column| { - column.semantic_type == SemanticType::Tag - && column.column_schema.name.as_str() == tag.name.as_ref() - && metadata.primary_key.contains(&tag.column_id) - }) + let mut logical_tag_count = 0; + for (column_id, name) in tags { + logical_tag_count += 1; + if !names.insert(name) + || !ids.insert(column_id) + || name == "__name__" + || is_internal_column(name) + || is_metric_engine_internal_column(name) + || ReservedColumnId::is_reserved(column_id) + || !key + .column(column_id) + .is_some_and(|column| column.name == name) { return InvalidSparsePrimaryKeySnafu { - reason: format!("invalid logical tag binding {}", tag.name), + reason: format!("invalid logical tag binding {name}"), } .fail(); } - let position = metadata.primary_key_index(tag.column_id).context( - InvalidSparsePrimaryKeySnafu { - reason: format!("logical tag {} is not a primary-key column", tag.name), - }, - )?; + let position = key + .position(column_id) + .context(InvalidSparsePrimaryKeySnafu { + reason: format!("logical tag {name} is not a primary-key column"), + })?; labels.push(BoundLabel { - id: tag.column_id, - name: tag.name.clone(), + id: column_id, + name: Arc::from(name), position, }); } @@ -143,14 +257,19 @@ impl PrimaryKeyLabelDecoder { let requested: HashSet<_> = requested.iter().map(AsRef::as_ref).collect(); labels.retain(|label| requested.contains(label.name.as_ref())); } - let layout = match metadata.primary_key_encoding { + let layout = match key.encoding { PrimaryKeyEncoding::Sparse => Layout::Sparse(SparseOffsetsCache::new()), PrimaryKeyEncoding::Dense => Layout::Dense { - codec: DensePrimaryKeyCodec::new(metadata), + codec: DensePrimaryKeyCodec::with_fields( + key.columns + .iter() + .map(|column| (column.id, SortField::new(column.data_type.clone()))) + .collect(), + ), offsets: Vec::new(), }, }; - let encoding = match metadata.primary_key_encoding { + let encoding = match key.encoding { PrimaryKeyEncoding::Sparse => "sparse", PrimaryKeyEncoding::Dense => "dense", }; @@ -158,13 +277,12 @@ impl PrimaryKeyLabelDecoder { debug!( encoding, projection_mode, - logical_tag_count = logical_tags.len(), + logical_tag_count, selected_tag_count = labels.len(), "Constructed primary key label decoder" ); Ok(Self { - metric_name: metric_name.clone(), - physical_metadata: physical_metadata.clone(), + key, complete, include_name, labels, @@ -182,8 +300,8 @@ impl PrimaryKeyLabelDecoder { let mut values = Vec::with_capacity(self.labels.len() + usize::from(self.include_name)); let Self { layout, + key, labels, - physical_metadata, complete, value_scratch, .. @@ -192,12 +310,7 @@ impl PrimaryKeyLabelDecoder { Layout::Sparse(cache) => { let mut view = SparsePrimaryKeyView::new(read_key, cache)?; view.visit_label_ids(|id| { - if *complete - && !physical_metadata.column_by_id(id).is_some_and(|column| { - column.semantic_type == SemanticType::Tag - && physical_metadata.primary_key.contains(&id) - }) - { + if *complete && key.column(id).is_none() { return InvalidSparsePrimaryKeySnafu { reason: format!("unrecognized sparse tag id {id}"), } @@ -240,7 +353,7 @@ impl PrimaryKeyLabelDecoder { } if *complete { Ok(DecodedLabels::Complete( - MetricName::try_new(Some(self.metric_name.clone()), values).map_err(|error| { + MetricName::try_new(Some(key.metric_name.clone()), values).map_err(|error| { InvalidSparsePrimaryKeySnafu { reason: error.to_string(), } @@ -251,7 +364,7 @@ impl PrimaryKeyLabelDecoder { if self.include_name { values.push(Label { name: Arc::from("__name__"), - value: self.metric_name.clone(), + value: self.key.metric_name.clone(), }); values.sort_unstable_by(|a, b| a.name.cmp(&b.name)); } diff --git a/src/mito-codec/src/labels/tests.rs b/src/mito-codec/src/labels/tests.rs index 99710e3d9fa..88651bd1c35 100644 --- a/src/mito-codec/src/labels/tests.rs +++ b/src/mito-codec/src/labels/tests.rs @@ -209,7 +209,7 @@ fn dense_and_sparse_decode_complete_or_projected_labels() { } #[test] -fn projection_is_bound_to_snapshot_and_distinguishes_missing_and_empty() { +fn projection_snapshots_its_layout_and_distinguishes_missing_and_empty() { let original = metadata( &[ (21, "old", ConcreteDataType::string_datatype()), @@ -250,7 +250,10 @@ fn projection_is_bound_to_snapshot_and_distinguishes_missing_and_empty() { let mut all = PrimaryKeyLabelDecoder::new(&old_source, LabelProjection::All).unwrap(); drop(old_source); drop(original); - assert!(weak_metadata.upgrade().is_some()); + assert!( + weak_metadata.upgrade().is_none(), + "the decoder keeps its own key layout (ids, positions, types), not the region metadata" + ); assert_eq!( decoder.decode(&key).unwrap(), DecodedLabels::Projected(vec![Label { @@ -601,3 +604,109 @@ fn dense_typed_values_keep_exact_integer_string_representation() { "9007199254740993" ); } + +/// One field is enough: the frame schema only carries ordered payload fields and metadata. +fn frame_series_schema() -> Arc { + Arc::new(datatypes::series::samples::SeriesSchema { + fields: vec![datatypes::series::samples::FieldSpec { + name: Arc::from("value"), + }] + .into(), + eval_grid: None, + }) +} + +/// A key-encoded frame's schema describes its own labels: recovering the description from the +/// schema and decoding with it matches decoding with the original bound source, for both primary +/// key encodings and both projections. +#[test] +fn frame_label_decode_description_decodes_like_the_bound_source() { + for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] { + let bound = source( + string_metadata(encoding), + "requests_total", + &[(8, "zeta"), (41, "alpha")], + ); + let description = common_query::series_frame::label_decode_spec(&bound) + .expect("a primary-key binding describes its key layout"); + let frame_schema = datatypes::series::frame::raw_series_frame_schema_with_label_decode( + &frame_series_schema(), + datatypes::series::frame::RawLabelMode::Key, + Some(&description), + ) + .unwrap(); + let recovered = datatypes::series::frame::raw_series_frame_label_decode(&frame_schema) + .unwrap() + .expect("the frame declares its decode description"); + assert_eq!(recovered, description); + + let key = encoded(encoding, &values()); + for projection in [ + LabelProjection::All, + LabelProjection::Only(Arc::from([Arc::::from("alpha")])), + LabelProjection::Only(Arc::from([])), + ] { + let mut direct = PrimaryKeyLabelDecoder::new(&bound, projection.clone()).unwrap(); + let mut from_frame = + PrimaryKeyLabelDecoder::from_decode_spec(&recovered, projection).unwrap(); + assert_eq!( + from_frame.decode(&key).unwrap(), + direct.decode(&key).unwrap() + ); + } + } +} + +/// A description this revision cannot honour is rejected when a decoder is bound, not decoded +/// best-effort. +#[test] +fn frame_label_decode_description_rejects_unsupported_version() { + let bound = source( + string_metadata(PrimaryKeyEncoding::Dense), + "requests_total", + &[(8, "zeta")], + ); + let description = common_query::series_frame::label_decode_spec(&bound).unwrap(); + let mut future = description.clone(); + future.version += 1; + assert!(PrimaryKeyLabelDecoder::from_decode_spec(&future, LabelProjection::All).is_err()); + + // The description must list the column a logical tag reads. + let mut unbound = description; + unbound.primary_key.retain(|column| column.column_id != 8); + assert!(PrimaryKeyLabelDecoder::from_decode_spec(&unbound, LabelProjection::All).is_err()); +} + +/// A primary-key column that is not a `Tag` column is rejected when a decoder is bound, because the +/// layout records no per-column role and decodes every listed key column as a tag. +/// +/// `RegionMetadata::validate` rejects such metadata, so this snapshot is built without validation to +/// exercise the check the decoder performs itself. +#[test] +fn binding_rejects_a_primary_key_column_that_is_not_a_tag() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(7, 1)); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("not_a_tag", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 8, + }); + builder.primary_key(vec![8]); + builder.primary_key_encoding(PrimaryKeyEncoding::Sparse); + let invalid = Arc::new(builder.build_without_validation().unwrap()); + + let bound = source(invalid, "m", &[(8, "not_a_tag")]); + let error = PrimaryKeyLabelDecoder::new(&bound, LabelProjection::All) + .err() + .expect("a non-tag primary-key column must be rejected") + .to_string(); + assert!(error.contains("not a tag column"), "{error}"); +} diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 68469e85324..be9a5f207ae 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -5118,7 +5118,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { }; let native_plan = Arc::new(RegionScanExec::new(native_scanner, request.clone(), None).unwrap()); let native_plan = native_plan - .try_to_series_scan(native_source.clone()) + .try_to_series_scan(native_source) .unwrap() .expect("Mito raw-series capability"); let native_schema = native_plan.series_schema().clone(); @@ -5133,8 +5133,7 @@ async fn test_mito_handle_query_raw_series_candidate_projection_and_mvcc() { frame_ordering, "the legacy producer must declare key-then-time frames" ); - let unpacked = - unpack_raw_series(native_frames, native_schema, native_source, frame_ordering).unwrap(); + let unpacked = unpack_raw_series(native_frames, Some(native_schema)).unwrap(); let native_schema = unpacked.schema().clone(); let mut native = unpacked.try_collect::>().await.unwrap(); assert_eq!( @@ -5370,14 +5369,14 @@ async fn test_two_phase_native_raw_stream_declares_contiguous_frames() { let schema = schema.expect("the scan must emit at least one partition stream"); let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( Arc::clone(&schema), - source.clone(), + source, SeriesOrdering::Contiguous, chunks, ))) .unwrap(); let ordering = raw_series_frame_ordering(packed.schema().as_ref()).unwrap(); assert_eq!(SeriesOrdering::Contiguous, ordering); - let validated = unpack_raw_series(packed, schema, source, ordering) + let validated = unpack_raw_series(packed, Some(schema)) .unwrap() .try_collect::>() .await @@ -5720,18 +5719,12 @@ async fn test_two_phase_raw_stream_stays_contiguous_across_disjoint_partition_as // not final. let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( Arc::clone(&schema), - source.clone(), + source, SeriesOrdering::Contiguous, complete.clone(), ))) .unwrap(); - let unpacked = unpack_raw_series( - packed, - Arc::clone(&schema), - source, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let unpacked = unpack_raw_series(packed, Some(schema)).unwrap(); let validated = unpacked.try_collect::>().await.unwrap(); assert_eq!( complete_rows, @@ -6064,18 +6057,12 @@ async fn test_two_phase_raw_same_partition_multi_round_assignment_stays_contiguo // the chunks this reader produced, including the chunks the second assignment added. let packed = pack_raw_series(Box::pin(ChainedRawSeriesStream::new( Arc::clone(&schema), - source.clone(), + source, SeriesOrdering::Contiguous, chunks.clone(), ))) .unwrap(); - let unpacked = unpack_raw_series( - packed, - Arc::clone(&schema), - source, - SeriesOrdering::Contiguous, - ) - .unwrap(); + let unpacked = unpack_raw_series(packed, Some(Arc::clone(&schema))).unwrap(); let validated = unpacked.try_collect::>().await.unwrap(); assert_eq!( stream_rows, @@ -6200,6 +6187,8 @@ fn primary_key_label_source( struct ChainedRawSeriesStream { schema: Arc, label_source: store_api::region_engine::RawLabelSource, + /// The description this stream declares, derived from `label_source` once, like a producer. + label_decode_spec: Option, ordering: datatypes::series::SeriesOrdering, chunks: std::vec::IntoIter, } @@ -6211,9 +6200,11 @@ impl ChainedRawSeriesStream { ordering: datatypes::series::SeriesOrdering, chunks: Vec, ) -> Self { + let label_decode_spec = common_query::series_frame::label_decode_spec(&label_source); Self { schema, label_source, + label_decode_spec, ordering, chunks: chunks.into_iter(), } @@ -6240,11 +6231,221 @@ impl store_api::region_engine::RawSeriesStream for ChainedRawSeriesStream { &self.label_source } + fn label_decode_spec(&self) -> Option<&datatypes::series::frame::LabelDecodeSpec> { + self.label_decode_spec.as_ref() + } + fn ordering(&self) -> datatypes::series::SeriesOrdering { self.ordering } } +/// Decodes each chunk's primary key through `description`, as a consumer holding only the stream +/// would, and asserts none of the chunks carries inline labels. +fn decode_declared_labels( + description: &datatypes::series::frame::LabelDecodeSpec, + chunks: &[datatypes::series::raw::RawSeriesChunk], +) -> Vec { + use mito_codec::labels::{DecodedLabels, LabelProjection, PrimaryKeyLabelDecoder}; + + let mut decoder = + PrimaryKeyLabelDecoder::from_decode_spec(description, LabelProjection::All).unwrap(); + chunks + .iter() + .map(|chunk| { + assert!( + chunk.metric.is_none(), + "a key-encoded chunk leaves its labels undecoded" + ); + match decoder.decode(chunk.primary_key.as_slice()).unwrap() { + DecodedLabels::Complete(metric) => metric, + DecodedLabels::Projected(_) => { + panic!("decoding every label must yield complete labels") + } + } + }) + .collect() +} + +/// Every key-encoded raw series stream declares its own label decode description before the first +/// poll, and that declaration is the one derived from the binding the scan read under. Dense and +/// sparse producers are both covered, plus a stream that ends without producing a chunk, so the +/// producer and the consumer are exercised against the same declaration. +#[tokio::test] +async fn test_mito_raw_series_stream_declares_its_label_decode_description() { + use store_api::region_engine::QueryScanContext; + + // Dense encoding: one flushed series under the default region layout. + let mut env = TestEnv::with_prefix("test_mito_raw_series_self_describing_dense").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + test_util::put_rows( + &engine, + region_id, + Rows { + schema, + rows: test_util::build_rows_for_key("host-a", 0, 2, 0), + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + let source = primary_key_label_source(&scanner.metadata()); + let expected = common_query::series_frame::label_decode_spec(&source) + .expect("a primary-key binding describes its key"); + let ranges = scanner.properties().partitions.clone(); + scanner + .prepare(PrepareRequest { + ranges: Some(ranges.clone()), + target_partitions: Some(ranges.len()), + ..Default::default() + }) + .unwrap(); + let metrics = ExecutionPlanMetricsSet::default(); + let mut chunks = Vec::new(); + for partition in 0..ranges.len() { + let stream = scanner + .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) + .unwrap() + .expect("a dense PerSeries scan produces raw series"); + // The producer states its layout up front, before anything is polled. + assert_eq!( + Some(&expected), + stream.label_decode_spec(), + "dense partition {partition} must declare its label decode description before polling" + ); + chunks.extend(stream.try_collect::>().await.unwrap()); + } + assert_eq!( + 2, + chunks.iter().map(|chunk| chunk.time.len()).sum::() + ); + let dense_metrics = decode_declared_labels(&expected, &chunks); + assert_eq!(dense_metrics.len(), chunks.len()); + for metric in &dense_metrics { + assert_eq!(metric.name(), Some("fixture_metric")); + assert_eq!(metric.get("tag_0"), Some("host-a")); + } + + // A statically empty stream declares the same layout while producing no chunk at all. + let empty_request = ScanRequest { + filters: vec![col("field_0").gt(lit(1.0e9_f64))], + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let mut empty_scanner = engine.handle_query(region_id, empty_request).await.unwrap(); + let empty_source = primary_key_label_source(&empty_scanner.metadata()); + let empty_expected = common_query::series_frame::label_decode_spec(&empty_source) + .expect("a primary-key binding describes its key"); + let empty_ranges = empty_scanner.properties().partitions.clone(); + assert!( + !empty_ranges.is_empty(), + "a filtered PerSeries scan still assigns its pinned input to a partition" + ); + empty_scanner + .prepare(PrepareRequest { + ranges: Some(empty_ranges.clone()), + target_partitions: Some(empty_ranges.len()), + ..Default::default() + }) + .unwrap(); + let mut empty_chunks = Vec::new(); + for partition in 0..empty_ranges.len() { + let stream = empty_scanner + .scan_series_partition( + &QueryScanContext::default(), + &metrics, + partition, + &empty_source, + ) + .unwrap() + .expect("a statically empty PerSeries scan still produces a raw series stream"); + assert_eq!( + Some(&empty_expected), + stream.label_decode_spec(), + "an empty partition {partition} must still declare its layout before polling" + ); + empty_chunks.extend(stream.try_collect::>().await.unwrap()); + } + assert!( + empty_chunks.is_empty(), + "the filtered-out stream produced no chunk" + ); + + // Sparse encoding: the two-phase producer encodes table id and TSID into the key too. + let mut env = TestEnv::with_prefix("test_mito_raw_series_self_describing_sparse").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let (metadata, row_schema) = create_two_phase_sparse_region(&engine, region_id).await; + let key = new_sparse_primary_key(&["host-a", "zone-1"], &metadata, 10, 0); + put_sparse_rows( + &engine, + region_id, + &row_schema, + vec![(key.clone(), 10, 1), (key, 20, 2)], + ) + .await; + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + preserve_pk_dictionary_encoding: true, + ..Default::default() + }; + let mut scanner = engine.handle_query(region_id, request).await.unwrap(); + let source = primary_key_label_source(&scanner.metadata()); + let expected = common_query::series_frame::label_decode_spec(&source) + .expect("a primary-key binding describes its key"); + let ranges = scanner.properties().partitions.clone(); + scanner + .prepare(PrepareRequest { + ranges: Some(ranges.clone()), + target_partitions: Some(ranges.len()), + ..Default::default() + }) + .unwrap(); + let mut chunks = Vec::new(); + for partition in 0..ranges.len() { + let stream = scanner + .scan_series_partition(&QueryScanContext::default(), &metrics, partition, &source) + .unwrap() + .expect("a sparse PerSeries scan produces raw series"); + assert_eq!( + Some(&expected), + stream.label_decode_spec(), + "sparse partition {partition} must declare its label decode description before polling" + ); + chunks.extend(stream.try_collect::>().await.unwrap()); + } + assert_eq!( + 2, + chunks.iter().map(|chunk| chunk.time.len()).sum::() + ); + let metrics_seen = decode_declared_labels(&expected, &chunks); + assert_eq!(metrics_seen.len(), chunks.len()); + for metric in &metrics_seen { + assert_eq!(metric.name(), Some("fixture_metric")); + assert_eq!(metric.get("tag_0"), Some("host-a")); + assert_eq!(metric.get("tag_1"), Some("zone-1")); + } +} + #[tokio::test] async fn test_mito_raw_series_int64_predicate_uses_original_precision() { use datatypes::series::samples::SampleRef; @@ -6388,7 +6589,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { }]), }; let native_plan = Arc::new(RegionScanExec::new(native_scanner, native_request, None).unwrap()) - .try_to_series_scan(native_source.clone()) + .try_to_series_scan(native_source) .unwrap() .expect("typed int64 gate is supported"); let native_schema = native_plan.series_schema().clone(); @@ -6402,7 +6603,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { ordering, "the legacy producer must declare key-then-time frames" ); - let chunks = unpack_raw_series(frames, native_schema, native_source, ordering) + let chunks = unpack_raw_series(frames, Some(native_schema)) .unwrap() .try_collect::>() .await @@ -6447,7 +6648,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { }]), }; let empty_plan = Arc::new(RegionScanExec::new(empty_scanner, empty_request, None).unwrap()) - .try_to_series_scan(empty_source.clone()) + .try_to_series_scan(empty_source) .unwrap() .expect("statically empty scan still has native capability"); let empty_schema = empty_plan.series_schema().clone(); @@ -6460,7 +6661,7 @@ async fn test_mito_raw_series_int64_predicate_uses_original_precision() { empty_ordering, "a statically empty legacy scan still declares key-then-time frames" ); - let empty_chunks = unpack_raw_series(empty_frames, empty_schema, empty_source, empty_ordering) + let empty_chunks = unpack_raw_series(empty_frames, Some(empty_schema)) .unwrap() .try_collect::>() .await diff --git a/src/mito2/src/read/series_scan/raw.rs b/src/mito2/src/read/series_scan/raw.rs index 4adac4263ae..00d6867921f 100644 --- a/src/mito2/src/read/series_scan/raw.rs +++ b/src/mito2/src/read/series_scan/raw.rs @@ -25,6 +25,7 @@ use common_query::error::Error as QueryError; use common_query::native_histogram::{ histogram_samples_from_arrow, is_native_histogram_value_type, }; +use common_query::series_frame::label_decode_spec; use common_recordbatch::filter; use datafusion::physical_plan::PhysicalExpr; use datatypes::arrow::array::{ @@ -37,6 +38,7 @@ use datatypes::arrow::datatypes::DataType; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::series::SeriesOrdering; +use datatypes::series::frame::LabelDecodeSpec; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{Samples, SeriesFields, SeriesSchema}; use futures::stream::BoxStream; @@ -128,6 +130,13 @@ fn chunk( struct RawStream { schema: Arc, label_source: RawLabelSource, + /// The label decode description this stream declares, derived from `label_source` once. + /// + /// Key-encoded chunks leave `metric` as `None`, so a consumer that holds only the stream needs + /// the description to decode them; deriving it here makes the producer self-describing instead + /// of leaving the derivation to whoever packs the chunks. An invalid binding is cached as + /// derived and rejected by the decoder downstream, never silently turned into `None`. + label_decode_spec: Option, ordering: SeriesOrdering, inner: BoxStream<'static, Result>, } @@ -146,6 +155,11 @@ impl RawSeriesStream for RawStream { fn label_source(&self) -> &RawLabelSource { &self.label_source } + /// Returns the description derived from this stream's binding before the first poll; chunks + /// carry key-encoded labels, so the description is what decoding needs. + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + self.label_decode_spec.as_ref() + } fn ordering(&self) -> SeriesOrdering { self.ordering } @@ -161,6 +175,7 @@ pub(crate) fn make_stream( ordering: SeriesOrdering, ) -> SendableRawSeriesStream { let output_schema = schema.clone(); + let label_decode_spec = label_decode_spec(&label_source); let stream = async_stream::stream! { let mut input = inner; let mut pending: Option = None; @@ -267,6 +282,7 @@ pub(crate) fn make_stream( Box::pin(RawStream { schema, label_source, + label_decode_spec, ordering, inner: Box::pin(stream), }) diff --git a/src/promql/src/extension_plan/instant_manipulate/native_exec.rs b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs index e2b41cb8aa4..95b7adc33ff 100644 --- a/src/promql/src/extension_plan/instant_manipulate/native_exec.rs +++ b/src/promql/src/extension_plan/instant_manipulate/native_exec.rs @@ -16,7 +16,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use common_query::series_frame::{raw_series_frame_ordering, unpack_raw_series}; +use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; use datafusion::arrow::array::Array; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -31,7 +31,6 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; -use datatypes::series::SeriesOrdering; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{DenseSamplesRef, Samples, SeriesFields, SeriesSchema}; use futures::{Stream, ready}; @@ -86,16 +85,16 @@ impl InstantManipulate { return Ok(None); } let declared_schema = input.schema(); - let ordering = raw_series_frame_ordering(declared_schema.as_ref())?; - // Validate the declared raw-frame schema without executing the child. + // Validate the declared raw-frame schema without executing the child. The frame states its + // own label layout, so the plan's binding is checked against that declaration rather than + // supplied to the decoder. + ensure_frame_matches_label_source(declared_schema.as_ref(), &label_source)?; unpack_raw_series( Box::pin(RecordBatchStreamAdapter::new( declared_schema.clone(), futures::stream::empty(), )), - series_schema.clone(), - label_source.clone(), - ordering, + Some(series_schema.clone()), )?; Ok(Some(Arc::new(NativeInstantManipulateExec::new( input, @@ -103,7 +102,6 @@ impl InstantManipulate { declared_schema, series_schema, label_source, - ordering, )))) } } @@ -114,7 +112,6 @@ struct NativeInstantManipulateExec { raw_schema: SchemaRef, series_schema: Arc, label_source: RawLabelSource, - ordering: SeriesOrdering, properties: Arc, } impl NativeInstantManipulateExec { @@ -124,7 +121,6 @@ impl NativeInstantManipulateExec { raw_schema: SchemaRef, series_schema: Arc, label_source: RawLabelSource, - ordering: SeriesOrdering, ) -> Self { let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(mapping.schema()), @@ -138,7 +134,6 @@ impl NativeInstantManipulateExec { raw_schema, series_schema, label_source, - ordering, properties, } } @@ -190,7 +185,6 @@ impl ExecutionPlan for NativeInstantManipulateExec { self.raw_schema.clone(), self.series_schema.clone(), self.label_source.clone(), - self.ordering, ))) } fn execute( @@ -204,12 +198,10 @@ impl ExecutionPlan for NativeInstantManipulateExec { )); } let frame = self.input.execute(partition, context.clone())?; - let source = unpack_raw_series( - frame, - self.series_schema.clone(), - self.label_source.clone(), - self.ordering, - )?; + // The child's frames state their own label layout; the plan's binding is an expectation + // checked against that declaration, not a decode input. + ensure_frame_matches_label_source(frame.schema().as_ref(), &self.label_source)?; + let source = unpack_raw_series(frame, Some(self.series_schema.clone()))?; let source = materialize_raw_labels(source) .map_err(|e| DataFusionError::Execution(e.to_string()))?; let reservation = diff --git a/src/promql/src/source.rs b/src/promql/src/source.rs index acd06605f19..03845419344 100644 --- a/src/promql/src/source.rs +++ b/src/promql/src/source.rs @@ -23,6 +23,7 @@ use std::task::{Context, Poll}; use common_error::ext::BoxedError; use datatypes::error::InvalidSeriesSnafu; use datatypes::series::SeriesOrdering; +use datatypes::series::frame::LabelDecodeSpec; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SeriesSchema; use futures::Stream; @@ -31,17 +32,20 @@ use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSerie pub use merge::merge_ordered_raw_series; -/// Materializes complete metric labels from a scan-bound raw-series source. +/// Materializes complete metric labels from a raw-series frame stream. /// -/// The caller supplies an already-bound source. Adjacent-chunk pack/unpack validation remains -/// the caller's responsibility, and nonadjacent key recurrence is governed by the producer -/// contract. Explicit labels are retained as supplied; their completeness cannot be inferred -/// from the schema. No authorization or cross-stream completion is performed here. +/// The labels are decoded from the stream's own label decode description, so a consumer that holds +/// only the frame needs no scan-bound binding; a stream that carries complete labels already, and +/// so declares no description, needs no decoder and is forwarded as supplied. Adjacent-chunk +/// pack/unpack validation remains the caller's responsibility, and nonadjacent key recurrence is +/// governed by the producer contract. Explicit labels are retained as supplied; their completeness +/// cannot be inferred from the schema. No authorization or cross-stream completion is performed +/// here. pub fn materialize_raw_labels( source: SendableRawSeriesStream, ) -> Result { let schema = Arc::clone(source.schema()); - let label_source = source.label_source().clone(); + let decode = source.label_decode_spec().cloned(); let ordering = source.ordering(); if schema.eval_grid.is_some() { return Err(BoxedError::new( @@ -59,12 +63,12 @@ pub fn materialize_raw_labels( .build(), )); } - let decoder = match &label_source { - RawLabelSource::PrimaryKey { .. } => Some( - PrimaryKeyLabelDecoder::new(&label_source, LabelProjection::All) + let decoder = match &decode { + Some(description) => Some( + PrimaryKeyLabelDecoder::from_decode_spec(description, LabelProjection::All) .map_err(BoxedError::new)?, ), - RawLabelSource::Explicit => None, + None => None, }; Ok(Box::pin(MaterializedRawSeriesStream { @@ -125,6 +129,12 @@ impl RawSeriesStream for MaterializedRawSeriesStream { &self.label_source } + /// Returns `None`: materialized labels are already complete, so this stream declares no label + /// decode description and a consumer needs no decoder. + fn label_decode_spec(&self) -> Option<&LabelDecodeSpec> { + None + } + fn ordering(&self) -> SeriesOrdering { self.ordering } @@ -150,10 +160,18 @@ impl MaterializedRawSeriesStream { } } (None, Some(metric)) => Some(metric), - _ => { + (Some(_), Some(_)) => { return Err(BoxedError::new( InvalidSeriesSnafu { - reason: "raw chunk label mode does not match source", + reason: "raw chunk already carries labels but the source describes a key", + } + .build(), + )); + } + (None, None) => { + return Err(BoxedError::new( + InvalidSeriesSnafu { + reason: "raw chunk carries no labels and the source describes no key", } .build(), )); diff --git a/src/promql/src/source/merge.rs b/src/promql/src/source/merge.rs index a734d63c51b..13f19145ea7 100644 --- a/src/promql/src/source/merge.rs +++ b/src/promql/src/source/merge.rs @@ -23,17 +23,25 @@ use std::task::{Context, Poll}; use common_error::ext::BoxedError; use datatypes::error::InvalidSeriesSnafu; use datatypes::series::SeriesOrdering; +use datatypes::series::frame::LabelDecodeSpec; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::{SeriesFields, SeriesSchema}; use futures::Stream; use mito_codec::labels::{LabelProjection, PrimaryKeyLabelDecoder}; -use store_api::region_engine::{RawLabelSource, RawSeriesStream, SendableRawSeriesStream}; +use store_api::region_engine::{ + NO_BOUND_LABEL_SOURCE, RawLabelSource, RawSeriesStream, SendableRawSeriesStream, +}; /// Merge compatible, explicitly supplied raw streams by complete read key and time. /// /// Callers provide authorized bindings and own snapshot/storage policy. This does not select an /// MVCC view, resolve overlapping physical copies, or roll back emitted prefixes after errors. /// `is_last` is global only over these inputs; there is no fallback, replay, or input sorting. +/// +/// The merge contract is the label layout, not a binding: every input must declare the same one — +/// from its frame if it arrived as a frame, otherwise from its own binding — and the merged stream +/// advertises that one declaration. A consumer therefore decodes the merged stream from the frame; +/// the identity it reports is none, see [`RawSeriesStream::label_source`]. pub fn merge_ordered_raw_series( inputs: Vec, ) -> Result { @@ -42,7 +50,7 @@ pub fn merge_ordered_raw_series( .next() .ok_or_else(|| invalid("raw merge requires at least one input"))?; let schema = Arc::clone(first.schema()); - let binding = first.label_source().clone(); + let decode = merged_decode_spec(first.as_ref().get_ref())?; let mut streams = vec![Input::new(first)]; for stream in inputs { streams.push(Input::new(stream)); @@ -52,13 +60,13 @@ pub fn merge_ordered_raw_series( .stream .as_ref() .ok_or_else(|| invalid("raw merge input disappeared"))?; - validate_binding(stream.label_source())?; - validate_input(stream.as_ref().get_ref(), &schema, &binding)?; + validate_input(stream.as_ref().get_ref(), &schema)?; + validate_decode(stream.as_ref().get_ref(), &decode)?; } Ok(Box::pin(OrderedRawMerge { streams, schema, - binding, + decode, candidate: None, })) } @@ -67,17 +75,50 @@ fn invalid(reason: &'static str) -> BoxedError { BoxedError::new(InvalidSeriesSnafu { reason }.build()) } -fn validate_binding(binding: &RawLabelSource) -> Result<(), BoxedError> { - PrimaryKeyLabelDecoder::new(binding, LabelProjection::All) - .map(|_| ()) - .map_err(BoxedError::new) +/// Returns the label layout a stream declares, with no binding required to state it. +/// +/// Every key-encoded stream declares its own description — a producer derives it from the binding it +/// read under and states it before the first poll, and a stream reassembled from a frame carries the +/// description its producer wrote — so the answer is exactly what decoding that stream's chunks +/// needs. A merge can therefore compare inputs that were framed and inputs that were not in the same +/// terms. `None` means the stream carries complete labels and declares no key layout at all. +fn declared_layout(stream: &dyn RawSeriesStream) -> Option { + stream.label_decode_spec().cloned() } -fn validate_input( +/// Returns the one label decode description the merged stream advertises for all inputs. +/// +/// The first input's declaration is forwarded unchanged, so a merge advertises exactly what its +/// inputs advertise, whether they are frame-decoded or straight from a producer. The declaration is +/// validated as a decode layout before it is advertised, so a downstream consumer holding only the +/// merged frame schema can decode it. +fn merged_decode_spec(stream: &dyn RawSeriesStream) -> Result { + let description = declared_layout(stream).ok_or_else(|| { + invalid("raw merge requires inputs that describe a primary-key label layout") + })?; + PrimaryKeyLabelDecoder::from_decode_spec(&description, LabelProjection::All) + .map_err(BoxedError::new)?; + Ok(description) +} + +/// Requires every input to declare the advertised label layout. +/// +/// The merged stream advertises one description, so an input declaring another would be decoded as +/// if it had the advertised encoding. Every input is compared through [`declared_layout`], which +/// reads the input's own declaration, so no input needs a binding to be present. +fn validate_decode( stream: &dyn RawSeriesStream, - schema: &SeriesSchema, - binding: &RawLabelSource, + decode: &LabelDecodeSpec, ) -> Result<(), BoxedError> { + if declared_layout(stream).as_ref() != Some(decode) { + return Err(invalid( + "raw merge inputs describe different label decode layouts", + )); + } + Ok(()) +} + +fn validate_input(stream: &dyn RawSeriesStream, schema: &SeriesSchema) -> Result<(), BoxedError> { let other = stream.schema(); if stream.ordering() != SeriesOrdering::KeyThenTime || other.eval_grid.is_some() @@ -88,58 +129,14 @@ fn validate_input( .iter() .zip(schema.fields.iter()) .any(|(a, b)| a.name != b.name) - || !bindings_compatible(stream.label_source(), binding) { return Err(invalid( - "raw merge requires equal raw schemas, primary-key bindings, and ReadKeyThenTime ordering", + "raw merge requires equal raw schemas and ReadKeyThenTime ordering", )); } Ok(()) } -fn bindings_compatible(left: &RawLabelSource, right: &RawLabelSource) -> bool { - let ( - RawLabelSource::PrimaryKey { - metric_name: lm, - physical_metadata: lmeta, - logical_tags: lt, - }, - RawLabelSource::PrimaryKey { - metric_name: rm, - physical_metadata: rmeta, - logical_tags: rt, - }, - ) = (left, right) - else { - return false; - }; - if lm != rm - || lt.len() != rt.len() - || lt.iter().any(|tag| { - !rt.iter() - .any(|other| tag.column_id == other.column_id && tag.name == other.name) - }) - { - return false; - } - let lm = lmeta.as_ref(); - let rm = rmeta.as_ref(); - lm.primary_key == rm.primary_key - && lm.primary_key_encoding == rm.primary_key_encoding - && lm - .primary_key - .iter() - .all(|id| match (lm.column_by_id(*id), rm.column_by_id(*id)) { - (Some(a), Some(b)) => { - a.column_id == b.column_id - && a.column_schema.name == b.column_schema.name - && a.column_schema.data_type == b.column_schema.data_type - && a.semantic_type == b.semantic_type - } - _ => false, - }) -} - struct Input { stream: Option, head: Option, @@ -223,7 +220,8 @@ impl Input { struct OrderedRawMerge { streams: Vec