Skip to content

Commit 5e551ce

Browse files
refactor(query-engine): execute native range query DAGs (#745)
* test(query-engine): characterize native range leaf execution * refactor(query-engine): validate native query DAG dependencies * refactor(query-engine): isolate range query store reads * refactor(query-engine): execute ranges from prepared reads * refactor(query-engine): execute native range DAGs * test(query-engine): cover native DAG execution order * refactor(query-engine): expose native PromQL execution errors * fix(query-engine): return native execution failures locally * fix(query-engine): return range execution failures locally * refactor(query-engine): stage native range DAG execution * test(query-engine): cover native DAG differential cases * test(query-engine): align off-grid rate fixture * test(query-engine): keep differential cases leaf-only * test(query-engine): keep temporal fixture planner-compatible * test(query-engine): restore temporal differential fixture * test(query-engine): add temporal DAG differential cases * test(query-engine): cover off-grid rate fallback * test(query-engine): isolate off-grid rate differential * test(query-engine): add native DAG differential suites * test(query-engine): add legacy range execution selector * test(query-engine): compare legacy and DAG range execution * test(query-engine): characterize sparse legacy DAG parity * test(query-engine): characterize keyed legacy DAG parity * test(query-engine): characterize topk legacy DAG parity * test(query-engine): cover malformed native plan errors * test(query-engine): cover native store failures * test(query-engine): characterize topk tie parity * test(query-engine): compare DeltaSet replay DAG execution * test(query-engine): cover local native execution failures * docs(query-engine): record native DAG follow-up work * fix(compliance): isolate differential compose runs * fix(compliance): seed differential data in the past * feat(query-engine): trace native DAG execution * fix(query-engine): preserve grouped quantile labels * refactor(query-engine): preserve typed DAG execution errors * test(compliance): drop removed sparse fixture * test(compliance): isolate Docker ports per matrix case * docs(query-engine): record native DAG matrix results * docs(query-engine): correct DAG matrix classification * fix(query-engine): preserve topk grouping * test(query-engine): cover grouped topk DAG parity * test(query-engine): exercise native DAG failure paths * docs(query-engine): record DAG validation matrix * fix(query-engine): support current Clippy lint
1 parent 4138a7f commit 5e551ce

20 files changed

Lines changed: 1718 additions & 185 deletions

‎asap-query-engine/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,4 @@ jemalloc = ["dep:tikv-jemallocator"]
9797
lock_profiling = []
9898
# Enable extra debugging output
9999
extra_debugging = []
100+
native_query_legacy_test_support = []

‎asap-query-engine/src/drivers/query/adapters/traits.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ impl std::error::Error for AdapterError {}
6060

6161
/// Trait for parsing incoming HTTP requests into internal query format
6262
/// Handles Axum extractors directly for different request types (GET/POST)
63+
#[allow(clippy::double_must_use)]
6364
#[async_trait]
6465
pub trait QueryRequestAdapter: Send + Sync {
6566
/// Parse a GET request with query parameters
@@ -106,6 +107,7 @@ pub trait QueryRequestAdapter: Send + Sync {
106107
}
107108

108109
/// Trait for formatting query results into protocol-specific HTTP responses
110+
#[allow(clippy::double_must_use)]
109111
#[async_trait]
110112
pub trait QueryResponseAdapter: Send + Sync {
111113
/// Format a successful query result into protocol response
@@ -135,6 +137,7 @@ pub trait QueryResponseAdapter: Send + Sync {
135137
/// define separate adapter traits.
136138
///
137139
/// Note: Fallback logic is handled separately via FallbackClient
140+
#[allow(clippy::double_must_use)]
138141
#[async_trait]
139142
pub trait HttpProtocolAdapter: QueryRequestAdapter + QueryResponseAdapter + Send + Sync {
140143
/// Get a descriptive name for this adapter (for logging/debugging)

‎asap-query-engine/src/drivers/query/fallback/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ impl IntoResponse for FallbackResponse {
3737
}
3838

3939
/// Client for forwarding unsupported queries to a fallback backend
40+
#[allow(clippy::double_must_use)]
4041
#[async_trait]
4142
pub trait FallbackClient: Send + Sync {
4243
/// Execute a query against the fallback backend

‎asap-query-engine/src/engines/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,7 @@ pub(crate) mod sliding_window_composition;
66
pub mod window_merger;
77

88
pub use query_result::{InstantVector, QueryResult, RangeVector, RangeVectorElement, Sample};
9+
#[cfg(feature = "native_query_legacy_test_support")]
10+
pub use simple_engine::NativeRangeExecutionMode;
911
pub use simple_engine::{QueryExecutionError, SimpleEngine};
1012
pub use window_merger::{create_window_merger, NaiveMerger, WindowMerger};

‎asap-query-engine/src/engines/query_plan.rs‎

Lines changed: 213 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22
33
use crate::engines::simple_engine::{RangeQueryExecutionContext, StoreQueryParams};
44
use asap_types::enums::WindowType;
5+
use promql_utilities::data_model::KeyByLabelNames;
56
use promql_utilities::query_logics::enums::Statistic;
7+
use tracing::debug;
68

79
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
810
pub(crate) struct NodeId(usize);
@@ -38,10 +40,12 @@ pub(crate) enum QueryPlanNode {
3840
LimitTopK {
3941
input: NodeId,
4042
k: String,
43+
grouping_labels: KeyByLabelNames,
4144
},
4245
Format {
4346
input: NodeId,
4447
include_metric_name: bool,
48+
metric: String,
4549
},
4650
}
4751

@@ -57,7 +61,44 @@ pub(crate) struct PlanOptions {
5761
pub format_output: bool,
5862
}
5963

64+
pub(crate) trait QueryPlanRuntime {
65+
type Output: Clone;
66+
type Error: std::fmt::Display;
67+
68+
fn execute_node(
69+
&self,
70+
id: NodeId,
71+
node: &QueryPlanNode,
72+
inputs: &[Self::Output],
73+
) -> Result<Self::Output, Self::Error>;
74+
}
75+
76+
#[derive(Debug)]
77+
pub(crate) enum QueryPlanExecutionError<E> {
78+
InvalidPlan(String),
79+
Node { id: NodeId, source: E },
80+
}
81+
82+
impl<E: std::fmt::Display> std::fmt::Display for QueryPlanExecutionError<E> {
83+
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84+
match self {
85+
Self::InvalidPlan(error) => write!(formatter, "invalid query plan: {error}"),
86+
Self::Node { id, source } => {
87+
write!(formatter, "Query plan node n{} failed: {source}", id.0)
88+
}
89+
}
90+
}
91+
}
92+
6093
impl QueryPlan {
94+
#[cfg(feature = "native_query_legacy_test_support")]
95+
pub(crate) fn malformed_for_test() -> Self {
96+
Self {
97+
nodes: Vec::new(),
98+
root: NodeId(0),
99+
}
100+
}
101+
61102
pub(crate) fn compile_range(
62103
context: &RangeQueryExecutionContext,
63104
options: PlanOptions,
@@ -114,18 +155,90 @@ impl QueryPlan {
114155
.ok_or_else(|| "Topk query is missing required `k` parameter".to_string())?;
115156
k.parse::<usize>()
116157
.map_err(|_| "Topk query has an invalid `k` parameter".to_string())?;
117-
root = Self::push(&mut nodes, QueryPlanNode::LimitTopK { input: root, k });
158+
root = Self::push(
159+
&mut nodes,
160+
QueryPlanNode::LimitTopK {
161+
input: root,
162+
k,
163+
grouping_labels: context.base.grouping_labels.clone(),
164+
},
165+
);
118166
}
119167
if options.format_output {
120168
root = Self::push(
121169
&mut nodes,
122170
QueryPlanNode::Format {
123171
input: root,
124-
include_metric_name: context.base.metadata.keep_metric_name,
172+
include_metric_name: context.base.metadata.statistic_to_compute
173+
== Statistic::Topk
174+
&& context.base.metadata.keep_metric_name,
175+
metric: context.base.metric.clone(),
125176
},
126177
);
127178
}
128-
Ok(Self { nodes, root })
179+
let plan = Self { nodes, root };
180+
plan.validate()?;
181+
Ok(plan)
182+
}
183+
184+
/// Rejects plans whose node dependencies cannot be executed safely.
185+
pub(crate) fn validate(&self) -> Result<(), String> {
186+
if self.nodes.is_empty() {
187+
return Err("Query plan has no nodes".to_string());
188+
}
189+
if self.root.0 != self.nodes.len() - 1 {
190+
return Err(format!(
191+
"Query plan root n{} does not include every node",
192+
self.root.0
193+
));
194+
}
195+
for (index, node) in self.nodes.iter().enumerate() {
196+
for input in node.inputs() {
197+
if input.0 >= index {
198+
return Err(format!(
199+
"Query plan node n{index} references unavailable input n{}",
200+
input.0
201+
));
202+
}
203+
}
204+
}
205+
Ok(())
206+
}
207+
208+
pub(crate) fn execute<R: QueryPlanRuntime>(
209+
&self,
210+
runtime: &R,
211+
) -> Result<R::Output, QueryPlanExecutionError<R::Error>> {
212+
self.validate()
213+
.map_err(QueryPlanExecutionError::InvalidPlan)?;
214+
let mut outputs: Vec<R::Output> = Vec::with_capacity(self.nodes.len());
215+
for (index, node) in self.nodes.iter().enumerate() {
216+
let inputs = node
217+
.inputs()
218+
.into_iter()
219+
.map(|input| outputs[input.0].clone())
220+
.collect::<Vec<_>>();
221+
debug!(
222+
node_id = index,
223+
node_kind = node.kind(),
224+
input_count = inputs.len(),
225+
"Executing native query plan node"
226+
);
227+
let output = runtime
228+
.execute_node(NodeId(index), node, &inputs)
229+
.map_err(|source| QueryPlanExecutionError::Node {
230+
id: NodeId(index),
231+
source,
232+
})?;
233+
debug!(
234+
node_id = index,
235+
node_kind = node.kind(),
236+
"Completed native query plan node"
237+
);
238+
outputs.push(output);
239+
}
240+
debug!(root_node_id = self.root.0, "Completed native query plan");
241+
Ok(outputs[self.root.0].clone())
129242
}
130243

131244
fn push(nodes: &mut Vec<QueryPlanNode>, node: QueryPlanNode) -> NodeId {
@@ -194,9 +307,11 @@ impl QueryPlan {
194307
kwargs.sort_unstable_by_key(|(key, _)| *key);
195308
format!("n{index} Estimate(n{}, {statistic}, {kwargs:?})", input.0)
196309
},
197-
QueryPlanNode::LimitTopK { input, k } => format!("n{index} LimitTopK(n{}, k={k})", input.0),
198-
QueryPlanNode::Format { input, include_metric_name } => format!(
199-
"n{index} Format(n{}, include_metric_name={include_metric_name})", input.0
310+
QueryPlanNode::LimitTopK { input, k, .. } => {
311+
format!("n{index} LimitTopK(n{}, k={k})", input.0)
312+
}
313+
QueryPlanNode::Format { input, include_metric_name, metric } => format!(
314+
"n{index} Format(n{}, include_metric_name={include_metric_name}) metric={metric}", input.0
200315
),
201316
};
202317
lines.push(line);
@@ -206,13 +321,43 @@ impl QueryPlan {
206321
}
207322
}
208323

324+
impl QueryPlanNode {
325+
fn kind(&self) -> &'static str {
326+
match self {
327+
Self::StoreRead { .. } => "StoreRead",
328+
Self::ComposeWindows { .. } => "ComposeWindows",
329+
Self::ResolveKeys { .. } => "ResolveKeys",
330+
Self::Estimate { .. } => "Estimate",
331+
Self::LimitTopK { .. } => "LimitTopK",
332+
Self::Format { .. } => "Format",
333+
}
334+
}
335+
336+
fn inputs(&self) -> Vec<NodeId> {
337+
match self {
338+
Self::StoreRead { .. } => Vec::new(),
339+
Self::ComposeWindows { input, .. }
340+
| Self::Estimate { input, .. }
341+
| Self::LimitTopK { input, .. }
342+
| Self::Format { input, .. } => vec![*input],
343+
Self::ResolveKeys { values, keys } => {
344+
keys.iter().copied().fold(vec![*values], |mut inputs, key| {
345+
inputs.push(key);
346+
inputs
347+
})
348+
}
349+
}
350+
}
351+
}
352+
209353
#[cfg(test)]
210354
mod tests {
211355
use super::*;
212356
use crate::data_model::AggregationIdInfo;
213357
use crate::engines::simple_engine::{QueryExecutionContext, QueryMetadata, StoreQueryPlan};
214358
use promql_utilities::data_model::KeyByLabelNames;
215359
use promql_utilities::query_logics::enums::AggregationType;
360+
use std::cell::RefCell;
216361
use std::collections::HashMap;
217362

218363
fn context() -> RangeQueryExecutionContext {
@@ -347,4 +492,66 @@ mod tests {
347492

348493
assert_eq!(error, "Topk query is missing required `k` parameter");
349494
}
495+
496+
#[test]
497+
fn rejects_a_node_that_references_a_later_node() {
498+
let plan = QueryPlan {
499+
nodes: vec![QueryPlanNode::Estimate {
500+
input: NodeId(1),
501+
statistic: Statistic::Sum,
502+
query_kwargs: HashMap::new(),
503+
}],
504+
root: NodeId(0),
505+
};
506+
507+
assert_eq!(
508+
plan.validate().expect_err("invalid plan must fail loudly"),
509+
"Query plan node n0 references unavailable input n1"
510+
);
511+
}
512+
513+
struct RecordingRuntime(RefCell<Vec<usize>>);
514+
515+
impl QueryPlanRuntime for RecordingRuntime {
516+
type Output = usize;
517+
type Error = std::convert::Infallible;
518+
519+
fn execute_node(
520+
&self,
521+
id: NodeId,
522+
_node: &QueryPlanNode,
523+
inputs: &[Self::Output],
524+
) -> Result<Self::Output, Self::Error> {
525+
self.0.borrow_mut().push(id.0);
526+
Ok(1 + inputs.iter().sum::<usize>())
527+
}
528+
}
529+
530+
#[test]
531+
fn executes_nodes_once_in_dependency_order() {
532+
let plan = QueryPlan {
533+
nodes: vec![
534+
QueryPlanNode::StoreRead {
535+
query: StoreQueryParams {
536+
metric: "requests".into(),
537+
aggregation_id: 7,
538+
start_timestamp: 0,
539+
end_timestamp: 1,
540+
},
541+
strategy: StoreReadStrategy::WindowGrid,
542+
},
543+
QueryPlanNode::ComposeWindows {
544+
input: NodeId(0),
545+
output_timestamps: vec![1],
546+
lookback_ms: 1,
547+
window_size_ms: 1,
548+
bucket_step_ms: 1,
549+
},
550+
],
551+
root: NodeId(1),
552+
};
553+
let runtime = RecordingRuntime(RefCell::new(Vec::new()));
554+
assert_eq!(plan.execute(&runtime).unwrap(), 2);
555+
assert_eq!(*runtime.0.borrow(), vec![0, 1]);
556+
}
350557
}

0 commit comments

Comments
 (0)