From 3f2bab7f48d789616efdf705a0f255883a78bc83 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 13:00:58 -0400 Subject: [PATCH] refactor(query-engine): dedup range-context arithmetic between usual and binary-expr range paths Extracts the lookback/step calc, extended_store_plan mutation, and RangeQueryExecutionContext construction shared by build_arm_range_context and build_range_query_execution_context_from_parsed into one finish_range_context helper. Closes ProjectASAP/ASAPQuery#578 Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/promql.rs | 107 +++++++----------- 1 file changed, 38 insertions(+), 69 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 8b37f84..bd51a88 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -540,28 +540,21 @@ impl SimpleEngine { } } - /// Recursively builds a range execution context for one arm of a binary - /// arithmetic expression. + /// Extends an instant `QueryExecutionContext` into a `RangeQueryExecutionContext`: + /// computes the lookback window from the aggregation's tumbling window size, + /// validates the range params, and widens the store plan to cover + /// `[start - lookback, end]` as a range (non-exact) fetch. /// - /// Leaf resolution (Paren-unwrap + structural config lookup) is shared - /// with `evaluate_binary_arm` 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( + /// Shared tail for `build_arm_range_context` (one binary-expr arm) and + /// `build_range_query_execution_context_from_parsed` (the whole query) — + /// they differ only in how `base_context` itself is obtained. + fn finish_range_context( &self, - arm_ast: &promql_parser::parser::Expr, + base_context: QueryExecutionContext, start: f64, end: f64, step: f64, - ) -> Option<(RangeQueryExecutionContext, Vec)> { - use promql_parser::parser::Expr; - - if matches!(arm_ast, Expr::NumberLiteral(_)) { - return None; // caller handles scalars - } - - let (base_context, label_names) = self.resolve_arm_leaf_context(arm_ast, end)?; - + ) -> Option { 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; @@ -575,7 +568,7 @@ impl SimpleEngine { self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms) .map_err(|e| { - warn!("Range arm query validation failed: {}", e); + warn!("Range query validation failed: {}", e); e }) .ok()?; @@ -591,7 +584,7 @@ impl SimpleEngine { extended_store_plan.values_query.end_timestamp = end_ms; extended_store_plan.values_query.is_exact_query = false; - let range_context = RangeQueryExecutionContext { + Some(RangeQueryExecutionContext { base: QueryExecutionContext { store_plan: extended_store_plan, ..base_context @@ -604,7 +597,31 @@ impl SimpleEngine { buckets_per_step, lookback_bucket_count, tumbling_window_ms, - }; + }) + } + + /// Recursively builds a range execution context for one arm of a binary + /// arithmetic expression. + /// + /// Leaf resolution (Paren-unwrap + structural config lookup) is shared + /// with `evaluate_binary_arm` 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, + start: f64, + end: f64, + step: f64, + ) -> Option<(RangeQueryExecutionContext, Vec)> { + use promql_parser::parser::Expr; + + if matches!(arm_ast, Expr::NumberLiteral(_)) { + return None; // caller handles scalars + } + + let (base_context, label_names) = self.resolve_arm_leaf_context(arm_ast, end)?; + let range_context = self.finish_range_context(base_context, start, end, step)?; Some((range_context, label_names)) } @@ -1185,55 +1202,7 @@ impl SimpleEngine { // Use 'end' as the reference time for parsing 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); - let end_ms = Self::convert_query_time_to_data_time(end); - let step_ms = (step * 1000.0) as u64; - - // Get window size - let tumbling_window_ms = self - .streaming_config - .read() - .unwrap() - .get_aggregation_config(base_context.agg_info.aggregation_id_for_value) - .map(|config| config.window_size_ms)?; - - // Validate parameters - self.validate_range_query_params(start_ms, end_ms, step_ms, tumbling_window_ms) - .map_err(|e| { - warn!("Range query validation failed: {}", e); - e - }) - .ok()?; - - // Calculate lookback from the base context's store plan - 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; - - // Modify the store plan to cover the entire range - 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; - // Range queries always use range fetch, not exact - extended_store_plan.values_query.is_exact_query = false; - - Some(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, - }) + self.finish_range_context(base_context, start, end, step) } /// Main entry point for range queries