rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(promql): serialize fused native prefixes through substrait

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-29 18:32:44 +08:00
committed by discord9
parent 7d113e2aa4
commit ecb76d7e30
4 changed files with 241 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/substrait/src/extension_serializer.rs
+10 -2
View File
@@ -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<Vec<u8>> {
match node.name() {
name if name == FusedSeries::name() => {
let fused = node
.as_any()
.downcast_ref::<FusedSeries>()
.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<Arc<dyn UserDefinedLogicalNode>> {
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))
src/promql/src/extension_plan/fused_series.rs
+172
View File
@@ -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<Vec<u8>> {
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<u8>, 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<Self> {
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<u8> {
Ok(self.take(1)?[0])
}
fn u32(&mut self) -> DataFusionResult<u32> {
Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
}
fn i64(&mut self) -> DataFusionResult<i64> {
Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
}
fn u64(&mut self) -> DataFusionResult<u64> {
Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
}
fn text(&mut self) -> DataFusionResult<String> {
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::<str>::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::<str>::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
src/promql/src/extension_plan/fused_series/tests.rs
+33
View File
@@ -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 { .. }) {
src/promql/src/native/wire.rs
+26 -23
View File
@@ -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<NativeNode, String> {
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<NativeNode, String> {
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"),