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"),