From 7d113e2aa4cf94b19c2c4f5c9dfe1ba2e8c7801c Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:53:28 +0800 Subject: [PATCH 1/5] feat(promql): add a versioned wire format for the native prefix Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/native/mod.rs | 1 + src/promql/src/native/wire.rs | 929 ++++++++++++++++++++++++++++++++++ 2 files changed, 930 insertions(+) create mode 100644 src/promql/src/native/wire.rs diff --git a/src/promql/src/native/mod.rs b/src/promql/src/native/mod.rs index fd04dcfda32..4fb62838b73 100644 --- a/src/promql/src/native/mod.rs +++ b/src/promql/src/native/mod.rs @@ -34,3 +34,4 @@ pub mod execute; pub mod metrics; pub mod plan; pub(crate) mod runs; +pub(crate) mod wire; diff --git a/src/promql/src/native/wire.rs b/src/promql/src/native/wire.rs new file mode 100644 index 00000000000..84baccfb925 --- /dev/null +++ b/src/promql/src/native/wire.rs @@ -0,0 +1,929 @@ +// 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. + +//! Native prefix wire format, protocol version 1. +//! +//! Layout: `u32 version`, then recursive node. Node tags are 1 Input, 2 Instant, 3 Window, +//! 4 Map; all integers are little-endian. Input is source index (u64), ordering tag (2 means +//! KeyThenTime), followed by a raw samples schema: field count (u32), each UTF-8 field name +//! (u32 byte length + bytes). Instant stores lookback i64, offset i64, optional-at marker + i64, +//! timestamp bool byte, then output grid. Window stores function tag and its arguments, range i64, +//! offset i64, optional-anchor marker + i64, modifier tag (0 none, 1 anchored, 2 smoothed) and +//! modifier lookback where present, input-grid marker + grid where present, output grid. Grid is +//! start_ms i64, step_ms i64, len u64. Map stores op count u32 and math operations. Math tags and +//! range-function tags below are wire-stable and MUST NOT follow Rust discriminants. Floating +//! values are raw `f64::to_bits()` u64s. Every payload ends exactly after the root node. +//! +//! Stable RangeFunction tags: 1 Rate, 2 Increase, 3 Delta, 4 Irate, 5 Idelta, 6 AvgOverTime, +//! 7 MinOverTime, 8 MaxOverTime, 9 SumOverTime, 10 CountOverTime, 11 LastOverTime, +//! 12 PresentOverTime, 13 StddevOverTime, 14 StdvarOverTime, 15 QuantileOverTime(f64), +//! 16 Resets, 17 Changes, 18 Deriv, 19 PredictLinear(f64), 20 DoubleExponentialSmoothing +//! (smoothing, trend), 21 FirstOverTime, 22 MadOverTime, 23 TsOfFirstOverTime, +//! 24 TsOfLastOverTime, 25 TsOfMinOverTime, 26 TsOfMaxOverTime. Stable MathOp tags: +//! 1 Abs, 2 Ceil, 3 Floor, 4 Sqrt, 5 Round(nearest), 6 Exp, 7 Ln, 8 Log2, 9 Log10, +//! 10 Sin, 11 Cos, 12 Tan, 13 Asin, 14 Acos, 15 Atan, 16 Sinh, 17 Cosh, 18 Tanh, +//! 19 Asinh, 20 Acosh, 21 Atanh, 22 Deg, 23 Rad, 24 Sgn, 25 Clamp(min,max), +//! 26 ClampMin(min), 27 ClampMax(max), 28-35 Date(DayOfMonth, DayOfWeek, DayOfYear, +//! DaysInMonth, Hour, Minute, Month, Year), in that order. + +//! Maximum payload is 1 MiB and maximum recursive node depth / operation count is 128. + +use crate::{ + map::MathOp, + map_time::DateFunction, + native::plan::{ + BoundInput, BoundInstant, BoundMap, BoundWindow, MapOp, NativeExpr, NativeNode, NativeTree, + NativeType, + }, + range_eval::{RangeFunction, RangeModifier, RangeWindow}, +}; +use datatypes::series::{ + SeriesOrdering, + samples::{FieldSpec, SeriesSchema}, + time::{EvalGrid, GridSpec}, +}; +use std::{convert::TryInto, sync::Arc}; + +pub(crate) const PROTOCOL_VERSION: u32 = 1; +const MAX_BYTES: usize = 1_048_576; +const MAX_COUNT: usize = 128; + +struct W(Vec); +impl W { + fn u8(&mut self, v: u8) { + self.0.push(v) + } + fn u32(&mut self, v: u32) { + self.0.extend(v.to_le_bytes()) + } + fn u64(&mut self, v: u64) { + self.0.extend(v.to_le_bytes()) + } + fn i64(&mut self, v: i64) { + self.0.extend(v.to_le_bytes()) + } + fn f(&mut self, v: f64) { + self.u64(v.to_bits()) + } + fn text(&mut self, s: &str) -> Result<(), String> { + self.u32(u32::try_from(s.len()).map_err(|_| "string too long")?); + self.0.extend(s.as_bytes()); + Ok(()) + } + fn grid(&mut self, g: &EvalGrid) { + let s = g.spec(); + self.i64(s.start_ms); + self.i64(s.step_ms); + self.u64(s.len) + } +} +struct R<'a> { + b: &'a [u8], + p: usize, +} +impl<'a> R<'a> { + fn take(&mut self, n: usize) -> Result<&'a [u8], String> { + let e = self.p.checked_add(n).ok_or("length overflow")?; + if e > self.b.len() { + return Err("truncated payload".into()); + } + let x = &self.b[self.p..e]; + self.p = e; + Ok(x) + } + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + fn i64(&mut self) -> Result { + Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + fn f(&mut self) -> Result { + Ok(f64::from_bits(self.u64()?)) + } + fn grid(&mut self) -> Result, String> { + let s = + GridSpec::try_new(self.i64()?, self.i64()?, self.u64()?).map_err(|e| e.to_string())?; + Ok(Arc::new( + EvalGrid::try_deferred(s).map_err(|e| e.to_string())?, + )) + } + fn text(&mut self) -> Result, String> { + let n = self.u32()? as usize; + Ok(std::str::from_utf8(self.take(n)?) + .map_err(|e| e.to_string())? + .into()) + } +} + +/// Encode only the documented Input→Instant/Window→Math-Map forms. +pub(crate) fn encode_prefix(tree: &NativeTree) -> Result, String> { + check_shape(tree.root())?; + let mut w = W(Vec::new()); + w.u32(PROTOCOL_VERSION); + enc_node(tree.root(), &mut w, 0)?; + if w.0.len() > MAX_BYTES { + return Err("payload exceeds MAX_BYTES".into()); + } + Ok(w.0) +} +/// Accepts exactly the supported shapes: a selection over raw input, optionally under one map. +/// +/// The map is optional because a bare selection is a prefix too — it is what a query like `x` +/// transports. When present there is one, and it holds only math operations: a label rewrite or a +/// histogram function would change the identity or the value type in ways this format does not +/// describe. +fn check_shape(n: &NativeNode) -> Result<(), String> { + let selection = match &n.expr { + NativeExpr::Map { input, op } => { + if op.ops.is_empty() || op.ops.iter().any(|x| !matches!(x, MapOp::Math(_))) { + return Err("unsupported Map operation (only Math)".into()); + } + input + } + NativeExpr::Instant { .. } | NativeExpr::Window { .. } => n, + other => return Err(format!("unsupported operation: {}", node_name(other))), + }; + let (NativeExpr::Instant { input: child, .. } | NativeExpr::Window { input: child, .. }) = + &selection.expr + else { + return Err(format!( + "unsupported operation: {}", + node_name(&selection.expr) + )); + }; + if !matches!(child.expr, NativeExpr::Input(_)) { + return Err(format!("unsupported operation: {}", node_name(&child.expr))); + } + Ok(()) +} +fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { + if d > MAX_COUNT { + return Err("nesting maximum exceeded".into()); + } + match &n.expr { + NativeExpr::Input(i) => { + w.u8(1); + w.u64(i.source_index as u64); + if i.ordering != SeriesOrdering::KeyThenTime { + return Err("unsupported Input ordering".into()); + } + w.u8(2); + let NativeType::Samples(s) = &i.output else { + return Err("unsupported Input output type".into()); + }; + if s.eval_grid.is_some() { + return Err("unsupported evaluated Input".into()); + } + w.u32(s.fields.len() as u32); + for f in s.fields.iter() { + w.text(&f.name)? + } + } + NativeExpr::Instant { input, op } => { + w.u8(2); + enc_node(input, w, d + 1)?; + w.i64(op.lookback_ms); + w.i64(op.offset_ms); + opt_i(w, op.at_ms); + w.u8(op.timestamp as u8); + w.grid(&op.output_grid) + } + NativeExpr::Window { input, op } => { + w.u8(3); + enc_node(input, w, d + 1)?; + enc_range(op.function, w)?; + w.i64(op.window.range_ms); + w.i64(op.window.offset_ms); + opt_i(w, op.window.anchor_ms); + match op.window.modifier { + None => w.u8(0), + Some(RangeModifier::Anchored { lookback_ms }) => { + w.u8(1); + w.i64(lookback_ms) + } + Some(RangeModifier::Smoothed { lookback_ms }) => { + w.u8(2); + w.i64(lookback_ms) + } + } + match &op.input_grid { + None => w.u8(0), + Some(g) => { + w.u8(1); + w.grid(g) + } + } + w.grid(&op.output_grid) + } + NativeExpr::Map { input, op } => { + w.u8(4); + enc_node(input, w, d + 1)?; + if op.ops.is_empty() || op.ops.len() > MAX_COUNT { + return Err("operation count exceeds supported maximum or is zero".into()); + } + w.u32(op.ops.len() as u32); + for x in op.ops.iter() { + if let MapOp::Math(m) = x { + enc_math(*m, w)? + } else { + return Err("unsupported Map operation (only Math)".into()); + } + } + } + _ => return Err(format!("unsupported operation: {}", node_name(&n.expr))), + } + Ok(()) +} +fn opt_i(w: &mut W, x: Option) { + if let Some(v) = x { + w.u8(1); + w.i64(v) + } else { + w.u8(0) + } +} +fn enc_math(m: MathOp, w: &mut W) -> Result<(), String> { + let (t, vals): (u8, &[f64]) = match m { + MathOp::Abs => (1, &[]), + MathOp::Ceil => (2, &[]), + MathOp::Floor => (3, &[]), + MathOp::Sqrt => (4, &[]), + MathOp::Round { nearest } => (5, &[nearest]), + MathOp::Exp => (6, &[]), + MathOp::Ln => (7, &[]), + MathOp::Log2 => (8, &[]), + MathOp::Log10 => (9, &[]), + MathOp::Sin => (10, &[]), + MathOp::Cos => (11, &[]), + MathOp::Tan => (12, &[]), + MathOp::Asin => (13, &[]), + MathOp::Acos => (14, &[]), + MathOp::Atan => (15, &[]), + MathOp::Sinh => (16, &[]), + MathOp::Cosh => (17, &[]), + MathOp::Tanh => (18, &[]), + MathOp::Asinh => (19, &[]), + MathOp::Acosh => (20, &[]), + MathOp::Atanh => (21, &[]), + MathOp::Deg => (22, &[]), + MathOp::Rad => (23, &[]), + MathOp::Sgn => (24, &[]), + MathOp::Clamp { min, max } => (25, &[min, max]), + MathOp::ClampMin { min } => (26, &[min]), + MathOp::ClampMax { max } => (27, &[max]), + MathOp::Date(f) => ( + 28 + match f { + DateFunction::DayOfMonth => 0, + DateFunction::DayOfWeek => 1, + DateFunction::DayOfYear => 2, + DateFunction::DaysInMonth => 3, + DateFunction::Hour => 4, + DateFunction::Minute => 5, + DateFunction::Month => 6, + DateFunction::Year => 7, + }, + &[], + ), + }; + w.u8(t); + for x in vals { + w.f(*x) + } + Ok(()) +} +fn enc_range(f: RangeFunction, w: &mut W) -> Result<(), String> { + let (tag, vs): (u8, &[f64]) = match f { + RangeFunction::Rate => (1, &[]), + RangeFunction::Increase => (2, &[]), + RangeFunction::Delta => (3, &[]), + RangeFunction::Irate => (4, &[]), + RangeFunction::Idelta => (5, &[]), + RangeFunction::AvgOverTime => (6, &[]), + RangeFunction::MinOverTime => (7, &[]), + RangeFunction::MaxOverTime => (8, &[]), + RangeFunction::SumOverTime => (9, &[]), + RangeFunction::CountOverTime => (10, &[]), + RangeFunction::LastOverTime => (11, &[]), + RangeFunction::PresentOverTime => (12, &[]), + RangeFunction::StddevOverTime => (13, &[]), + RangeFunction::StdvarOverTime => (14, &[]), + RangeFunction::QuantileOverTime(x) => (15, &[x]), + RangeFunction::Resets => (16, &[]), + RangeFunction::Changes => (17, &[]), + RangeFunction::Deriv => (18, &[]), + RangeFunction::PredictLinear(x) => (19, &[x]), + RangeFunction::DoubleExponentialSmoothing { smoothing, trend } => (20, &[smoothing, trend]), + RangeFunction::FirstOverTime => (21, &[]), + RangeFunction::MadOverTime => (22, &[]), + RangeFunction::TsOfFirstOverTime => (23, &[]), + RangeFunction::TsOfLastOverTime => (24, &[]), + RangeFunction::TsOfMinOverTime => (25, &[]), + RangeFunction::TsOfMaxOverTime => (26, &[]), + }; + w.u8(tag); + for x in vs { + w.f(*x) + } + Ok(()) +} +fn node_name(e: &NativeExpr) -> &'static str { + match e { + NativeExpr::Binary { .. } => "Binary", + NativeExpr::Aggregate { .. } => "Aggregate", + NativeExpr::Vector { .. } => "Vector", + NativeExpr::Scalar(_) => "Scalar", + NativeExpr::Timestamp { .. } => "Timestamp", + _ => "unsupported", + } +} +/// Decode and validate a supported prefix tree. +pub(crate) fn decode_prefix(bytes: &[u8]) -> Result { + if bytes.len() > MAX_BYTES { + return Err("payload exceeds MAX_BYTES".into()); + } + let mut r = R { b: bytes, p: 0 }; + let v = r.u32()?; + if v != PROTOCOL_VERSION { + return Err(format!("unknown version {v}")); + } + let n = dec_node(&mut r, 0)?; + check_shape(&n)?; + if r.p != bytes.len() { + return Err("trailing bytes after complete payload".into()); + } + NativeTree::try_new(n).map_err(|e| e.to_string()) +} +fn dec_node(r: &mut R<'_>, d: usize) -> Result { + if d > MAX_COUNT { + return Err("nesting maximum exceeded".into()); + } + match r.u8()? { + 1 => { + let ix = usize::try_from(r.u64()?).map_err(|_| "source index overflow")?; + if r.u8()? != 2 { + return Err("unsupported Input ordering".into()); + } + let n = r.u32()? as usize; + if n > MAX_COUNT { + return Err("operation count exceeds maximum".into()); + } + let mut fs = Vec::with_capacity(n); + for _ in 0..n { + fs.push(FieldSpec { name: r.text()? }) + } + let ty = NativeType::Samples(Arc::new(SeriesSchema { + fields: fs.into(), + eval_grid: None, + })); + Ok(NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: ix, + output: ty.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: ty, + }) + } + 2 => { + let input = Box::new(dec_node(r, d + 1)?); + let lookback_ms = r.i64()?; + let offset_ms = r.i64()?; + let at_ms = read_opt(r)?; + let timestamp = match r.u8()? { + 0 => false, + 1 => true, + _ => return Err("invalid timestamp flag".into()), + }; + let output_grid = r.grid()?; + let schema = match &input.output { + NativeType::Samples(s) => Arc::new(SeriesSchema { + fields: s.fields.clone(), + eval_grid: Some(output_grid.clone()), + }), + _ => return Err("unsupported Instant input type".into()), + }; + let output = NativeType::Samples(schema); + Ok(NativeNode { + expr: NativeExpr::Instant { + input, + op: BoundInstant { + lookback_ms, + offset_ms, + at_ms, + timestamp, + output_grid, + }, + }, + output, + }) + } + 3 => { + let input = Box::new(dec_node(r, d + 1)?); + let function = dec_range(r)?; + let range_ms = r.i64()?; + let offset_ms = r.i64()?; + let anchor_ms = read_opt(r)?; + let modifier = match r.u8()? { + 0 => None, + 1 => Some(RangeModifier::Anchored { + lookback_ms: r.i64()?, + }), + 2 => Some(RangeModifier::Smoothed { + lookback_ms: r.i64()?, + }), + _ => return Err("unknown range modifier".into()), + }; + let input_grid = match r.u8()? { + 0 => None, + 1 => Some(r.grid()?), + _ => return Err("invalid input grid marker".into()), + }; + let output_grid = r.grid()?; + let fields = match &input.output { + NativeType::Samples(s) => s.fields.clone(), + _ => return Err("unsupported Window input type".into()), + }; + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields, + eval_grid: Some(output_grid.clone()), + })); + Ok(NativeNode { + expr: NativeExpr::Window { + input, + op: BoundWindow { + function, + window: RangeWindow { + range_ms, + offset_ms, + anchor_ms, + modifier, + }, + input_grid, + output_grid, + }, + }, + output, + }) + } + 4 => { + let input = Box::new(dec_node(r, d + 1)?); + let n = r.u32()? as usize; + if n == 0 || n > MAX_COUNT { + return Err("operation count exceeds maximum".into()); + } + let mut ops = Vec::new(); + for _ in 0..n { + ops.push(MapOp::Math(dec_math(r)?)) + } + let output = match &input.output { + NativeType::Samples(s) => NativeType::Samples(s.clone()), + _ => return Err("unsupported Map input type".into()), + }; + Ok(NativeNode { + expr: NativeExpr::Map { + input, + op: BoundMap { ops: ops.into() }, + }, + output, + }) + } + x => Err(format!("unsupported operation tag {x}")), + } +} +fn read_opt(r: &mut R<'_>) -> Result, String> { + match r.u8()? { + 0 => Ok(None), + 1 => Ok(Some(r.i64()?)), + _ => Err("invalid optional integer marker".into()), + } +} +fn dec_range(r: &mut R<'_>) -> Result { + Ok(match r.u8()? { + 1 => RangeFunction::Rate, + 2 => RangeFunction::Increase, + 3 => RangeFunction::Delta, + 4 => RangeFunction::Irate, + 5 => RangeFunction::Idelta, + 6 => RangeFunction::AvgOverTime, + 7 => RangeFunction::MinOverTime, + 8 => RangeFunction::MaxOverTime, + 9 => RangeFunction::SumOverTime, + 10 => RangeFunction::CountOverTime, + 11 => RangeFunction::LastOverTime, + 12 => RangeFunction::PresentOverTime, + 13 => RangeFunction::StddevOverTime, + 14 => RangeFunction::StdvarOverTime, + 15 => RangeFunction::QuantileOverTime(r.f()?), + 16 => RangeFunction::Resets, + 17 => RangeFunction::Changes, + 18 => RangeFunction::Deriv, + 19 => RangeFunction::PredictLinear(r.f()?), + 20 => RangeFunction::DoubleExponentialSmoothing { + smoothing: r.f()?, + trend: r.f()?, + }, + 21 => RangeFunction::FirstOverTime, + 22 => RangeFunction::MadOverTime, + 23 => RangeFunction::TsOfFirstOverTime, + 24 => RangeFunction::TsOfLastOverTime, + 25 => RangeFunction::TsOfMinOverTime, + 26 => RangeFunction::TsOfMaxOverTime, + x => return Err(format!("unsupported RangeFunction tag {x}")), + }) +} +fn dec_math(r: &mut R<'_>) -> Result { + Ok(match r.u8()? { + 1 => MathOp::Abs, + 2 => MathOp::Ceil, + 3 => MathOp::Floor, + 4 => MathOp::Sqrt, + 5 => MathOp::Round { nearest: r.f()? }, + 6 => MathOp::Exp, + 7 => MathOp::Ln, + 8 => MathOp::Log2, + 9 => MathOp::Log10, + 10 => MathOp::Sin, + 11 => MathOp::Cos, + 12 => MathOp::Tan, + 13 => MathOp::Asin, + 14 => MathOp::Acos, + 15 => MathOp::Atan, + 16 => MathOp::Sinh, + 17 => MathOp::Cosh, + 18 => MathOp::Tanh, + 19 => MathOp::Asinh, + 20 => MathOp::Acosh, + 21 => MathOp::Atanh, + 22 => MathOp::Deg, + 23 => MathOp::Rad, + 24 => MathOp::Sgn, + 25 => MathOp::Clamp { + min: r.f()?, + max: r.f()?, + }, + 26 => MathOp::ClampMin { min: r.f()? }, + 27 => MathOp::ClampMax { max: r.f()? }, + 28 => MathOp::Date(DateFunction::DayOfMonth), + 29 => MathOp::Date(DateFunction::DayOfWeek), + 30 => MathOp::Date(DateFunction::DayOfYear), + 31 => MathOp::Date(DateFunction::DaysInMonth), + 32 => MathOp::Date(DateFunction::Hour), + 33 => MathOp::Date(DateFunction::Minute), + 34 => MathOp::Date(DateFunction::Month), + 35 => MathOp::Date(DateFunction::Year), + x => return Err(format!("unsupported MathOp tag {x}")), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + fn input() -> NativeNode { + let t = NativeType::Samples(Arc::new(SeriesSchema { + fields: vec![FieldSpec { name: "v".into() }].into(), + eval_grid: None, + })); + NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: 0, + output: t.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: t, + } + } + fn grid() -> Arc { + Arc::new(EvalGrid::try_deferred(GridSpec::try_new(100, 10, 3).unwrap()).unwrap()) + } + /// `rate` over the raw input on the fixture grid: the simplest window a prefix transports. + fn rate_window() -> NativeTree { + let input = Box::new(input()); + let g = grid(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields: match &input.output { + NativeType::Samples(s) => s.fields.clone(), + _ => unreachable!(), + }, + eval_grid: Some(g.clone()), + })); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input, + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 3_000, + offset_ms: 0, + anchor_ms: None, + modifier: None, + }, + input_grid: None, + output_grid: g, + }, + }, + output, + }) + .unwrap() + } + + fn instant() -> NativeTree { + let input = Box::new(input()); + let g = grid(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields: match &input.output { + NativeType::Samples(s) => s.fields.clone(), + _ => unreachable!(), + }, + eval_grid: Some(g.clone()), + })); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input, + op: BoundInstant { + lookback_ms: 300_000, + offset_ms: 7, + at_ms: Some(88), + timestamp: true, + output_grid: g, + }, + }, + output, + }) + .unwrap() + } + fn window() -> NativeTree { + let input = Box::new(input()); + let g = grid(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields: match &input.output { + NativeType::Samples(s) => s.fields.clone(), + _ => unreachable!(), + }, + eval_grid: Some(g.clone()), + })); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input, + op: BoundWindow { + function: RangeFunction::QuantileOverTime(f64::from_bits( + 0x7ff8_0000_0000_0042, + )), + window: RangeWindow { + range_ms: 500, + offset_ms: 9, + anchor_ms: Some(-4), + modifier: Some(RangeModifier::Smoothed { lookback_ms: 99 }), + }, + input_grid: None, + output_grid: g, + }, + }, + output, + }) + .unwrap() + } + fn mapped() -> NativeTree { + let instant = instant(); + let input = Box::new(instant.root().clone()); + let output = input.output.clone(); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![ + MapOp::Math(MathOp::Round { nearest: 0.5 }), + MapOp::Math(MathOp::Clamp { + min: -1.0, + max: 2.0, + }), + MapOp::Math(MathOp::Abs), + ] + .into(), + }, + }, + output, + }) + .unwrap() + } + #[test] + fn instant_roundtrip_and_fields() { + let a = instant(); + let b = encode_prefix(&a).unwrap(); + let d = decode_prefix(&b).unwrap(); + assert_eq!(encode_prefix(&d).unwrap(), b); + let NativeExpr::Instant { op, .. } = &d.root().expr else { + panic!() + }; + assert_eq!(op.lookback_ms, 300_000); + assert_eq!(op.offset_ms, 7); + assert_eq!(op.at_ms, Some(88)); + assert!(op.timestamp); + assert_eq!(op.output_grid.spec(), grid().spec()); + } + #[test] + fn window_roundtrip_and_fields() { + let a = window(); + let b = encode_prefix(&a).unwrap(); + let d = decode_prefix(&b).unwrap(); + assert_eq!(encode_prefix(&d).unwrap(), b); + let NativeExpr::Window { op, .. } = &d.root().expr else { + panic!() + }; + assert!( + matches!(op.function,RangeFunction::QuantileOverTime(x) if x.to_bits()==0x7ff8_0000_0000_0042) + ); + assert_eq!(op.window.range_ms, 500); + assert_eq!(op.output_grid.spec(), grid().spec()); + } + #[test] + fn map_roundtrip_and_ops() { + let a = mapped(); + let b = encode_prefix(&a).unwrap(); + let d = decode_prefix(&b).unwrap(); + assert_eq!(encode_prefix(&d).unwrap(), b); + let NativeExpr::Map { op, .. } = &d.root().expr else { + panic!() + }; + assert_eq!(op.ops.len(), 3); + assert!( + matches!(op.ops[0],MapOp::Math(MathOp::Round{nearest}) if nearest.to_bits()==0.5f64.to_bits()) + ); + } + #[test] + fn floats_preserve_bits() { + let values = [ + f64::from_bits(0x7ff8_0000_0000_0042), + -0.0, + 0.0, + f64::INFINITY, + f64::NEG_INFINITY, + f64::from_bits(1), + ]; + for x in values { + let mut w = W(Vec::new()); + w.f(x); + let mut r = R { b: &w.0, p: 0 }; + assert_eq!(r.f().unwrap().to_bits(), x.to_bits()); + } + } + fn labels() -> Option { + datatypes::series::SeriesLabels::try_new( + Some("m".into()), + vec![datatypes::series::Label { + name: "job".into(), + value: "x".into(), + }], + ) + .ok() + } + + /// The decoded tree must compute what the original computes, not merely encode back to the + /// same bytes. + /// + /// Both trees are driven through the real executor over the same raw input, and their answers + /// are compared. A codec that dropped a parameter — a lookback, a window offset, a clamp bound + /// — would still round-trip its own bytes, so only executing them can tell. + #[test] + fn decoded_tree_computes_the_same_answer() { + use crate::native::execute::{InputId, NativeExecutionState, NativeExecutor}; + use datafusion::execution::memory_pool::UnboundedMemoryPool; + use datatypes::arrow::array::Float64Array; + use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; + use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; + use datatypes::series::samples::{Samples, SeriesFields}; + + // A plain rate window over raw input: a legal tree the executor accepts. The float-bits + // fixture carries a `QuantileOverTime(NaN)` under a smoothed modifier, which the executor + // rejects on purpose, so it cannot be the tree this test runs. + let tree = rate_window(); + let decoded = decode_prefix(&encode_prefix(&tree).unwrap()).unwrap(); + + let answer = |tree: &NativeTree| { + let mut execution = NativeExecutor::new( + NativeExecutionState::try_new(tree, Arc::new(UnboundedMemoryPool::default())) + .unwrap(), + ); + let schema = tree + .root() + .output + .grid() + .map(|_| ()) + .and_then(|()| match &tree.root().expr { + NativeExpr::Window { input, .. } => match &input.output { + NativeType::Samples(schema) => Some(Arc::clone(schema)), + other => panic!("a window reads samples, not {other:?}"), + }, + other => panic!("the fixture is a window, not {other:?}"), + }) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from_slice_ref(b"series-a"), + labels(), + // Samples before, inside and after the fixture grid (100, 110, 120) so the + // window has a counter to measure across every evaluation point. + RawTimestamps::try_new(ScalarBuffer::from(vec![60i64, 100, 110, 120, 140])) + .unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + 0.0, 10.0, 20.0, 30.0, 50.0, + ]))), + true, + ) + .unwrap(); + execution.push_raw(InputId::new(0), chunk).unwrap(); + execution.end_of_input().unwrap(); + let mut values = Vec::new(); + while let Some(chunk) = execution.next_chunk().unwrap() { + for sample in chunk.series.fields.get(0).unwrap().iter() { + match sample { + datatypes::series::samples::SampleRef::Float(value) => { + values.push(Some(value)) + } + _ => values.push(None), + } + } + } + values + }; + + let original = answer(&tree); + let round_tripped = answer(&decoded); + assert_eq!( + original, round_tripped, + "the decoded tree answered differently from the original" + ); + assert!( + original.iter().any(Option::is_some), + "the fixture produced no values, so it proves nothing" + ); + } + + #[test] + fn rejects_unknown_version() { + let mut b = encode_prefix(&instant()).unwrap(); + b[0] = 2; + assert!(decode_prefix(&b).unwrap_err().contains("unknown version")); + } + #[test] + fn rejects_unsupported_node_kind() { + let n = NativeNode { + expr: NativeExpr::Scalar(crate::native::plan::BoundScalar { value: 1.0 }), + output: NativeType::Scalar(grid()), + }; + let t = NativeTree::try_new(n).unwrap(); + assert!(encode_prefix(&t).unwrap_err().contains("Scalar")); + } + #[test] + fn rejects_truncated_payload() { + let mut b = encode_prefix(&instant()).unwrap(); + b.pop(); + assert!(decode_prefix(&b).unwrap_err().contains("truncated")); + } + #[test] + fn rejects_trailing_bytes() { + let mut b = encode_prefix(&instant()).unwrap(); + b.push(0); + assert!(decode_prefix(&b).unwrap_err().contains("trailing")); + } + #[test] + fn rejects_oversized_payload() { + let b = vec![0; MAX_BYTES + 1]; + assert!(decode_prefix(&b).unwrap_err().contains("MAX_BYTES")); + } + #[test] + fn rejects_excess_operation_count() { + let mut w = W(Vec::new()); + w.u32(PROTOCOL_VERSION); + w.u8(4); + w.u8(1); + w.u64(0); + w.u8(2); + w.u32(1); + w.text("v").unwrap(); + w.u8(2); + w.i64(1); + w.i64(0); + opt_i(&mut w, None); + w.u8(0); + w.grid(&grid()); + w.u32((MAX_COUNT + 1) as u32); + assert!(decode_prefix(&w.0).unwrap_err().contains("operation count")); + } +} From ecb76d7e3007d4080ce442dba9fc7291ec5ee5b2 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:04:20 +0800 Subject: [PATCH 2/5] feat(promql): serialize fused native prefixes through substrait Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../substrait/src/extension_serializer.rs | 12 +- src/promql/src/extension_plan/fused_series.rs | 172 ++++++++++++++++++ .../src/extension_plan/fused_series/tests.rs | 33 ++++ src/promql/src/native/wire.rs | 49 ++--- 4 files changed, 241 insertions(+), 25 deletions(-) diff --git a/src/common/substrait/src/extension_serializer.rs b/src/common/substrait/src/extension_serializer.rs index 5e94a1a57ae..e14cf627515 100644 --- a/src/common/substrait/src/extension_serializer.rs +++ b/src/common/substrait/src/extension_serializer.rs @@ -19,8 +19,8 @@ use datafusion::execution::registry::SerializerRegistry; use datafusion_common::DataFusionError; use datafusion_expr::UserDefinedLogicalNode; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate, - SeriesDivide, SeriesNormalize, UnionDistinctOn, + Absent, EmptyMetric, FusedSeries, HistogramFold, InstantManipulate, RangeManipulate, + ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, }; #[derive(Debug)] @@ -31,6 +31,13 @@ impl SerializerRegistry for ExtensionSerializer { /// input plans. fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result> { match node.name() { + name if name == FusedSeries::name() => { + let fused = node + .as_any() + .downcast_ref::() + .expect("Failed to downcast to FusedSeries"); + fused.serialize() + } name if name == InstantManipulate::name() => { let instant_manipulate = node .as_any() @@ -105,6 +112,7 @@ impl SerializerRegistry for ExtensionSerializer { bytes: &[u8], ) -> Result> { match name { + name if name == FusedSeries::name() => Ok(Arc::new(FusedSeries::deserialize(bytes)?)), name if name == InstantManipulate::name() => { let instant_manipulate = InstantManipulate::deserialize(bytes)?; Ok(Arc::new(instant_manipulate)) diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index f36d63f8839..f02c09e0f5d 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -232,6 +232,173 @@ pub struct FusedSeries { } impl FusedSeries { + /// Serializes the native tree and the binding required to reconstruct this plan node. + /// The PromQL text is included for diagnostics only and never drives execution; children + /// remain ordinary Substrait inputs. Version 1 supports only SeriesFrame without partial folds: + /// decoding must rebuild its output schema, and only SeriesFrame derives that schema from the tree. + pub fn serialize(&self) -> DataFusionResult> { + if self.format != FusedOutputFormat::SeriesFrame { + return Err(DataFusionError::NotImplemented( + "FusedSeries serialization does not support PlainRows output schema".into(), + )); + } + if self.partial_fold.is_some() { + return Err(DataFusionError::NotImplemented( + "FusedSeries serialization does not support partial_fold".into(), + )); + } + fn text(out: &mut Vec, s: &str) -> DataFusionResult<()> { + let n = u32::try_from(s.len()) + .map_err(|_| DataFusionError::Plan("FusedSeries string too long".into()))?; + out.extend(n.to_le_bytes()); + out.extend(s.as_bytes()); + Ok(()) + } + let tree = crate::native::wire::encode_prefix(&self.tree) + .map_err(|e| DataFusionError::Plan(format!("cannot encode FusedSeries tree: {e}")))?; + let mut out = Vec::new(); + out.extend(1u32.to_le_bytes()); + out.extend( + u32::try_from(tree.len()) + .map_err(|_| DataFusionError::Plan("FusedSeries tree too large".into()))? + .to_le_bytes(), + ); + out.extend(tree); + let spec = self.binding.grid.spec(); + out.extend(spec.start_ms.to_le_bytes()); + out.extend(spec.step_ms.to_le_bytes()); + out.extend(spec.len.to_le_bytes()); + match &self.binding.metric_name { + Some(name) => { + out.push(1); + text(&mut out, name)?; + } + None => out.push(0), + } + text(&mut out, &self.binding.time_index)?; + out.extend( + u32::try_from(self.binding.tags.len()) + .map_err(|_| DataFusionError::Plan("too many FusedSeries tags".into()))? + .to_le_bytes(), + ); + for tag in &self.binding.tags { + text(&mut out, tag)?; + } + out.push(1); // explicit stable tag: SeriesFrame + text(&mut out, &self.promql)?; + Ok(out) + } + + /// Rebuilds a serialized SeriesFrame node. Substrait attaches its child plans separately. + pub fn deserialize(bytes: &[u8]) -> DataFusionResult { + struct Reader<'a> { + b: &'a [u8], + p: usize, + } + impl<'a> Reader<'a> { + fn take(&mut self, n: usize) -> DataFusionResult<&'a [u8]> { + let e = self + .p + .checked_add(n) + .ok_or_else(|| DataFusionError::Plan("FusedSeries length overflow".into()))?; + if e > self.b.len() { + return Err(DataFusionError::Plan( + "truncated FusedSeries payload".into(), + )); + } + let r = &self.b[self.p..e]; + self.p = e; + Ok(r) + } + fn u8(&mut self) -> DataFusionResult { + Ok(self.take(1)?[0]) + } + fn u32(&mut self) -> DataFusionResult { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + fn i64(&mut self) -> DataFusionResult { + Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + fn u64(&mut self) -> DataFusionResult { + Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + fn text(&mut self) -> DataFusionResult { + let n = self.u32()? as usize; + String::from_utf8(self.take(n)?.to_vec()) + .map_err(|e| DataFusionError::Plan(format!("invalid FusedSeries UTF-8: {e}"))) + } + } + let mut r = Reader { b: bytes, p: 0 }; + let version = r.u32()?; + if version != 1 { + return Err(DataFusionError::Plan(format!( + "unsupported FusedSeries envelope version {version}" + ))); + } + let n = r.u32()? as usize; + let tree = crate::native::wire::decode_prefix(r.take(n)?) + .map_err(|e| DataFusionError::Plan(format!("cannot decode FusedSeries tree: {e}")))?; + let grid = Arc::new( + EvalGrid::try_deferred( + datatypes::series::time::GridSpec::try_new(r.i64()?, r.i64()?, r.u64()?) + .map_err(|e| DataFusionError::Plan(format!("invalid FusedSeries grid: {e}")))?, + ) + .map_err(|e| DataFusionError::Plan(format!("invalid FusedSeries grid: {e}")))?, + ); + let metric_name = match r.u8()? { + 0 => None, + 1 => Some(Arc::::from(r.text()?)), + x => { + return Err(DataFusionError::Plan(format!( + "invalid FusedSeries metric-name tag {x}" + ))); + } + }; + let time_index = r.text()?; + let count = r.u32()? as usize; + if count > 1024 { + return Err(DataFusionError::Plan("too many FusedSeries tags".into())); + } + let mut tags = Vec::with_capacity(count); + for _ in 0..count { + tags.push(r.text()?) + } + let format = match r.u8()? { + 1 => FusedOutputFormat::SeriesFrame, + x => { + return Err(DataFusionError::Plan(format!( + "unsupported FusedSeries output format tag {x}" + ))); + } + }; + let promql = Arc::::from(r.text()?); + if r.p != bytes.len() { + return Err(DataFusionError::Plan( + "trailing bytes in FusedSeries payload".into(), + )); + } + let children = vec![ + LogicalPlan::EmptyRelation(datafusion::logical_expr::EmptyRelation { + produce_one_row: false, + schema: Arc::new(datafusion::common::DFSchema::empty()) + }); + collect_source_schemas(&tree)?.len() + ]; + Self::try_new_with_format( + promql, + tree, + FusedSeriesBinding { + grid, + metric_name, + time_index, + tags, + }, + children, + Arc::new(datafusion::common::DFSchema::empty()), + format, + ) + } + /// The name this node is shown under. pub const fn name() -> &'static str { "FusedSeries" @@ -334,6 +501,11 @@ impl FusedSeries { &self.promql } + /// The selected output representation. + pub fn format(&self) -> FusedOutputFormat { + self.format + } + /// The compiled tree this node runs. pub fn tree(&self) -> &NativeTree { &self.tree diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index ccc7c037615..eb1ed2ad769 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -924,6 +924,39 @@ impl Drop for StagedStream { } } +#[test] +fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() { + let tree = rate_tree(); + let original_binding = binding(&tree); + let node = FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree, + original_binding.clone(), + vec![input_plan(frame_schema())], + output_schema().to_dfschema_ref().unwrap(), + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let bytes = node.serialize().unwrap(); + let decoded = FusedSeries::deserialize(&bytes).unwrap(); + assert_eq!( + crate::native::wire::encode_prefix(decoded.tree()).unwrap(), + crate::native::wire::encode_prefix(node.tree()).unwrap() + ); + assert_eq!(decoded.binding(), node.binding()); + assert_eq!(decoded.schema(), node.schema()); + assert!(fused(rate_tree()).serialize().is_err()); + let mut bad_version = bytes.clone(); + bad_version[..4].copy_from_slice(&99u32.to_le_bytes()); + assert!( + FusedSeries::deserialize(&bad_version) + .unwrap_err() + .to_string() + .contains("version") + ); + assert!(FusedSeries::deserialize(&bytes[..5]).is_err()); +} + fn fused(tree: NativeTree) -> FusedSeries { let binding = binding(&tree); let output_schema = if matches!(&tree.root().expr, NativeExpr::Binary { .. }) { diff --git a/src/promql/src/native/wire.rs b/src/promql/src/native/wire.rs index 84baccfb925..4b29406885b 100644 --- a/src/promql/src/native/wire.rs +++ b/src/promql/src/native/wire.rs @@ -15,8 +15,8 @@ //! Native prefix wire format, protocol version 1. //! //! Layout: `u32 version`, then recursive node. Node tags are 1 Input, 2 Instant, 3 Window, -//! 4 Map; all integers are little-endian. Input is source index (u64), ordering tag (2 means -//! KeyThenTime), followed by a raw samples schema: field count (u32), each UTF-8 field name +//! 4 Map; all integers are little-endian. Input is source index (u64), ordering tag (0 +//! Interleaved, 1 Contiguous, 2 KeyThenTime), followed by a raw samples schema: field count (u32), each UTF-8 field name //! (u32 byte length + bytes). Instant stores lookback i64, offset i64, optional-at marker + i64, //! timestamp bool byte, then output grid. Window stores function tag and its arguments, range i64, //! offset i64, optional-anchor marker + i64, modifier tag (0 none, 1 anchored, 2 smoothed) and @@ -181,10 +181,14 @@ fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { NativeExpr::Input(i) => { w.u8(1); w.u64(i.source_index as u64); - if i.ordering != SeriesOrdering::KeyThenTime { - return Err("unsupported Input ordering".into()); - } - w.u8(2); + // The ordering is data, not a policy of this format: every value is carried so the + // tree decides on decode whether it is legal for its consumers (a window refuses only + // `Interleaved`). Refusing one here would make the wire narrower than the node. + w.u8(match i.ordering { + SeriesOrdering::Interleaved => 0, + SeriesOrdering::Contiguous => 1, + SeriesOrdering::KeyThenTime => 2, + }); let NativeType::Samples(s) = &i.output else { return Err("unsupported Input output type".into()); }; @@ -377,9 +381,12 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { match r.u8()? { 1 => { let ix = usize::try_from(r.u64()?).map_err(|_| "source index overflow")?; - if r.u8()? != 2 { - return Err("unsupported Input ordering".into()); - } + let ordering = match r.u8()? { + 0 => SeriesOrdering::Interleaved, + 1 => SeriesOrdering::Contiguous, + 2 => SeriesOrdering::KeyThenTime, + x => return Err(format!("unsupported Input ordering tag {x}")), + }; let n = r.u32()? as usize; if n > MAX_COUNT { return Err("operation count exceeds maximum".into()); @@ -396,7 +403,7 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { expr: NativeExpr::Input(BoundInput { source_index: ix, output: ty.clone(), - ordering: SeriesOrdering::KeyThenTime, + ordering, }), output: ty, }) @@ -820,19 +827,15 @@ mod tests { NativeExecutionState::try_new(tree, Arc::new(UnboundedMemoryPool::default())) .unwrap(), ); - let schema = tree - .root() - .output - .grid() - .map(|_| ()) - .and_then(|()| match &tree.root().expr { - NativeExpr::Window { input, .. } => match &input.output { - NativeType::Samples(schema) => Some(Arc::clone(schema)), - other => panic!("a window reads samples, not {other:?}"), - }, - other => panic!("the fixture is a window, not {other:?}"), - }) - .unwrap(); + // The fixture is a window over raw input, so the schema its samples are addressed on + // comes from the window's input, not from the window's own evaluated output. + let NativeExpr::Window { input, .. } = &tree.root().expr else { + panic!("the fixture is a window"); + }; + let NativeType::Samples(schema) = &input.output else { + panic!("a window reads samples"); + }; + let schema = Arc::clone(schema); let chunk = RawSeriesChunk::try_new( &schema, Buffer::from_slice_ref(b"series-a"), From 1d4aae73717ffcd4c77bea579590c20467002c58 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 28 Oct 2026 13:26:36 +0800 Subject: [PATCH 3/5] fix(promql): make the native prefix codec agree with itself MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A payload the encoder accepts has to be one this version can decode. Three places disagreed. An `Input`'s field count was bounded only on the way in, so a 129-field `Input → Instant` encoded to 983 bytes and then failed to decode. The decoder's limit was `MAX_COUNT`, a constant that also bounds nesting depth and `Map` operation count, which is why the missing half was easy to miss; each limit now has its own name and both sides check it. `FusedSeries::serialize` only checked that the tag count fit in a `u32` while `deserialize` bounded it at 1024. Both sides now share `MAX_TAG_COUNT`. `FusedSeries::deserialize` passed an empty schema to `try_new_with_format`, which checks a frame node's declared schema against the layout derived from its tree rather than overwriting it. Every serialized frame node was therefore impossible to rebuild. The schema is derived from the decoded tree now, which is where it comes from in the first place, so nothing is added to the wire. Encoding refuses an over-limit count instead of writing a payload it cannot read back, and a field-count overflow no longer reports "operation count exceeds maximum". `rejects_excess_operation_count` hand-built a payload whose operation count was read from the wrong offset: it placed an `Instant` where the decoder expects the count, so it exercised the limit with 258 while claiming 129. It now starts from a tree the codec accepts, asserts the bytes it edits are the count it means, and checks the error names both the operation limit and the value. The same is added for field count, tag count and nesting. The integration test moves to its own file: it was written against `tests-integration/tests/raw_series.rs`, which the raw-fan-in line creates and this one does not carry. Signed-off-by: discord9 --- src/promql/src/extension_plan/fused_series.rs | 29 +- .../src/extension_plan/fused_series/tests.rs | 51 +++- src/promql/src/native/wire.rs | 272 +++++++++++++++--- tests-integration/tests/main.rs | 1 + tests-integration/tests/native_prefix.rs | 265 +++++++++++++++++ 5 files changed, 580 insertions(+), 38 deletions(-) create mode 100644 tests-integration/tests/native_prefix.rs diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index f02c09e0f5d..2c8ba8f6f85 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -99,6 +99,10 @@ const METRIC_NAME: &str = "__name__"; /// again, so nothing is lost by stopping early. const POLL_BUDGET: usize = 64; +/// Maximum number of tags the envelope carries, shared by [`FusedSeries::serialize`] and +/// [`FusedSeries::deserialize`] so a serialized node is always one this version can rebuild. +const MAX_TAG_COUNT: usize = 1024; + /// The statement grid and output naming that travel with a compiled tree. /// /// Source schemas belong to the tree's input nodes. This binding supplies the grid that addresses @@ -276,6 +280,12 @@ impl FusedSeries { None => out.push(0), } text(&mut out, &self.binding.time_index)?; + if self.binding.tags.len() > MAX_TAG_COUNT { + return Err(DataFusionError::Plan(format!( + "too many FusedSeries tags: {} exceeds maximum {MAX_TAG_COUNT}", + self.binding.tags.len() + ))); + } out.extend( u32::try_from(self.binding.tags.len()) .map_err(|_| DataFusionError::Plan("too many FusedSeries tags".into()))? @@ -356,8 +366,10 @@ impl FusedSeries { }; let time_index = r.text()?; let count = r.u32()? as usize; - if count > 1024 { - return Err(DataFusionError::Plan("too many FusedSeries tags".into())); + if count > MAX_TAG_COUNT { + return Err(DataFusionError::Plan(format!( + "too many FusedSeries tags: {count} exceeds maximum {MAX_TAG_COUNT}" + ))); } let mut tags = Vec::with_capacity(count); for _ in 0..count { @@ -384,6 +396,17 @@ impl FusedSeries { }); collect_source_schemas(&tree)?.len() ]; + // A frame output declares the layout of the tree's evaluated root, which the tree itself + // carries. `try_new_with_format` compares the two, so passing an empty schema here would + // make every serialized frame node impossible to rebuild. + let output_schema = if format == FusedOutputFormat::SeriesFrame { + let evaluated = evaluated_output(&tree)?; + Arc::new(datafusion::common::DFSchema::try_from( + series_frame_output_schema(evaluated.as_ref())?, + )?) + } else { + Arc::new(datafusion::common::DFSchema::empty()) + }; Self::try_new_with_format( promql, tree, @@ -394,7 +417,7 @@ impl FusedSeries { tags, }, children, - Arc::new(datafusion::common::DFSchema::empty()), + output_schema, format, ) } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index eb1ed2ad769..9e0fa848d22 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -80,7 +80,7 @@ use store_api::region_engine::RawLabelSource; mod frame_output; use super::{ - FusedOutputFormat, FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, + FusedOutputFormat, FusedSeries, FusedSeriesBinding, MAX_TAG_COUNT, OutputColumn, POLL_BUDGET, series_to_record_batch, }; use crate::evaluated::EvaluatedChunk; @@ -928,12 +928,15 @@ impl Drop for StagedStream { fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() { let tree = rate_tree(); let original_binding = binding(&tree); + // A frame-output node is checked against the codec's own frame layout, so it is built with + // `canonical_frame_schema` rather than with the plain row `output_schema`. + let declared = canonical_frame_schema(&tree); let node = FusedSeries::try_new_with_format( Arc::from(PROMQL), tree, original_binding.clone(), vec![input_plan(frame_schema())], - output_schema().to_dfschema_ref().unwrap(), + declared, FusedOutputFormat::SeriesFrame, ) .unwrap(); @@ -957,6 +960,50 @@ fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() { assert!(FusedSeries::deserialize(&bytes[..5]).is_err()); } +/// The `SeriesFrame` node of [`rate_tree`] declaring `tags` as its labels. +/// +/// A `SeriesFrame` node is checked against the codec's own frame layout, so it is built with +/// [`canonical_frame_schema`] rather than with the plain row [`output_schema`]. +fn envelope_node_with_tags(tags: Vec) -> FusedSeries { + let tree = rate_tree(); + let binding = binding_of(&tree, raw_schema(), tags); + let output_schema = canonical_frame_schema(&tree); + FusedSeries::try_new_with_format( + Arc::from(PROMQL), + tree, + binding, + vec![input_plan(frame_schema())], + output_schema, + FusedOutputFormat::SeriesFrame, + ) + .expect("the frame layout matches the evaluated root") +} + +/// The largest tag count the envelope supports is written whole. +#[test] +fn fused_series_envelope_serializes_the_maximum_tag_count() { + let tags: Vec = (0..MAX_TAG_COUNT).map(|i| format!("t{i}")).collect(); + let node = envelope_node_with_tags(tags); + assert!( + node.serialize().is_ok(), + "the maximum tag count must still serialize" + ); +} + +/// The encoder must refuse the tag count, so it never writes a payload its own decoder rejects. +#[test] +fn fused_series_envelope_rejects_excess_tag_count() { + let tags: Vec = (0..=MAX_TAG_COUNT).map(|i| format!("t{i}")).collect(); + let node = envelope_node_with_tags(tags); + let message = node.serialize().unwrap_err().to_string(); + assert!(message.contains("too many FusedSeries tags"), "{message}"); + assert!(message.contains(&MAX_TAG_COUNT.to_string()), "{message}"); + assert!( + message.contains(&(MAX_TAG_COUNT + 1).to_string()), + "{message}" + ); +} + fn fused(tree: NativeTree) -> FusedSeries { let binding = binding(&tree); let output_schema = if matches!(&tree.root().expr, NativeExpr::Binary { .. }) { diff --git a/src/promql/src/native/wire.rs b/src/promql/src/native/wire.rs index 4b29406885b..d70fb2fcaa2 100644 --- a/src/promql/src/native/wire.rs +++ b/src/promql/src/native/wire.rs @@ -18,7 +18,7 @@ //! 4 Map; all integers are little-endian. Input is source index (u64), ordering tag (0 //! Interleaved, 1 Contiguous, 2 KeyThenTime), followed by a raw samples schema: field count (u32), each UTF-8 field name //! (u32 byte length + bytes). Instant stores lookback i64, offset i64, optional-at marker + i64, -//! timestamp bool byte, then output grid. Window stores function tag and its arguments, range i64, +//! value tag (0 sample, 1 timestamp, 2 info timestamp), then output grid. Window stores function tag and its arguments, range i64, //! offset i64, optional-anchor marker + i64, modifier tag (0 none, 1 anchored, 2 smoothed) and //! modifier lookback where present, input-grid marker + grid where present, output grid. Grid is //! start_ms i64, step_ms i64, len u64. Map stores op count u32 and math operations. Math tags and @@ -37,14 +37,16 @@ //! 26 ClampMin(min), 27 ClampMax(max), 28-35 Date(DayOfMonth, DayOfWeek, DayOfYear, //! DaysInMonth, Hour, Minute, Month, Year), in that order. -//! Maximum payload is 1 MiB and maximum recursive node depth / operation count is 128. +//! Maximum payload is 1 MiB, maximum recursive node depth is 128, maximum field count of one raw +//! samples schema is 128, and maximum operation count of one Map is 128. The encoder refuses every +//! payload that exceeds one of these, so a payload it produces is always decodable by this version. use crate::{ map::MathOp, map_time::DateFunction, native::plan::{ - BoundInput, BoundInstant, BoundMap, BoundWindow, MapOp, NativeExpr, NativeNode, NativeTree, - NativeType, + BoundInput, BoundInstant, BoundMap, BoundWindow, InstantValue, MapOp, NativeExpr, + NativeNode, NativeTree, NativeType, }, range_eval::{RangeFunction, RangeModifier, RangeWindow}, }; @@ -57,7 +59,12 @@ use std::{convert::TryInto, sync::Arc}; pub(crate) const PROTOCOL_VERSION: u32 = 1; const MAX_BYTES: usize = 1_048_576; -const MAX_COUNT: usize = 128; +/// Maximum recursive nesting depth of a node tree, counted from the root as depth 0. +const MAX_DEPTH: usize = 128; +/// Maximum number of fields in one raw samples schema (an [`NativeExpr::Input`]'s output). +const MAX_FIELD_COUNT: usize = 128; +/// Maximum number of operations in one [`NativeExpr::Map`]. +const MAX_OPERATION_COUNT: usize = 128; struct W(Vec); impl W { @@ -174,8 +181,8 @@ fn check_shape(n: &NativeNode) -> Result<(), String> { Ok(()) } fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { - if d > MAX_COUNT { - return Err("nesting maximum exceeded".into()); + if d > MAX_DEPTH { + return Err(format!("nesting maximum exceeded: {d} > {MAX_DEPTH}")); } match &n.expr { NativeExpr::Input(i) => { @@ -195,6 +202,12 @@ fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { if s.eval_grid.is_some() { return Err("unsupported evaluated Input".into()); } + if s.fields.len() > MAX_FIELD_COUNT { + return Err(format!( + "field count {} exceeds maximum {MAX_FIELD_COUNT}", + s.fields.len() + )); + } w.u32(s.fields.len() as u32); for f in s.fields.iter() { w.text(&f.name)? @@ -206,7 +219,7 @@ fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { w.i64(op.lookback_ms); w.i64(op.offset_ms); opt_i(w, op.at_ms); - w.u8(op.timestamp as u8); + w.u8(instant_value_tag(op.value)); w.grid(&op.output_grid) } NativeExpr::Window { input, op } => { @@ -239,8 +252,11 @@ fn enc_node(n: &NativeNode, w: &mut W, d: usize) -> Result<(), String> { NativeExpr::Map { input, op } => { w.u8(4); enc_node(input, w, d + 1)?; - if op.ops.is_empty() || op.ops.len() > MAX_COUNT { - return Err("operation count exceeds supported maximum or is zero".into()); + if op.ops.is_empty() || op.ops.len() > MAX_OPERATION_COUNT { + return Err(format!( + "operation count {} exceeds maximum {MAX_OPERATION_COUNT} or is zero", + op.ops.len() + )); } w.u32(op.ops.len() as u32); for x in op.ops.iter() { @@ -263,6 +279,14 @@ fn opt_i(w: &mut W, x: Option) { w.u8(0) } } +/// The tag of an [`InstantValue`], wire-stable and independent of the Rust discriminants. +fn instant_value_tag(value: InstantValue) -> u8 { + match value { + InstantValue::Sample => 0, + InstantValue::Timestamp => 1, + InstantValue::InfoTimestamp => 2, + } +} fn enc_math(m: MathOp, w: &mut W) -> Result<(), String> { let (t, vals): (u8, &[f64]) = match m { MathOp::Abs => (1, &[]), @@ -375,8 +399,8 @@ pub(crate) fn decode_prefix(bytes: &[u8]) -> Result { NativeTree::try_new(n).map_err(|e| e.to_string()) } fn dec_node(r: &mut R<'_>, d: usize) -> Result { - if d > MAX_COUNT { - return Err("nesting maximum exceeded".into()); + if d > MAX_DEPTH { + return Err(format!("nesting maximum exceeded: {d} > {MAX_DEPTH}")); } match r.u8()? { 1 => { @@ -388,8 +412,8 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { x => return Err(format!("unsupported Input ordering tag {x}")), }; let n = r.u32()? as usize; - if n > MAX_COUNT { - return Err("operation count exceeds maximum".into()); + if n > MAX_FIELD_COUNT { + return Err(format!("field count {n} exceeds maximum {MAX_FIELD_COUNT}")); } let mut fs = Vec::with_capacity(n); for _ in 0..n { @@ -413,10 +437,11 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { let lookback_ms = r.i64()?; let offset_ms = r.i64()?; let at_ms = read_opt(r)?; - let timestamp = match r.u8()? { - 0 => false, - 1 => true, - _ => return Err("invalid timestamp flag".into()), + let value = match r.u8()? { + 0 => InstantValue::Sample, + 1 => InstantValue::Timestamp, + 2 => InstantValue::InfoTimestamp, + _ => return Err("invalid instant value tag".into()), }; let output_grid = r.grid()?; let schema = match &input.output { @@ -434,7 +459,7 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { lookback_ms, offset_ms, at_ms, - timestamp, + value, output_grid, }, }, @@ -492,8 +517,10 @@ fn dec_node(r: &mut R<'_>, d: usize) -> Result { 4 => { let input = Box::new(dec_node(r, d + 1)?); let n = r.u32()? as usize; - if n == 0 || n > MAX_COUNT { - return Err("operation count exceeds maximum".into()); + if n == 0 || n > MAX_OPERATION_COUNT { + return Err(format!( + "operation count {n} exceeds maximum {MAX_OPERATION_COUNT} or is zero" + )); } let mut ops = Vec::new(); for _ in 0..n { @@ -667,7 +694,7 @@ mod tests { lookback_ms: 300_000, offset_ms: 7, at_ms: Some(88), - timestamp: true, + value: InstantValue::Timestamp, output_grid: g, }, }, @@ -741,7 +768,7 @@ mod tests { assert_eq!(op.lookback_ms, 300_000); assert_eq!(op.offset_ms, 7); assert_eq!(op.at_ms, Some(88)); - assert!(op.timestamp); + assert_eq!(op.value, InstantValue::Timestamp); assert_eq!(op.output_grid.spec(), grid().spec()); } #[test] @@ -910,23 +937,202 @@ mod tests { let b = vec![0; MAX_BYTES + 1]; assert!(decode_prefix(&b).unwrap_err().contains("MAX_BYTES")); } - #[test] - fn rejects_excess_operation_count() { + /// A tree with `n` fields: [`instant`]'s shape over a raw input of `n` float fields. + /// + /// The input node and its binding share one output `Arc`, which is what `NativeTree` + /// validation compares by identity. + fn instant_with_fields(n: usize) -> NativeTree { + let g = grid(); + let fields: Arc<[FieldSpec]> = (0..n) + .map(|i| FieldSpec { + name: Arc::from(format!("f{i}").as_str()), + }) + .collect::>() + .into(); + let raw = NativeType::Samples(Arc::new(SeriesSchema { + fields: fields.clone(), + eval_grid: None, + })); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields, + eval_grid: Some(g.clone()), + })); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: 0, + ordering: SeriesOrdering::KeyThenTime, + output: raw.clone(), + }), + output: raw, + }), + op: BoundInstant { + lookback_ms: 300_000, + offset_ms: 0, + at_ms: None, + value: InstantValue::Sample, + output_grid: g, + }, + }, + output, + }) + .unwrap() + } + + /// A tree with `n` math operations: one [`NativeExpr::Map`] of `Abs` over [`instant`]. + /// + /// `Abs` has a one-byte wire tag, so the payload ends with exactly `n` tag bytes. + fn mapped_with_ops(n: usize) -> NativeTree { + let instant = instant(); + let input = Box::new(instant.root().clone()); + let output = input.output.clone(); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![MapOp::Math(MathOp::Abs); n].into(), + }, + }, + output, + }) + .unwrap() + } + + /// A payload of `depth` nested `Instant` nodes over one raw input. + /// + /// No supported shape accepts the result, but the decoder walks the nodes before it checks the + /// shape, so this is how the recursion limit is reached. The encoder's own depth guard is + /// unreachable through [`encode_prefix`], which rejects a deeper shape before encoding it. + fn nested_instants(depth: usize) -> Vec { let mut w = W(Vec::new()); w.u32(PROTOCOL_VERSION); - w.u8(4); + for _ in 0..depth { + w.u8(2); + } w.u8(1); w.u64(0); w.u8(2); w.u32(1); w.text("v").unwrap(); + for _ in 0..depth { + w.i64(300_000); + w.i64(0); + opt_i(&mut w, None); + w.u8(0); + w.grid(&grid()); + } + w.0 + } + + /// The largest field count this version supports is transported whole. + #[test] + fn field_count_at_the_maximum_round_trips() { + let a = instant_with_fields(MAX_FIELD_COUNT); + let b = encode_prefix(&a).unwrap(); + let d = decode_prefix(&b).unwrap(); + assert_eq!(encode_prefix(&d).unwrap(), b); + let NativeType::Samples(s) = &d.root().output else { + panic!() + }; + assert_eq!(s.fields.len(), MAX_FIELD_COUNT); + assert_eq!(s.fields[MAX_FIELD_COUNT - 1].name.as_ref(), "f127"); + } + + /// The encoder must refuse the field count, not write a payload its own decoder rejects. + #[test] + fn rejects_excess_field_count_on_encode() { + let err = encode_prefix(&instant_with_fields(MAX_FIELD_COUNT + 1)).unwrap_err(); + assert!(err.contains("field count"), "{err}"); + assert!(err.contains(&MAX_FIELD_COUNT.to_string()), "{err}"); + assert!( + !err.contains("operation") && !err.contains("nesting"), + "a field-count limit must name its own limit: {err}" + ); + } + + #[test] + fn rejects_excess_field_count_on_decode() { + // A raw input at the root is not a supported shape, but the count is checked while the node + // is read, before any shape check, so the field names need not follow. + let mut w = W(Vec::new()); + w.u32(PROTOCOL_VERSION); + w.u8(1); + w.u64(0); w.u8(2); - w.i64(1); - w.i64(0); - opt_i(&mut w, None); - w.u8(0); - w.grid(&grid()); - w.u32((MAX_COUNT + 1) as u32); - assert!(decode_prefix(&w.0).unwrap_err().contains("operation count")); + w.u32((MAX_FIELD_COUNT + 1) as u32); + let err = decode_prefix(&w.0).unwrap_err(); + assert!(err.contains("field count"), "{err}"); + assert!(err.contains(&MAX_FIELD_COUNT.to_string()), "{err}"); + assert!(!err.contains("operation"), "{err}"); + } + + /// The largest operation count this version supports is transported whole. + #[test] + fn operation_count_at_the_maximum_round_trips() { + let a = mapped_with_ops(MAX_OPERATION_COUNT); + let b = encode_prefix(&a).unwrap(); + let d = decode_prefix(&b).unwrap(); + assert_eq!(encode_prefix(&d).unwrap(), b); + let NativeExpr::Map { op, .. } = &d.root().expr else { + panic!() + }; + assert_eq!(op.ops.len(), MAX_OPERATION_COUNT); + } + + #[test] + fn rejects_excess_operation_count() { + // Start from a tree the codec accepts, so the bytes this test edits are the ones the + // encoder actually writes for `MAX_OPERATION_COUNT` operations. + let tree = mapped_with_ops(MAX_OPERATION_COUNT); + let mut bytes = encode_prefix(&tree).unwrap(); + assert!(decode_prefix(&bytes).is_ok()); + + // `Map` writes its input first, then the operation count, then the operations, and nothing + // follows the root node. Every operation here is `Abs` (wire tag `1`), so the last + // `MAX_OPERATION_COUNT` bytes are those tags and the four bytes before them hold the count. + let at = bytes.len() - MAX_OPERATION_COUNT - 4; + assert!( + bytes[at + 4..].iter().all(|b| *b == 1), + "the payload does not end with one-byte Abs tags" + ); + assert_eq!( + u32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()), + MAX_OPERATION_COUNT as u32, + "this is not the operation count field" + ); + bytes[at..at + 4].copy_from_slice(&(MAX_OPERATION_COUNT as u32 + 1).to_le_bytes()); + + // The operations themselves are untouched, so the only reason to refuse this payload is the + // count it now declares. + let err = decode_prefix(&bytes).unwrap_err(); + assert!(err.contains("operation count"), "{err}"); + assert!(err.contains(&MAX_OPERATION_COUNT.to_string()), "{err}"); + assert!(!err.contains("field count"), "{err}"); + assert!(!err.contains("nesting"), "{err}"); + assert!( + err.contains(&(MAX_OPERATION_COUNT + 1).to_string()), + "{err}" + ); + } + + #[test] + fn rejects_excess_operation_count_on_encode() { + let err = encode_prefix(&mapped_with_ops(MAX_OPERATION_COUNT + 1)).unwrap_err(); + assert!(err.contains("operation count"), "{err}"); + assert!(err.contains(&MAX_OPERATION_COUNT.to_string()), "{err}"); + assert!(!err.contains("field count"), "{err}"); + } + + /// Depth 128 is still walked; depth 129 is refused, and by the depth limit rather than the + /// shape check that also rejects every payload this deep. + #[test] + fn nesting_boundary_is_128() { + let allowed = decode_prefix(&nested_instants(MAX_DEPTH)).unwrap_err(); + assert!(!allowed.contains("nesting"), "{allowed}"); + let refused = decode_prefix(&nested_instants(MAX_DEPTH + 1)).unwrap_err(); + assert!(refused.contains("nesting"), "{refused}"); + assert!(refused.contains(&MAX_DEPTH.to_string()), "{refused}"); + assert!(refused.contains(&(MAX_DEPTH + 1).to_string()), "{refused}"); } } diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 72d0a1850bd..75a475c64e7 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -39,6 +39,7 @@ mod repartition_event; mod repartition_expr_version; mod export_logical_tables; mod mysql; +mod native_prefix; mod object_store_wal; mod reconciliation_event; mod view_ddl_event; diff --git a/tests-integration/tests/native_prefix.rs b/tests-integration/tests/native_prefix.rs new file mode 100644 index 00000000000..f35cff67dd2 --- /dev/null +++ b/tests-integration/tests/native_prefix.rs @@ -0,0 +1,265 @@ +// 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. + +//! A serialized native PromQL prefix executed on a datanode over its own region scan. + +use std::sync::Arc; + +use api::v1::region::{QueryRequest, RegionRequestHeader}; +use common_recordbatch::RecordBatches; +use common_recordbatch::adapter::DfRecordBatchStreamAdapter; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datatypes::series::SeriesOrdering; +use datatypes::series::samples::SampleRef; +use datatypes::series::samples::{FieldSpec, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpec}; +use futures::{StreamExt, TryStreamExt}; +use promql::extension_plan::{FusedOutputFormat, FusedSeries, FusedSeriesBinding}; +use promql::native::plan::{ + BoundInput, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, +}; +use promql::range_eval::{RangeFunction, RangeWindow}; +use query::parser::QueryLanguageParser; +use query::query_engine::DefaultSerializer; +use session::context::QueryContext; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::execute_sql; + +/// The logical frame schema a `FusedOutputFormat::SeriesFrame` node declares: the shared evaluated +/// packer's layout for the tree's evaluated root, which `try_new_with_format` compares against. +fn evaluated_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef { + let NativeType::Samples(samples) = &tree.root().output else { + panic!("the prefix root evaluates to samples"); + }; + let stream = common_query::series_frame::pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(samples), + ) + .expect("an evaluated root declares a frame"); + Arc::new( + datafusion::common::DFSchema::try_from(stream.schema()) + .expect("the frame schema converts to a logical schema"), + ) +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_prefix_executes_over_real_mito_region() { + let instance = GreptimeDbStandaloneBuilder::new("native_prefix_region_scan") + .with_auto_create_table(false) + .build() + .await; + let fe = instance.fe_instance(); + execute_sql(fe, "CREATE TABLE native_prefix (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING, PRIMARY KEY(host))").await; + execute_sql(fe, "INSERT INTO native_prefix(host, val, ts) VALUES ('a', 0.0, 1000), ('a', 10.0, 2000), ('a', 20.0, 3000), ('a', 30.0, 4000)").await; + execute_sql(fe, "ADMIN FLUSH_TABLE('native_prefix')").await; + + let ctx = QueryContext::arc(); + let statement = QueryLanguageParser::parse_sql("SELECT * FROM native_prefix", &ctx).unwrap(); + let planned = fe + .query_engine() + .planner() + .plan(&statement, ctx.clone()) + .await + .unwrap(); + let datafusion::logical_expr::LogicalPlan::Projection(projection) = planned else { + panic!("expected projection: {planned:?}"); + }; + let datafusion::logical_expr::LogicalPlan::TableScan(scan) = projection.input.as_ref() else { + panic!("expected table scan: {projection:?}"); + }; + let scan = scan.clone(); + let child = LogicalPlan::TableScan(datafusion::logical_expr::logical_plan::TableScan { + projection: Some((0..scan.projected_schema.fields().len()).collect()), + ..scan + }); + let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(3000, 1000, 2).unwrap()).unwrap()); + let raw = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("val"), + }]), + eval_grid: None, + }); + let evaluated = Arc::new(SeriesSchema { + fields: Arc::clone(&raw.fields), + eval_grid: Some(Arc::clone(&grid)), + }); + // The input node's output must be the same binding as its `BoundInput.output`: the tree + // compares those two by allocation identity, so an equal-but-separate schema is refused. + let raw_type = NativeType::Samples(raw); + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: 0, + output: raw_type.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw_type, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 3000, + offset_ms: 0, + anchor_ms: None, + modifier: None, + }, + input_grid: None, + output_grid: Arc::clone(&grid), + }, + }, + output: NativeType::Samples(evaluated), + }) + .unwrap(); + let binding = FusedSeriesBinding { + grid, + metric_name: Some(Arc::from("native_prefix")), + time_index: "ts".into(), + tags: vec!["host".into()], + }; + // A frame-output node declares the frame layout of its evaluated root, not the projection of + // the scan below it: `try_new_with_format` compares the two and refuses a mismatch. + let frame_schema = evaluated_frame_schema(&tree); + let fused = FusedSeries::try_new_with_format( + Arc::from("rate(native_prefix[3s])"), + tree, + binding, + vec![child], + frame_schema, + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let fused_payload = fused.serialize().unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(fused), + }); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap() + .to_vec(); + let region = instance + .mito_engine + .regions() + .into_iter() + .find(|r| { + r.metadata() + .column_metadatas + .iter() + .any(|c| c.column_schema.name == "host") + }) + .unwrap(); + let metadata = region.metadata(); + let server = instance.datanode.region_server(); + let mut stream = server + .handle_remote_read( + QueryRequest { + header: Some(RegionRequestHeader { + query_context: Some((&*ctx).into()), + ..Default::default() + }), + region_id: metadata.region_id.as_u64(), + plan: bytes.clone(), + }, + ctx.clone(), + ) + .await + .unwrap(); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + let schema = batches.first().expect("evaluated frames").schema.clone(); + let batches = RecordBatches::try_new(schema, batches).unwrap(); + let series = common_query::series_frame::unpack_evaluated_series( + Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())), + None, + ) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(series.len(), 1); + let (series, _) = &series[0]; + // `rate` drops `__name__`, as Prometheus does for every range function except + // `last_over_time` and `first_over_time`; the series' own tags survive. + assert_eq!(series.metric.name(), None); + assert_eq!( + series + .metric + .get("host") + .map(ToString::to_string) + .as_deref(), + Some("a") + ); + assert_eq!(series.time.timestamps().unwrap().as_ref(), &[3000, 4000]); + // At 3000 ms the (0, 3000] window has samples 0, 10, 20: increase 20 over + // 2 observed seconds. Prometheus extrapolates the 1s left-boundary gap and + // 0s right-boundary gap to the full 3s range: 20 * (3 / 2) / 3 = 10/s. + // At 4000 ms, (1000, 4000] has 10, 20, 30: increase 20 over 2 seconds, + // likewise extrapolated to 10/s. + let values = series.fields.get(0).unwrap(); + // At 3000 ms the window holds 0, 10, 20: an increase of 20 over 2 observed + // seconds. Prometheus would keep the 1 s leading gap (below the 1.1x average + // threshold), but a counter whose first sample is 0 cannot be extrapolated + // backwards at all: `duration_to_zero = sampled * first/value = 0` clamps it, + // so the rate is 20 * (2 + 0) / 2 / 3 = 20/3. + // + // At 4000 ms the window holds 10, 20, 30: also 20 over 2 s, but now + // `duration_to_zero = 2 * (10/20) = 1 s` does not shorten the 1 s leading gap, + // so the rate is 20 * (2 + 1 + 0) / 2 / 3 = 10. + let expected = [20.0 / 3.0, 10.0]; + for (index, expected) in expected.into_iter().enumerate() { + let SampleRef::Float(value) = values.sample(index) else { + panic!("the rate field is float at position {index}"); + }; + assert!( + (value - expected).abs() < 1e-12, + "value at position {index} is {value}, not {expected}" + ); + } + let mut tampered = bytes; + // The version is the first four bytes of the embedded FusedSeries payload; flipping the + // Substrait envelope is intentionally avoided, so locate the native protocol version. + let native_prefix_version = 1u32.to_le_bytes(); + let embedded = tampered + .windows(fused_payload.len()) + .position(|window| window == fused_payload.as_slice()) + .expect("FusedSeries payload in Substrait plan"); + let prefix_version = embedded + 8; + assert_eq!( + &tampered[prefix_version..prefix_version + 4], + &native_prefix_version + ); + tampered[prefix_version] = 2; + let error = server + .handle_remote_read( + QueryRequest { + header: Some(RegionRequestHeader { + query_context: Some((&*ctx).into()), + ..Default::default() + }), + region_id: metadata.region_id.as_u64(), + plan: tampered, + }, + ctx, + ) + .await + .err() + .expect("invalid prefix rejected"); + assert!( + format!("{error:?}").contains("unknown version"), + "{error:?}" + ); +} From dfd052605d760707cc2f869bdd05a3bb0eaddeb4 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 29 Oct 2026 18:52:18 +0800 Subject: [PATCH 4/5] feat(promql): carry nested math maps in the native prefix The wire format accepted at most one Math map above the selection, so abs(ceil(x)) - which compiles to one map over another - could not be transported even though a map is transparent to shape: it keeps its input samples type and grid, and only the operation kinds are restricted. Walk every map before the selection check instead, and keep rejecting non-math operations at any depth. The tests cover a nested round trip that would notice flattening, a non-math operation under a valid outer map, three distinct levels, and the depth boundary in both directions. Signed-off-by: discord9 --- src/promql/src/native/wire.rs | 224 +++++++++++++++++++++++++++++++--- 1 file changed, 205 insertions(+), 19 deletions(-) diff --git a/src/promql/src/native/wire.rs b/src/promql/src/native/wire.rs index d70fb2fcaa2..2b79f1b35b7 100644 --- a/src/promql/src/native/wire.rs +++ b/src/promql/src/native/wire.rs @@ -57,6 +57,8 @@ use datatypes::series::{ }; use std::{convert::TryInto, sync::Arc}; +// This experimental branch has no deployed readers or compatibility obligation; the unreleased +// format can widen to nested math maps without changing its version. pub(crate) const PROTOCOL_VERSION: u32 = 1; const MAX_BYTES: usize = 1_048_576; /// Maximum recursive nesting depth of a node tree, counted from the root as depth 0. @@ -150,23 +152,23 @@ pub(crate) fn encode_prefix(tree: &NativeTree) -> Result, String> { } Ok(w.0) } -/// Accepts exactly the supported shapes: a selection over raw input, optionally under one map. +/// Accepts exactly the supported shapes: a selection over raw input, under any number of math maps. /// -/// The map is optional because a bare selection is a prefix too — it is what a query like `x` -/// transports. When present there is one, and it holds only math operations: a label rewrite or a -/// histogram function would change the identity or the value type in ways this format does not -/// describe. +/// The maps are optional because a bare selection is a prefix too — it is what a query like `x` +/// transports. When present they hold only math operations: a label rewrite or a histogram +/// function would change the identity or the value type in ways this format does not describe. +/// +/// A map is transparent to the shape. It keeps its input's samples type and grid, so `abs(ceil(x))` +/// — which the compiler builds as one map over another — is the same shape as `abs(x)`, and only +/// the operation kinds are restricted. [`MAX_DEPTH`] still bounds how deep the payload can go. fn check_shape(n: &NativeNode) -> Result<(), String> { - let selection = match &n.expr { - NativeExpr::Map { input, op } => { - if op.ops.is_empty() || op.ops.iter().any(|x| !matches!(x, MapOp::Math(_))) { - return Err("unsupported Map operation (only Math)".into()); - } - input + let mut selection = n; + while let NativeExpr::Map { input, op } = &selection.expr { + if op.ops.is_empty() || op.ops.iter().any(|x| !matches!(x, MapOp::Math(_))) { + return Err("unsupported Map operation (only Math)".into()); } - NativeExpr::Instant { .. } | NativeExpr::Window { .. } => n, - other => return Err(format!("unsupported operation: {}", node_name(other))), - }; + selection = input; + } let (NativeExpr::Instant { input: child, .. } | NativeExpr::Window { input: child, .. }) = &selection.expr else { @@ -629,6 +631,7 @@ fn dec_math(r: &mut R<'_>) -> Result { #[cfg(test)] mod tests { use super::*; + use crate::histogram_eval::NativeHistogramFunction; fn input() -> NativeNode { let t = NativeType::Samples(Arc::new(SeriesSchema { fields: vec![FieldSpec { name: "v".into() }].into(), @@ -800,6 +803,191 @@ mod tests { matches!(op.ops[0],MapOp::Math(MathOp::Round{nearest}) if nearest.to_bits()==0.5f64.to_bits()) ); } + /// `abs(ceil(x))` compiles to one map over another. A map keeps its input's type and grid, so + /// that is the same shape as one map and has to survive the round trip. + #[test] + fn nested_maps_roundtrip() { + let inner = mapped(); + let input = Box::new(inner.root().clone()); + let output = input.output.clone(); + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![MapOp::Math(MathOp::Abs)].into(), + }, + }, + output, + }) + .unwrap(); + let bytes = encode_prefix(&tree).unwrap(); + let decoded = decode_prefix(&bytes).unwrap(); + assert_eq!(encode_prefix(&decoded).unwrap(), bytes); + let NativeExpr::Map { input, op } = &decoded.root().expr else { + panic!("the outer map is not a map") + }; + assert!(matches!(op.ops[0], MapOp::Math(MathOp::Abs))); + let NativeExpr::Map { op, .. } = &input.expr else { + panic!("the inner map was flattened away: {:?}", input.expr) + }; + assert_eq!(op.ops.len(), 3); + let NativeType::Samples(decoded_schema) = &decoded.root().output else { + panic!("decoded output is not samples") + }; + let NativeType::Samples(original_schema) = &tree.root().output else { + panic!("original output is not samples") + }; + assert_eq!( + decoded_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>(), + original_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .collect::>() + ); + assert_eq!( + decoded_schema.eval_grid.as_ref().map(|grid| grid.spec()), + original_schema.eval_grid.as_ref().map(|grid| grid.spec()) + ); + } + /// The shape validator must inspect inner maps, not only the root map. + #[test] + fn nested_maps_still_refuse_non_math_operations() { + let inner_input = Box::new(instant().root().clone()); + let inner = NativeNode { + output: inner_input.output.clone(), + expr: NativeExpr::Map { + input: inner_input, + op: BoundMap { + ops: vec![MapOp::Histogram(NativeHistogramFunction::Count)].into(), + }, + }, + }; + let input = Box::new(inner); + let tree = NativeTree::try_new(NativeNode { + output: input.output.clone(), + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![MapOp::Math(MathOp::Abs)].into(), + }, + }, + }) + .unwrap(); + assert_eq!( + encode_prefix(&tree).unwrap_err(), + "unsupported Map operation (only Math)" + ); + } + #[test] + fn three_nested_maps_roundtrip() { + let mut tree = mapped(); + for op in [MathOp::Ceil, MathOp::Abs] { + let input = Box::new(tree.root().clone()); + tree = NativeTree::try_new(NativeNode { + output: input.output.clone(), + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![MapOp::Math(op)].into(), + }, + }, + }) + .unwrap(); + } + let bytes = encode_prefix(&tree).unwrap(); + let decoded = decode_prefix(&bytes).unwrap(); + assert_eq!(encode_prefix(&decoded).unwrap(), bytes); + let NativeExpr::Map { input: middle, .. } = &decoded.root().expr else { + panic!("outer map missing") + }; + let NativeExpr::Map { input: inner, .. } = &middle.expr else { + panic!("middle map missing") + }; + assert!( + matches!(inner.expr, NativeExpr::Map { .. }), + "inner map missing" + ); + } + #[test] + fn nested_inner_non_math_map_is_rejected() { + let inner_input = Box::new(mapped().root().clone()); + let inner = NativeNode { + output: inner_input.output.clone(), + expr: NativeExpr::Map { + input: inner_input, + op: BoundMap { + ops: vec![MapOp::Histogram(NativeHistogramFunction::Count)].into(), + }, + }, + }; + let outer_input = Box::new(inner); + let tree = NativeTree::try_new(NativeNode { + output: outer_input.output.clone(), + expr: NativeExpr::Map { + input: outer_input, + op: BoundMap { + ops: vec![MapOp::Math(MathOp::Abs)].into(), + }, + }, + }) + .unwrap(); + assert_eq!( + encode_prefix(&tree).unwrap_err(), + "unsupported Map operation (only Math)" + ); + } + fn make_map_chain(map_count: usize) -> NativeTree { + let mut node = instant().root().clone(); + for _ in 0..map_count { + let input = Box::new(node); + node = NativeNode { + output: input.output.clone(), + expr: NativeExpr::Map { + input, + op: BoundMap { + ops: vec![MapOp::Math(MathOp::Abs)].into(), + }, + }, + }; + } + NativeTree::try_new(node).unwrap() + } + #[test] + fn map_chain_depth_boundary_is_enforced_on_encode_and_decode() { + // 127 maps plus the selection and input fit the protocol's depth boundary. + let allowed = encode_prefix(&make_map_chain(MAX_DEPTH - 1)).unwrap(); + let decoded = decode_prefix(&allowed).unwrap(); + assert_eq!(encode_prefix(&decoded).unwrap(), allowed); + + let too_deep = make_map_chain(MAX_DEPTH); + let error = encode_prefix(&too_deep).unwrap_err(); + assert!(error.starts_with("nesting maximum exceeded:"), "{error}"); + assert!(error.contains(&MAX_DEPTH.to_string()), "{error}"); + + // Decode the same over-depth shape from bytes and verify it fails at the depth guard. + let mut bytes = W(Vec::new()); + bytes.u32(PROTOCOL_VERSION); + for _ in 0..MAX_DEPTH { + bytes.u8(4); + } + bytes.u8(2); + bytes.u8(1); + bytes.u64(0); + bytes.u8(2); + bytes.u32(1); + bytes.text("v").unwrap(); + for _ in 0..MAX_DEPTH { + bytes.u32(1); + bytes.u8(1); + } + let error = decode_prefix(&bytes.0).unwrap_err(); + assert!(error.starts_with("nesting maximum exceeded:"), "{error}"); + } #[test] fn floats_preserve_bits() { let values = [ @@ -1001,9 +1189,8 @@ mod tests { /// A payload of `depth` nested `Instant` nodes over one raw input. /// - /// No supported shape accepts the result, but the decoder walks the nodes before it checks the - /// shape, so this is how the recursion limit is reached. The encoder's own depth guard is - /// unreachable through [`encode_prefix`], which rejects a deeper shape before encoding it. + /// The decoder walks nodes before shape validation, so unsupported deep trees exercise its + /// recursion guard. Supported map chains exercise the encoder's guard separately. fn nested_instants(depth: usize) -> Vec { let mut w = W(Vec::new()); w.u32(PROTOCOL_VERSION); @@ -1124,8 +1311,7 @@ mod tests { assert!(!err.contains("field count"), "{err}"); } - /// Depth 128 is still walked; depth 129 is refused, and by the depth limit rather than the - /// shape check that also rejects every payload this deep. + /// The decoder walks depth 128, then refuses depth 129 at its recursion limit. #[test] fn nesting_boundary_is_128() { let allowed = decode_prefix(&nested_instants(MAX_DEPTH)).unwrap_err(); From 6f99b36e444823b0c3a6baeacebd24497316c207 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 29 Oct 2026 20:17:00 +0800 Subject: [PATCH 5/5] fix(promql): bound the prefix depth and name what a refusal names Lower MAX_DEPTH to 32. The old 128 was validated recursively, and a debug build overflows a default test-thread stack well before reaching it, so the depth guard itself could abort instead of refusing. Report the refused node by its own name instead of "unsupported", which is what every kind used to say. Add the multi-partition datanode case. It asserts the failure the review reproduced - a series chain cut by CoalescePartitionsExec no longer decodes - so the limitation stays visible and turns red when the region merge becomes chain-atomic. Reconcile the retained-rate comment with the values the test asserts. Signed-off-by: discord9 --- src/promql/src/native/wire.rs | 41 ++++-- tests-integration/tests/native_prefix.rs | 172 ++++++++++++++++++++++- 2 files changed, 199 insertions(+), 14 deletions(-) diff --git a/src/promql/src/native/wire.rs b/src/promql/src/native/wire.rs index 2b79f1b35b7..e3773dbeb3a 100644 --- a/src/promql/src/native/wire.rs +++ b/src/promql/src/native/wire.rs @@ -37,7 +37,7 @@ //! 26 ClampMin(min), 27 ClampMax(max), 28-35 Date(DayOfMonth, DayOfWeek, DayOfYear, //! DaysInMonth, Hour, Minute, Month, Year), in that order. -//! Maximum payload is 1 MiB, maximum recursive node depth is 128, maximum field count of one raw +//! Maximum payload is 1 MiB, maximum recursive node depth is 32, maximum field count of one raw //! samples schema is 128, and maximum operation count of one Map is 128. The encoder refuses every //! payload that exceeds one of these, so a payload it produces is always decodable by this version. @@ -62,7 +62,8 @@ use std::{convert::TryInto, sync::Arc}; pub(crate) const PROTOCOL_VERSION: u32 = 1; const MAX_BYTES: usize = 1_048_576; /// Maximum recursive nesting depth of a node tree, counted from the root as depth 0. -const MAX_DEPTH: usize = 128; +/// Kept low to leave ample stack headroom in debug builds on default-sized test threads. +const MAX_DEPTH: usize = 32; /// Maximum number of fields in one raw samples schema (an [`NativeExpr::Input`]'s output). const MAX_FIELD_COUNT: usize = 128; /// Maximum number of operations in one [`NativeExpr::Map`]. @@ -375,12 +376,21 @@ fn enc_range(f: RangeFunction, w: &mut W) -> Result<(), String> { } fn node_name(e: &NativeExpr) -> &'static str { match e { - NativeExpr::Binary { .. } => "Binary", + NativeExpr::Input(_) => "Input", + NativeExpr::Scalar(_) => "Scalar", + NativeExpr::Instant { .. } => "Instant", + NativeExpr::Window { .. } => "Window", + NativeExpr::Merge { .. } => "Merge", + NativeExpr::Map { .. } => "Map", + NativeExpr::Timestamp { .. } => "Timestamp", NativeExpr::Aggregate { .. } => "Aggregate", NativeExpr::Vector { .. } => "Vector", - NativeExpr::Scalar(_) => "Scalar", - NativeExpr::Timestamp { .. } => "Timestamp", - _ => "unsupported", + NativeExpr::Binary { .. } => "Binary", + NativeExpr::EvalTime { .. } => "EvalTime", + NativeExpr::VectorOf { .. } => "VectorOf", + NativeExpr::ParameterizedMap { .. } => "ParameterizedMap", + NativeExpr::Set { .. } => "Set", + NativeExpr::Info { .. } => "Info", } } /// Decode and validate a supported prefix tree. @@ -959,7 +969,7 @@ mod tests { } #[test] fn map_chain_depth_boundary_is_enforced_on_encode_and_decode() { - // 127 maps plus the selection and input fit the protocol's depth boundary. + // Maps up to the configured node-depth boundary fit alongside the selection and input. let allowed = encode_prefix(&make_map_chain(MAX_DEPTH - 1)).unwrap(); let decoded = decode_prefix(&allowed).unwrap(); assert_eq!(encode_prefix(&decoded).unwrap(), allowed); @@ -1311,9 +1321,22 @@ mod tests { assert!(!err.contains("field count"), "{err}"); } - /// The decoder walks depth 128, then refuses depth 129 at its recursion limit. #[test] - fn nesting_boundary_is_128() { + fn unsupported_operation_names_node_kind() { + let error = encode_prefix( + &NativeTree::try_new(NativeNode { + expr: NativeExpr::Scalar(crate::native::plan::BoundScalar { value: 1.0 }), + output: NativeType::Scalar(grid()), + }) + .unwrap(), + ) + .unwrap_err(); + assert_eq!(error, "unsupported operation: Scalar"); + } + + /// The decoder walks to the configured depth, then refuses one level beyond it. + #[test] + fn nesting_boundary_is_enforced() { let allowed = decode_prefix(&nested_instants(MAX_DEPTH)).unwrap_err(); assert!(!allowed.contains("nesting"), "{allowed}"); let refused = decode_prefix(&nested_instants(MAX_DEPTH + 1)).unwrap_err(); diff --git a/tests-integration/tests/native_prefix.rs b/tests-integration/tests/native_prefix.rs index f35cff67dd2..9b7685e793c 100644 --- a/tests-integration/tests/native_prefix.rs +++ b/tests-integration/tests/native_prefix.rs @@ -30,6 +30,7 @@ use promql::native::plan::{ BoundInput, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, }; use promql::range_eval::{RangeFunction, RangeWindow}; +use query::options::QueryOptions; use query::parser::QueryLanguageParser; use query::query_engine::DefaultSerializer; use session::context::QueryContext; @@ -54,6 +55,172 @@ fn evaluated_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef ) } +// Reproduction from the review: 16 hosts × 20k samples, `rate[10s]`, and 16 DN +// partitions produces four batches per chain and fails while unpacking because +// `CoalescePartitionsExec` interleaves their batches, so a chain is cut and the frames no +// longer decode. This test asserts that failure, which is why it runs rather than being +// ignored: it keeps the boundary honest until the frontend region merge becomes +// chain-atomic, and it is the test to change when that fix lands. The parallelism-one +// path is exercised by `native_prefix_executes_over_real_mito_region`. +#[tokio::test(flavor = "multi_thread")] +async fn native_prefix_multi_partition_series_chains() { + let instance = GreptimeDbStandaloneBuilder::new("native_prefix_partitioned_chains") + .with_query_options(QueryOptions { + parallelism: 16, + ..Default::default() + }) + .with_auto_create_table(false) + .build() + .await; + let fe = instance.fe_instance(); + execute_sql(fe, "CREATE TABLE native_prefix_partitioned (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING, PRIMARY KEY(host))").await; + let mut values = String::new(); + for host in 0..16 { + for sample in 0..20_000 { + let timestamp = sample as u64 * 1_000; + let value = sample as f64; + values.push_str(&format!("('host{host}', {value}, {timestamp}),")); + } + } + values.pop(); + execute_sql( + fe, + &format!("INSERT INTO native_prefix_partitioned(host, val, ts) VALUES {values}"), + ) + .await; + execute_sql(fe, "ADMIN FLUSH_TABLE('native_prefix_partitioned')").await; + + let ctx = QueryContext::arc(); + let statement = + QueryLanguageParser::parse_sql("SELECT * FROM native_prefix_partitioned", &ctx).unwrap(); + let planned = fe + .query_engine() + .planner() + .plan(&statement, ctx.clone()) + .await + .unwrap(); + let LogicalPlan::Projection(projection) = planned else { + panic!("expected projection"); + }; + let LogicalPlan::TableScan(scan) = projection.input.as_ref() else { + panic!("expected table scan"); + }; + let scan = scan.clone(); + let child = LogicalPlan::TableScan(datafusion::logical_expr::logical_plan::TableScan { + projection: Some((0..scan.projected_schema.fields().len()).collect()), + ..scan + }); + let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(20_000, 1_000, 1).unwrap()).unwrap()); + let raw = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("val"), + }]), + eval_grid: None, + }); + let evaluated = Arc::new(SeriesSchema { + fields: Arc::clone(&raw.fields), + eval_grid: Some(Arc::clone(&grid)), + }); + let raw_type = NativeType::Samples(raw); + let tree = NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + source_index: 0, + output: raw_type.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw_type, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 10_000, + offset_ms: 0, + anchor_ms: None, + modifier: None, + }, + input_grid: None, + output_grid: grid.clone(), + }, + }, + output: NativeType::Samples(evaluated), + }) + .unwrap(); + let binding = FusedSeriesBinding { + grid, + metric_name: Some(Arc::from("native_prefix_partitioned")), + time_index: "ts".into(), + tags: vec!["host".into()], + }; + let frame_schema = evaluated_frame_schema(&tree); + let fused = FusedSeries::try_new_with_format( + Arc::from("rate(native_prefix_partitioned[10s])"), + tree, + binding, + vec![child], + frame_schema, + FusedOutputFormat::SeriesFrame, + ) + .unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(fused), + }); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap() + .to_vec(); + let region = instance + .mito_engine + .regions() + .into_iter() + .find(|region| { + region + .metadata() + .column_metadatas + .iter() + .any(|column| column.column_schema.name == "host") + }) + .unwrap(); + let metadata = region.metadata(); + let mut stream = instance + .datanode + .region_server() + .handle_remote_read( + QueryRequest { + header: Some(RegionRequestHeader { + query_context: Some((&*ctx).into()), + ..Default::default() + }), + region_id: metadata.region_id.as_u64(), + plan: bytes, + }, + ctx, + ) + .await + .unwrap(); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + let schema = batches.first().expect("evaluated frames").schema.clone(); + let batches = RecordBatches::try_new(schema, batches).unwrap(); + let error = common_query::series_frame::unpack_evaluated_series( + Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())), + None, + ) + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("evaluated series identity changed before its chain closed"), + "{error}" + ); +} + #[tokio::test(flavor = "multi_thread")] async fn native_prefix_executes_over_real_mito_region() { let instance = GreptimeDbStandaloneBuilder::new("native_prefix_region_scan") @@ -204,11 +371,6 @@ async fn native_prefix_executes_over_real_mito_region() { Some("a") ); assert_eq!(series.time.timestamps().unwrap().as_ref(), &[3000, 4000]); - // At 3000 ms the (0, 3000] window has samples 0, 10, 20: increase 20 over - // 2 observed seconds. Prometheus extrapolates the 1s left-boundary gap and - // 0s right-boundary gap to the full 3s range: 20 * (3 / 2) / 3 = 10/s. - // At 4000 ms, (1000, 4000] has 10, 20, 30: increase 20 over 2 seconds, - // likewise extrapolated to 10/s. let values = series.fields.get(0).unwrap(); // At 3000 ms the window holds 0, 10, 20: an increase of 20 over 2 observed // seconds. Prometheus would keep the 1 s leading gap (below the 1.1x average