Skip to content

refactor(query-engine): dedup instant/range PromQL arm-building and parsing in promql.rs #517

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

Merged
milindsrivastava1997 merged 1 commit into from
Jul 8, 2026
Merged

refactor(query-engine): dedup instant/range PromQL arm-building and parsing in promql.rs #517

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
282 changes: 178 additions & 104 deletions asap-query-engine/src/engines/simple_engine/promql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,22 @@ use std::collections::HashMap;
use std::time::Instant;
use tracing::{debug, warn};

/// Detects whether either side of a PromQL binary expression is a scalar
/// (numeric literal), returning the scalar value, the other (vector) arm,
/// and whether the scalar was on the left. Shared by instant and range
/// binary-expression handling.
fn detect_scalar_arm<'a>(
lhs: &'a promql_parser::parser::Expr,
rhs: &'a promql_parser::parser::Expr,
) -> Option<(f64, &'a promql_parser::parser::Expr, bool)> {
use promql_parser::parser::Expr;
match (lhs, rhs) {
(_, Expr::NumberLiteral(nl)) => Some((nl.val, lhs, false)),
(Expr::NumberLiteral(nl), _) => Some((nl.val, rhs, true)),
_ => None,
}
}

impl SimpleEngine {
/// Calculates start timestamp for PromQL queries
fn calculate_start_timestamp_promql(
Expand Down Expand Up @@ -346,11 +362,37 @@ impl SimpleEngine {
})
}

/// Recursively unwraps `Paren`, then structurally resolves a leaf PromQL
/// arm (i.e. not `Binary` or `NumberLiteral`) to its `QueryConfig` and
/// base `QueryExecutionContext`. Shared leaf-resolution step for both
/// `build_arm_logical_plan` (instant) and `build_arm_range_context` (range).
///
/// Returns `None` for `Binary` arms (caller handles recursion) and
/// `NumberLiteral` arms (caller handles scalars).
fn resolve_arm_leaf_context(
&self,
arm_ast: &promql_parser::parser::Expr,
time: f64,
) -> Option<(QueryExecutionContext, Vec<String>)> {
use promql_parser::parser::Expr;

match arm_ast {
Expr::NumberLiteral(_) | Expr::Binary(_) => None,
Expr::Paren(paren) => self.resolve_arm_leaf_context(&paren.expr, time),
other => {
let config = self.find_query_config_promql_structural(other)?;
let ctx = self.build_query_execution_context_from_ast(other, &config, time)?;
let label_names = ctx.metadata.query_output_labels.labels.clone();
Some((ctx, label_names))
}
}
}

/// Recursively builds a DataFusion logical plan for one arm of a binary
/// arithmetic expression.
///
/// - Leaf arm (supported PromQL pattern): look up config structurally, build
/// context, return its `to_logical_plan()` together with the output label names.
/// - Leaf arm (supported PromQL pattern): resolved via `resolve_arm_leaf_context`,
/// returning its `to_logical_plan()` together with the output label names.
/// - Binary arm: recursively build both sub-arms and combine with
/// `build_binary_vector_plan`.
/// - Scalar literal: returns `None` (handled by the caller separately).
Expand All @@ -374,11 +416,8 @@ impl SimpleEngine {
.ok()?;
Some((combined, lhs_labels))
}
other => {
// Leaf pattern: structural config lookup + context + plan
let config = self.find_query_config_promql_structural(other)?;
let ctx = self.build_query_execution_context_from_ast(other, &config, time)?;
let label_names = ctx.metadata.query_output_labels.labels.clone();
_ => {
let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?;
let plan = ctx.to_logical_plan().ok()?;
Some((plan, label_names))
}
Expand Down Expand Up @@ -408,13 +447,7 @@ impl SimpleEngine {
let rhs = binary.rhs.as_ref();
let op = &binary.op;

// Scalar case: either side may be a numeric literal
let scalar_case: Option<(f64, &Expr, bool)> = match (lhs, rhs) {
(_, Expr::NumberLiteral(nl)) => Some((nl.val, lhs, false)),
(Expr::NumberLiteral(nl), _) => Some((nl.val, rhs, true)),
_ => None,
};
if let Some((scalar, vector_arm, scalar_on_left)) = scalar_case {
if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) {
let (vector_plan, label_names) = self.build_arm_logical_plan(vector_arm, time)?;
let combined =
build_scalar_plan(vector_plan, scalar, op, scalar_on_left, label_names.clone())
Expand Down Expand Up @@ -469,7 +502,13 @@ impl SimpleEngine {
}
}

/// Recursively builds a range execution context for one arm of a binary arithmetic expression.
/// Recursively builds a range execution context for one arm of a binary
/// arithmetic expression.
///
/// Leaf resolution (Paren-unwrap + structural config lookup) is shared
/// with `build_arm_logical_plan` via `resolve_arm_leaf_context`. Note this
/// does not support nested `Binary` arms (e.g. `(a+b)*c` over a range) —
/// tracked separately in #516.
fn build_arm_range_context(
&self,
arm_ast: &promql_parser::parser::Expr,
Expand All @@ -479,63 +518,57 @@ impl SimpleEngine {
) -> Option<(RangeQueryExecutionContext, Vec<String>)> {
use promql_parser::parser::Expr;

match arm_ast {
Expr::NumberLiteral(_) => None, // caller handles scalars
Expr::Paren(paren) => self.build_arm_range_context(&paren.expr, start, end, step),
other => {
let config = self.find_query_config_promql_structural(other)?;
let base_context =
self.build_query_execution_context_from_ast(other, &config, end)?;
let label_names = base_context.metadata.query_output_labels.labels.clone();
if matches!(arm_ast, Expr::NumberLiteral(_)) {
return None; // caller handles scalars
}

let start_ms = Self::convert_query_time_to_data_time(start);
let end_ms = Self::convert_query_time_to_data_time(end);
let step_ms = (step * 1000.0) as u64;
let (base_context, label_names) = self.resolve_arm_leaf_context(arm_ast, end)?;

let tumbling_window_ms = self
.streaming_config
.read()
.unwrap()
.get_aggregation_config(base_context.agg_info.aggregation_id_for_value)
.map(|c| c.window_size_ms)?;

self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms)
.map_err(|e| {
warn!("Range arm query validation failed: {}", e);
e
})
.ok()?;
let start_ms = Self::convert_query_time_to_data_time(start);
let end_ms = Self::convert_query_time_to_data_time(end);
let step_ms = (step * 1000.0) as u64;

let lookback_ms = base_context.store_plan.values_query.end_timestamp
- base_context.store_plan.values_query.start_timestamp;

let buckets_per_step = (step_ms / tumbling_window_ms) as usize;
let lookback_bucket_count = (lookback_ms / tumbling_window_ms) as usize;

let mut extended_store_plan = base_context.store_plan.clone();
extended_store_plan.values_query.start_timestamp =
start_ms.saturating_sub(lookback_ms);
extended_store_plan.values_query.end_timestamp = end_ms;
extended_store_plan.values_query.is_exact_query = false;

let range_context = RangeQueryExecutionContext {
base: QueryExecutionContext {
store_plan: extended_store_plan,
..base_context
},
range_params: RangeQueryParams {
start: start_ms,
end: end_ms,
step: step_ms,
},
buckets_per_step,
lookback_bucket_count,
tumbling_window_ms,
};

Some((range_context, label_names))
}
}
let tumbling_window_ms = self
.streaming_config
.read()
.unwrap()
.get_aggregation_config(base_context.agg_info.aggregation_id_for_value)
.map(|c| c.window_size_ms)?;

self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms)
.map_err(|e| {
warn!("Range arm query validation failed: {}", e);
e
})
.ok()?;

let lookback_ms = base_context.store_plan.values_query.end_timestamp
- base_context.store_plan.values_query.start_timestamp;

let buckets_per_step = (step_ms / tumbling_window_ms) as usize;
let lookback_bucket_count = (lookback_ms / tumbling_window_ms) as usize;

let mut extended_store_plan = base_context.store_plan.clone();
extended_store_plan.values_query.start_timestamp = start_ms.saturating_sub(lookback_ms);
extended_store_plan.values_query.end_timestamp = end_ms;
extended_store_plan.values_query.is_exact_query = false;

let range_context = RangeQueryExecutionContext {
base: QueryExecutionContext {
store_plan: extended_store_plan,
..base_context
},
range_params: RangeQueryParams {
start: start_ms,
end: end_ms,
step: step_ms,
},
buckets_per_step,
lookback_bucket_count,
tumbling_window_ms,
};

Some((range_context, label_names))
}

/// Handles a binary arithmetic PromQL expression for range queries.
Expand All @@ -561,13 +594,7 @@ impl SimpleEngine {
let rhs = binary.rhs.as_ref();
let op = &binary.op;

// Scalar case: either side may be a numeric literal
let scalar_case: Option<(f64, &Expr, bool)> = match (lhs, rhs) {
(_, Expr::NumberLiteral(nl)) => Some((nl.val, lhs, false)),
(Expr::NumberLiteral(nl), _) => Some((nl.val, rhs, true)),
_ => None,
};
if let Some((scalar, vector_arm, scalar_on_left)) = scalar_case {
if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) {
let (ctx, labels) = self.build_arm_range_context(vector_arm, start, end, step)?;
let results = self.execute_range_query_pipeline(&ctx).ok()?;
let combined: Vec<RangeVectorElement> = results
Expand Down Expand Up @@ -869,21 +896,27 @@ impl SimpleEngine {
let query_start_time = Instant::now();
debug!("Handling query: {} at time {}", query, time);

let ast = match promql_parser::parser::parse(&query) {
Ok(ast) => ast,
Err(e) => {
warn!("Failed to parse PromQL query '{}': {}", query, e);
return None;
}
};

// Check for binary arithmetic before attempting single-query dispatch.
// Binary expressions won't have a matching query_config, so we handle them here.
if let Ok(ast) = promql_parser::parser::parse(&query) {
if matches!(&ast, promql_parser::parser::Expr::Binary(_)) {
let result = self.handle_binary_expr_promql(&ast, time);
let total_query_duration = query_start_time.elapsed();
debug!(
"Binary arithmetic query handling took: {:.2}ms",
total_query_duration.as_secs_f64() * 1000.0
);
return result;
}
if matches!(&ast, promql_parser::parser::Expr::Binary(_)) {
let result = self.handle_binary_expr_promql(&ast, time);
let total_query_duration = query_start_time.elapsed();
debug!(
"Binary arithmetic query handling took: {:.2}ms",
total_query_duration.as_secs_f64() * 1000.0
);
return result;
}

let context = self.build_query_execution_context_promql(query, time)?;
let context = self.build_query_execution_context_from_parsed(&ast, &query, time)?;

debug!(
"Querying store for metric: {}, aggregation_id: {}, range: [{}, {}]",
Expand Down Expand Up @@ -974,8 +1007,6 @@ impl SimpleEngine {
query: String,
time: f64,
) -> Option<QueryExecutionContext> {
let query_time = Self::convert_query_time_to_data_time(time);

// Parse PromQL AST using promql-parser crate
let parse_start_time = Instant::now();
let ast = match promql_parser::parser::parse(&query) {
Expand All @@ -993,9 +1024,23 @@ impl SimpleEngine {
}
};

self.build_query_execution_context_from_parsed(&ast, &query, time)
}

/// Variant of `build_query_execution_context_promql` that accepts an
/// already-parsed AST, avoiding a redundant re-parse when the caller
/// (e.g. `handle_query_promql`) has already parsed `query` once.
fn build_query_execution_context_from_parsed(
&self,
ast: &promql_parser::parser::Expr,
query: &str,
time: f64,
) -> Option<QueryExecutionContext> {
let query_time = Self::convert_query_time_to_data_time(time);

let pattern_match_start_time = Instant::now();

let found_match = self.find_matching_controller_pattern(&ast, &query);
let found_match = self.find_matching_controller_pattern(ast, query);

let (query_pattern_type, match_result) = match found_match {
Some((pt, result)) => {
Expand All @@ -1017,7 +1062,7 @@ impl SimpleEngine {
let query_context_start_time = Instant::now();

// Resolve aggregation: try pre-configured query_configs first, fall back to capability matching.
let agg_info: AggregationIdInfo = if let Some(config) = self.find_query_config(&query) {
let agg_info: AggregationIdInfo = if let Some(config) = self.find_query_config(query) {
self.get_aggregation_id_info(&config)
.map_err(|e| {
warn!("{}", e);
Expand All @@ -1036,7 +1081,7 @@ impl SimpleEngine {
_ => &empty_schema,
};
let requirements = build_query_requirements_promql(
&query,
query,
&match_result,
query_pattern_type,
metric_schema,
Expand Down Expand Up @@ -1072,10 +1117,32 @@ impl SimpleEngine {
start: f64,
end: f64,
step: f64,
) -> Option<RangeQueryExecutionContext> {
let ast = match promql_parser::parser::parse(&query) {
Ok(ast) => ast,
Err(e) => {
warn!("Failed to parse PromQL query '{}': {}", query, e);
return None;
}
};

self.build_range_query_execution_context_from_parsed(&ast, &query, start, end, step)
}

/// Variant of `build_range_query_execution_context_promql` that accepts an
/// already-parsed AST, avoiding a redundant re-parse when the caller
/// (e.g. `handle_range_query_promql`) has already parsed `query` once.
fn build_range_query_execution_context_from_parsed(
&self,
ast: &promql_parser::parser::Expr,
query: &str,
start: f64,
end: f64,
step: f64,
) -> Option<RangeQueryExecutionContext> {
// First, build the base instant query context (reuse existing logic)
// Use 'end' as the reference time for parsing
let base_context = self.build_query_execution_context_promql(query, end)?;
let base_context = self.build_query_execution_context_from_parsed(ast, query, end)?;

// Convert to milliseconds
let start_ms = Self::convert_query_time_to_data_time(start);
Expand Down Expand Up @@ -1142,20 +1209,27 @@ impl SimpleEngine {
query, start, end, step
);

// Check for binary arithmetic before attempting single-query dispatch.
if let Ok(ast) = promql_parser::parser::parse(&query) {
if matches!(&ast, promql_parser::parser::Expr::Binary(_)) {
let result = self.handle_binary_expr_range_promql(&ast, start, end, step);
let total_duration = query_start_time.elapsed();
debug!(
"Binary arithmetic range query handling took: {:.2}ms",
total_duration.as_secs_f64() * 1000.0
);
return result;
let ast = match promql_parser::parser::parse(&query) {
Ok(ast) => ast,
Err(e) => {
warn!("Failed to parse PromQL query '{}': {}", query, e);
return None;
}
};

// Check for binary arithmetic before attempting single-query dispatch.
if matches!(&ast, promql_parser::parser::Expr::Binary(_)) {
let result = self.handle_binary_expr_range_promql(&ast, start, end, step);
let total_duration = query_start_time.elapsed();
debug!(
"Binary arithmetic range query handling took: {:.2}ms",
total_duration.as_secs_f64() * 1000.0
);
return result;
}

let context = self.build_range_query_execution_context_promql(query, start, end, step)?;
let context =
self.build_range_query_execution_context_from_parsed(&ast, &query, start, end, step)?;

// Execute range query pipeline
let results: Vec<RangeVectorElement> = self
Expand Down
Loading
Loading