mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
feat(promql): serialize the native prefix and execute it on a datanode (#10029)
I hereby agree to the terms of the [GreptimeDB CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md). ## Refer to a related PR or issue link (optional) This is the wire-format half of the closed #9993 line. The raw fan-in, aggregate split and remote-source work that #10002 also carried moves to a separate PR on top of the public execution framework. ## What's changed and what's your intention? ### The problem A native PromQL plan could not leave the frontend. `FusedSeries` carries a `NativeTree`, and that tree had no wire representation, so a datanode could only ever be asked for raw samples. ### The wire format `src/promql/src/native/wire.rs` defines protocol version 1 for the prefix shapes the remote route can actually produce — `Input → Instant|Window → Map(Math*)+`, i.e. a selection under any number of math maps — as an explicit byte layout documented in the module: - every node kind and every operation has a **fixed tag**; the tags are declared wire-stable and deliberately do not follow Rust discriminants, so reordering an enum cannot change the protocol; - floats travel as `f64::to_bits()`, so NaN payloads, signed zero, infinities and subnormals round-trip exactly; - the input ordering is carried as data (0 Interleaved, 1 Contiguous, 2 KeyThenTime) rather than filtered by the encoder: legality belongs to `NativeTree::try_new` on decode, not to the transport; - decoding rebuilds grids through their validating constructor and ends with `NativeTree::try_new`, so a structurally invalid tree is refused at decode; compatibility between a window modifier and its function, or a negative `lookback_ms`, is refused later, when the datanode builds the physical plan; - the payload is bounded (1 MiB) and rejects unknown versions, unknown tags, truncation and trailing bytes instead of guessing. `Instant`'s value selector is carried as a tag (0 sample, 1 timestamp, 2 info timestamp) rather than a bool, because `InstantValue` has three cases since `info()` became native. Anything outside the supported shapes fails to encode *and* to decode with an error naming the unsupported operation. `label_replace`/`label_join` are not encoded because admission already declines them. ### Carrying it `FusedSeries` gains `serialize`/`deserialize` (a small versioned envelope around the tree plus its binding) and registers arms in the existing Substrait extension serializer, so a datanode's `DefaultPlanDecoder` already understands the node. No new transport, no protocol change, no `greptime-proto` bump: extension payloads are opaque bytes. Serialization is restricted to the `SeriesFrame` output with no `partial_fold`, because only that shape derives its output schema from the tree and can therefore rebuild itself after decoding. ### Executing it on a datanode The existing `PromExtensionPlanner::plan_fused_series` already knows how to build a decoded `FusedSeries` over a `RegionScanExec`. The integration test `native_prefix_executes_over_real_mito_region` uses a real standalone Mito region to prove the whole path: a flushed table, an encoded `Input → Window(rate)` prefix, `encode` → `handle_remote_read` on the real region server, decode and execute, then unpack the returned evaluated frame and check that the tag `host` survives, `__name__` is dropped as `rate` requires, the grid timestamps are `3000`/`4000`, and the rates are `20/3` and `10`. Those two numbers are not interchangeable: a counter whose first sample is `0` cannot be extrapolated backwards at all, while the second window's sample pattern allows the full 1 s gap. The test also tampers with the embedded version byte and asserts the datanode fails rather than returning a plausible frame. ### Encode/decode agreement A payload the encoder accepts must be one this version can decode. Three places disagreed and are fixed here: - **Field count.** The decoder limited an `Input`'s field count to 128 while the encoder wrote any count, so a 129-field `Input → Instant` encoded to 983 bytes and failed to decode. Both sides now use `MAX_FIELD_COUNT`. - **Tag count.** `serialize` only checked that the tag count fit in a `u32` while `deserialize` limited it to 1024. Both sides now use `MAX_TAG_COUNT`. - **Frame schema.** `deserialize` passed an empty schema to `try_new_with_format`, which since #9968 *checks* a frame node's declared schema against the layout derived from its tree instead of overwriting it. Every serialized frame node was therefore impossible to rebuild. `deserialize` now derives the schema from the decoded tree, which is where it comes from in the first place. Counts were left at their previously published values (depth 128, operations 128, fields 128, tags 1024); only the missing encoder-side checks and the wrong schema were added. `MAX_COUNT` is split into `MAX_DEPTH` / `MAX_FIELD_COUNT` / `MAX_OPERATION_COUNT` so each limit names what it bounds, and a field-count overflow no longer reports "operation count exceeds maximum". ### Limitations - **The frontend does not route queries through this yet.** It is a plan-serialization and datanode-execution boundary, not a query-plan change. Wiring the frontend to push the prefix and consume evaluated frames is the next increment. - Only one input source and the prefix shapes listed above are transported; binary, vector-wide, aggregate-state and subquery trees are not. - No performance claim: a datanode that executes the prefix returns evaluated frames instead of raw samples, but the frontend still gathers every selected region into one partition. - A selection may sit under any number of math maps, because a map is transparent to shape: it keeps its input's samples type and grid, and only the operation kinds are restricted. Non-math operations are still refused at every level. The payload stays bounded by `MAX_DEPTH` (128 node levels, so a selection plus 126 maps), and both the encoder and the decoder enforce it. ### Known limitation: the datanode boundary is not chain-atomic The region-read merge point wraps a non-streaming plan in `CoalescePartitionsExec`, which coalesces **per batch**. A series frame declares `Contiguous` and one series chain spans several batches (the packer bounds rows and flushes on `Pending`), so on a multi-partition datanode the batches of one chain are interleaved with another's and decoding fails with `evaluated series identity changed before its chain closed`. This is reproduced by `native_prefix_multi_partition_series_chains` (16 hosts × 20k samples, `rate[10s]`, 16 partitions). That test asserts the failure, so it passes today and will turn red when the merge is fixed — which is the signal to change it. The `parallelism: 1` path is covered by `native_prefix_executes_over_real_mito_region`. The fix does not belong here: the merge point should become chain-atomic in the frontend-integration change, for example by reusing the public `SeriesRepartitionExec`. Until then, **a multi-partition datanode read of a series frame is not correct** — this PR is not evidence that the datanode path is usable. ## PR Checklist Please convert it to a draft if some of the following conditions are not met. - [x] I have written the necessary rustdoc comments. - [x] I have added the necessary unit tests and integration tests. - [ ] This PR requires documentation updates. - [x] API changes are backward compatible. - [x] Schema or data changes are backward compatible. - [ ] This PR needs to be backported to release branches, and I have added the `backport-<target>` labels (e.g. `backport-v1.3` targets `release/v1.3`).
This commit is contained in:
discord9
@@ -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))
|
||||
|
||||
@@ -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
|
||||
@@ -232,6 +236,192 @@ 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)?;
|
||||
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()))?
|
||||
.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 > 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 {
|
||||
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()
|
||||
];
|
||||
// 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,
|
||||
FusedSeriesBinding {
|
||||
grid,
|
||||
metric_name,
|
||||
time_index,
|
||||
tags,
|
||||
},
|
||||
children,
|
||||
output_schema,
|
||||
format,
|
||||
)
|
||||
}
|
||||
|
||||
/// The name this node is shown under.
|
||||
pub const fn name() -> &'static str {
|
||||
"FusedSeries"
|
||||
@@ -334,6 +524,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
|
||||
|
||||
@@ -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;
|
||||
@@ -924,6 +924,86 @@ impl Drop for StagedStream {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
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())],
|
||||
declared,
|
||||
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());
|
||||
}
|
||||
|
||||
/// 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<String>) -> 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<String> = (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<String> = (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 { .. }) {
|
||||
|
||||
@@ -34,3 +34,4 @@ pub mod execute;
|
||||
pub mod metrics;
|
||||
pub mod plan;
|
||||
pub(crate) mod runs;
|
||||
pub(crate) mod wire;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
|
||||
@@ -0,0 +1,427 @@
|
||||
// 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::options::QueryOptions;
|
||||
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"),
|
||||
)
|
||||
}
|
||||
|
||||
// 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::<Vec<_>>()
|
||||
.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")
|
||||
.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::<Vec<_>>()
|
||||
.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]);
|
||||
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:?}"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user