Skip to content
32 changes: 18 additions & 14 deletions asap-common/dependencies/rs/asap_types/src/query_requirements.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use promql_utilities::ast_matching::PromQLMatchResult;
use promql_utilities::data_model::KeyByLabelNames;
use promql_utilities::query_logics::enums::{QueryPatternType, Statistic};
use promql_utilities::query_logics::enums::Statistic;
use promql_utilities::query_logics::parsing::{
get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute,
};
Expand Down Expand Up @@ -46,13 +46,12 @@ pub struct QueryRequirements {
pub fn build_query_requirements_promql(
query: &str,
match_result: &PromQLMatchResult,
pattern_type: QueryPatternType,
metric_schema: &PromQLSchema,
data_ingestion_interval_ms: u64,
) -> Option<QueryRequirements> {
let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);

let statistics = get_statistics_to_compute(pattern_type, match_result)
let statistics = get_statistics_to_compute(match_result)
.map_err(|err| {
warn!(
query = %query,
Expand All @@ -63,28 +62,33 @@ pub fn build_query_requirements_promql(
})
.ok()?;

let data_range_ms = match pattern_type {
QueryPatternType::OnlySpatial => data_ingestion_interval_ms,
let has_temporal_function = match_result.tokens.contains_key("function");
let has_aggregation = match_result.tokens.contains_key("aggregation");

let data_range_ms = if has_temporal_function {
// promql-parser supports a literal `ms` duration suffix (e.g. `[500ms]`),
// so .num_seconds() would truncate sub-second ranges to 0.
_ => match_result
match_result
.get_range_duration()
.map(|d| d.num_milliseconds() as u64)?,
.map(|d| d.num_milliseconds() as u64)?
} else {
// OnlySpatial (no temporal component): the query has no range of its
// own, so its data range is exactly one scrape interval.
data_ingestion_interval_ms
};

let all_labels = metric_schema
.get_labels(&metric)
.cloned()
.unwrap_or_else(KeyByLabelNames::empty);

let grouping_labels = match pattern_type {
let grouping_labels = if has_aggregation {
// OnlySpatial and (collapsable, see #508) OneTemporalOneSpatial encode
// their output labels in the AST's `by (...)` / `without (...)` clause.
get_spatial_aggregation_output_labels(match_result, &all_labels)
} else {
// OnlyTemporal preserves all labels.
QueryPatternType::OnlyTemporal => all_labels,
// OnlySpatial and OneTemporalOneSpatial encode their output labels in
// the AST's `by (...)` / `without (...)` clause.
QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => {
get_spatial_aggregation_output_labels(match_result, &all_labels)
}
all_labels
};

Some(QueryRequirements {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,24 +3,6 @@ use std::fmt;
use std::str::FromStr;
use tracing::debug;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum QueryPatternType {
OnlyTemporal,
OnlySpatial,
OneTemporalOneSpatial,
}

impl std::fmt::Display for QueryPatternType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
debug!("Formatting QueryPatternType: {:?}", self);
match self {
QueryPatternType::OnlyTemporal => write!(f, "only_temporal"),
QueryPatternType::OnlySpatial => write!(f, "only_spatial"),
QueryPatternType::OneTemporalOneSpatial => write!(f, "one_temporal_one_spatial"),
}
}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum QueryTreatmentType {
Exact,
Expand Down Expand Up @@ -69,6 +51,22 @@ impl std::fmt::Display for Statistic {

#[allow(clippy::should_implement_trait)]
impl Statistic {
/// Returns `true` for statistics whose result requires approximate
/// pre-aggregation, regardless of whether they were reached via a
/// temporal function (`PromQLFunction::is_approximate`) or a spatial
/// aggregation operator (`AggregationOperator::is_approximate`) — for
/// every `Statistic` reachable from either origin, the two origins agree.
pub fn is_approximate(self) -> bool {
matches!(
self,
Statistic::Count
| Statistic::Sum
| Statistic::Cardinality
| Statistic::Quantile
| Statistic::Topk
)
}

pub fn from_str(s: &str) -> Option<Self> {
debug!("Parsing Statistic from string: {}", s);
match s.to_lowercase().as_str() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,22 @@ use tracing::debug;
use crate::ast_matching::promql_pattern::AggregationModifierType;
use crate::ast_matching::PromQLMatchResult;
use crate::data_model::KeyByLabelNames;
use crate::query_logics::enums::{AggregationOperator, QueryPatternType, Statistic};
use crate::query_logics::enums::{AggregationOperator, PromQLFunction, Statistic};
use crate::query_logics::logics::get_is_collapsable;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StatisticExtractionError {
MissingStatistic { pattern_type: QueryPatternType },
MissingStatistic,
UnsupportedStatistic { statistic: String },
}

impl std::fmt::Display for StatisticExtractionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingStatistic { pattern_type } => {
Self::MissingStatistic => {
write!(
f,
"No statistic found for query pattern type {pattern_type:?}"
"No temporal function or aggregation operation found in match result"
)
}
Self::UnsupportedStatistic { statistic } => {
Expand Down Expand Up @@ -72,28 +73,66 @@ pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (Strin
(metric_name, spatial_filter)
}

/// Get statistics to compute based on pattern type and tokens.
/// Get statistics to compute from a matched query's tokens.
///
/// Explicitly handles the three reachable shapes:
/// - Only a temporal function (`"function"` token, no `"aggregation"`): the
/// statistic comes from the function name.
/// - Only a spatial aggregation (`"aggregation"` token, no `"function"`): the
/// statistic comes from the aggregation operator.
/// - Both (a spatial aggregation wrapping a temporal function): only ever
/// reachable via a pattern already narrowed to a collapsable `(function,
/// op)` pair (see `get_is_collapsable`, and #508's pattern-narrowing fix
/// that makes non-collapsable combinations fail to match at all) — asserted
/// below rather than silently trusted. The statistic still comes from the
/// *function*, never the outer op: e.g. `count_over_time` + `sum` needs a
/// `Count` accumulator, not a `Sum` one — summing per-series counts gives
/// the group's total count, so the outer op only describes how per-series
/// results combine, never which statistic must be precomputed.
///
/// Returns a typed error if the matched statistic/function name is not
/// recognized, so callers can decide whether to skip or fail the query.
pub fn get_statistics_to_compute(
pattern_type: QueryPatternType,
match_result: &PromQLMatchResult,
) -> Result<Vec<Statistic>, StatisticExtractionError> {
debug!("Computing statistics for pattern type {:?}", pattern_type);
let statistic_to_compute: Option<String> = match pattern_type {
QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => {
match_result.get_function_name().map(|function_name| {
let name = function_name.to_lowercase();
name.split('_').next().unwrap_or(&name).to_string()
})
}
QueryPatternType::OnlySpatial => match_result
let has_function = match_result.tokens.contains_key("function");
let has_aggregation = match_result.tokens.contains_key("aggregation");
debug!("Computing statistics (has_function={has_function}, has_aggregation={has_aggregation})");

let function_statistic = |match_result: &PromQLMatchResult| {
match_result.get_function_name().map(|function_name| {
let name = function_name.to_lowercase();
name.split('_').next().unwrap_or(&name).to_string()
})
};

let statistic_to_compute: Option<String> = if has_function && has_aggregation {
debug_assert!(
match_result
.get_function_name()
.and_then(|f| f.parse::<PromQLFunction>().ok())
.zip(
match_result
.get_aggregation_op()
.and_then(|o| o.parse::<AggregationOperator>().ok())
)
.is_some_and(|(f, o)| get_is_collapsable(f, o)),
"a match with both function and aggregation tokens must be collapsable \
(patterns are narrowed to only collapsable pairs, see #508)"
);
function_statistic(match_result)
} else if has_function {
function_statistic(match_result)
} else if has_aggregation {
match_result
.get_aggregation_op()
.map(|agg| agg.to_lowercase()),
.map(|agg| agg.to_lowercase())
} else {
None
};

let Some(statistic_to_compute) = statistic_to_compute else {
return Err(StatisticExtractionError::MissingStatistic { pattern_type });
return Err(StatisticExtractionError::MissingStatistic);
};

debug!("Found statistic to compute: {}", statistic_to_compute);
Expand Down Expand Up @@ -194,9 +233,7 @@ mod tests {

#[test]
fn unsupported_matched_temporal_statistic_returns_typed_error() {
let err =
get_statistics_to_compute(QueryPatternType::OnlyTemporal, &temporal_match("stddev"))
.unwrap_err();
let err = get_statistics_to_compute(&temporal_match("stddev")).unwrap_err();

assert_eq!(
err,
Expand All @@ -208,17 +245,9 @@ mod tests {

#[test]
fn missing_statistic_returns_typed_error() {
let err = get_statistics_to_compute(
QueryPatternType::OnlyTemporal,
&PromQLMatchResult::with_tokens(HashMap::new()),
)
.unwrap_err();
let err =
get_statistics_to_compute(&PromQLMatchResult::with_tokens(HashMap::new())).unwrap_err();

assert_eq!(
err,
StatisticExtractionError::MissingStatistic {
pattern_type: QueryPatternType::OnlyTemporal
}
);
assert_eq!(err, StatisticExtractionError::MissingStatistic);
}
}
5 changes: 2 additions & 3 deletions asap-planner-rs/src/optimizer/aqe_extractor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,10 +170,10 @@ fn extract_requirements(
let ast = promql_parser::parser::parse(query).ok()?;
let patterns = build_patterns();

let (pattern_type, match_result) = patterns.iter().find_map(|(pt, pat)| {
let match_result = patterns.iter().find_map(|pat| {
let r = pat.matches(&ast);
if r.matches {
Some((*pt, r))
Some(r)
} else {
None
}
Expand All @@ -182,7 +182,6 @@ fn extract_requirements(
build_query_requirements_promql(
query,
&match_result,
pattern_type,
metric_schema,
data_ingestion_interval_ms,
)
Expand Down
2 changes: 1 addition & 1 deletion asap-planner-rs/src/optimizer/greedy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ mod tests {
);

let mut seen_ids: StdHashMap<u64, ()> = StdHashMap::new();
for (id, _) in solution.deployed_configs.iter() {
for id in solution.deployed_configs.keys() {
assert!(
seen_ids.insert(*id, ()).is_none(),
"duplicate aggregation_id"
Expand Down
Loading
Loading