Skip to content

refactor(query-engine): make native DAG execution self-contained #788

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
milindsrivastava1997 wants to merge 4 commits into
base: main
Choose a base branch
from
Open

refactor(query-engine): make native DAG execution self-contained #788

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
337ae1f
refactor(query-engine): make native DAG execution self-contained
milindsrivastava1997 Oct 5, 2026
b53ad94
refactor(query-engine): make range plan specs cohesive
milindsrivastava1997 Oct 8, 2026
3539309
fix(query-engine): validate native range plan invariants
milindsrivastava1997 Oct 11, 2026
659bff8
refactor(query-engine): derive key reads from plan spec
milindsrivastava1997 Oct 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
238 changes: 167 additions & 71 deletions asap-query-engine/src/engines/query_plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,30 +5,78 @@ use crate::engines::simple_engine::{RangeQueryExecutionContext, StoreQueryParams
use asap_types::enums::WindowType;
use asap_types::query_config::QueryTimeAggregation;
use promql_utilities::data_model::KeyByLabelNames;
use promql_utilities::query_logics::enums::Statistic;
use promql_utilities::query_logics::enums::{AggregationType, Statistic};
use tracing::debug;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct NodeId(usize);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum StoreReadStrategy {
WindowGrid,
SlidingExactCover,
SlidingExactCover {
output_timestamps: Vec<u64>,
lookback_ms: u64,
window_size_ms: u64,
bucket_step_ms: u64,
},
}

impl StoreReadStrategy {
fn for_window(
window_type: WindowType,
output_timestamps: &[u64],
lookback_ms: u64,
window_size_ms: u64,
bucket_step_ms: u64,
) -> Self {
match window_type {
WindowType::Tumbling => Self::WindowGrid,
WindowType::Sliding => Self::SlidingExactCover {
output_timestamps: output_timestamps.to_vec(),
lookback_ms,
window_size_ms,
bucket_step_ms,
},
}
}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum StoreReadRole {
Values,
Keys,
}

#[derive(Debug, Clone)]
pub(crate) struct RangeEstimateSpec {

@milindsrivastava1997 milindsrivastava1997 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#754: "Do not copy the existing context wholesale into a node: extract cohesive plan-owned specs with explicit invariants."

This struct is close to a field-for-field copy of RangeQueryExecutionContext, including four parallel keys_* Options that must be all-Some or all-None. Nothing enforces that, so the estimator .expect()s each one separately and the compiler falls back with unwrap_or.

Suggest a single keys: Option<KeyWindowSpec { window_type, window_size_ms, lookback_ms, bucket_step_ms }> (and the same for the value window). Build it once in compile_range, failing if any field is missing, and use it for both the key StoreRead strategy and the estimator. That makes the half-set state unrepresentable, and the read and estimate can't disagree.

pub output_timestamps: Vec<u64>,
pub query_range_ms: u64,
pub buckets_per_step: usize,
pub lookback_bucket_count: usize,
pub tumbling_window_ms: u64,
pub window_type: WindowType,
pub window_size_ms: u64,
pub keys_window_type: Option<WindowType>,

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The keys-window settings are stored twice. RangeEstimateSpec copies keys_window_type, keys_window_size_ms, keys_lookback_ms, keys_tumbling_window_ms and output_timestamps, which the keys StoreRead strategy already holds.

The read resolves these with unwrap_or defaults, while the estimator calls .expect() on its own copies. If one side changes and the other does not, the read fetches one set of windows and the estimator looks for buckets from another, giving silently incomplete covers.

pub keys_window_size_ms: Option<u64>,
pub keys_lookback_ms: Option<u64>,
pub keys_tumbling_window_ms: Option<u64>,
pub value_aggregation_type: AggregationType,
pub key_aggregation_type: AggregationType,
pub grouping_labels: KeyByLabelNames,
pub aggregated_labels: KeyByLabelNames,
pub row_label_order: KeyByLabelNames,
}

#[derive(Debug, Clone)]
pub(crate) enum QueryPlanNode {
StoreRead {
query: StoreQueryParams,
role: StoreReadRole,
strategy: StoreReadStrategy,
},
ComposeWindows {
PrepareBuckets {
input: NodeId,
output_timestamps: Vec<u64>,
lookback_ms: u64,
window_size_ms: u64,
bucket_step_ms: u64,
},
ResolveKeys {
values: NodeId,
Expand All @@ -39,6 +87,7 @@ pub(crate) enum QueryPlanNode {
statistic: Statistic,
query_kwargs: std::collections::HashMap<String, String>,
output_labels: KeyByLabelNames,
spec: RangeEstimateSpec,
},
AggregateVector {
input: NodeId,
Expand All @@ -49,6 +98,7 @@ pub(crate) enum QueryPlanNode {
input: NodeId,
k: String,
grouping_labels: KeyByLabelNames,
row_label_order: KeyByLabelNames,
},
Format {
input: NodeId,
Expand Down Expand Up @@ -113,46 +163,53 @@ impl QueryPlan {
query_time_aggregations: &[QueryTimeAggregation],
) -> Result<Self, String> {
let mut nodes = Vec::new();
let values_lookback_ms =

@milindsrivastava1997 milindsrivastava1997 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lookback_bucket_count * tumbling_window_ms is now computed here, in estimate_range_query (mod.rs:2721), and in legacy read_range_query_inputs. The read strategy and the estimator each hold their own copy of lookback, window size and step.

If one copy changes (e.g. the bucket-lookback rounding this PR just fixed), the reads fetch a different cover than the estimator expects, and steps get silently skipped as "incomplete Sliding value cover". This works against #754's "each node consumes the fields that define its semantics": let the estimator take the window spec from the plan instead of recomputing it.

(context.lookback_bucket_count as u64) * context.tumbling_window_ms;
let values_read = Self::push_read(
&mut nodes,
&context.base.store_plan.values_query,
context.window_type,
);
let values = Self::push_compose(
&mut nodes,
values_read,
&context.output_timestamps,
context.query_range_ms,
context.window_size_ms,
context.tumbling_window_ms,
StoreReadRole::Values,
StoreReadStrategy::for_window(
context.window_type,
&context.output_timestamps,
values_lookback_ms,
context.window_size_ms,
context.tumbling_window_ms,
),
);
let values = Self::push_prepare_buckets(&mut nodes, values_read);
let keys = context.base.store_plan.keys_query.as_ref().map(|query| {
let keys_lookback_ms = context.keys_lookback_ms.unwrap_or(context.query_range_ms);

@milindsrivastava1997 milindsrivastava1997 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Silent defaults for required key-window config. keys_lookback_ms.unwrap_or(query_range_ms), keys_window_size_ms.unwrap_or(window_size_ms), keys_tumbling_window_ms.unwrap_or(tumbling_window_ms) and (below) keys_window_type.unwrap_or(window_type) replace the explicit errors the legacy reader returned ("Sliding keys query is missing its lookback", etc.).

Both production builders (simple_engine/mod.rs:870-920, promql.rs:655-702) set all four keys_* fields together, so this can't be reached today. But if a future builder leaves one unset, the DAG reads the key aggregation with the value aggregation's W/S/lookback instead of failing. It may also pick SlidingExactCover where legacy would have used WindowGrid, and estimate_range_query would then panic on keys_lookback_ms.expect(...). This breaks the "fail loud" rule (code-design-review §1: "config resolution that silently picks a default when a required value is missing").

Suggested fix: see the comment on RangeEstimateSpec. A single Option<KeyWindowSpec> built once at compile time removes these fallbacks and the downstream .expects.

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing keys settings now fall back silently. The keys StoreRead resolves keys_lookback_ms, keys_window_size_ms and keys_tumbling_window_ms with unwrap_or fallbacks to the values query's settings. The removed read_range_query_inputs returned explicit errors in these cases.

Example: with keys_window_type=Some(Sliding) and keys_lookback_ms=None, the old code failed with "Sliding keys query is missing its lookback". Now the read plans a cover from the values' query_range_ms, window size and step, so it fetches the wrong windows. estimate_range_query then panics on keys_lookback_ms.expect(...). These fallbacks used to sit only in the explain-only ComposeWindows node; now they decide real store reads. This breaks the fail-loud rule in the code-design-review checklist.

let keys_window_size_ms = context
.keys_window_size_ms
.unwrap_or(context.window_size_ms);
let keys_bucket_step_ms = context
.keys_tumbling_window_ms
.unwrap_or(context.tumbling_window_ms);
let read = Self::push_read(
&mut nodes,
query,
context.keys_window_type.unwrap_or(context.window_type),
StoreReadRole::Keys,
StoreReadStrategy::for_window(
context.keys_window_type.unwrap_or(context.window_type),

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The DAG and the legacy oracle choose different keys read strategies. Here the strategy comes from keys_window_type.unwrap_or(context.window_type); the legacy oracle checks keys_window_type == Some(Sliding).

When values are Sliding, keys_query is Some and keys_window_type is None, the DAG does a SlidingExactCover read and the oracle does a WindowGrid read. They read different buckets, so the differential oracle no longer checks the same logic. Both should derive the strategy from one shared function.

&context.output_timestamps,
keys_lookback_ms,
keys_window_size_ms,
keys_bucket_step_ms,
),
);
Self::push_compose(
&mut nodes,
read,
&context.output_timestamps,
context.keys_lookback_ms.unwrap_or(context.query_range_ms),
context
.keys_window_size_ms
.unwrap_or(context.window_size_ms),
context
.keys_tumbling_window_ms
.unwrap_or(context.tumbling_window_ms),
)
Self::push_prepare_buckets(&mut nodes, read)
});
let resolved = Self::push(&mut nodes, QueryPlanNode::ResolveKeys { values, keys });
let estimate_spec = RangeEstimateSpec::from(context);
let mut root = Self::push(
&mut nodes,
QueryPlanNode::Estimate {
input: resolved,
statistic: context.base.metadata.statistic_to_compute,
query_kwargs: context.base.metadata.query_kwargs.clone(),
output_labels: context.base.metadata.query_output_labels.clone(),
spec: estimate_spec.clone(),

@milindsrivastava1997 milindsrivastava1997 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: estimate_spec can be moved here once row_label_order is taken out first. Also, output_timestamps is now cloned into each SlidingExactCover strategy and into the spec, and explain() prints the full vector up to three times. On long range queries with debug logging on, debug!(plan = %plan.explain()) gets very long. Consider summarizing the timestamps in explain() (count, first, last) or sharing them behind an Arc.

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The timestamp and row-order lists are copied several times. compile_range copies output_timestamps into every Sliding StoreRead strategy and into the Estimate spec, and copies row_label_order into both the spec and LimitTopK. estimate_range_query then clones spec.row_label_order again (~line 2903), where a borrow would do.

A query with 10k steps ends up holding three copies of the timestamp vector per plan. Sharing them through an Arc<[u64]> or borrowing avoids this.

},
);
if options.limit_topk && context.base.metadata.statistic_to_compute == Statistic::Topk {
Expand All @@ -171,6 +228,7 @@ impl QueryPlan {
input: root,
k,
grouping_labels: context.base.grouping_labels.clone(),
row_label_order: estimate_spec.row_label_order.clone(),
},
);
}
Expand Down Expand Up @@ -285,62 +343,52 @@ impl QueryPlan {
fn push_read(
nodes: &mut Vec<QueryPlanNode>,
query: &StoreQueryParams,
window_type: WindowType,
role: StoreReadRole,
strategy: StoreReadStrategy,
) -> NodeId {
let strategy = match window_type {
WindowType::Tumbling => StoreReadStrategy::WindowGrid,
WindowType::Sliding => StoreReadStrategy::SlidingExactCover,
};
Self::push(
nodes,
QueryPlanNode::StoreRead {
query: query.clone(),
role,
strategy,
},
)
}

fn push_compose(
nodes: &mut Vec<QueryPlanNode>,
input: NodeId,
output_timestamps: &[u64],
lookback_ms: u64,
window_size_ms: u64,
bucket_step_ms: u64,
) -> NodeId {
Self::push(
nodes,
QueryPlanNode::ComposeWindows {
input,
output_timestamps: output_timestamps.to_vec(),
lookback_ms,
window_size_ms,
bucket_step_ms,
},
)
fn push_prepare_buckets(nodes: &mut Vec<QueryPlanNode>, input: NodeId) -> NodeId {
Self::push(nodes, QueryPlanNode::PrepareBuckets { input })
}

pub(crate) fn explain(&self) -> String {
let mut lines = Vec::with_capacity(self.nodes.len() + 1);
for (index, node) in self.nodes.iter().enumerate() {
let line = match node {
QueryPlanNode::StoreRead { query, strategy } => format!(
"n{index} StoreRead({strategy:?}, {}#{}, [{}, {}])",
QueryPlanNode::StoreRead { query, role, strategy } => format!(
"n{index} StoreRead({strategy:?}, role={role:?}, {}#{}, [{}, {}])",
query.metric, query.aggregation_id, query.start_timestamp, query.end_timestamp
),
QueryPlanNode::ComposeWindows { input, output_timestamps, lookback_ms, window_size_ms, bucket_step_ms } => format!(
"n{index} ComposeWindows(n{}, outputs={:?}, lookback={lookback_ms}ms, window={window_size_ms}ms, step={bucket_step_ms}ms)",
input.0, output_timestamps
),
QueryPlanNode::PrepareBuckets { input } => {
format!("n{index} PrepareBuckets(n{})", input.0)
}
QueryPlanNode::ResolveKeys { values, keys } => format!(
"n{index} ResolveKeys(values=n{}, keys={})",
values.0,
keys.map(|id| format!("n{}", id.0)).unwrap_or_else(|| "self".to_string())
),
QueryPlanNode::Estimate { input, statistic, query_kwargs, .. } => {
QueryPlanNode::Estimate {
input,
statistic,
query_kwargs,
spec,
..
} => {
let mut kwargs: Vec<_> = query_kwargs.iter().collect();
kwargs.sort_unstable_by_key(|(key, _)| *key);
format!("n{index} Estimate(n{}, {statistic}, {kwargs:?})", input.0)
format!(
"n{index} Estimate(n{}, {statistic}, {kwargs:?}, outputs={:?})",

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

explain() prints the full timestamp list. It prints all of output_timestamps for each Sliding StoreRead (through the Debug output of SlidingExactCover) and again for Estimate, and the native path logs the plan at debug level.

For a query with thousands of steps, one debug log line holds the list up to three times. Suggest printing a summary instead: the count plus the first and last timestamp.

input.0, spec.output_timestamps
)
},
QueryPlanNode::LimitTopK { input, k, .. } => {
format!("n{index} LimitTopK(n{}, k={k})", input.0)
Expand All @@ -359,11 +407,38 @@ impl QueryPlan {
}
}

impl From<&RangeQueryExecutionContext> for RangeEstimateSpec {
fn from(context: &RangeQueryExecutionContext) -> Self {
Self {
output_timestamps: context.output_timestamps.clone(),
query_range_ms: context.query_range_ms,
buckets_per_step: context.buckets_per_step,
lookback_bucket_count: context.lookback_bucket_count,
tumbling_window_ms: context.tumbling_window_ms,
window_type: context.window_type,
window_size_ms: context.window_size_ms,
keys_window_type: context.keys_window_type,
keys_window_size_ms: context.keys_window_size_ms,
keys_lookback_ms: context.keys_lookback_ms,
keys_tumbling_window_ms: context.keys_tumbling_window_ms,
value_aggregation_type: context.base.agg_info.aggregation_type_for_value,
key_aggregation_type: context.base.agg_info.aggregation_type_for_key,
grouping_labels: context.base.grouping_labels.clone(),
aggregated_labels: context.base.aggregated_labels.clone(),
row_label_order: crate::engines::simple_engine::SimpleEngine::topk_row_label_order(
&context.base.metadata,
&context.base.grouping_labels,
&context.base.aggregated_labels,
),
}
}
}

impl QueryPlanNode {
fn kind(&self) -> &'static str {
match self {
Self::StoreRead { .. } => "StoreRead",
Self::ComposeWindows { .. } => "ComposeWindows",
Self::PrepareBuckets { .. } => "PrepareBuckets",
Self::ResolveKeys { .. } => "ResolveKeys",
Self::Estimate { .. } => "Estimate",
Self::AggregateVector { .. } => "AggregateVector",
Expand All @@ -375,7 +450,7 @@ impl QueryPlanNode {
fn inputs(&self) -> Vec<NodeId> {
match self {
Self::StoreRead { .. } => Vec::new(),
Self::ComposeWindows { input, .. }
Self::PrepareBuckets { input, .. }
| Self::Estimate { input, .. }
| Self::AggregateVector { input, .. }
| Self::LimitTopK { input, .. }
Expand Down Expand Up @@ -473,7 +548,9 @@ mod tests {
.explain();

assert!(explanation.contains("n4 ResolveKeys(values=n1, keys=n3)"));
assert!(explanation.contains("n2 StoreRead(SlidingExactCover, requests#8"));
assert!(explanation.contains("role=Values"));
assert!(explanation.contains("role=Keys"));
assert!(explanation.contains("n2 StoreRead(SlidingExactCover {"));

@milindsrivastava1997 milindsrivastava1997 Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assertion doesn't check the key read's parameters (lookback, window, step). The fixture (L537) sets keys_window_type = Some(Sliding) but leaves keys_lookback_ms and the other key fields as None, which production builders never produce and legacy execution rejected. So the test passes on the silent unwrap_or fallbacks and locks in that invalid configuration. Suggest filling in all key fields in the fixture with values distinct from the value window, and asserting them in the key StoreRead.

assert!(explanation.ends_with("root: n5"));
}

Expand Down Expand Up @@ -531,6 +608,29 @@ mod tests {
assert!(explanation.contains("outputs=[1000, 2000, 3000]"));
}

#[test]

@milindsrivastava1997 milindsrivastava1997 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The lookback test is brittle, and the empty-read behavior is untested. sliding_value_read_uses_bucket_lookback checks the lookback only by substring-matching the Debug output of explain(), so a formatting change can break or pass it for unrelated reasons.

No native-path test checks that an empty Keys read succeeds while an empty Values read returns NoLocalData. The PR description says this behavior is preserved, but a regression in the role == StoreReadRole::Values check would go uncaught.

fn sliding_value_read_uses_bucket_lookback() {
let mut context = context();
context.window_type = WindowType::Sliding;
context.query_range_ms = 1_500;
context.lookback_bucket_count = 1;

let explanation = QueryPlan::compile_range(
&context,
PlanOptions {
limit_topk: false,
format_output: false,
},
&[],
)
.unwrap()
.explain();

// Store reads must match the estimator's whole-bucket lookback.
assert!(explanation.contains("lookback_ms: 1000"));
assert!(!explanation.contains("lookback_ms: 1500"));
}

#[test]
fn topk_formatting_is_the_plan_root() {
let mut context = context();
Expand Down Expand Up @@ -584,6 +684,7 @@ mod tests {
statistic: Statistic::Sum,
query_kwargs: HashMap::new(),
output_labels: KeyByLabelNames::empty(),
spec: RangeEstimateSpec::from(&context()),
}],
root: NodeId(0),
};
Expand Down Expand Up @@ -622,15 +723,10 @@ mod tests {
start_timestamp: 0,
end_timestamp: 1,
},
role: StoreReadRole::Values,
strategy: StoreReadStrategy::WindowGrid,
},
QueryPlanNode::ComposeWindows {
input: NodeId(0),
output_timestamps: vec![1],
lookback_ms: 1,
window_size_ms: 1,
bucket_step_ms: 1,
},
QueryPlanNode::PrepareBuckets { input: NodeId(0) },
],
root: NodeId(1),
};
Expand Down
Loading
Loading