feat(promql): carry the native plan with serde instead of a hand-written codec (#10072)

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)

Replaces the serialization introduced by #10029 (merged into `narthex`);
this replaces rather than layers on top of that change.

## What's changed and what's your intention?

### Summary

This replaces the hand-written `NativeTree` byte codec and shape
whitelist with derived serde serialization encoded by postcard. The
versioned `FusedSeries` envelope now serializes a payload containing the
native tree and its plan binding, including the output Arrow `Schema`;
Substrait continues to carry children separately. The former explicit
`NativeTree` codec is removed.

The output schema is serialized directly using Arrow Schema's serde
representation and postcard; it is not inferred from the tree on decode.
The payload no longer transports the separate `output_tags`,
`output_columns`, or `plain_rows_schema` representations. Carrying Arrow
`Schema` directly preserves column order, schema and field metadata, and
custom scalar column names through the round trip. Decode reconstructs a
`FusedSeries` node and checks format-specific schema constraints:
`SeriesFrame` must match the evaluated frame layout; `PlainRows` must
pass output-schema classification and consistency checks; partial-fold
output must match the expected state schema, including tag order. These
checks cover the time column, tags, value columns, and scalar-root
constraints.

### Why

The old codec duplicated native node kinds and admitted only a
hand-maintained set of tree shapes. That made transport capability
diverge from what the compiler could build. Using serde over the tree
avoids a separate per-variant wire codec and lets native plan validation
continue to define legal trees.

The envelope has a version and payload length. Decode checks the version
before deserializing, rejects length/truncation/trailing-byte errors,
and validates the decoded plan's binding and output representation.
`NativeTree` construction checks nesting and validates the root; its
custom `Deserialize` likewise checks nesting and validates the decoded
root. `EvalGrid`, `SeriesLabels`, and `LabelRewrite` retain custom serde
handling for their invariants (deferred grid reconstruction, label
validation, and regex recompilation).

Postcard preserves floating-point bit patterns; tests cover a
non-canonical NaN payload, signed zero, infinities, and a subnormal. The
dependency adds `cobs` and `embedded-io`; postcard default features are
disabled.

### Structural equality and binding

Native plan equality and hashing use structural content rather than
display strings; float-bearing leaves use `f64::to_bits()`. Binding
comparison uses ordered field specifications and optional
evaluation-grid specifications rather than `Arc` pointer identity. The
repeated `validate_binding` native-tree depth check was removed: checked
construction and custom deserialization already enforce the tree-depth
invariant.

### Depth safety

The limit is a maximum **nesting depth of 64 native nodes**, not a
total-node limit. Checked native-tree construction and custom
deserialization enforce the boundary; deserialization also uses a depth
guard while recursively reading nodes. Over-limit input is rejected,
although postcard reports the custom serde failure as its unit
`SerdeDeCustom` variant. In the combined admission test, 63 selectors
fit within the depth limit; 64 are rejected by the native route, and a
comparison against the legacy fallback passes. This is not a broad
fallback guarantee: a measured 70-term `time()` expression overflows the
stack on the legacy path.

### Known limitation: schema depth is not bounded on decode

The native tree is depth-bounded while it deserializes. The carried
Arrow `Schema` is not: `DataType` is independently recursive, and a
payload whose schema field nests `List` thousands deep aborts the
decoding thread before any validation runs. A reviewer reproduced this
with a hand-built 80 KB payload on an 8 MiB thread.

This schema cannot grow that way from any query. Production builds every
`FusedSeries` at two call sites in `planner/native.rs`, both taking
their schema from `fused_output_schema`, which hardcodes
`Timestamp(Millisecond, None)`, `Float64` values and `Utf8` tags and
never constructs a nested type. The partial-fold form uses the fixed
`state_schema`, and the frame form is checked against
`series_frame_output_schema`, a constant layout. Only a payload written
outside this planner can deepen it, and anything able to supply that
payload can supply any other plan.

It is still a real difference from the tree, so it is recorded rather
than glossed: this transport is not depth-safe against crafted input.
Bounding it would mean a depth guard inside the schema's own
`Deserialize`, not a second schema representation.

### Verification

- Serde branch at `7e8f31ab1bc`: targeted `fused_series` tests plus the
named-stack `NativeTree` lifecycle test, 77 passed (not the full
`promql` suite).
- Combined with the admission branch at `209cb81ddae`: query and promql
suites, 2,099 passed and 2 skipped, on the default stack.
- The depth-fallback regression uses 64 selector terms: 63 fit the
limit, 64 are rejected with `NotImplemented` and fall back to the legacy
plan. It is run at `RUST_MIN_STACK=8388608` because the legacy planner,
not the native route, overflows the default 8 MiB stack on a deep
`time()` chain.
- `PlanRewriter::new` now encodes the root once instead of twice; the
`dist_plan` suite, 207 tests, passed after that change.
- `cargo fmt --all -- --check` and `git diff --check`: passed.

All Rust runs used a local `greptime-proto` path patch, which is why a
`Cargo.lock` difference appears locally; it is not part of this PR.

The earlier datatypes, integration, and clippy counts were not rerun at
the current head and are not claimed as current verification. No
`#10055` integration result is claimed.

## 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.
- [ ] 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
2026-11-03 19:30:36 +08:00
committed by GitHub
28 changed files with 1685 additions and 1822 deletions
Generated
+35
View File
@@ -2253,6 +2253,15 @@ version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a"
[[package]]
name = "cobs"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0fa961b519f0b462e3a3b4a34b64d119eeaca1d59af726fe450bbba07a9fc0a1"
dependencies = [
"thiserror 2.0.17",
]
[[package]]
name = "codespan-reporting"
version = "0.12.0"
@@ -5020,6 +5029,18 @@ dependencies = [
"serde",
]
[[package]]
name = "embedded-io"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced"
[[package]]
name = "embedded-io"
version = "0.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d"
[[package]]
name = "ena"
version = "0.14.3"
@@ -11056,6 +11077,18 @@ dependencies = [
"portable-atomic",
]
[[package]]
name = "postcard"
version = "1.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6764c3b5dd454e283a30e6dfe78e9b31096d9e32036b5d1eaac7a6119ccb9a24"
dependencies = [
"cobs",
"embedded-io 0.4.0",
"embedded-io 0.6.1",
"serde",
]
[[package]]
name = "postgres-protocol"
version = "0.6.12"
@@ -11373,10 +11406,12 @@ dependencies = [
"greptime-proto",
"lazy_static",
"mito-codec",
"postcard",
"prometheus 0.14.0",
"prost 0.14.1",
"rand 0.9.4",
"regex",
"serde",
"snafu 0.8.6",
"store-api",
"table",
+1
View File
@@ -211,6 +211,7 @@ parquet-variant-json = "=59.2.0"
paste = "1.0"
percent-encoding = "2.3"
pin-project = "1.0"
postcard = { version = "1.1.3", default-features = false, features = ["alloc", "use-std"] }
pretty_assertions = "1.4.0"
prometheus = { version = "0.14", features = ["process"] }
promql-parser = { version = "0.10.0", features = ["ser"] }
+1 -1
View File
@@ -33,7 +33,7 @@ parquet-variant-compute.workspace = true
parquet-variant-json.workspace = true
paste.workspace = true
regex.workspace = true
serde.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json.workspace = true
snafu.workspace = true
sqlparser.workspace = true
+21 -3
View File
@@ -36,7 +36,7 @@ use crate::error::{InvalidSeriesSnafu, Result};
/// evaluated streams use current name/label identity, not storage primary keys or object address/hash.
/// This ordering
/// is not a query-global ordering, partition-ownership rule, or finality guarantee.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum SeriesOrdering {
/// Series chunks may interleave; positions within each series remain strictly increasing.
Interleaved,
@@ -47,7 +47,7 @@ pub enum SeriesOrdering {
}
/// A label with shared UTF-8 name and value strings.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub struct Label {
/// Label name.
pub name: Arc<str>,
@@ -61,12 +61,30 @@ pub struct Label {
/// The metric name is only the `__name__` value (see [`Self::name`]); this type is
/// the whole label set a PromQL series shows. It is neither a canonical
/// matching/grouping identity nor a storage identity (read key / primary key).
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct SeriesLabels {
name: Option<Arc<str>>,
tags: Arc<[Label]>,
}
impl serde::Serialize for SeriesLabels {
fn serialize<S: serde::Serializer>(
&self,
serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
(&self.name, self.tags.as_ref()).serialize(serializer)
}
}
impl<'de> serde::Deserialize<'de> for SeriesLabels {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
let (name, tags): (Option<Arc<str>>, Vec<Label>) =
serde::Deserialize::deserialize(deserializer)?;
Self::try_new(name, tags).map_err(serde::de::Error::custom)
}
}
impl SeriesLabels {
/// Creates a metric name, sorting labels and rejecting duplicates or `__name__` tags.
pub fn try_new(name: Option<Arc<str>>, mut tags: Vec<Label>) -> Result<Self> {
+2 -2
View File
@@ -664,12 +664,12 @@ impl SeriesFields {
}
/// Shared name and position of each output field.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct FieldSpec {
pub name: Arc<str>,
}
/// Shared schema describing series fields.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct SeriesSchema {
pub fields: Arc<[FieldSpec]>,
/// Shared grid for regular series; its `GridSpec` must match each series grid axis. Explicit axes require `None`.
+30
View File
@@ -90,6 +90,36 @@ pub struct EvalGrid {
spec: GridSpec,
timestamps: OnceCell<ScalarBuffer<i64>>,
}
impl serde::Serialize for EvalGrid {
fn serialize<S: serde::Serializer>(
&self,
serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
self.spec.serialize(serializer)
}
}
impl<'de> serde::Deserialize<'de> for EvalGrid {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
let spec = GridSpec::deserialize(deserializer)?;
Self::try_deferred(spec).map_err(serde::de::Error::custom)
}
}
impl PartialEq for EvalGrid {
fn eq(&self, other: &Self) -> bool {
self.spec == other.spec
}
}
impl Eq for EvalGrid {}
impl std::hash::Hash for EvalGrid {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.spec.start_ms.hash(state);
self.spec.step_ms.hash(state);
self.spec.len.hash(state)
}
}
impl EvalGrid {
/// Allocates the grid timestamps once for sharing across result series.
+2
View File
@@ -24,9 +24,11 @@ futures.workspace = true
greptime-proto.workspace = true
lazy_static.workspace = true
mito-codec.workspace = true
postcard.workspace = true
prometheus.workspace = true
prost.workspace = true
regex.workspace = true
serde = { workspace = true, features = ["rc"] }
snafu.workspace = true
store-api.workspace = true
table.workspace = true
@@ -226,10 +226,7 @@ impl AggregateFinalize {
)?))
}
/// The comparable form of the split this node carries, for its plan line and its identity.
///
/// [`BoundAggregate`] carries no `Eq`/`Hash` - a quantile's parameter is an `f64` - so its
/// `Debug` rendering stands in, the way [`FusedSeries`] renders the tree it runs.
/// The readable form of the split, for its plan line only.
fn key(&self) -> String {
format!("{:?} {:?} {:?}", self.op.kind, self.op.group, self.op.stage)
}
@@ -268,14 +265,21 @@ impl UserDefinedLogicalNodeCore for AggregateFinalize {
"AggregateFinalize should have exactly one input".to_string(),
));
}
// The aggregate, the binding and the suffix are this node's own, so a rewritten input changes
// where the state comes from and nothing about how it is merged.
let input = inputs.swap_remove(0);
let expected = state_schema(&self.binding.tags);
if input.schema().as_arrow() != expected.as_ref() {
return Err(DataFusionError::Plan(format!(
"AggregateFinalize input must export aggregate state with schema {expected:?}, got {:?}",
input.schema().as_arrow()
)));
}
// A rewritten input changes where the correctly-shaped state comes from, not how it merges.
Ok(Self {
promql: Arc::clone(&self.promql),
op: self.op.clone(),
binding: self.binding.clone(),
suffix: Arc::clone(&self.suffix),
input: inputs.swap_remove(0),
input,
output_schema: Arc::clone(&self.output_schema),
series_schema: Arc::clone(&self.series_schema),
})
@@ -284,12 +288,10 @@ impl UserDefinedLogicalNodeCore for AggregateFinalize {
impl PartialEq for AggregateFinalize {
fn eq(&self, other: &Self) -> bool {
// `NativeTree` carries no structural equality, so the rendered suffix stands in for it, as
// [`FusedSeries`] compares the tree it runs.
self.promql == other.promql
&& self.key() == other.key()
&& self.op == other.op
&& self.binding == other.binding
&& self.suffix.to_string() == other.suffix.to_string()
&& self.suffix == other.suffix
&& self.input == other.input
&& self.output_schema == other.output_schema
}
@@ -300,9 +302,9 @@ impl Eq for AggregateFinalize {}
impl Hash for AggregateFinalize {
fn hash<H: Hasher>(&self, state: &mut H) {
self.promql.hash(state);
self.key().hash(state);
self.op.hash(state);
self.binding.hash(state);
self.suffix.to_string().hash(state);
self.suffix.hash(state);
self.input.hash(state);
self.output_schema.hash(state);
}
@@ -310,16 +312,15 @@ impl Hash for AggregateFinalize {
impl PartialOrd for AggregateFinalize {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
// The rendered forms stand in for the trees, as they do in equality.
let mine = (self.promql.to_string(), self.key(), self.suffix.to_string());
let theirs = (
other.promql.to_string(),
other.key(),
other.suffix.to_string(),
);
match mine.partial_cmp(&theirs) {
Some(Ordering::Equal) => {}
ord => return ord,
match self.promql.cmp(&other.promql) {
Ordering::Equal => {}
ord => return Some(ord),
}
if self.op != other.op {
return None;
}
if self.suffix != other.suffix {
return None;
}
match self.binding.partial_cmp(&other.binding) {
Some(Ordering::Equal) => {}
+192 -227
View File
@@ -71,7 +71,7 @@ use datafusion::physical_plan::{
use datatypes::series::SeriesOrdering;
use datatypes::series::raw::RawSeriesChunk;
use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesSchema};
use datatypes::series::time::EvalGrid;
use datatypes::series::time::{EvalGrid, GridSpec};
use futures::Stream;
use store_api::region_engine::RawLabelSource;
@@ -102,10 +102,6 @@ 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
@@ -180,7 +176,9 @@ impl PartialOrd for FusedSeriesBinding {
}
/// The output representation fixed when a fused plan is built.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[derive(
Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
pub enum FusedOutputFormat {
/// The existing tabular query result.
PlainRows,
@@ -238,191 +236,142 @@ pub struct FusedSeries {
format: FusedOutputFormat,
}
#[derive(serde::Serialize, serde::Deserialize)]
struct FusedSeriesPayload {
tree: NativeTree,
grid: GridSpec,
metric_name: Option<Arc<str>>,
time_index: String,
tags: Vec<String>,
format: FusedOutputFormat,
promql: Arc<str>,
partial_fold: Option<BoundAggregate>,
output_schema: datafusion::arrow::datatypes::Schema,
}
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.
/// Serializes the native tree and binding required to reconstruct this plan node.
/// Substrait carries children separately; the envelope contains only a version, payload
/// length and the single derived payload.
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)?;
let payload = FusedSeriesPayload {
tree: self.tree.clone(),
grid: self.binding.grid.spec(),
metric_name: self.binding.metric_name.clone(),
time_index: self.binding.time_index.clone(),
tags: self.binding.tags.clone(),
format: self.format,
promql: Arc::clone(&self.promql),
partial_fold: self.partial_fold.clone(),
output_schema: self.output_schema.as_arrow().clone(),
};
let encoded = postcard::to_stdvec(&payload).map_err(|e| {
DataFusionError::Plan(format!("cannot serialize FusedSeries payload: {e}"))
})?;
let len = u32::try_from(encoded.len())
.map_err(|_| DataFusionError::Plan("FusedSeries payload too large".into()))?;
let mut out = Vec::with_capacity(8 + encoded.len());
out.extend(3u32.to_le_bytes());
out.extend(len.to_le_bytes());
out.extend(encoded);
Ok(out)
}
/// Rebuilds a serialized SeriesFrame node. Substrait attaches its child plans separately.
/// Rebuilds a serialized FusedSeries node, validating its carried output schema against the
/// execution contract. Substrait attaches children separately.
pub fn deserialize(bytes: &[u8]) -> DataFusionResult<Self> {
struct Reader<'a> {
b: &'a [u8],
p: usize,
if bytes.len() < 8 {
return Err(DataFusionError::Plan(
"truncated FusedSeries envelope".into(),
));
}
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 {
let version = u32::from_le_bytes(bytes[..4].try_into().unwrap());
if version != 3 {
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 len = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
let end = 8usize
.checked_add(len)
.ok_or_else(|| DataFusionError::Plan("FusedSeries length overflow".into()))?;
if end != bytes.len() {
return Err(DataFusionError::Plan(if end > bytes.len() {
"truncated FusedSeries payload".into()
} else {
"trailing bytes in FusedSeries payload".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() {
let (payload, remainder) = postcard::take_from_bytes::<FusedSeriesPayload>(&bytes[8..])
.map_err(|e| {
DataFusionError::Plan(format!("cannot deserialize FusedSeries payload: {e}"))
})?;
if !remainder.is_empty() {
return Err(DataFusionError::Plan(
"trailing bytes in FusedSeries payload".into(),
));
}
let FusedSeriesPayload {
tree,
grid,
metric_name,
time_index,
tags,
format,
promql,
partial_fold,
output_schema,
} = payload;
let grid = Arc::new(
EvalGrid::try_deferred(grid)
.map_err(|e| DataFusionError::Plan(format!("invalid FusedSeries grid: {e}")))?,
);
let children = vec![
LogicalPlan::EmptyRelation(datafusion::logical_expr::EmptyRelation {
produce_one_row: false,
schema: Arc::new(datafusion::common::DFSchema::empty())
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())
let binding = FusedSeriesBinding {
grid,
metric_name,
time_index,
tags,
};
Self::try_new_with_format(
promql,
tree,
FusedSeriesBinding {
grid,
metric_name,
time_index,
tags,
},
children,
output_schema,
format,
)
if binding.tags.iter().any(|tag| tag == &binding.time_index) {
return Err(DataFusionError::Plan(
"tag column cannot be the time column".into(),
));
}
let output_schema = Arc::new(datafusion::common::DFSchema::try_from(output_schema)?);
if let Some(op) = partial_fold {
if format != FusedOutputFormat::PlainRows {
return Err(DataFusionError::Plan(
"partial FusedSeries must use PlainRows".into(),
));
}
let partial = Self::try_new_partial(promql, tree, binding, children, op)?;
if partial.schema() != &output_schema {
return Err(DataFusionError::Plan(format!(
"partial FusedSeries schema does not match state schema {:?}",
state_schema(&partial.binding.tags)
)));
}
return Ok(partial);
}
let node =
Self::try_new_with_format(promql, tree, binding, children, output_schema, format)?;
if format == FusedOutputFormat::PlainRows {
let schema = Arc::new(node.output_schema.as_arrow().clone());
validate_plain_output_schema(
&schema,
&node.tree,
&node.output_tags(),
&node.binding.time_index,
)?;
}
Ok(node)
}
/// The name this node is shown under.
@@ -581,27 +530,30 @@ impl FusedSeries {
)?))
}
/// The label columns of the output: every column that is neither the time index nor a field of
/// the root. A scalar root has no labels; its one other column is the value.
///
/// The planner decides them from the whole tree - a match narrows them, `or` joins both sides'
/// and a wrapper around either keeps them - and the output schema is where it records that.
/// The label columns of the output, independent of the raw-input binding.
fn output_tags(&self) -> Vec<String> {
let Some(root) = self.tree.root().output.samples() else {
return Vec::new();
};
let field_names = self
.tree
.root()
.output
.samples()
.map(|schema| {
schema
.fields
.iter()
.map(|field| field.name.as_ref())
.collect::<std::collections::HashSet<_>>()
})
.unwrap_or_default();
self.output_schema
.fields()
.iter()
.map(|field| field.name())
.filter(|name| {
*name != &self.binding.time_index
&& !root
.fields
.iter()
.any(|field| field.name.as_ref() == name.as_str())
.filter(|field| {
field.name() != &self.binding.time_index
&& !field_names.contains(field.name().as_str())
&& !matches!(self.tree.root().output, NativeType::Scalar(_))
})
.cloned()
.map(|field| field.name().clone())
.collect()
}
}
@@ -660,17 +612,12 @@ impl UserDefinedLogicalNodeCore for FusedSeries {
impl PartialEq for FusedSeries {
fn eq(&self, other: &Self) -> bool {
// `NativeTree` carries no structural equality, so the rendered tree stands in for it. The
// renderer names every node with its bound operation and grids, which is what makes two
// nodes the same computation; comparing the text alone would call two different trees
// equal whenever their PromQL happens to read the same.
self.promql == other.promql
&& self.tree.to_string() == other.tree.to_string()
&& tree_input_fields(&self.tree) == tree_input_fields(&other.tree)
&& self.tree == other.tree
&& self.binding == other.binding
&& self.children == other.children
&& self.output_schema == other.output_schema
&& partial_key(self.partial_fold.as_ref()) == partial_key(other.partial_fold.as_ref())
&& self.partial_fold == other.partial_fold
&& self.format == other.format
}
}
@@ -683,7 +630,8 @@ impl Hash for FusedSeries {
self.binding.hash(state);
self.children.hash(state);
self.output_schema.hash(state);
partial_key(self.partial_fold.as_ref()).hash(state);
self.tree.hash(state);
self.partial_fold.hash(state);
self.format.hash(state);
}
}
@@ -702,24 +650,16 @@ impl PartialOrd for FusedSeries {
Some(Ordering::Equal) => {}
ord => return ord,
}
match self.tree.to_string().partial_cmp(&other.tree.to_string()) {
Some(Ordering::Equal) => {}
ord => return ord,
if self.tree != other.tree {
return None;
}
match partial_key(self.partial_fold.as_ref())
.partial_cmp(&partial_key(other.partial_fold.as_ref()))
{
Some(Ordering::Equal) => {}
ord => return ord,
if self.partial_fold != other.partial_fold {
return None;
}
match self.format.cmp(&other.format) {
Ordering::Equal => {}
ord => return Some(ord),
}
match tree_input_fields(&self.tree).cmp(&tree_input_fields(&other.tree)) {
Ordering::Equal => {}
ord => return Some(ord),
}
if self.output_schema == other.output_schema {
Some(Ordering::Equal)
} else {
@@ -741,29 +681,6 @@ fn render_partial_fold(op: &BoundAggregate) -> String {
format!("{kind} {group} {:?}", op.stage)
}
/// A stable comparable representation for an optional partial fold.
fn partial_key(partial: Option<&BoundAggregate>) -> Option<String> {
partial.map(render_partial_fold)
}
/// The raw field names bound to each input leaf, retaining source association and leaf order.
fn tree_input_fields(tree: &NativeTree) -> Vec<(usize, Vec<String>)> {
tree_inputs(tree)
.into_iter()
.map(|input| {
let fields = match &input.output {
NativeType::Samples(schema) => schema
.fields
.iter()
.map(|field| field.name.to_string())
.collect(),
NativeType::Scalar(_) | NativeType::AggregateState(_) => Vec::new(),
};
(input.source_index, fields)
})
.collect()
}
/// The schemas read by distinct raw sources, in dense source-index order.
fn collect_source_schemas(tree: &NativeTree) -> DataFusionResult<Vec<Arc<SeriesSchema>>> {
let mut schemas: Vec<(usize, Arc<SeriesSchema>)> = Vec::new();
@@ -907,6 +824,38 @@ pub(crate) fn tree_source_fields(tree: &NativeTree) -> DataFusionResult<Vec<Vec<
.collect())
}
fn validate_plain_output_schema(
schema: &SchemaRef,
tree: &NativeTree,
tags: &[String],
time: &str,
) -> DataFusionResult<()> {
let fields = match &tree.root().output {
NativeType::Samples(s) => s.fields.to_vec(),
NativeType::Scalar(_) => {
let non_time_columns = schema
.fields()
.iter()
.filter(|field| field.name().as_str() != time)
.count();
if non_time_columns != 1 {
return Err(DataFusionError::Plan(format!(
"FusedSeriesExec scalar root cannot produce {} tag columns",
non_time_columns.saturating_sub(1)
)));
}
scalar_root_fields(schema, &[], time)?
}
NativeType::AggregateState(_) => {
return Err(DataFusionError::Plan(
"PlainRows does not support aggregate-state roots".into(),
));
}
};
classify_output(schema, &fields, tags, time)?;
Ok(())
}
fn validate_physical_tree(tree: &NativeTree) -> DataFusionResult<Vec<Arc<SeriesSchema>>> {
let schemas = collect_source_schemas(tree)?;
if tree.root().output.grid().is_none() {
@@ -966,6 +915,17 @@ pub(crate) fn classify_output(
tags: &[String],
time_column: &str,
) -> DataFusionResult<Vec<OutputColumn>> {
if output_schema
.fields()
.iter()
.filter(|field| field.name().as_str() == time_column)
.count()
> 1
{
return Err(DataFusionError::Plan(format!(
"duplicate time column {time_column:?} in FusedSeriesExec output schema"
)));
}
for (index, tag) in tags.iter().enumerate() {
if tags[..index].contains(tag) {
return Err(DataFusionError::Plan(format!(
@@ -1086,6 +1046,11 @@ pub(crate) fn classify_output(
)));
}
}
if output_schema.fields().len() != 1 + series_fields.len() + tags.len() {
return Err(DataFusionError::Plan(
"FusedSeriesExec output schema has an unexpected number of columns".into(),
));
}
Ok(classification)
}
@@ -80,7 +80,7 @@ use store_api::region_engine::RawLabelSource;
mod frame_output;
use super::{
FusedOutputFormat, FusedSeries, FusedSeriesBinding, MAX_TAG_COUNT, OutputColumn, POLL_BUDGET,
FusedOutputFormat, FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET,
series_to_record_batch,
};
use crate::evaluated::EvaluatedChunk;
@@ -925,7 +925,7 @@ impl Drop for StagedStream {
}
#[test]
fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() {
fn fused_series_ordinary_encoding_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
@@ -942,13 +942,60 @@ fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() {
.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.tree(), node.tree());
assert_eq!(decoded.binding(), node.binding());
assert_eq!(decoded.schema(), node.schema());
assert!(fused(rate_tree()).serialize().is_err());
let rows = fused(rate_tree());
let row_decoded = FusedSeries::deserialize(&rows.serialize().unwrap()).unwrap();
assert_eq!(row_decoded.schema(), rows.schema());
assert_eq!(row_decoded.tree(), rows.tree());
// A legal noncanonical order must survive the wire unchanged.
let canonical = rows.schema().as_arrow();
let reordered_schema = Arc::new(Schema::new(vec![
canonical.field_with_name(VALUE).unwrap().clone(),
canonical.field_with_name(HOST).unwrap().clone(),
canonical.field_with_name(NAME).unwrap().clone(),
canonical.field_with_name(TIME).unwrap().clone(),
]));
let reordered = FusedSeries::try_new(
Arc::from(PROMQL),
rows.tree().clone(),
rows.binding().clone(),
rows.children().to_vec(),
Arc::new(
datafusion::common::DFSchema::try_from(reordered_schema.as_ref().clone()).unwrap(),
),
)
.unwrap();
let reordered_decoded = FusedSeries::deserialize(&reordered.serialize().unwrap()).unwrap();
assert_eq!(reordered_decoded.schema(), reordered.schema());
// The output schema owns the label layout; the source binding may have no tags or different
// labels after a tree operation such as label rewriting.
let empty_binding = binding_of(rows.tree(), raw_schema(), Vec::new());
let empty_tags = FusedSeries::try_new(
Arc::from(PROMQL),
rate_tree(),
empty_binding,
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
)
.unwrap();
let decoded = FusedSeries::deserialize(&empty_tags.serialize().unwrap()).unwrap();
assert_eq!(decoded.binding().tags, empty_tags.binding().tags);
assert_eq!(decoded.schema(), empty_tags.schema());
let projected = fused(binary_tree(raw_schema(), raw_schema()));
let projected_bytes = projected.serialize().unwrap();
let projected_decoded = FusedSeries::deserialize(&projected_bytes).unwrap();
assert_eq!(projected_decoded.schema(), projected.schema());
let partial = partial_node(AggregateKind::Sum);
let partial_decoded = FusedSeries::deserialize(&partial.serialize().unwrap()).unwrap();
assert_eq!(partial_decoded.partial_fold(), partial.partial_fold());
assert_eq!(partial_decoded.binding().tags, partial.binding().tags);
assert_eq!(partial_decoded.schema(), partial.schema());
let mut bad_version = bytes.clone();
bad_version[..4].copy_from_slice(&99u32.to_le_bytes());
assert!(
@@ -958,6 +1005,15 @@ fn fused_series_envelope_round_trip_and_rejects_unsupported_payloads() {
.contains("version")
);
assert!(FusedSeries::deserialize(&bytes[..5]).is_err());
let mut trailing_payload = bytes.clone();
let payload_len = u32::from_le_bytes(trailing_payload[4..8].try_into().unwrap()) as usize;
trailing_payload.insert(8 + payload_len, 0);
trailing_payload[4..8].copy_from_slice(&((payload_len + 1) as u32).to_le_bytes());
assert!(FusedSeries::deserialize(&trailing_payload).is_err());
let mut trailing = bytes.clone();
trailing.push(0);
assert!(FusedSeries::deserialize(&trailing).is_err());
}
/// The `SeriesFrame` node of [`rate_tree`] declaring `tags` as its labels.
@@ -979,29 +1035,218 @@ fn envelope_node_with_tags(tags: Vec<String>) -> FusedSeries {
.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);
fn mutate_envelope_payload(
bytes: &[u8],
mutate: impl FnOnce(&mut super::FusedSeriesPayload),
) -> Vec<u8> {
let (mut payload, remainder) =
postcard::take_from_bytes::<super::FusedSeriesPayload>(&bytes[8..]).unwrap();
assert!(remainder.is_empty());
mutate(&mut payload);
let changed = postcard::to_stdvec(&payload).unwrap();
let mut result = Vec::with_capacity(8 + changed.len());
result.extend(3u32.to_le_bytes());
result.extend((changed.len() as u32).to_le_bytes());
result.extend(changed);
result
}
fn mutate_schema(bytes: &[u8], schema: Schema) -> Vec<u8> {
mutate_envelope_payload(bytes, |payload| payload.output_schema = schema)
}
fn scalar_node(tree: NativeTree, value_name: &str) -> FusedSeries {
let schema = Arc::new(Schema::new(vec![
Field::new(
TIME,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new(value_name, DataType::Float64, true),
]));
FusedSeries::try_new(
Arc::from(PROMQL),
tree,
FusedSeriesBinding {
grid: grid(),
time_index: TIME.into(),
tags: Vec::new(),
metric_name: None,
},
Vec::new(),
Arc::new(datafusion::common::DFSchema::try_from(schema.as_ref().clone()).unwrap()),
)
.unwrap()
}
fn assert_decode_error(bytes: &[u8], expected: &str) {
let error = FusedSeries::deserialize(bytes).unwrap_err().to_string();
assert!(
node.serialize().is_ok(),
"the maximum tag count must still serialize"
error.contains(expected),
"expected {expected:?}, got {error:?}"
);
}
/// 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 plain_rows_decode_rejects_invalid_schemas() {
let node = fused(rate_tree());
let bytes = node.serialize().unwrap();
let scalar = scalar_node(scalar_tree(grid(), 1.0), "clock_value");
let scalar_bytes = scalar.serialize().unwrap();
let scalar_with_tag = mutate_envelope_payload(&scalar_bytes, |payload| {
let mut fields = payload.output_schema.fields().to_vec();
fields.push(Arc::new(Field::new(HOST, DataType::Utf8, true)));
payload.output_schema = Schema::new(fields);
});
assert_decode_error(&scalar_with_tag, "cannot produce 1 tag columns");
let malformed_schemas = [
(
Schema::new(vec![Field::new(TIME, DataType::Utf8, false)]),
"must be Timestamp",
),
(
Schema::new(vec![
Field::new(
TIME,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new(VALUE, DataType::Float64, true),
Field::new(VALUE, DataType::Float64, true),
]),
"duplicate",
),
(
Schema::new(vec![
Field::new(
TIME,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new(HOST, DataType::Utf8, true),
Field::new(NAME, DataType::Utf8, true),
]),
"missing field column",
),
(
Schema::new(vec![
Field::new(
TIME,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new(VALUE, DataType::Float64, true),
Field::new(HOST, DataType::Float64, true),
Field::new(NAME, DataType::Utf8, true),
]),
"must be Utf8",
),
];
for (schema, expected) in malformed_schemas {
assert_decode_error(&mutate_schema(&bytes, schema), expected);
}
let scalar_decoded = FusedSeries::deserialize(&scalar_bytes).unwrap();
assert_eq!(scalar_decoded.schema(), scalar.schema());
// Binding tags describe the source; the carried schema independently describes outputs.
let empty_binding = binding_of(node.tree(), raw_schema(), Vec::new());
let cross_table = FusedSeries::try_new(
Arc::from(PROMQL),
node.tree().clone(),
empty_binding,
node.children().to_vec(),
output_schema().to_dfschema_ref().unwrap(),
)
.unwrap();
let decoded = FusedSeries::deserialize(&cross_table.serialize().unwrap()).unwrap();
assert_eq!(decoded.schema(), cross_table.schema());
assert!(decoded.binding().tags.is_empty());
}
#[test]
fn a_non_default_column_order_survives_the_round_trip() {
let tree = rate_tree();
let binding = binding(&tree);
let reordered = Schema::new_with_metadata(
vec![
Field::new(VALUE, DataType::Float64, true),
Field::new(
TIME,
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new(HOST, DataType::Utf8, true)
.with_metadata([("origin".to_string(), "carried".to_string())].into()),
Field::new(NAME, DataType::Utf8, true),
],
[("schema-origin".to_string(), "carried".to_string())].into(),
);
let node = FusedSeries::try_new(
Arc::from(PROMQL),
tree,
binding,
vec![input_plan(frame_schema())],
Arc::new(datafusion::common::DFSchema::try_from(reordered.clone()).unwrap()),
)
.unwrap();
let decoded = FusedSeries::deserialize(&node.serialize().unwrap()).unwrap();
assert_eq!(decoded.schema(), node.schema());
assert_eq!(decoded.schema().as_arrow().metadata(), reordered.metadata());
assert_eq!(
decoded
.schema()
.field_with_name(None, HOST)
.unwrap()
.metadata(),
reordered.field_with_name(HOST).unwrap().metadata()
);
}
#[test]
fn series_frame_decode_rejects_schema_mismatch() {
let node = envelope_node_with_tags(vec![HOST.into(), NAME.into()]);
let bytes = node.serialize().unwrap();
let wrong = mutate_envelope_payload(&bytes, |payload| {
let mut fields = payload.output_schema.fields().to_vec();
fields.push(Arc::new(Field::new("extra", DataType::Utf8, true)));
payload.output_schema = Schema::new(fields);
});
assert_decode_error(&wrong, "does not match the evaluated frame layout");
}
#[test]
fn partial_decode_rejects_format_schema_and_tag_order_mismatch() {
let tree = rate_tree();
let node = FusedSeries::try_new_partial(
Arc::from(PROMQL),
tree.clone(),
partial_binding_of(&tree, &[HOST, JOB]),
vec![input_plan(frame_schema())],
partial_op(AggregateKind::Sum, &[HOST, JOB]),
)
.unwrap();
let bytes = node.serialize().unwrap();
let wrong_schema = mutate_envelope_payload(&bytes, |payload| {
let mut fields = payload.output_schema.fields().to_vec();
fields.swap(0, 1);
payload.output_schema = Schema::new(fields);
});
assert_decode_error(&wrong_schema, "does not match state schema");
let wrong_format = mutate_envelope_payload(&bytes, |payload| {
payload.format = FusedOutputFormat::SeriesFrame;
});
assert_decode_error(&wrong_format, "partial FusedSeries must use PlainRows");
}
#[test]
fn fused_series_encoding_round_trips_large_tag_lists() {
let tags: Vec<String> = (0..1025).map(|i| format!("t{i}")).collect();
let node = envelope_node_with_tags(tags);
let decoded = FusedSeries::deserialize(&node.serialize().unwrap()).unwrap();
assert_eq!(decoded.binding().tags, node.binding().tags);
}
fn fused(tree: NativeTree) -> FusedSeries {
+30 -2
View File
@@ -50,7 +50,7 @@ use crate::map::validate_grid_series;
const BUCKET_LABEL: &str = "le";
/// `histogram_quantile` or `histogram_fraction` with its constant arguments.
#[derive(Debug, Clone, Copy, PartialEq)]
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
pub enum HistogramFunction {
/// `histogram_quantile(q, v)`.
Quantile(f64),
@@ -273,7 +273,7 @@ pub(crate) fn histogram_string(h: &HistogramView<'_>) -> Result<String> {
}
/// A function of one native histogram.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum NativeHistogramFunction {
Count,
Sum,
@@ -346,3 +346,31 @@ impl NativeHistogramFunction {
})
}
}
impl PartialEq for HistogramFunction {
fn eq(&self, o: &Self) -> bool {
match (self, o) {
(Self::Quantile(a), Self::Quantile(b)) => a.to_bits() == b.to_bits(),
(Self::Fraction { lower: a, upper: b }, Self::Fraction { lower: c, upper: d }) => {
a.to_bits() == c.to_bits() && b.to_bits() == d.to_bits()
}
_ => false,
}
}
}
impl Eq for HistogramFunction {}
impl std::hash::Hash for HistogramFunction {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
match self {
Self::Quantile(x) => {
0u8.hash(h);
x.to_bits().hash(h)
}
Self::Fraction { lower, upper } => {
1u8.hash(h);
lower.to_bits().hash(h);
upper.to_bits().hash(h)
}
}
}
}
+1 -1
View File
@@ -43,7 +43,7 @@ use crate::vector_match::{Budget, HISTOGRAM_ROW_BYTES, owned_histogram_bytes};
const IDENTIFYING_LABELS: [&str; 2] = ["instance", "job"];
/// The second argument of `info`, split the way Prometheus uses it.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct InfoSelector {
/// The `__name__` matchers, which choose the info metrics; none means `target_info`.
pub name_matchers: Vec<(MatchOp, Arc<str>)>,
+33 -1
View File
@@ -46,7 +46,7 @@ use crate::grid_series::GridSeriesBuilder;
use crate::map_time::DateFunction;
/// Float math supported by the evaluated-series kernel.
#[derive(Clone, Copy, Debug)]
#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
pub enum MathOp {
/// Absolute value.
Abs,
@@ -452,3 +452,35 @@ fn float_sample(sample: SampleRef<'_>) -> Option<f64> {
_ => unreachable!("validated as float"),
}
}
impl PartialEq for MathOp {
fn eq(&self, o: &Self) -> bool {
match (self, o) {
(Self::Round { nearest: a }, Self::Round { nearest: b }) => a.to_bits() == b.to_bits(),
(Self::Clamp { min: a, max: b }, Self::Clamp { min: c, max: d }) => {
a.to_bits() == c.to_bits() && b.to_bits() == d.to_bits()
}
(Self::ClampMin { min: a }, Self::ClampMin { min: b })
| (Self::ClampMax { max: a }, Self::ClampMax { max: b }) => a.to_bits() == b.to_bits(),
(Self::Date(a), Self::Date(b)) => a == b,
_ => std::mem::discriminant(self) == std::mem::discriminant(o),
}
}
}
impl Eq for MathOp {}
impl std::hash::Hash for MathOp {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
std::mem::discriminant(self).hash(h);
match self {
Self::Round { nearest: x } | Self::ClampMin { min: x } | Self::ClampMax { max: x } => {
x.to_bits().hash(h)
}
Self::Clamp { min, max } => {
min.to_bits().hash(h);
max.to_bits().hash(h)
}
Self::Date(d) => d.hash(h),
_ => (),
}
}
}
+149
View File
@@ -80,6 +80,89 @@ fn rewrite_all(
/// label names are checked.
#[derive(Debug)]
pub struct LabelRewrite(Change);
impl serde::Serialize for LabelRewrite {
fn serialize<S: serde::Serializer>(
&self,
serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
#[derive(serde::Serialize)]
enum Wire<'a> {
Replace {
dst: &'a str,
replacement: &'a str,
src: &'a str,
pattern: &'a str,
},
Join {
dst: &'a str,
separator: &'a str,
src: Vec<&'a str>,
},
}
let wire = match &self.0 {
Change::Replace {
dst,
replacement,
src,
pattern,
..
} => Wire::Replace {
dst,
replacement,
src,
pattern,
},
Change::Join {
dst,
separator,
src,
} => Wire::Join {
dst,
separator,
src: src.iter().map(AsRef::as_ref).collect(),
},
};
serde::Serialize::serialize(&wire, serializer)
}
}
impl<'de> serde::Deserialize<'de> for LabelRewrite {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
#[derive(serde::Deserialize)]
enum Wire {
Replace {
dst: String,
replacement: String,
src: String,
pattern: String,
},
Join {
dst: String,
separator: String,
src: Vec<String>,
},
}
match <Wire as serde::Deserialize>::deserialize(deserializer)? {
Wire::Replace {
dst,
replacement,
src,
pattern,
} => {
Self::replace(&dst, &replacement, &src, &pattern).map_err(serde::de::Error::custom)
}
Wire::Join {
dst,
separator,
src,
} => {
let refs: Vec<_> = src.iter().map(String::as_str).collect();
Self::join(&dst, &separator, &refs).map_err(serde::de::Error::custom)
}
}
}
}
#[derive(Debug)]
enum Change {
@@ -98,6 +181,72 @@ enum Change {
},
}
impl PartialEq for LabelRewrite {
fn eq(&self, other: &Self) -> bool {
match (&self.0, &other.0) {
(
Change::Replace {
dst: a,
replacement: b,
src: c,
pattern: d,
..
},
Change::Replace {
dst: w,
replacement: x,
src: y,
pattern: z,
..
},
) => (a, b, c, d) == (w, x, y, z),
(
Change::Join {
dst: a,
separator: b,
src: c,
},
Change::Join {
dst: w,
separator: x,
src: y,
},
) => (a, b, c) == (w, x, y),
_ => false,
}
}
}
impl Eq for LabelRewrite {}
impl std::hash::Hash for LabelRewrite {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
match &self.0 {
Change::Replace {
dst,
replacement,
src,
pattern,
..
} => {
0u8.hash(state);
dst.hash(state);
replacement.hash(state);
src.hash(state);
pattern.hash(state);
}
Change::Join {
dst,
separator,
src,
} => {
1u8.hash(state);
dst.hash(state);
separator.hash(state);
src.hash(state);
}
}
}
}
impl LabelRewrite {
/// `label_replace`'s change, failing on an invalid expression or an empty destination.
pub fn replace(dst: &str, replacement: &str, src: &str, regex: &str) -> Result<Self> {
+1 -1
View File
@@ -27,7 +27,7 @@ use crate::grid_series::external;
mod tests;
/// A PromQL date function: `minute(v)`, `year()` and so on.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum DateFunction {
DayOfMonth,
DayOfWeek,
-1
View File
@@ -34,4 +34,3 @@ pub mod execute;
pub mod metrics;
pub mod plan;
pub(crate) mod runs;
pub(crate) mod wire;
+237 -28
View File
@@ -26,8 +26,36 @@
//! What the tree cannot express yet is deliberate for this stage: `scalar(v)` has no node. The
//! legality below is therefore about the operations that exist, not about PromQL's whole surface.
use std::cell::Cell;
use std::hash::Hash;
use std::sync::Arc;
const MAX_NATIVE_NODE_DEPTH: usize = 64;
thread_local! {
static NATIVE_NODE_DEPTH: Cell<usize> = const { Cell::new(0) };
}
struct NativeNodeDepthGuard;
impl NativeNodeDepthGuard {
fn enter<E: serde::de::Error>() -> std::result::Result<Self, E> {
NATIVE_NODE_DEPTH.with(|depth| {
let current = depth.get();
if current >= MAX_NATIVE_NODE_DEPTH {
return Err(E::custom(format!(
"native tree nesting exceeds {MAX_NATIVE_NODE_DEPTH}"
)));
}
depth.set(current + 1);
Ok(Self)
})
}
}
impl Drop for NativeNodeDepthGuard {
fn drop(&mut self) {
NATIVE_NODE_DEPTH.with(|depth| depth.set(depth.get() - 1));
}
}
use datafusion::error::{DataFusionError, Result};
use datatypes::series::samples::{FieldSpec, SeriesSchema};
use datatypes::series::time::EvalGrid;
@@ -48,7 +76,7 @@ use crate::vector_match::{BinaryMatching, MatchingLabels, SetOp};
/// a grid it is raw samples carrying their own timestamps. A scalar is one value per grid point
/// and no labels. Intermediate aggregate state is a different type altogether, so an operation
/// that cannot consume it is rejected by type rather than by convention.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub enum NativeType {
/// Samples, evaluated when the schema carries a grid and raw otherwise.
Samples(Arc<SeriesSchema>),
@@ -95,14 +123,12 @@ impl NativeType {
/// Whether `self` binds the same thing as `other`: the same samples or the same state.
///
/// Bindings are compared by allocation identity, because a compiler that meant one schema or
/// one state has no reason to allocate it twice. A scalar is produced by a node and never
/// handed to one by a source, so it never matches. This is about the binding itself: grids are
/// aligned by [`EvalGrid::spec`] everywhere else, here and in the kernels.
/// Bindings compare their semantic content, not process-local allocation identity. Scalars are
/// produced by a node and never handed to one by a source, so they never match.
fn same_binding(&self, other: &NativeType) -> bool {
match (self, other) {
(Self::Samples(mine), Self::Samples(theirs)) => Arc::ptr_eq(mine, theirs),
(Self::AggregateState(mine), Self::AggregateState(theirs)) => Arc::ptr_eq(mine, theirs),
(Self::Samples(mine), Self::Samples(theirs)) => **mine == **theirs,
(Self::AggregateState(mine), Self::AggregateState(theirs)) => **mine == **theirs,
_ => false,
}
}
@@ -112,7 +138,7 @@ impl NativeType {
///
/// State is indexed by `(group labels, absolute grid index)`, so it carries the grid it was
/// accumulated on. It is not a query value and does not pretend to be one.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct AggregateStateSchema {
/// How the groups its rows are indexed by are identified.
pub group: Grouping,
@@ -127,7 +153,7 @@ pub struct AggregateStateSchema {
/// The binding records what the input was created for: the output a consumer type-checks against,
/// and the ordering it guarantees. An instant selection and a window both read the ordering,
/// rejecting an interleaved input; no other consumer reads it yet.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct BoundInput {
/// Child index in the DataFusion plan. The compiler initially assigns selector-occurrence
/// indices; the caller maps them to child indices when constructing the plan.
@@ -146,14 +172,14 @@ pub struct BoundInput {
/// A scalar that varies per grid point is the same [`NativeType::Scalar`] as this constant: what
/// keeps a binary operation from label-matching it is its type, not its constness. The varying
/// ones are [`NativeExpr::EvalTime`] and a binary operation between scalars.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct BoundScalar {
/// The constant value.
pub value: f64,
}
/// An instant selection: the last sample at or before each grid point, within a lookback.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct BoundInstant {
/// How far back a sample may lie before it counts as stale.
pub lookback_ms: i64,
@@ -168,7 +194,7 @@ pub struct BoundInstant {
}
/// What an instant selection reports at a grid point for the sample it selected there.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum InstantValue {
/// The sample itself.
Sample,
@@ -183,7 +209,7 @@ pub enum InstantValue {
///
/// The input and output grids are bound separately: a subquery evaluates its inner expression on
/// its own grid and consumes the result directly, without resampling it onto the outer grid.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct BoundWindow {
/// The function this node runs.
pub function: RangeFunction,
@@ -196,7 +222,7 @@ pub struct BoundWindow {
}
/// A chain of point-wise operations applied to a whole evaluated value.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct BoundMap {
/// The operations, applied in order.
pub ops: Arc<[MapOp]>,
@@ -205,7 +231,7 @@ pub struct BoundMap {
/// A math function whose scalar parameters are bound as nodes, because at least one of them varies
/// per grid point; Prometheus reads each parameter at every step. With constant parameters the same
/// function is a [`MathOp`] in a [`BoundMap`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum ParameterizedOp {
/// `round(v, to_nearest)`.
Round,
@@ -230,7 +256,7 @@ impl ParameterizedOp {
/// One operation of a map, applied to each series on its own: math over its float values, a
/// function of its native histograms, or a change of its labels. Each can give two series one label
/// set, which is why all live in a map.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub enum MapOp {
Math(MathOp),
Histogram(NativeHistogramFunction),
@@ -251,7 +277,7 @@ pub enum MapOp {
/// [`RangeFunction::QuantileOverTime`] does.
///
/// [`RangeFunction::QuantileOverTime`]: crate::range_eval::RangeFunction::QuantileOverTime
#[derive(Clone, Copy, Debug, PartialEq)]
#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
pub enum AggregateKind {
/// `sum`: the total of each group's values at each grid point.
Sum,
@@ -278,7 +304,7 @@ pub enum AggregateKind {
}
/// Where an aggregate sits in the execution, which decides whether its output is a query value.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum AggregateStage {
/// Every contribution to each group lies within this execution, so the output is complete.
Complete,
@@ -296,7 +322,7 @@ pub enum AggregateStage {
/// `without` written that way would group a query by labels it meant to remove.
///
/// [`aggregate::Grouping`]: crate::aggregate::Grouping
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum Grouping {
/// The labels a group is identified by; every other label is dropped.
By(Arc<[Arc<str>]>),
@@ -305,7 +331,7 @@ pub enum Grouping {
}
/// An aggregate bound to its kind, grouping and stage.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct BoundAggregate {
/// Which aggregate to compute.
pub kind: AggregateKind,
@@ -321,7 +347,7 @@ pub struct BoundAggregate {
///
/// Its constants are bound here, as a window's are: a `k` or a ratio only known at run time has no
/// field.
#[derive(Clone, Debug, PartialEq)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub enum BoundVectorOp {
/// `topk` (`top`) or `bottomk`: the `k` largest or smallest values of each group.
TopK { k: f64, top: bool, group: Grouping },
@@ -349,7 +375,7 @@ pub enum BoundVectorOp {
///
/// It has no cardinality, `bool` or fill: PromQL matches a set operator many-to-many and refuses
/// the other modifiers on it.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub struct BoundSetOp {
/// `and`, `or` or `unless`.
pub op: SetOp,
@@ -358,7 +384,7 @@ pub struct BoundSetOp {
}
/// What a node computes. Every variant holds a bound operation, never a name to resolve.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
pub enum NativeExpr {
/// A scan or an exchange boundary.
Input(BoundInput),
@@ -434,7 +460,7 @@ pub enum NativeExpr {
}
/// One node of a bound tree.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, PartialEq, Eq, Hash)]
pub struct NativeNode {
/// What this node computes, including its bound operands.
pub expr: NativeExpr,
@@ -504,7 +530,7 @@ fn merge_shape_supported(node: &NativeNode) -> bool {
///
/// The tree itself carries no run-time state: every execution builds its own, so one compiled plan
/// can serve every partition.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, serde::Serialize, PartialEq, Eq, Hash)]
pub struct NativeTree {
root: NativeNode,
}
@@ -512,10 +538,10 @@ pub struct NativeTree {
impl NativeTree {
/// Creates a tree over `root`, checking that every node's contract holds.
///
/// A rejected tree is a compiler defect rather than a query error: the compiler only builds a
/// node after deciding that this combination is supported, so the error is
/// [`DataFusionError::Internal`] and not a statement the client wrote wrong.
/// A structurally inconsistent tree is a compiler defect; a tree beyond the native depth
/// limit is unsupported by this route and can fall back to the existing planner.
pub fn try_new(root: NativeNode) -> Result<Self> {
ensure_native_node_depth(&root)?;
root.validate()?;
Ok(Self { root })
}
@@ -526,6 +552,78 @@ impl NativeTree {
}
}
impl<'de> serde::Deserialize<'de> for NativeNode {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
let _guard = NativeNodeDepthGuard::enter::<D::Error>()?;
#[derive(serde::Deserialize)]
struct Wire {
expr: NativeExpr,
output: NativeType,
}
let wire = Wire::deserialize(deserializer)?;
Ok(Self {
expr: wire.expr,
output: wire.output,
})
}
}
impl<'de> serde::Deserialize<'de> for NativeTree {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
#[derive(serde::Deserialize)]
struct Wire {
root: NativeNode,
}
let wire = Wire::deserialize(deserializer)?;
ensure_native_node_depth(&wire.root).map_err(serde::de::Error::custom)?;
wire.root.validate().map_err(serde::de::Error::custom)?;
Ok(Self { root: wire.root })
}
}
fn ensure_native_node_depth(root: &NativeNode) -> Result<()> {
fn visit(node: &NativeNode, depth: usize) -> Result<()> {
if depth > MAX_NATIVE_NODE_DEPTH {
return Err(DataFusionError::NotImplemented(format!(
"native tree nesting exceeds {MAX_NATIVE_NODE_DEPTH}"
)));
}
match &node.expr {
NativeExpr::VectorOf { input }
| NativeExpr::Instant { input, .. }
| NativeExpr::Merge { input }
| NativeExpr::Window { input, .. }
| NativeExpr::Map { input, .. }
| NativeExpr::Timestamp { input }
| NativeExpr::Aggregate { input, .. }
| NativeExpr::Vector { input, .. } => visit(input, depth + 1),
NativeExpr::ParameterizedMap {
input, parameters, ..
} => {
visit(input, depth + 1)?;
for parameter in parameters {
visit(parameter, depth + 1)?;
}
Ok(())
}
NativeExpr::Binary { left, right, .. } | NativeExpr::Set { left, right, .. } => {
visit(left, depth + 1)?;
visit(right, depth + 1)
}
NativeExpr::Info { base, info, .. } => {
visit(base, depth + 1)?;
visit(info, depth + 1)
}
NativeExpr::Input(_) | NativeExpr::Scalar(_) | NativeExpr::EvalTime { .. } => Ok(()),
}
}
visit(root, 1)
}
impl NativeNode {
/// Checks this node and its inputs against the contract of its operation.
pub fn validate(&self) -> Result<()> {
@@ -1025,3 +1123,114 @@ fn invalid(reason: &str) -> DataFusionError {
#[cfg(test)]
mod tests;
impl PartialEq for BoundScalar {
fn eq(&self, o: &Self) -> bool {
self.value.to_bits() == o.value.to_bits()
}
}
impl Eq for BoundScalar {}
impl std::hash::Hash for BoundScalar {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
self.value.to_bits().hash(h);
}
}
impl PartialEq for AggregateKind {
fn eq(&self, o: &Self) -> bool {
match (self, o) {
(Self::Quantile { quantile: a }, Self::Quantile { quantile: b }) => {
a.to_bits() == b.to_bits()
}
(Self::Sum, Self::Sum)
| (Self::Avg, Self::Avg)
| (Self::Count, Self::Count)
| (Self::Min, Self::Min)
| (Self::Max, Self::Max)
| (Self::Group, Self::Group)
| (Self::Stddev, Self::Stddev)
| (Self::Stdvar, Self::Stdvar) => true,
_ => false,
}
}
}
impl Eq for AggregateKind {}
impl std::hash::Hash for AggregateKind {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
std::mem::discriminant(self).hash(h);
if let Self::Quantile { quantile } = self {
quantile.to_bits().hash(h);
}
}
}
impl PartialEq for BoundVectorOp {
fn eq(&self, o: &Self) -> bool {
match (self, o) {
(
Self::TopK {
k: a,
top: b,
group: c,
},
Self::TopK {
k: x,
top: y,
group: z,
},
) => a.to_bits() == x.to_bits() && b == y && c == z,
(Self::LimitK { k: a, group: b }, Self::LimitK { k: x, group: y }) => {
a.to_bits() == x.to_bits() && b == y
}
(Self::LimitRatio { ratio: a }, Self::LimitRatio { ratio: b }) => {
a.to_bits() == b.to_bits()
}
(Self::Sort(a), Self::Sort(b)) => a == b,
(Self::Absent { labels: a }, Self::Absent { labels: b }) => a == b,
(Self::Histogram(a), Self::Histogram(b)) => a == b,
(
Self::CountValues { label: a, group: b },
Self::CountValues { label: x, group: y },
) => a == x && b == y,
_ => false,
}
}
}
impl std::hash::Hash for BoundVectorOp {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
match self {
Self::TopK { k, top, group } => {
0u8.hash(h);
k.to_bits().hash(h);
top.hash(h);
group.hash(h)
}
Self::LimitK { k, group } => {
1u8.hash(h);
k.to_bits().hash(h);
group.hash(h)
}
Self::LimitRatio { ratio } => {
2u8.hash(h);
ratio.to_bits().hash(h)
}
Self::Sort(x) => {
3u8.hash(h);
x.hash(h)
}
Self::Absent { labels } => {
4u8.hash(h);
labels.hash(h)
}
Self::Histogram(x) => {
5u8.hash(h);
x.hash(h)
}
Self::CountValues { label, group } => {
6u8.hash(h);
label.hash(h);
group.hash(h)
}
}
}
}
impl Eq for BoundVectorOp {}
+260 -11
View File
@@ -14,7 +14,7 @@
use std::sync::Arc;
use datafusion::common::DFSchema;
use datafusion::common::{DFSchema, DataFusionError};
use datatypes::series::SeriesOrdering;
use datatypes::series::samples::{FieldSpec, SeriesSchema};
use datatypes::series::time::{EvalGrid, GridSpec};
@@ -28,7 +28,9 @@ use crate::extension_plan::{FusedSeries, FusedSeriesBinding};
use crate::info::InfoSelector;
use crate::map::MathOp;
use crate::range_eval::{RangeFunction, RangeWindow};
use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels};
use crate::vector_match::{
BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels, SetOp,
};
fn grid(start: i64, step: i64, len: u64) -> Arc<EvalGrid> {
Arc::new(EvalGrid::try_new(GridSpec::try_new(start, step, len).unwrap()).unwrap())
@@ -1345,7 +1347,10 @@ fn a_scalar_node_must_produce_a_scalar() {
fn an_input_must_report_the_binding_it_was_bound_to() {
let node = NativeNode {
expr: NativeExpr::Input(BoundInput {
output: raw(),
output: NativeType::Samples(Arc::new(SeriesSchema {
fields: fields(1),
eval_grid: None,
})),
source_index: 0,
ordering: SeriesOrdering::KeyThenTime,
}),
@@ -1372,17 +1377,105 @@ fn an_input_can_bind_aggregate_state() {
// A `Final` aggregate can then consume it, which is the shape of a distributed aggregate.
assert!(final_aggregate(bound).validate().is_ok());
// A fresh state allocation is a different binding, so it cannot be reported as the bound one.
let other = NativeNode {
// Equivalent independent allocations are accepted; field, order, grid, and group changes remain invalid.
let evaluated_a = evaluated(fields(2), grid(0, 60_000, 4));
let equivalent = NativeNode {
expr: NativeExpr::Input(BoundInput {
output: state.output.clone(),
output: evaluated_a.clone(),
source_index: 0,
ordering: SeriesOrdering::Contiguous,
}),
output: partial_aggregate(window_over(input(), None)).output,
output: evaluated(fields(2), grid(0, 60_000, 4)),
};
let error = other.validate().unwrap_err();
assert!(error.to_string().contains("not the binding"), "{error}");
assert!(equivalent.validate().is_ok());
let mismatches = [
evaluated(fields(1), grid(0, 60_000, 4)),
evaluated(
Arc::from([
FieldSpec {
name: Arc::from("f1"),
},
FieldSpec {
name: Arc::from("f0"),
},
]),
grid(0, 60_000, 4),
),
evaluated(fields(2), grid(60_000, 60_000, 4)),
];
for output in mismatches {
let mismatch = NativeNode {
expr: NativeExpr::Input(BoundInput {
output: evaluated_a.clone(),
source_index: 0,
ordering: SeriesOrdering::Contiguous,
}),
output,
};
assert!(mismatch.validate().is_err());
}
let state_a = partial_aggregate(window_over(input(), None)).output;
let NativeType::AggregateState(schema_a) = &state_a else {
unreachable!()
};
let state_group_mismatch = NativeType::AggregateState(Arc::new(super::AggregateStateSchema {
group: Grouping::Without(Arc::new([Arc::from("job")])),
accumulators: schema_a.accumulators.clone(),
grid: schema_a.grid.clone(),
}));
let mismatch = NativeNode {
expr: NativeExpr::Input(BoundInput {
output: state_a,
source_index: 0,
ordering: SeriesOrdering::Contiguous,
}),
output: state_group_mismatch,
};
assert!(mismatch.validate().is_err());
let state_grid = grid(0, 60_000, 4);
let state = NativeType::AggregateState(Arc::new(super::AggregateStateSchema {
group: Grouping::By(Arc::new([Arc::from("job")])),
accumulators: fields(2),
grid: Arc::clone(&state_grid),
}));
let equivalent_state = NativeType::AggregateState(Arc::new(super::AggregateStateSchema {
group: Grouping::By(Arc::new([Arc::from("job")])),
accumulators: fields(2),
grid: Arc::clone(&state_grid),
}));
let state_input = |bound: NativeType, output: NativeType| NativeNode {
expr: NativeExpr::Input(BoundInput {
output: bound,
source_index: 0,
ordering: SeriesOrdering::Contiguous,
}),
output,
};
assert!(
state_input(state.clone(), equivalent_state)
.validate()
.is_ok()
);
let reversed_accumulators = NativeType::AggregateState(Arc::new(super::AggregateStateSchema {
group: Grouping::By(Arc::new([Arc::from("job")])),
accumulators: Arc::from([
FieldSpec {
name: Arc::from("f1"),
},
FieldSpec {
name: Arc::from("f0"),
},
]),
grid: state_grid,
}));
assert!(
state_input(state.clone(), reversed_accumulators)
.validate()
.is_err()
);
}
/// A vector operation reads the values at each grid point, so raw samples, which have no shared points,
@@ -1400,14 +1493,14 @@ fn a_vector_operation_cannot_consume_raw_samples() {
assert!(error.to_string().contains("an evaluated vector"), "{error}");
}
/// A node that reports samples while holding no input node is refused where the node is built.
/// A root that reports samples while its expression is a constant breaks its own contract.
///
/// The tree is constructed directly here — a child module of `native::plan` may — because the node
/// is one the public constructors decline: a constant does not produce samples, and a zero-input
/// tree has no scan that would refuse it later. Every public constructor therefore stays a checked
/// one.
#[test]
fn a_samples_root_without_an_input_is_refused() {
fn a_scalar_root_claiming_samples_is_refused() {
let statement_grid = grid(0, 60_000, 4);
let tree = NativeTree {
root: NativeNode {
@@ -1437,6 +1530,27 @@ fn a_samples_root_without_an_input_is_refused() {
);
}
/// A samples root without an input is refused by FusedSeries.
#[test]
fn a_samples_root_without_an_input_is_refused() {
let tree = NativeTree::try_new(input()).unwrap();
let binding = FusedSeriesBinding {
grid: grid(0, 60_000, 4),
time_index: "time".to_string(),
tags: vec![],
metric_name: None,
};
let error = FusedSeries::try_new(
Arc::from("some_metric"),
tree,
binding,
Vec::new(),
Arc::new(DFSchema::empty()),
)
.expect_err("a samples root without an input is refused");
assert!(error.to_string().contains("source indexes"), "{error}");
}
/// `info` reads its info side's values as sample times, so it accepts only the instant selection
/// that reports them: one that reports the samples themselves would compare every info series as
/// equally new.
@@ -1476,3 +1590,138 @@ fn info_reads_the_sample_times_of_one_field() {
assert!(error.to_string().contains(expected), "{error}");
}
}
#[test]
fn native_tree_depth_limit_holds_across_lifecycle_on_named_worker_stack() {
use super::MAX_NATIVE_NODE_DEPTH;
fn nested_maps(map_count: usize) -> NativeTree {
let mut node = input_on(grid(0, 60_000, 1));
let output = node.output.clone();
for _ in 0..map_count {
node = NativeNode {
expr: NativeExpr::Map {
input: Box::new(node),
op: BoundMap {
ops: Arc::new([MapOp::Math(MathOp::Abs)]),
},
},
output: output.clone(),
};
}
NativeTree { root: node }
}
// A leaf counts as one node; test exactly the cap and rejection on a named 2 MiB worker.
std::thread::Builder::new()
.name("native-tree-depth-boundary".to_string())
.stack_size(2 * 1024 * 1024)
.spawn(|| {
let at_limit = nested_maps(MAX_NATIVE_NODE_DEPTH - 1);
let checked = NativeTree::try_new(at_limit.root.clone()).unwrap();
let bytes = postcard::to_stdvec(&checked).unwrap();
let decoded: NativeTree = postcard::from_bytes(&bytes).unwrap();
decoded.root().validate().unwrap();
drop((decoded, bytes, checked, at_limit));
// Construct wire data beyond the cap without passing through try_new.
let over_limit = nested_maps(MAX_NATIVE_NODE_DEPTH);
let bytes = postcard::to_stdvec(&over_limit).unwrap();
assert_eq!(
postcard::from_bytes::<NativeTree>(&bytes).unwrap_err(),
postcard::Error::SerdeDeCustom
);
assert!(matches!(
NativeTree::try_new(over_limit.root.clone()),
Err(DataFusionError::NotImplemented(_))
));
drop(over_limit);
})
.unwrap()
.join()
.unwrap();
}
#[test]
fn native_tree_postcard_float_fidelity_round_trip() {
for bits in [
0x7ff8_0000_0000_0001,
0,
0x8000_0000_0000_0000,
f64::INFINITY.to_bits(),
f64::NEG_INFINITY.to_bits(),
1,
] {
let value = f64::from_bits(bits);
let tree = NativeTree::try_new(scalar_on(value, grid(0, 1, 1))).unwrap();
let decoded: NativeTree =
postcard::from_bytes(&postcard::to_stdvec(&tree).unwrap()).unwrap();
let NativeExpr::Scalar(value) = &decoded.root().expr else {
panic!("scalar shape changed")
};
assert_eq!(value.value.to_bits(), bits);
}
}
#[test]
fn native_tree_postcard_rejects_trailing_bytes() {
let tree = NativeTree::try_new(input()).unwrap();
let mut bytes = postcard::to_stdvec(&tree).unwrap();
bytes.push(0);
let (_, remainder) = postcard::take_from_bytes::<NativeTree>(&bytes).unwrap();
assert!(!remainder.is_empty());
}
#[test]
fn native_tree_non_prefix_shapes_round_trip() {
let grid = grid(0, 60_000, 4);
let lhs = window_on(input(), None, Arc::clone(&grid));
let rhs = window_on(input_with_fields(0), None, Arc::clone(&grid));
let binary = NativeNode {
expr: NativeExpr::Binary {
left: Box::new(lhs.clone()),
right: Box::new(rhs.clone()),
matching: add(),
},
output: evaluated(Arc::new([]), Arc::clone(&grid)),
};
let aggregate = partial_aggregate(window_on(input(), None, Arc::clone(&grid)));
let set = NativeNode {
expr: NativeExpr::Set {
left: Box::new(lhs),
right: Box::new(rhs),
op: super::BoundSetOp {
op: SetOp::Or,
labels: MatchingLabels::Ignoring(Vec::new()),
},
},
output: evaluated(Arc::new([]), Arc::clone(&grid)),
};
let trees = vec![
NativeTree::try_new(input()).unwrap(),
NativeTree::try_new(input_on(Arc::clone(&grid))).unwrap(),
NativeTree::try_new(scalar_on(3.0, Arc::clone(&grid))).unwrap(),
NativeTree::try_new(binary).unwrap(),
NativeTree::try_new(aggregate).unwrap(),
NativeTree::try_new(set).unwrap(),
];
for tree in trees {
let decoded: NativeTree =
postcard::from_bytes(&postcard::to_stdvec(&tree).unwrap()).unwrap();
assert_eq!(decoded, tree);
}
}
#[test]
fn independently_allocated_binding_mismatch_is_rejected() {
let node = input_with_fields(1);
let mut mismatch = node.clone();
let NativeExpr::Input(binding) = &mut mismatch.expr else {
unreachable!()
};
binding.output = NativeType::Samples(Arc::new(SeriesSchema {
fields: fields(2),
eval_grid: None,
}));
assert!(mismatch.validate().is_err());
}
File diff suppressed because it is too large Load Diff
+2 -2
View File
@@ -45,7 +45,7 @@ pub(crate) use function::quantile;
/// `(E - range_ms, E]`, where `E` is `T - offset_ms`, or `anchor_ms - offset_ms` with `@`. A
/// modifier widens that: `anchored` also reads the lookback before `E - range_ms`, `smoothed` the
/// lookback on both sides, and a scan feeding the window has to read that much more.
#[derive(Clone, Copy, Debug)]
#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize, PartialEq, Hash, Eq)]
pub struct RangeWindow {
/// Positive window length in milliseconds.
pub range_ms: i64,
@@ -62,7 +62,7 @@ pub struct RangeWindow {
/// The experimental range selector modifiers of Prometheus, which read samples beyond the
/// window to pin its boundary values. Not to be confused with an `@` anchor, which fixes the
/// window's end.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum RangeModifier {
/// `anchored`: the window also reads the last sample at or before its start, up to the
/// statement's lookback before it, and takes it as the value at the start.
+35 -1
View File
@@ -23,7 +23,7 @@ use crate::go_sort;
use crate::range_eval::{BatchKernel, FloatWindow, FloatWindows, RangeModifier, WindowKernel};
/// A PromQL range-vector function over float and native histogram samples.
#[derive(Clone, Copy, Debug, PartialEq)]
#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
pub enum RangeFunction {
Rate,
Increase,
@@ -611,3 +611,37 @@ fn median_absolute_deviation(
deviations.extend(values.iter().map(|value| (value - median).abs()));
quantile(0.5, deviations, sorted)
}
impl PartialEq for RangeFunction {
fn eq(&self, o: &Self) -> bool {
match (self, o) {
(Self::QuantileOverTime(a), Self::QuantileOverTime(b))
| (Self::PredictLinear(a), Self::PredictLinear(b)) => a.to_bits() == b.to_bits(),
(
Self::DoubleExponentialSmoothing {
smoothing: a,
trend: b,
},
Self::DoubleExponentialSmoothing {
smoothing: c,
trend: d,
},
) => a.to_bits() == c.to_bits() && b.to_bits() == d.to_bits(),
_ => std::mem::discriminant(self) == std::mem::discriminant(o),
}
}
}
impl Eq for RangeFunction {}
impl std::hash::Hash for RangeFunction {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
std::mem::discriminant(self).hash(h);
match self {
Self::QuantileOverTime(x) | Self::PredictLinear(x) => x.to_bits().hash(h),
Self::DoubleExponentialSmoothing { smoothing, trend } => {
smoothing.to_bits().hash(h);
trend.to_bits().hash(h)
}
_ => (),
}
}
}
+2 -2
View File
@@ -44,7 +44,7 @@ pub enum Grouping {
}
/// The order `sort_instant` arranges an instant result in.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum SortOrder {
/// `sort`: ascending values, NaN last.
Asc,
@@ -371,7 +371,7 @@ pub struct LabelMatcher {
}
/// The operator of a [`LabelMatcher`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum MatchOp {
Equal,
NotEqual,
+12 -6
View File
@@ -68,7 +68,7 @@ pub(crate) fn visited_pairs() -> usize {
}
/// A PromQL binary operator between two instant vectors.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum BinaryOp {
Add,
Sub,
@@ -114,7 +114,7 @@ impl BinaryOp {
}
/// The labels two series are matched by.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum MatchingLabels {
/// `on(...)`: only these labels, `__name__` included when listed.
On(Vec<Arc<str>>),
@@ -124,7 +124,7 @@ pub enum MatchingLabels {
}
/// How many series of each side may share a match group.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum Cardinality {
OneToOne,
/// `group_left(...)`: the right side is the "one" side; the listed labels are copied from it.
@@ -134,7 +134,7 @@ pub enum Cardinality {
}
/// A vector–vector binary operation with its modifiers.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub struct BinaryMatching {
pub op: BinaryOp,
/// `bool`: comparisons return 1 or 0 instead of filtering.
@@ -147,7 +147,7 @@ pub struct BinaryMatching {
/// The values of `fill`, `fill_left` and `fill_right`, which stand in for a sample that has no
/// match on the other side: `lhs` for a missing left-hand sample and `rhs` for a missing
/// right-hand one, as the query writes them, `group_right` included (prometheus/prometheus#19939).
#[derive(Clone, Copy, Debug, Default)]
#[derive(Clone, Copy, Debug, Default, serde::Serialize, serde::Deserialize)]
pub struct FillValues {
pub lhs: Option<f64>,
pub rhs: Option<f64>,
@@ -162,9 +162,15 @@ impl PartialEq for FillValues {
}
impl Eq for FillValues {}
impl std::hash::Hash for FillValues {
fn hash<H: std::hash::Hasher>(&self, h: &mut H) {
std::hash::Hash::hash(&self.lhs.map(f64::to_bits), h);
std::hash::Hash::hash(&self.rhs.map(f64::to_bits), h);
}
}
/// A PromQL set operator between two instant vectors.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Hash)]
pub enum SetOp {
And,
Or,
+17 -26
View File
@@ -445,15 +445,12 @@ enum RewriterStatus {
#[derive(Debug, Default)]
struct PlanRewriter {
/// Whether the whole plan this rewriter walks can be encoded to Substrait.
/// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check.
///
/// Used for the ordinary standalone fast path when the root plan is encodable. It is not
/// needed when the root itself is the in-process `FusedSeries` exception.
/// Whether the whole plan encodes, allowing per-node checks to be skipped.
whole_plan_encodable: bool,
/// In-process transport skips the encodability gate only when the rewritten root is the native
/// route's whole plan: a `FusedSeries` node, or the `AggregateFinalize` merge of one's state.
skip_encodability_gate: bool,
/// In-process root shapes can contain AggregateFinalize, which has no Substrait serializer.
skip_in_process_encodability: bool,
/// Whether an in-process root fused plan was recognized for placement-gate evaluation.
in_process_fused_root: bool,
/// Current level in the tree
level: usize,
/// Simulated stack for the `rewrite` recursion
@@ -620,15 +617,15 @@ impl PlanRewriter {
LogicalPlan::Analyze(analyze) => analyze.input.as_ref(),
_ => plan,
};
let skip_encodability_gate = in_process_region_transport && is_in_process_plan(native_root);
let in_process_fused_root = in_process_region_transport && is_in_process_plan(native_root);
let whole_plan_encodable = DFLogicalSubstraitConvertor
.encode(plan, DefaultSerializer)
.is_ok();
let skip_in_process_encodability = in_process_fused_root && !whole_plan_encodable;
Self {
// Preserve the ordinary whole-plan fast path unless this is the in-process root
// exception: no node in that tree needs an encoding check.
skip_encodability_gate,
whole_plan_encodable: !skip_encodability_gate
&& DFLogicalSubstraitConvertor
.encode(plan, DefaultSerializer)
.is_ok(),
in_process_fused_root,
skip_in_process_encodability,
whole_plan_encodable,
..Default::default()
}
}
@@ -661,13 +658,7 @@ impl PlanRewriter {
// When the root does not encode, the plan holds at least one node that has to stay
// on the frontend and only the per-node check locates it.
//
// For the in-process root `FusedSeries` exception, encodability says nothing about where
// a node may run, and asking it would cut the plan at the first node with no serializer
// branch before `check_plan` could decide. Skipping it only removes a reason to expand: every
// check below still runs, and a node they call `NonCommutative`/`Unimplemented`/
// `Unsupported` still expands, so a node no serializer lists only rides to a region when
// the categorizer places it there.
if !self.skip_encodability_gate
if !self.skip_in_process_encodability
&& !self.whole_plan_encodable
&& let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer)
{
@@ -926,9 +917,9 @@ impl PlanRewriter {
// if there is a new child plan, use it as the new root
on_node = new_child_plan;
}
let bypass_multi_input_gate =
self.skip_encodability_gate && is_admitted_pair_plan(&on_node);
if bypass_multi_input_gate {
let has_fused_pair_placement_proof =
self.in_process_fused_root && is_admitted_pair_plan(&on_node);
if has_fused_pair_placement_proof {
self.column_requirements.clear();
} else {
let mut rewriter = EnforceDistRequirementRewriter::new(
+2 -2
View File
@@ -98,11 +98,11 @@ fn ordinary_encodable_plan_keeps_in_process_encodability_gate() {
let in_process_rewriter = PlanRewriter::new(&plan, true);
assert!(in_process_rewriter.whole_plan_encodable);
assert!(!in_process_rewriter.skip_encodability_gate);
assert!(!in_process_rewriter.skip_in_process_encodability);
let wire_rewriter = PlanRewriter::new(&plan, false);
assert!(wire_rewriter.whole_plan_encodable);
assert!(!wire_rewriter.skip_encodability_gate);
assert!(!wire_rewriter.skip_in_process_encodability);
}
struct MergeScanRemoteDynFilterProducerIdCollector<'a> {
+54 -97
View File
@@ -990,49 +990,44 @@ mod tests {
);
}
/// The encoding gate is what kept the node on the frontend, and it only exists for a plan that
/// has to be encoded.
///
/// `PlanRewriter::should_expand` asked two questions about a node before it let the descent
/// continue, and the Substrait one came first: unless the *whole* root plan encoded, every
/// subtree that failed to encode was expanded on, whichever way `check_plan` answered. A
/// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers `NotImplemented`
/// for every name it does not list — so over the wire the boundary is still drawn at the node
/// and the region receives the bare scan.
///
/// A plan handed to `RegionServer::handle_read` as a `LogicalPlan` is never encoded, so there
/// the gate has nothing to say and the commutativity arm decides: the node rides inside the
/// merge scan.
/// Ordinary fused producers are serializable; placement follows the in-process routing rules.
#[test]
fn fused_series_rides_to_the_region_only_when_regions_are_reached_in_process() {
fn fused_series_encoding_and_placement_are_independent() {
let plan = fused_series_plan(table_scan_plan());
let encoded = DFLogicalSubstraitConvertor
.encode(&plan, DefaultSerializer)
.expect("FusedSeries uses its ordinary serializer");
assert!(!encoded.is_empty());
assert!(
DFLogicalSubstraitConvertor
.encode(&plan, DefaultSerializer)
.is_err(),
"the fused node encodes now, so the encoding gate decides nothing"
);
assert!(
DFLogicalSubstraitConvertor
.encode(&table_scan_plan(), DefaultSerializer)
.is_ok(),
"the scan under the fused node encodes, so the fused node is the only reason this plan does not"
);
let in_process = analyze(DistPlannerAnalyzer::new(true), plan.clone());
assert!(any_remote_input_holds(&in_process, holds_fused_series));
// Over the wire the plan has to encode, so the boundary is drawn at the node.
let encoded = analyze(DistPlannerAnalyzer::new(false), plan.clone());
assert!(
!any_remote_input_holds(&encoded, holds_fused_series),
"the analyzer placed the node although its plan has to encode: {encoded}"
);
// In process nothing is encoded, so the categorizer places the node.
let in_process = analyze(DistPlannerAnalyzer::new(true), plan);
assert!(
any_remote_input_holds(&in_process, holds_fused_series),
"the analyzer kept the node on the frontend although regions are reached in process: {in_process}"
);
// A partial producer exports aggregate state; its merge remains frontend-side.
let split = aggregate_finalize_plan(table_scan_plan());
let LogicalPlan::Extension(split_extension) = &split else {
unreachable!()
};
let finalize = split_extension
.node
.as_any()
.downcast_ref::<AggregateFinalize>()
.unwrap();
let inputs = datafusion_expr::UserDefinedLogicalNodeCore::inputs(finalize);
let [producer] = inputs.as_slice() else {
unreachable!()
};
DFLogicalSubstraitConvertor
.encode(*producer, DefaultSerializer)
.expect("partial FusedSeries producer uses its ordinary serializer");
let placed = analyze(DistPlannerAnalyzer::new(true), split);
assert!(as_aggregate_finalize(&placed).is_some());
let child = placed.inputs();
let [child] = child.as_slice() else {
panic!("expected finalize input: {placed}")
};
let scan = first_merge_scan(child).expect("merge scan below finalize");
let producer = first_fused_series(scan.input()).expect("partial producer in region input");
assert!(producer.partial_fold().is_some());
}
/// A split aggregate's merge, over a scan and bound to `sum(some_metric)`.
@@ -1125,20 +1120,10 @@ mod tests {
));
}
/// The merge stays on the frontend and the producer rides inside the merge scan, so the plan a
/// region is handed is exactly the per-series prefix over its own scan.
///
/// Over the wire nothing changes: the merge has no serializer branch either, so the boundary is
/// still drawn at the producer, which is the same place the existing route draws it.
/// The merge stays on the frontend while its partial producer rides inside the merge scan.
#[test]
fn a_split_aggregate_merges_above_the_region_only_in_process() {
fn a_split_aggregate_merges_above_the_region_with_partial_producer() {
let plan = aggregate_finalize_plan(table_scan_plan());
assert!(
DFLogicalSubstraitConvertor
.encode(&plan, DefaultSerializer)
.is_err(),
"the merge encodes now, so the encoding gate decides nothing"
);
let in_process = analyze(DistPlannerAnalyzer::new(true), plan.clone());
assert!(
@@ -1165,48 +1150,15 @@ mod tests {
"the producer still reads its own scan in: {in_process}"
);
// Over the wire the plan has to encode, so the boundary is drawn at the producer.
// The merge itself has no Substrait serializer, so its frontend placement is unchanged.
let encoded = analyze(DistPlannerAnalyzer::new(false), plan);
assert!(
!any_remote_input_holds(&encoded, holds_fused_series),
"the analyzer placed the producer although its plan has to encode: {encoded}"
);
assert!(as_aggregate_finalize(&encoded).is_some());
}
/// The merge of `plan`, with its partition-local producer replaced by `producer`.
///
/// A swap like this is what a node has to tolerate: `with_exprs_and_inputs` rebuilds it from
/// whatever input an optimizer holds, so the analyzer cannot take a merge over a fused node to be
/// the one `AggregateFinalize::try_new` built.
fn merge_over(plan: LogicalPlan, producer: LogicalPlan) -> LogicalPlan {
let LogicalPlan::Extension(extension) = &plan else {
panic!("a split aggregate is an extension node: {plan}");
};
let finalize = extension
.node
.as_any()
.downcast_ref::<AggregateFinalize>()
.expect("a split aggregate is a merge");
// Qualified: both `UserDefinedLogicalNode` and `UserDefinedLogicalNodeCore` declare this
// method, and only the latter's answer is the node itself.
let swapped =
UserDefinedLogicalNodeCore::with_exprs_and_inputs(finalize, Vec::new(), vec![producer])
.expect("a merge holds the one input it is given");
LogicalPlan::Extension(Extension {
node: Arc::new(swapped),
})
}
/// The in-process exception admits the split the route builds, not the shape it looks like: a
/// merge whose child is an ordinary `FusedSeries` - one that answers with computed values - over
/// the same scan is refused the exception, so the encoding gate still draws the boundary at the
/// scan and the merge stays on the frontend.
///
/// Skipping the gate for it would let a region's values be merged as if they were that
/// partition's contribution to the group, which is the shape no constructor check can rule out:
/// `with_exprs_and_inputs` accepts the replacement.
/// A merge admits only a producer that exports partial aggregate state. A computed-values
/// producer must be refused by `with_exprs_and_inputs`.
#[test]
fn a_merge_over_an_ordinary_fused_series_keeps_the_encoding_gate() {
fn a_merge_requires_a_partial_state_producer_not_computed_values() {
let scan = table_scan_plan();
let split = aggregate_finalize_plan(scan.clone());
// The control: the split's own producer exports state, so it is the shape the exception
@@ -1218,16 +1170,21 @@ mod tests {
);
// The same merge over the same scan, with a producer that answers with values.
let LogicalPlan::Extension(extension) = &split else {
unreachable!()
};
let finalize = extension
.node
.as_any()
.downcast_ref::<AggregateFinalize>()
.unwrap();
let values = fused_series_plan(scan);
let swapped = merge_over(split, values);
let analyzed = analyze(DistPlannerAnalyzer::new(true), swapped);
let error =
UserDefinedLogicalNodeCore::with_exprs_and_inputs(finalize, Vec::new(), vec![values])
.unwrap_err();
assert!(
!any_remote_input_holds(&analyzed, holds_fused_series),
"an ordinary fused node exports no state, so the exception does not apply: {analyzed}"
);
assert!(
any_remote_input_holds(&analyzed, holds_table_scan),
"the region still reads the scan the producer reads: {analyzed}"
error.to_string().contains("must export aggregate state"),
"{error}"
);
}
+49
View File
@@ -13,6 +13,7 @@
// limitations under the License.
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use catalog::RegisterTableRequest;
@@ -2596,6 +2597,54 @@ async fn native_fused_whole_vector_operation_falls_back() {
}
}
#[tokio::test]
async fn native_fused_depth_limit_falls_back_to_legacy_plan() {
// At 63 selector terms the tree reaches depth 64; the 64th exceeds the limit.
let schema = Arc::new(datatypes::series::samples::SeriesSchema {
fields: Arc::from([datatypes::series::samples::FieldSpec { name: "v0".into() }]),
eval_grid: None,
});
let compile = |terms| {
let query = std::iter::repeat_n("some_metric", terms)
.collect::<Vec<_>>()
.join(" + ");
crate::promql::native_compile::compile_selectors(
&build_eval_stmt(&query),
&|_| Ok(Arc::clone(&schema)),
SeriesOrdering::KeyThenTime,
)
};
compile(63).expect("63 selector terms must fit native tree depth 64");
let error = compile(64).expect_err("64 selector terms must exceed native tree depth 64");
assert!(
matches!(
&error,
datafusion::error::DataFusionError::NotImplemented(_)
),
"expected a NotImplemented depth-limit rejection, got: {error}"
);
assert!(
error.to_string().contains("native tree nesting exceeds 64"),
"unexpected compiler rejection: {error}"
);
let query = std::iter::repeat_n("some_metric", 64)
.collect::<Vec<_>>()
.join(" + ");
let native = build_native_fused_plan(&query, true).await;
let legacy = build_native_fused_plan(&query, false).await;
let display = native.display_indent().to_string();
assert!(
!display.contains("FusedSeries"),
"64 selector terms exceed the native tree depth cap: {display}"
);
assert_eq!(
display,
legacy.display_indent().to_string(),
"a native depth-limit rejection must preserve the legacy plan"
);
}
#[tokio::test]
async fn native_fused_unsupported_shape_falls_back() {
for query in [
+224 -14
View File
@@ -65,7 +65,7 @@ fn evaluated_frame_schema(tree: &NativeTree) -> datafusion::common::DFSchemaRef
// a possible empty closing chunk, so it does not exercise a chain spanning many batches. 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() {
async fn native_prefix_multi_partition_single_position_chains() {
let instance = GreptimeDbStandaloneBuilder::new("native_prefix_partitioned_chains")
.with_query_options(QueryOptions {
parallelism: 16,
@@ -234,9 +234,8 @@ async fn native_prefix_multi_partition_series_chains() {
entry
.1
.extend(field.iter().filter_map(|sample| match sample {
datatypes::series::samples::SampleRef::Float(value) => Some(value),
datatypes::series::samples::SampleRef::Histogram(_)
| datatypes::series::samples::SampleRef::Missing => None,
SampleRef::Float(value) => Some(value),
SampleRef::Histogram(_) | SampleRef::Missing => None,
}));
}
}
@@ -251,6 +250,209 @@ async fn native_prefix_multi_partition_series_chains() {
}
}
// The one-position case above cannot exercise a chain that spans chunks. This 8,192-position grid
// across 16 region-scan partitions does: whatever chunk boundaries the merge produces, each host's
// chain must arrive complete, in time order, and close exactly once.
#[tokio::test(flavor = "multi_thread")]
async fn native_prefix_multi_partition_multi_chunk_chains() {
let instance = GreptimeDbStandaloneBuilder::new("native_prefix_partitioned_multi_chunk_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(10_000, 1_000, 8_192).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 decoded = common_query::series_frame::unpack_evaluated_series(
Box::pin(DfRecordBatchStreamAdapter::new(batches.as_stream())),
None,
)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut per_host = std::collections::BTreeMap::<String, (usize, Vec<i64>, Vec<f64>)>::new();
for (series, is_last) in decoded {
let host = series
.metric
.get("host")
.unwrap_or_else(|| panic!("a decoded series keeps its host label: {:?}", series.metric))
.to_string();
let entry = per_host.entry(host).or_default();
if is_last {
entry.0 += 1;
}
if !series.time.is_empty() {
entry
.1
.extend_from_slice(series.time.timestamps().unwrap().as_ref());
}
for field in series.fields.iter() {
entry
.2
.extend(field.iter().filter_map(|sample| match sample {
SampleRef::Float(value) => Some(value),
SampleRef::Histogram(_) | SampleRef::Missing => None,
}));
}
}
assert_eq!(16, per_host.len(), "one chain per host: {per_host:?}");
let expected_timestamps = (0..8_192).map(|i| 10_000 + i * 1_000).collect::<Vec<_>>();
for (host, (closes, timestamps, values)) in &per_host {
assert_eq!(1, *closes, "`{host}` closes its chain exactly once");
assert_eq!(
8_192,
timestamps.len(),
"`{host}` keeps every grid position across chunk boundaries"
);
assert_eq!(
&expected_timestamps, timestamps,
"`{host}` arrives in order over its whole contiguous grid"
);
assert_eq!(8_192, values.len(), "`{host}` keeps every value");
for value in values {
assert!(
(value - 1.0).abs() < 1e-12,
"`{host}` reports the rate of its own samples, which a cut chain would lose: {value}"
);
}
}
}
#[tokio::test(flavor = "multi_thread")]
async fn native_prefix_executes_over_real_mito_region() {
let instance = GreptimeDbStandaloneBuilder::new("native_prefix_region_scan")
@@ -422,19 +624,27 @@ async fn native_prefix_executes_over_real_mito_region() {
);
}
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
// Locate the complete FusedSeries envelope in the Substrait bytes. Its version is the first
// word and its payload length is the second; no payload-field offsets are needed here.
let envelope_at = tampered
.windows(fused_payload.len())
.position(|window| window == fused_payload.as_slice())
.expect("FusedSeries payload in Substrait plan");
let prefix_version = embedded + 8;
.expect("FusedSeries envelope in Substrait plan");
let version_bytes: [u8; 4] = tampered[envelope_at..envelope_at + 4]
.try_into()
.unwrap_or_else(|_| panic!("missing version bytes at {envelope_at}"));
let version = u32::from_le_bytes(version_bytes);
assert_eq!(version, 3, "observed version bytes: {version_bytes:?}");
let length_bytes: [u8; 4] = tampered[envelope_at + 4..envelope_at + 8]
.try_into()
.unwrap_or_else(|_| panic!("missing length bytes at {}", envelope_at + 4));
let payload_length = u32::from_le_bytes(length_bytes);
assert_eq!(
&tampered[prefix_version..prefix_version + 4],
&native_prefix_version
payload_length as usize,
fused_payload.len() - 8,
"observed length bytes: {length_bytes:?}"
);
tampered[prefix_version] = 2;
tampered[envelope_at] = 99;
let error = server
.handle_remote_read(
QueryRequest {
@@ -451,7 +661,7 @@ async fn native_prefix_executes_over_real_mito_region() {
.err()
.expect("invalid prefix rejected");
assert!(
format!("{error:?}").contains("unknown version"),
format!("{error:?}").contains("unsupported FusedSeries envelope version"),
"{error:?}"
);
}